report#

ray.train.report(metrics: Dict[str, Any], checkpoint: Checkpoint | None = None, checkpoint_dir_name: str | None = None, checkpoint_upload_mode: CheckpointUploadMode = CheckpointUploadMode.SYNC, delete_local_checkpoint_after_upload: bool | None = None, checkpoint_upload_fn: Callable[[Checkpoint, str], Checkpoint] | None = None, validation: bool | ValidationTaskConfig = False)#

Save a checkpoint of the training state so that training can recover from failures.

The checkpoint is persisted to storage. If workers fail and Ray Train restarts them, the training function can load the latest reported checkpoint with ray.train.get_checkpoint() and resume from it. See the fault tolerance guide.

Ray Train attaches the metrics to the checkpoint, uses them to decide which checkpoints to keep, and returns them in Result. See CheckpointConfig. Ray Train doesn’t persist metrics reported without a checkpoint. To track metrics over the course of training, log them from the workers to an experiment tracking tool instead. See Experiment tracking.

If this is called in multiple distributed training workers:

  • Only the metrics reported by the rank 0 worker will be attached to the checkpoint.

  • A checkpoint will be registered as long as one or more workers reports checkpoint that is not None. See the checkpointing guide.

  • Checkpoints from multiple workers will be merged into one directory in persistent storage. See the distributed checkpointing guide.

Warning

All workers must call ray.train.report the same number of times so that Ray Train can properly synchronize the training state across workers. This method acts as a barrier across all workers, so be sure that every worker reaches this method.

Example

import tempfile

import ray.train
from ray.train.torch import TorchTrainer


def train_func(config):
    start_epoch = 0

    # Resume from the latest reported checkpoint after a failure.
    checkpoint = ray.train.get_checkpoint()
    if checkpoint:
        with checkpoint.as_directory() as checkpoint_dir:
            # Load the model and optimizer state...
            # start_epoch = ... + 1
            pass

    for epoch in range(start_epoch, config.get("num_epochs", 10)):
        # Do training...

        metrics = {"loss": ...}

        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
           # Save the checkpoint...
           # torch.save(...)

            checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)

            # Example: Only the rank 0 worker uploads the checkpoint.
            if ray.train.get_context().get_world_rank() == 0:
                ray.train.report(metrics, checkpoint=checkpoint)
            else:
                ray.train.report(metrics, checkpoint=None)

trainer = TorchTrainer(
    train_func, scaling_config=ray.train.ScalingConfig(num_workers=2)
)
Parameters:

  • metrics (Dict[str, Any]) – The metrics to attach to the checkpoint, such as the validation loss. Ray Train uses them to rank checkpoints.

  • checkpoint (Checkpoint | None) – The checkpoint to persist. Workers that don’t hold any checkpoint data, such as non-zero ranks in data-parallel training, pass None.

  • checkpoint_dir_name (str | None) – Custom name for the checkpoint directory. If not provided, a unique directory name will be automatically generated. If provided, it must be unique across all checkpoints per worker to avoid naming collisions. Consider including identifiers such as the epoch or batch index in the name.

  • checkpoint_upload_mode (CheckpointUploadMode) – The manner in which we want to upload the checkpoint. Defaults to uploading the checkpoint synchronously. This works when no checkpoint is provided but is not useful in that case.

  • delete_local_checkpoint_after_upload (bool | None) – Whether to delete the checkpoint after it is uploaded.

  • checkpoint_upload_fn (Callable[[Checkpoint, str], Checkpoint] | None) – A user defined function that will be called with the checkpoint to upload it. If not provided, defaults to using the pyarrow.fs.copy_files utility for copying to the destination storage_path.

  • validation (bool | ValidationTaskConfig) – [Alpha] If True, triggers validation with default kwargs from validation_config. If a ValidationTaskConfig, validation is run using fn_kwargs merged with validation_config defaults, with fn_kwargs taking precedence on conflicts. If False, no validation.