[train] Fix broken tune tests and support ray storage

[justinvyu]

Why are these changes needed?

This PR re-introduces support for ray storage ray.init(storage="s3://...") and fixes a broken tune controller test.

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 :(

[justinvyu]

this is pretty confusing -- we do pause + save checkpoint as 2 steps:

1. first time we try to schedule a pause, we schedule a save instead. then we CACHE a pause decision.
2. then, once the save finishes, we pop the cached decision and execute it, which should be a PAUSE again, in which case we enter the should_checkpoint=False condition and just stop the trial and set the status to PAUSED.

In between these 2 steps, if the scheduler outputs some decision that's not PAUSED (ex: NOOP), then this thing will just hang forever. So we need to make sure that the cached PAUSE decision is not overriden by something random while the save is happening.

I allow STOP decisions to override, but unclear if this ever happens..

[krfricke]

Hm, the _cached_trial_decision is only updated on saves. The scheduler actions only affect _queued_trial_decision. (I know the naming is confusing... happy to rename it).

Does this still come up? Is this like a double triggered save? In that case, should we have this line:

            if trial.temporary_state.saving_to:
                # If a save is already in progress, don't schedule another one.
                return trial.temporary_state.saving_to

in the new storage path _schedule_trial_save as well?

[justinvyu]

I see, you're right. The problem is a little different from what I described above. Here's what's actually happening to cause that unit test to fail:

1. The test defines a scheduler that manually calls tune_controller.pause_trial(trial) during the on_trial_result hook
2. should_checkpoint=True by default, so this will schedule a SAVE and set the cached trial decision here:

ray/python/ray/tune/execution/tune_controller.py

Line 1631 in 3984b85

self._cached_trial_decisions[trial.trial_id] = TrialScheduler.PAUSE

3. At this point, we return a NOOP trial scheduler decision (since we paused manually) and end up here:

ray/python/ray/tune/execution/tune_controller.py

Lines 1767 to 1768 in 3984b85

	decision = self._scheduler_alg.on_trial_result(
	self._wrapped(), trial, flat_result

4. The trial IS SAVING at this point, so we enter this block:

ray/python/ray/tune/execution/tune_controller.py

Lines 1806 to 1820 in 3984b85

	if trial.is_saving:
	logger.debug(f"Caching trial decision for trial {trial}: {decision}")
	# Cache decision to execute on after the save is processed.
	# This prevents changing the trial's state or kicking off
	# another training step prematurely.
	# If a decision is already cached, don't override it for CONTINUE/NOOP
	# decisions. Only escalate the cached decision to a STOP/PAUSE if requested.
	# This is only very relevant for pausing trials, since we cache a PAUSE
	# decision to happen after a save finishes.
	# We need to make sure that we don't override it in the
	# time between the save operation starting and finishing.
	if decision in [TrialScheduler.STOP, TrialScheduler.PAUSE]:
	self._cached_trial_decisions[trial.trial_id] = decision
	return None

5. We overwrite the PAUSE cached decision with the NOOP, leading to an infinite hang.

[justinvyu]

The conclusion: Calling pause_trial(should_checkpoint=True) directly inside a scheduler's on_trial_result leads to a hang.

[justinvyu]

PBT doesn't run into this problem because it calls pause_trial(should_checkpoint=False)

[justinvyu]

this thing errors since it's still RUNNING (even if you put it after pause_trial)

[justinvyu]

I've reverted the tune controller changes and updated the test to pass, but calling pause_trial within on_trial_result is still an issue.