[AIR/Train] Add Trainer.restore API for train experiment-level fault tolerance

[justinvyu]

Why are these changes needed?

This PR introduces a Trainer.restore API that resumes a Train experiment that crashed/got interrupted. Previously, resume_from_checkpoint only allowed starting a new experiment, which writes to a completely new log directory and a new set of results.

Context

The restoration API is inconsistent between Train and Tune, since Tuner.restore. It’s not clear how to restore an interrupted Train experiment, and the existing resume_from_checkpoint API has issues as reported by users (see here [github, discourse, github]).

Added Documentation

• FAQ on restoring a Train experiment
• API reference
• Docstrings

TODO before merging

• Merge Jun's PR, overwrite the param space, and do the shutdown/init in the unit test to make sure that it doesn't hang/crash.
• Add to docs
• Fix restore for with user-defined Trainer subclasses that are defined in the same scope as the trainer.fit call.

Future Todos

Dealing with ActorHandle

TL;DR: Object references don't immediately throw an exception when unpickled from a new ray cluster, but actor handles do.

If an ActorHandle gets captured in the scope of a training function/config and gets pickled along with the trainer, then it will be impossible to load the pickled trainer from a new ray cluster. You will run into an error that looks like:

Traceback (most recent call last):
  File "train_example_nfs.py", line 85, in <module>
    trainer = DataParallelTrainer.restore(
  File "/Users/justin/Developer/justinvyu-dev/python/ray/train/data_parallel_trainer.py", line 315, in restore
    return super(DataParallelTrainer, cls).restore(
  File "/Users/justin/Developer/justinvyu-dev/python/ray/train/base_trainer.py", line 281, in restore
    original_trainer = pickle.load(fp)
  File "/Users/justin/Developer/justinvyu-dev/python/ray/_private/serialization.py", line 89, in _actor_handle_deserializer
    return ray.actor.ActorHandle._deserialization_helper(serialized_obj, outer_id)
  File "/Users/justin/Developer/justinvyu-dev/python/ray/actor.py", line 1280, in _deserialization_helper
    return worker.core_worker.deserialize_and_register_actor_handle(
  File "python/ray/_raylet.pyx", line 2310, in ray._raylet.CoreWorker.deserialize_and_register_actor_handle
  File "python/ray/_raylet.pyx", line 2279, in ray._raylet.CoreWorker.make_actor_handle
  File "/Users/justin/Developer/justinvyu-dev/python/ray/_private/function_manager.py", line 522, in load_actor_class
    actor_class = self._load_actor_class_from_gcs(
  File "/Users/justin/Developer/justinvyu-dev/python/ray/_private/function_manager.py", line 617, in _load_actor_class_from_gcs
    class_name = ensure_str(class_name)
  File "/Users/justin/Developer/justinvyu-dev/python/ray/_private/utils.py", line 293, in ensure_str
    assert isinstance(s, bytes)

We can't really do anything about this, and this de-serialization in a new ray cluster behavior needs to be fixed by Core.

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

[gjoliver]

it kinda bothers me a lot now that we have to create a custom restore() for every single Trainer we have.
nobody would understand how to contribute a new Trainer if we do this ...

[justinvyu]

What about:
Allow user to re-specify everything in __init__ except for the run config. We can inspect the arguments of cls.__init__ to see what's allowed. Then just overwrite everything that's passed in. Now, the subclasses don't need to re-implement restore.

One downside is that validation will be harder, and the signature of restore will be something like def restore(cls, **kwargs): which is very opaque.

[justinvyu]

Latest iteration: User subclass doesn't need to set any attributes themselves. They just need to define what attributes from their constructor can get re-specified. Then, they just pass these through as kwargs to the BaseTrainer.restore.

class MyTrainerSubclass(BaseTrainer):
    def restore(cls, path, my_arg, ...):
        return super().restore(path, my_arg=my_arg, ...)
        # No more setting things afterwards.

[gjoliver]

this is nice!

[krfricke]

Thanks! Please ping me when tests pass