[tune] Fix trial runner/controller whitelist attributes

python/ray/tune/execution/trial_runner.py

@@ -1314,8 +1314,20 @@ def _wrapped(self):
  return TrialRunnerWrapper(
  self,
  self.trial_executor,
- runner_whitelist_attr={"search_alg", "get_trials", "_set_trial_status"},
- executor_whitelist_attr={"has_resources_for_trial", "pause_trial", "save"},
+ runner_whitelist_attr={
+ "search_alg",
+ "get_trials",
+ "get_live_trials",
+ "_set_trial_status",
+ "pause_trial",
+ "stop_trial",
+ },
+ executor_whitelist_attr={
+ "has_resources_for_trial",
+ "pause_trial",
+ "save",
+ "_resource_updater",
+ },
  )
 
  def update_max_pending_trials(self, max_pending_trials: Optional[int] = None):
@@ -1545,6 +1557,9 @@ def __init__(
  self._trial_executor = trial_executor
  self._whitelist_attr = whitelist_attr or set()
 
+ for attr in self._whitelist_attr:
+ assert hasattr(self._trial_executor, attr)
+
  def __getattr__(self, attr):
  if attr not in self._whitelist_attr:
  if log_once("restrict_accessing_trial_executor"):
@@ -1584,6 +1599,9 @@ def __init__(
  )
  self._runner_whitelist_attr = runner_whitelist_attr or set()
 
+ for attr in self._runner_whitelist_attr:
+ assert hasattr(self, attr)
+
  def __getattr__(self, attr):
  if attr == self._EXECUTOR_ATTR:
  return self._trial_executor

python/ray/tune/execution/tune_controller.py

@@ -154,8 +154,20 @@ def _wrapped(self):
  return TrialRunnerWrapper(
  self,
  trial_executor=_FakeRayTrialExecutor(self),
- runner_whitelist_attr={"search_alg", "get_trials", "_set_trial_status"},
- executor_whitelist_attr={"has_resources_for_trial", "pause_trial", "save"},
+ runner_whitelist_attr={
+ "search_alg",
+ "get_trials",
+ "get_live_trials",
+ "_set_trial_status",
+ "pause_trial",
+ "stop_trial",
+ },
+ executor_whitelist_attr={
+ "has_resources_for_trial",
+ "pause_trial",
+ "save",
+ "_resource_updater",
+ },
  )
 
  def _used_resources_string(self) -> str:

python/ray/tune/schedulers/hb_bohb.py

@@ -109,16 +109,16 @@ def on_trial_result(
  # as intended.
  # There should be a better API for this.
  # TODO(team-ml): Refactor alongside HyperBandForBOHB
- trial_runner._search_alg.searcher.on_pause(trial.trial_id)
+ trial_runner.search_alg.searcher.on_pause(trial.trial_id)
  return TrialScheduler.PAUSE
  action = self._process_bracket(trial_runner, bracket)
  if action == TrialScheduler.PAUSE:
- trial_runner._search_alg.searcher.on_pause(trial.trial_id)
+ trial_runner.search_alg.searcher.on_pause(trial.trial_id)
  return action
 
  def _unpause_trial(self, trial_runner: "trial_runner.TrialRunner", trial: Trial):
  # Hack. See comment in on_trial_result
- trial_runner._search_alg.searcher.on_unpause(trial.trial_id)
+ trial_runner.search_alg.searcher.on_unpause(trial.trial_id)
 
  def choose_trial_to_run(
  self, trial_runner: "trial_runner.TrialRunner", allow_recurse: bool = True

python/ray/tune/tests/test_trial_scheduler.py

@@ -733,8 +733,8 @@ def result(score, ts):
  metric="episode_reward_mean", mode="max", max_t=3, reduction_factor=3
  )
  runner = _MockTrialRunner(sched)
- runner._search_alg = MagicMock()
- runner._search_alg.searcher = MagicMock()
+ runner.search_alg = MagicMock()
+ runner.search_alg.searcher = MagicMock()
  trials = [Trial("__fake") for i in range(3)]
  for t in trials:
  runner.add_trial(t)
@@ -748,8 +748,8 @@ def result(score, ts):
  decision = sched.on_trial_result(runner, trials[-1], spy_result)
  self.assertEqual(decision, TrialScheduler.STOP)
  sched.choose_trial_to_run(runner)
- self.assertEqual(runner._search_alg.searcher.on_pause.call_count, 2)
- self.assertEqual(runner._search_alg.searcher.on_unpause.call_count, 1)
+ self.assertEqual(runner.search_alg.searcher.on_pause.call_count, 2)
+ self.assertEqual(runner.search_alg.searcher.on_unpause.call_count, 1)
  self.assertTrue("hyperband_info" in spy_result)
  self.assertEqual(spy_result["hyperband_info"]["budget"], 1)
 
@@ -761,8 +761,8 @@ def result(score, ts):
  metric="episode_reward_mean", mode="min", max_t=3, reduction_factor=3
  )
  runner = _MockTrialRunner(sched)
- runner._search_alg = MagicMock()
- runner._search_alg.searcher = MagicMock()
+ runner.search_alg = MagicMock()
+ runner.search_alg.searcher = MagicMock()
  trials = [Trial("__fake") for i in range(3)]
  for t in trials:
  runner.add_trial(t)
@@ -776,7 +776,7 @@ def result(score, ts):
  decision = sched.on_trial_result(runner, trials[-1], spy_result)
  self.assertEqual(decision, TrialScheduler.CONTINUE)
  sched.choose_trial_to_run(runner)
- self.assertEqual(runner._search_alg.searcher.on_pause.call_count, 2)
+ self.assertEqual(runner.search_alg.searcher.on_pause.call_count, 2)
  self.assertTrue("hyperband_info" in spy_result)
  self.assertEqual(spy_result["hyperband_info"]["budget"], 1)
 
@@ -788,8 +788,8 @@ def result(score, ts):
  metric="episode_reward_mean", mode="min", max_t=10, reduction_factor=3
  )
  runner = _MockTrialRunner(sched)
- runner._search_alg = MagicMock()
- runner._search_alg.searcher = MagicMock()
+ runner.search_alg = MagicMock()
+ runner.search_alg.searcher = MagicMock()
  trials = [Trial("__fake") for i in range(3)]
  for t in trials:
  runner.add_trial(t)
@@ -2344,8 +2344,8 @@ def testBOHBNanInf(self):
  scheduler = HyperBandForBOHB(metric="episode_reward_mean", mode="max")
 
  runner = _MockTrialRunner(scheduler)
- runner._search_alg = MagicMock()
- runner._search_alg.searcher = MagicMock()
+ runner.search_alg = MagicMock()
+ runner.search_alg.searcher = MagicMock()
 
  t1, t2, t3 = self.nanInfSetup(scheduler, runner)
  # skip trial complete in this mock setting