[tune] Fix trial runner/controller whitelist attributes (ray-project#…

…35769)

With PBT and BOHB, we currently can get these warnings:

```
2023-05-24 04:03:59,962 WARNING trial_runner.py:1543 -- You are trying to access pause_trial interface of TrialRunner in TrialScheduler, which is being restricted. If you believe it is reasonable for your scheduler to access this TrialRunner API, please reach out to Ray team on GitHub. A more strict API access pattern would be enforced starting 1.12s.0
```

While technically we should have seen those warnings before (in the old execution path), it looks like they only come up ever since we activated the new execution path.

This PR adds the rest of those attributes we access in our tune-provided schedulers to the whitelist to get rid of the warnings.

Signed-off-by: Kai Fricke <kai@anyscale.com>
Signed-off-by: e428265 <arvind.chandramouli@lmco.com>

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