How to enable fault tolerance in Ray Tune#

In distributed machine learning experiments, fault tolerance can help mitigate the impact of node failures caused by out-of-memory and out-of-disk issues.

Fault tolerance can preserve training progress when a node fails, which saves time and resources. In a distributed setting, it also gives you access to the cost savings of preemptible spot instance nodes.

Experiment-level fault tolerance in Tune#

At the experiment level, Tuner.restore resumes a previously interrupted experiment from where it left off.

Use Tuner.restore in the following cases:

The driver script that calls Tuner.fit() errors out, for example because the head node runs out of memory or disk.
You manually interrupt the experiment with Ctrl+C.
The entire cluster crashes, taking the experiment with it, because of an ephemeral error such as the network going down or Ray object store memory filling up.

Note

Tuner.restore isn’t meant for resuming a terminated experiment with modified hyperparameter search spaces or stopping criteria. Experiment restoration is meant to resume and complete the exact job that you previously submitted through Tuner.fit.

For example, consider a Tune experiment configured to run for 10 training iterations, where all trials have already completed. You can’t use Tuner.restore to restore the experiment, change the number of training iterations to 20, and then continue training.

Instead, start a new experiment and initialize your model weights with a checkpoint from the previous experiment. For an example, see the FAQ entry on iterative experimentation.

Note

Restoration can’t fix bugs in your training loop. For restoration to help, the issue that crashed the experiment needs to be ephemeral, so that the retry after restoring can succeed.

Restore a Tune experiment#

Suppose you configure your initial Tune experiment as follows. The training loop is for demonstration purposes only. The important detail is that the trainable saves and loads checkpoints.

import json
import os
import tempfile

from ray import tune


def trainable(config):
    # Checkpoint loading
    checkpoint = tune.get_checkpoint()
    start = 1
    if checkpoint:
        with checkpoint.as_directory() as checkpoint_dir:
            with open(os.path.join(checkpoint_dir, "checkpoint.json"), "r") as f:
                state = json.load(f)
        start = state["epoch"] + 1

    for epoch in range(start, config["num_epochs"]):
        # Do some training...

        # Checkpoint saving
        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
            with open(os.path.join(temp_checkpoint_dir, "checkpoint.json"), "w") as f:
                json.dump({"epoch": epoch}, f)
            tune.report(
                {"epoch": epoch},
                checkpoint=tune.Checkpoint.from_directory(temp_checkpoint_dir),
            )


tuner = tune.Tuner(
    trainable,
    param_space={"num_epochs": 10},
    run_config=tune.RunConfig(
        storage_path=os.path.expanduser("~/ray_results"),
        name="tune_fault_tolerance_guide",
    ),
)
result_grid = tuner.fit()

Tune saves the experiment’s results and checkpoints to ~/ray_results/tune_fault_tolerance_guide, as configured by RunConfig. If the experiment stops for one of the reasons listed earlier, resume it from this path:

tuner = tune.Tuner.restore(
    os.path.expanduser("~/ray_results/tune_fault_tolerance_guide"),
    trainable=trainable,
    resume_errored=True,
)
tuner.fit()

Tip

You can also restore the experiment from a cloud bucket path:

tuner = tune.Tuner.restore(
    path="s3://cloud-bucket/tune_fault_tolerance_guide", trainable=trainable
)

See Configure persistent storage in Ray Tune.

Configure which trials to resume#

You can configure which trials Tune resumes, based on each trial’s status when the experiment stopped:

By default, Tune resumes unfinished trials left in the RUNNING state.
You can resume ERRORED trials or retry them from scratch.
You can’t resume TERMINATED trials.

tuner = tune.Tuner.restore(
    os.path.expanduser("~/ray_results/tune_fault_tolerance_guide"),
    trainable=trainable,
    resume_errored=True,
    restart_errored=False,
    resume_unfinished=True,
)

Auto-resume#

In a production setting, you might want a single script that launches the initial training run, or restores the experiment if the initial run already happened.

To do this, use the Tuner.can_restore utility:

import os
from ray import tune

storage_path = os.path.expanduser("~/ray_results")
exp_name = "tune_fault_tolerance_guide"
path = os.path.join(storage_path, exp_name)

if tune.Tuner.can_restore(path):
    tuner = tune.Tuner.restore(path, trainable=trainable, resume_errored=True)
else:
    tuner = tune.Tuner(
        trainable,
        param_space={"num_epochs": 10},
        run_config=tune.RunConfig(storage_path=storage_path, name=exp_name),
    )
tuner.fit()

The first time you run this script, it launches the initial training run. The second time, it attempts to resume from the outputs of the first run.

Advanced: Restore a Tune experiment that uses Ray object references#

Experiment restoration often happens in a different Ray session than the original run. In that case, Ray automatically garbage collects the original object references. If the experiment state includes object references, for example in each trial’s configuration dictionary, retrieving those objects doesn’t work after restoration, because the objects the references point to no longer exist.

To work around this, re-create the objects, put them in the Ray object store, and pass the new object references to Tune.

Example: Re-create object references to pre-trained models#

Suppose you have a large pre-trained model that you want to use in your training loop. For example, it could be an image classification model that calculates an Inception Score to evaluate the quality of a generative model. You might have multiple models to tune over, where each trial samples one of the models.

import ray
from ray import tune


class LargeModel:
    def __init__(self, model_id):
        self.model_id = model_id
        # Load weights based on the `model_id`...


def train_fn(config):
    # Retrieve the model from the object store.
    model = ray.get(config["model_ref"])
    print(model.model_id)


# These models may be large, so `ray.put` them in the Ray Object Store
# to share the models between trials.
model_refs = [ray.put(LargeModel(1)), ray.put(LargeModel(2))]

tuner = tune.Tuner(
    train_fn,
    # Tune over the object references!
    param_space={"model_ref": tune.grid_search(model_refs)},
    run_config=tune.RunConfig(
        storage_path=os.path.expanduser("~/ray_results"), name="restore_object_refs"
    ),
)
tuner.fit()

To restore, re-specify the param_space through Tuner.restore:

# Re-create the objects and put them in the object store.
param_space = {
    "model_ref": tune.grid_search([ray.put(LargeModel(1)), ray.put(LargeModel(2))])
}

tuner = tune.Tuner.restore(
    os.path.expanduser("~/ray_results/restore_object_refs"),
    trainable=train_fn,
    # Re-specify the `param_space` to update the object references.
    param_space=param_space,
    resume_errored=True,
)
tuner.fit()

Note

If you’re tuning over Ray Data datasets, you also need to re-specify them in the param_space. Ray Data datasets can contain object references, so the same problems apply.

The following example shows how:

ds_1 = ray.data.from_items([{"x": i, "y": 2 * i} for i in range(128)])
ds_2 = ray.data.from_items([{"x": i, "y": 3 * i} for i in range(128)])

param_space = {
    "datasets": {"train": tune.grid_search([ds_1, ds_2])},
}

tuner = tune.Tuner.restore(..., param_space=param_space)

Trial-level fault tolerance in Tune#

Trial-level fault tolerance handles the failure of individual trials in the cluster. Causes of these failures include the following:

Preemptible spot instances
Ephemeral network connection issues
Nodes running out of memory or out of disk space

Configure failure handling for individual trials with FailureConfig.

The following example configures FailureConfig for the trainable from the previous example, which implements saving and loading trial checkpoints:

from ray import tune

tuner = tune.Tuner(
    trainable,
    param_space={"num_epochs": 10},
    run_config=tune.RunConfig(
        storage_path=os.path.expanduser("~/ray_results"),
        name="trial_fault_tolerance",
        failure_config=tune.FailureConfig(max_failures=3),
    ),
)
tuner.fit()

With this configuration, when a trial encounters a runtime error, Tune reschedules that trial up to max_failures=3 times.

Similarly, if node X fails, for example because it’s preempted or loses its connection, Tune reschedules all trials that ran on node X up to three times.

Summary#

This guide covered how to enable experiment-level and trial-level fault tolerance in Ray Tune.

See the following resources for more information: