[Tune] bohb_example patch fix

[justinvyu]

Why are these changes needed?

The bohb_example is super flaky after #35338 modified the example to run more trials. This is just a patch fix to deflake the example. There is still some underlying root cause to investigate. The other BOHB example (the Jupyter notebook) is also flaky, but for a different reason. See #35428.

In general, I unsure about the current hyperband implementation. See comments below for parts that I don't really understand.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[krfricke]

Should this be

Suggested change

	elif t.status in Trial.RUNNING:
	elif t.status in {Trial.RUNNING, Trial.PENDING}:

?

[justinvyu]

I initially did this, but it started giving some errors on the execution side.

  File "/Users/justin/Developer/justinvyu-dev/python/ray/tune/execution/trial_runner.py", line 748, in _process_trial_results
    decision = self._process_trial_result(trial, result)
  File "/Users/justin/Developer/justinvyu-dev/python/ray/tune/execution/trial_runner.py", line 791, in _process_trial_result
    decision = self._scheduler_alg.on_trial_result(
  File "/Users/justin/Developer/justinvyu-dev/python/ray/tune/schedulers/hb_bohb.py", line 88, in on_trial_result
    bracket.update_trial_stats(trial, result)
  File "/Users/justin/Developer/justinvyu-dev/python/ray/tune/schedulers/hyperband.py", line 470, in update_trial_stats
    assert trial in self._live_trials
AssertionError

The trial is pending, gets cleaned up in the scheduler state, then starts running and returns a result. Then, the scheduler is confused because it already cleaned up the trial.

Basically, we should be calling stop_trial for bad trials, but that didn't really work when I tried it -- the example started to hang after a certain number of trials (see comment above).

[justinvyu]

Why do we stop_trial on a paused trial? It's already stopped (doesn't have a live actor). Doesn't this just change the status from TERMINATED to ERROR?

Why do we not stop_trial on a running or pending trial? Why do we cleanup the bracket instead? Doesn't this cause us to not track a running trial?
Instead, we return STOP as the action to take if any of the bad trials is running, even though the returned action is not even for the running trial that we are inspecting in the loop. It's some arbitrary trial that got a result and started this process_bracket logic.

[justinvyu]

I initially did this, but it started giving some errors on the execution side.

  File "/Users/justin/Developer/justinvyu-dev/python/ray/tune/execution/trial_runner.py", line 748, in _process_trial_results
    decision = self._process_trial_result(trial, result)
  File "/Users/justin/Developer/justinvyu-dev/python/ray/tune/execution/trial_runner.py", line 791, in _process_trial_result
    decision = self._scheduler_alg.on_trial_result(
  File "/Users/justin/Developer/justinvyu-dev/python/ray/tune/schedulers/hb_bohb.py", line 88, in on_trial_result
    bracket.update_trial_stats(trial, result)
  File "/Users/justin/Developer/justinvyu-dev/python/ray/tune/schedulers/hyperband.py", line 470, in update_trial_stats
    assert trial in self._live_trials
AssertionError

The trial is pending, gets cleaned up in the scheduler state, then starts running and returns a result. Then, the scheduler is confused because it already cleaned up the trial.

Basically, we should be calling stop_trial for bad trials, but that didn't really work when I tried it -- the example started to hang after a certain number of trials (see comment above).