classmethod BaseTrainer.restore(path: str, datasets: Optional[Dict[str, Union[Dataset, Callable[[], Dataset]]]] = None, preprocessor: Optional[Preprocessor] = None, scaling_config: Optional[ray.air.config.ScalingConfig] = None, **kwargs) BaseTrainer[source]#

Restores a Train experiment from a previously interrupted/failed run.

Restore should be used for experiment-level fault tolerance in the event that the head node crashes (e.g., OOM or some other runtime error) or the entire cluster goes down (e.g., network error affecting all nodes).

The following example can be paired with implementing job retry using Ray Jobs to produce a Train experiment that will attempt to resume on both experiment-level and trial-level failures:

import os
from ray import air
from ray.data.preprocessors import BatchMapper
from ray.train.trainer import BaseTrainer

experiment_name = "unique_experiment_name"
local_dir = "~/ray_results"
experiment_dir = os.path.join(local_dir, experiment_name)

# Define some dummy inputs for demonstration purposes
datasets = {"train": ray.data.from_items([{"a": i} for i in range(10)])}
preprocessor = BatchMapper(lambda x: x, batch_format="numpy")

class CustomTrainer(BaseTrainer):
    def training_loop(self):

if CustomTrainer.can_restore(experiment_dir):
    trainer = CustomTrainer.restore(
    trainer = CustomTrainer(
            # Tip: You can also enable retries on failure for
            # worker-level fault tolerance

result = trainer.fit()
  • path – The path to the experiment directory of the training run to restore. This can be a local path or a remote URI if the experiment was uploaded to the cloud.

  • datasets – Re-specified datasets used in the original training run. This must include all the datasets that were passed in the original trainer constructor.

  • preprocessor – Optionally re-specified preprocessor that was passed in the original trainer constructor. This should be used to re-supply the preprocessor if it is not restorable in a new Ray cluster. This preprocessor will be fit at the start before resuming training. If no preprocessor is passed in restore, then the old preprocessor will be loaded from the latest checkpoint and will not be re-fit.

  • scaling_config – Optionally re-specified scaling config. This can be modified to be different from the original spec.

  • **kwargs – Other optionally re-specified arguments, passed in by subclasses.


ValueError – If all datasets were not re-supplied on restore.


A restored instance of the class that is calling this method.

Return type