[air/execution] Fix new execution backend for BOHB (#34828)

The new execution backend did not work with BOHB. The reason was a faulty check when we start actors. We should not eagerly start paused trials, as the scheduler may want to keep them paused for synchronous training. However, a leftover code piece allowed for trials that were once pending but then paused to be started. This PR fixes the bug by requiring trials to be strictly pending, and also makes sure that the `trial_to_run` is set to PENDING when chosen by the scheduler. BOHB tests are now enabled for the new execution backend. Signed-off-by: Kai Fricke <[email protected]>
ray-project · Apr 28, 2023 · 92f4fe7 · 92f4fe7
1 parent 922e2c8
commit 92f4fe7
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 10 deletions.
diff --git a/python/ray/tune/execution/tune_controller.py b/python/ray/tune/execution/tune_controller.py
@@ -437,6 +437,7 @@ def _maybe_add_actors(self) -> None:
                 and trial_to_run not in self._trial_to_actor
             ):
                 logger.debug(f"Staging trial to run: {trial_to_run}")
+                self._set_trial_status(trial_to_run, Trial.PENDING)
                 self._staged_trials.add(trial_to_run)
                 self._actor_cache.increase_max(trial_to_run.placement_group_factory)
                 # schedule_trial_actor also potentially uses cached actors
@@ -460,7 +461,7 @@ def _maybe_add_actors(candidates: List[Trial]):
                 # If the trial is part of the list, but not of the set,
                 # we just ignore it. Removing it from the list on status
                 # change is too expensive.
-                if trial not in (self._pending_trials | self._paused_trials):
+                if trial not in self._pending_trials:
                     continue
 
                 if trial in self._trial_to_actor:
@@ -541,7 +542,7 @@ def _schedule_trial_actor(self, trial: Trial):
         """
         logger.debug(f"Trying to schedule new ACTOR for trial {trial}")
 
-        self._set_trial_status(trial, Trial.PENDING)
+        assert trial.status == Trial.PENDING
 
         trial.init_logdir()
         # We checkpoint metadata here to try mitigating logdir duplication

diff --git a/python/ray/tune/tests/test_trial_scheduler.py b/python/ray/tune/tests/test_trial_scheduler.py
@@ -807,10 +807,6 @@ def result(score, ts):
             [t.status for t in trials], [Trial.PAUSED, Trial.PENDING, Trial.PAUSED]
         )
 
-    @pytest.mark.skipif(
-        os.environ.get("TUNE_NEW_EXECUTION") == "1",
-        reason="BOHB does not currently work with the new execution backend.",
-    )
     def testNonstopBOHB(self):
         from ray.tune.search.bohb import TuneBOHB
 

diff --git a/python/ray/tune/tests/test_tune_restore_warm_start.py b/python/ray/tune/tests/test_tune_restore_warm_start.py
@@ -470,10 +470,6 @@ def cost(space, reporter):
         return search_alg, cost
 
 
-@pytest.mark.skipif(
-    os.environ.get("TUNE_NEW_EXECUTION") == "1",
-    reason="BOHB does not currently work with the new execution backend.",
-)
 class BOHBWarmStartTest(AbstractWarmStartTest, unittest.TestCase):
     def set_basic_conf(self):
         space = {"width": tune.uniform(0, 20), "height": tune.uniform(-100, 100)}