Train + Tune API Revamp

[justinvyu]

Summary

Ray Tune and Ray Train have been tightly coupled since Ray 2.0, when Ray Tune became the common execution engine for both libraries.

Ray Train execution invokes Tune’s execution logic under the hood, which leads to a complex, layered system.
The original intention behind this was to increase the interoperability of the two libraries, but the dependency of Ray Train on Ray Tune has led to many usability and stability issues, and it has stalled feature development.

The goal of these changes is to improve:

• Usability by mostly keeping feature parity, while introducing more intuitive APIs in place of inherited Ray Tune APIs that do not fit in the context of Ray Train.
• Extensibility by introducing more modular execution components that can be more easily customized.
• Interoperability between Ray Tune and Ray Train so that using both together is more natural and still provides separation between the two libraries.
• Testability with proper unit tests rather than hundreds of mini end to end tests.

[hongpeng-guo]

This is legendary! Thanks for putting everything together.

Will approve before getting public discussion feedbacks

[robertnishihara]

A few questions

1. If there is nothing to restore, e.g., this is the first time the job is submitted, then unique_run_name will be None and the TorchTrainer will see that there is no checkpoint to load from and will not load anything?
2. What's the reason for removing the explicit restore call?
3. Why do I need the storage info in RunConfig? Couldn't that all just go into get_checkpoint?

[justinvyu]

1. That's right.
2. This is actually something I was debating back and forth about. This section outlines some of the issues with the existing TorchTrainer.restore API.
   • Re-implementing with exact API parity would require us to pickle the user code and python objects again.
   • The explicit Trainer.restore API also had some UX friction (needing the user to manually concatenate their storage_path with the run name.
   • I took this chance to just improve the usability by pulling the "restore path" from the RunConfig directly and avoiding all pickling. So, the RunConfig(storage_path, name, storage_filesystem) tuple is what uniquely defines a Ray Train run, and users should pass in a unique name per job rather than re-use the same name for multiple runs.
   • Using a colliding name runs into other problems such as checkpoints being uploaded to the same directory with possible overwriting, so it's already unsupported and undefined behavior.
   • We can also definitely add a restore_run_if_exists flag to make the job restoration behavior explicit if users find this behavior is too magical.
3. The storage info sets the location where checkpoints and Ray Train's driver state gets saved. Without it, we can't load the state if restoring the run from a new Ray cluster, and we don't know what checkpoint to populate get_checkpoint with.

[robertnishihara]

This looks awesome. A few questions.
Training health monitoring is becoming more and more important and can be quite complex and compute intensive (not just looking at lots of metrics, but also running evals). Where does this logic live?

If the health monitoring takes a while to run but eventually decides that we need to restart from an earlier checkpoint, how does that get implemented?

[justinvyu]

Where does this logic live?

This logic lives in the Train driver right now, where the controller actor periodically pings the worker actor tasks and handles any errors raised by those health-check tasks.

If the health monitoring takes a while to run but eventually decides that we need to restart from an earlier checkpoint, how does that get implemented?

Health-checks are every few seconds at the moment, so the decisions happen pretty quickly after Ray detects that a node/actor has died. We also handle more extreme edge-cases where the underlying RPC times out randomly, in which case we'll wait for some configurable time until restarting from the latest checkpoint.

(not just looking at lots of metrics, but also running evals)

Oh, maybe this part was a bit misleading. At the moment, the "health and training progress" just refers to the underlying actor/node health and reported checkpoints, rather than the health of the training jobs (we do NOT implement anything to handle how well things are converging, how validation metrics are doing, etc).

[robertnishihara]

Minor, but isn't config.get("resume_from_checkpoint") enough? Why do you also need ray.train.get_checkpoint()?

[justinvyu]

ray.train.get_checkpoint is still needed if you want to support fault tolerance, in addition to a checkpoint to start fine-tuning from initially. Otherwise, if a node fails, the training progress always gets reset to the initial checkpoint, throwing away all finetuning progress.