Callbacks

These callbacks are used internally by the pipelines to report metrics and checkpoints to Ray Train at the end of each Lightning training epoch. You only need to interact with them directly if you want to customise the checkpointing strategy (see the tutorial).

TrainerReportOnIntervalCallback

class TrainerReportOnIntervalCallback(*args, **kwargs)

Bases: Callback

PyTorch Lightning callback that reports metrics to Ray Train every N epochs.

At each epoch whose index is divisible by interval a full checkpoint is saved and reported to Ray; all other epochs report metrics only. The per-epoch temporary directory is removed after reporting to keep disk usage bounded.

Variables:

• CHECKPOINT_NAME (str) – Fixed filename used when saving the Lightning checkpoint inside the temporary directory.

• trial_name (str) – Name of the current Ray Train trial, retrieved from the train context.

• local_rank (int) – Local rank of the current worker inside the trial.

• tmpdir_prefix (str) – Root path under which per-epoch checkpoint directories are created.

• interval (int) – How often (in epochs) a checkpoint is included in the Ray report.

• step (int) – Internal counter incremented after each on_train_epoch_end call.

CHECKPOINT_NAME = 'checkpoint.ckpt'

__init__(interval=1)

Parameters:

interval (int, optional) – Number of epochs between checkpoint saves. A value of 1 saves a checkpoint every epoch; 2 saves every other epoch, etc. Defaults to 1.

on_train_epoch_end(trainer, pl_module)

Report metrics (and optionally a checkpoint) to Ray Train.

Called automatically by the Lightning Trainer at the end of every training epoch. A checkpoint is included in the report only when self.step % self.interval == 0; otherwise only metrics are sent. The temporary epoch directory is deleted on rank 0 after reporting.

Parameters:

• trainer (lightning.pytorch.Trainer) – The active Lightning trainer instance.

• pl_module (lightning.pytorch.LightningModule) – The model being trained (unused directly; provided by the callback protocol).

TrainerReportKeepOnlyLastCallback

class TrainerReportKeepOnlyLastCallback(*args, **kwargs)

Bases: Callback

PyTorch Lightning callback that always reports the most recent checkpoint.

Every epoch overwrites the single last/ checkpoint directory so only the latest weights are kept on disk. This minimises storage at the cost of not being able to restore from earlier epochs.

Variables:

• CHECKPOINT_NAME (str) – Fixed filename used when saving the Lightning checkpoint inside the temporary directory.

• trial_name (str) – Name of the current Ray Train trial, retrieved from the train context.

• local_rank (int) – Local rank of the current worker inside the trial.

• tmpdir_prefix (str) – Root path under which the last/ checkpoint directory is created.

CHECKPOINT_NAME = 'checkpoint.ckpt'

__init__()

Initialise the callback and clean up any leftover trial directory.

on_train_epoch_end(trainer, pl_module)

Overwrite the last checkpoint and report metrics to Ray Train.

Called automatically by the Lightning Trainer at the end of every training epoch. The previous epoch’s checkpoint directory is removed before saving the new one, so at most one checkpoint exists at any time. The temporary directory is deleted on rank 0 after reporting.

Parameters:

• trainer (lightning.pytorch.Trainer) – The active Lightning trainer instance.

• pl_module (lightning.pytorch.LightningModule) – The model being trained (unused directly; provided by the callback protocol).