[train] Storage refactor: Support PBT and BOHB

[krfricke]

Why are these changes needed?

Previously PBT and BOHB do not work with the new storage context. This is because they directly affect the Ray Tune control flow and rely on specific behavior.

In particular, PBT and BOHB made heavy use of pausing trials, and PBT also saved and restored from "objects".

However, the new storage backend does not support memory checkpoints anymore. Saving and restoring from objects was also removed. This means PBT and BOHB have to be adjusted and the pausing logic has to be revamped.

For BOHB/Hypberband in general, we now avoid controlling trial status directly, and instead rely on choose_trial_to_run to unpause trials. This means the tune control loop has greater control over the pausing logic.

In the tune control loop, we now detect double saves. If a save() future is scheduled while another one is in-flight (which can happen in PBT), we don't schedule another save.

In PBT, we now schedule persistent checkpoints instead of memory checkpoints. The main difference here is that persistent checkpoints may get deleted before the exploiting trial gets a chance to load from it. For this reason, we detect too small num_to_keep values and log a warning.

Related issue number

Closes #38569

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]

Leaving some review guidance

[krfricke]

We now may schedule multiple saves subsequently (e.g. save is scheduled, trial is instructed to pause which also schedules a save, or maybe PBT schedules an additional save).

Previously, these additional saves were memory checkpoints that were not registered. We want to safeguard here against the same checkpoint being registered multiple times.

[krfricke]

This wrapper is only used by PBT and it serves a similar purpose as the previous future in _TrackedCheckpoint.

It's not used anywhere else (except being created in the tune controller).

[krfricke]

This just changes the API to the new Checkpoint API

[krfricke]

Same here, this just changes the API to the new Checkpoint API

[krfricke]

Trials now may be PENDING and failing - e.g. when a save resolves late

[justinvyu]

This is if a trial is paused, schedules a save, then gets unpaused (set to pending) and then the save fails?

[krfricke]

That's right!

[krfricke]

This also de-dupes checkpoints

[justinvyu]

old codepath

[justinvyu]

Can we also add a typehint to saving_to in the _TemporaryTrialState class? It's always a _FutureTrainingResult now

[krfricke]

We schedule some restores directly from the object ref

[justinvyu]

old codepath

[krfricke]

This is a change in the hyperband - instead of directly setting the trial status to PENDING, we add them to the trials_to_unpause set. Then the scheduler will selet them in choose_trial_to_run

[justinvyu]

nice I like this.

[krfricke]

Here we actually synchronously wait for the save to resolve. This may impact performance slightly but let's see if it's actually a problem. We only wait for trials we actually exploit and we do need the attached result.

[justinvyu]

I think it's good to go.

I have one suggestion for clarity that I'd like to get in but not 100% needed.

Basically, let's make PBTTrialState.last_checkpoint always either a _FutureTrainingResult or a _TrainingResult. Right now it can be those as well as a Checkpoint.

Basically here, just set clone_state.last_checkpoint = training_result:

ray/python/ray/tune/schedulers/pbt.py

Lines 698 to 712 in c3bf12b

	last_checkpoint = clone_state.last_checkpoint
	logger.debug(
	f"Trial {trial} is in lower quantile. "
	f"Exploiting trial {trial_to_clone}."
	)
	if isinstance(last_checkpoint, _FutureTrainingResult):
	training_result = last_checkpoint.resolve()
	if training_result:
	clone_state.last_result = training_result.metrics
	clone_state.last_checkpoint = training_result.checkpoint
	last_checkpoint = clone_state.last_checkpoint
	else:

Then, no need to construct a new "training result" on exploit.

[krfricke]

It makes sense, but it's a bit confusing (why is a checkpoint a training result). We also use trial.checkpoint above, which returns the actual checkpoint and not the result.

Let's keep it as is for now (the whole situation is very confusing right now anyway - two different checkpoint, three checkpoint managers, training results, tracked checkpoints...). But I'd like to specifically make time to clean up the leftovers from the old code path.

[justinvyu]

Ok, good with me. Here's what should remain after everything is all cleaned up:

train._internal.checkpoint_manager.CheckpointManager
train.Checkpoint
train._internal.session.TrainingResult
train._internal.session.FutureTrainingResult (hopefully temporary)

[zhe-thoughts]

Necessary change for Train