How to save and load trial checkpoints

Trial checkpoints are one of the three types of data stored by Tune. You define their contents, and they’re meant to snapshot your training progress.

You save trial-level checkpoints through the trainable API. The trainable is where you define your custom training logic, and it’s also where you define which trial state to checkpoint. This guide shows how to save and load checkpoints with the function API and the class API, and walks you through the configuration options.

Save and load checkpoints with the function API

With the Ray Tune function API, save and load checkpoints as the following example shows. To create a checkpoint, use the from_directory() API.

import os
import tempfile

from ray import tune
from ray.tune import Checkpoint


def train_func(config):
    start = 1
    my_model = MyModel()

    checkpoint = tune.get_checkpoint()
    if checkpoint:
        with checkpoint.as_directory() as checkpoint_dir:
            checkpoint_dict = torch.load(os.path.join(checkpoint_dir, "checkpoint.pt"))
            start = checkpoint_dict["epoch"] + 1
            my_model.load_state_dict(checkpoint_dict["model_state"])

    for epoch in range(start, config["epochs"] + 1):
        # Model training here
        # ...

        metrics = {"metric": 1}
        with tempfile.TemporaryDirectory() as tempdir:
            torch.save(
                {"epoch": epoch, "model_state": my_model.state_dict()},
                os.path.join(tempdir, "checkpoint.pt"),
            )
            tune.report(metrics=metrics, checkpoint=Checkpoint.from_directory(tempdir))


tuner = tune.Tuner(train_func, param_space={"epochs": 5})
result_grid = tuner.fit()

The preceding code does the following:

• It saves a checkpoint with tune.report(..., checkpoint=checkpoint). You must report every checkpoint alongside a set of metrics, so that Tune can order checkpoints by a specified metric.

• During training iteration epoch, it saves the checkpoint to the path <storage_path>/<exp_name>/<trial_name>/checkpoint_<epoch> on the node where training happens. Depending on the storage configuration, Tune can then sync the checkpoint to a consolidated storage location.

• It loads a checkpoint with tune.get_checkpoint(). Whenever Tune restores a trial, this call returns the trial’s latest checkpoint. Tune restores a trial when the trial is configured to retry after a failure, when the experiment is restored, or when the trial resumes after a pause, such as with PBT.

Note

checkpoint_frequency and checkpoint_at_end don’t work with function API checkpointing. With a function trainable, you control checkpointing manually instead. For example, to checkpoint every three epochs, use the following code:

NUM_EPOCHS = 12
# checkpoint every three epochs.
CHECKPOINT_FREQ = 3


def train_func(config):
    for epoch in range(1, config["epochs"] + 1):
        # Model training here
        # ...

        # Report metrics and save a checkpoint
        metrics = {"metric": "my_metric"}
        if epoch % CHECKPOINT_FREQ == 0:
            with tempfile.TemporaryDirectory() as tempdir:
                # Save a checkpoint in tempdir.
                tune.report(metrics, checkpoint=Checkpoint.from_directory(tempdir))
        else:
            tune.report(metrics)


tuner = tune.Tuner(train_func, param_space={"epochs": NUM_EPOCHS})
result_grid = tuner.fit()

For more information on creating checkpoints, see Checkpoint.

Save and load checkpoints with the class API

You can also implement checkpointing and restoring with the class API:

import os
import torch
from torch import nn

from ray import tune


class MyTrainableClass(tune.Trainable):
    def setup(self, config):
        self.model = nn.Sequential(
            nn.Linear(config.get("input_size", 32), 32), nn.ReLU(), nn.Linear(32, 10)
        )

    def step(self):
        return {}

    def save_checkpoint(self, tmp_checkpoint_dir):
        checkpoint_path = os.path.join(tmp_checkpoint_dir, "model.pth")
        torch.save(self.model.state_dict(), checkpoint_path)
        return tmp_checkpoint_dir

    def load_checkpoint(self, tmp_checkpoint_dir):
        checkpoint_path = os.path.join(tmp_checkpoint_dir, "model.pth")
        self.model.load_state_dict(torch.load(checkpoint_path))


tuner = tune.Tuner(
    MyTrainableClass,
    param_space={"input_size": 64},
    run_config=tune.RunConfig(
        stop={"training_iteration": 2},
        checkpoint_config=tune.CheckpointConfig(checkpoint_frequency=2),
    ),
)
tuner.fit()

You can checkpoint manually, periodically, or at termination.

Checkpoint manually from the trainable

To trigger checkpointing manually from a custom Trainable, return should_checkpoint: True, or tune.result.SHOULD_CHECKPOINT: True, in the result dictionary of step. This can be especially helpful on spot instances:

import random


# to be implemented by user.
def detect_instance_preemption():
    choice = random.randint(1, 100)
    # simulating a 1% chance of preemption.
    return choice <= 1


def train_func(self):
    # training code
    result = {"mean_accuracy": "my_accuracy"}
    if detect_instance_preemption():
        result.update(should_checkpoint=True)
    return result

In the preceding example, when detect_instance_preemption returns True, the trainable sets should_checkpoint=True in its result to trigger a manual checkpoint.

Checkpoint manually from a Tuner callback

As with manual checkpointing from the trainable, you can trigger checkpointing from Tuner Callback methods. In the on_trial_result() method of your custom callback, set the result["should_checkpoint"] = True flag, or result[tune.result.SHOULD_CHECKPOINT] = True. Unlike checkpointing in the class API, this approach decouples checkpointing logic from training logic. It also gives you access to all Trial instances, so you can build more complex checkpointing strategies.

from ray import tune
from ray.tune.experiment import Trial
from ray.tune.result import SHOULD_CHECKPOINT, TRAINING_ITERATION


class CheckpointByStepsTaken(tune.Callback):
    def __init__(self, iterations_per_checkpoint: int):
        self.steps_per_checkpoint = iterations_per_checkpoint
        self._trials_last_checkpoint = {}

    def on_trial_result(
        self, iteration: int, trials: list[Trial], trial: Trial, result: dict, **info
    ):
        current_iteration = result[TRAINING_ITERATION]
        if (
            current_iteration - self._trials_last_checkpoint.get(trial, -1)
            >= self.steps_per_checkpoint
        ):
            result[SHOULD_CHECKPOINT] = True
            self._trials_last_checkpoint[trial] = current_iteration

Checkpoint periodically

To checkpoint trials every N iterations, set checkpoint_frequency=N, as in the following example:

tuner = tune.Tuner(
    MyTrainableClass,
    run_config=tune.RunConfig(
        stop={"training_iteration": 2},
        checkpoint_config=tune.CheckpointConfig(checkpoint_frequency=10),
    ),
)
tuner.fit()

Checkpoint at termination

The checkpoint_frequency might not coincide with the exact end of a trial. To also create a checkpoint at the end of a trial, set checkpoint_at_end=True:

tuner = tune.Tuner(
    MyTrainableClass,
    run_config=tune.RunConfig(
        stop={"training_iteration": 2},
        checkpoint_config=tune.CheckpointConfig(
            checkpoint_frequency=10, checkpoint_at_end=True
        ),
    ),
)
tuner.fit()

Configure checkpointing

Configure checkpointing with CheckpointConfig. Some of its options don’t apply to the function API, because you determine checkpointing frequency manually in your own training loop. The following compatibility matrix shows which options each API supports.

	Class API	Function API
num_to_keep	✅	✅
checkpoint_score_attribute	✅	✅
checkpoint_score_order	✅	✅
checkpoint_frequency	✅	❌
checkpoint_at_end	✅	❌

Summary

This guide covered how to save and load trial checkpoints in Tune. After you enable checkpointing, see one of the following guides:

• Extract checkpoints from Tune experiment results

• Configure persistent storage options for a distributed Tune experiment

Appendix: Types of data stored by Tune

Tune stores three types of data, which the following sections describe.

Experiment checkpoints

Experiment-level checkpoints save the experiment state, which includes the following:

• The state of the search algorithm.

• The list of trials and their statuses, such as PENDING, RUNNING, TERMINATED, and ERROR.

• Metadata for each trial, such as its hyperparameter configuration and some derived trial results, including the minimum, maximum, and last values.

The driver on the head node saves the experiment-level checkpoint periodically. By default, Tune adjusts the save frequency automatically so that it spends at most 5% of the time saving experiment checkpoints and the remaining time handling training results and scheduling. You can also adjust this time with the TUNE_GLOBAL_CHECKPOINT_S environment variable.

Trial checkpoints

Trial-level checkpoints capture the per-trial state, which often includes the model and optimizer states. Trial checkpoints have several uses, including the following:

• If a trial is interrupted for some reason, such as on a spot instance, it can resume from its last state, so no training time is lost.

• Some search algorithms and schedulers pause trials to free up resources for other trials to train in the meantime. Pausing only makes sense if the trials can then continue training from their latest state.

• You can use the checkpoint later for downstream tasks such as batch inference.

To learn how to save and load trial checkpoints, see How to save and load trial checkpoints.

Trial results

Tune saves and logs the metrics that trials report to their respective trial directories. It stores this data in CSV, JSON, or TensorBoard format. TensorBoard files match events.out.tfevents.*. You can inspect this data with TensorBoard and use it for post-experiment analysis.