CheckpointManager
CheckpointManager(output_directory: str | PathLike, configuration: CheckpointConfiguration, file_types: Iterable[CheckpointFileRegistration] = ()) -> NoneInitialize the checkpoint manager.
Parameters:
output_directory(str | PathLike) – The experiment's output directory (checkpoints live in acheckpointssubdirectory).configuration(CheckpointConfiguration) – The checkpoint retention configuration.file_types(Iterable[CheckpointFileRegistration]) – The checkpoint file types this manager can decode.
Properties
checkpoints_directory
checkpoints_directory: Pathlatest_checkpoint_number
latest_checkpoint_number: int | NoneGet the number of the latest checkpoint, or None if no checkpoints exist.
Returns:
int | None– The latest checkpoint number or None.
Methods
load_latest
load_latest() -> Checkpointmove_checkpoints_to_rollback_directory
move_checkpoints_to_rollback_directory(numbers: list[int]) -> Path | NoneMove checkpoint directories into a new rollback_N directory.
Creates a rollback_1 directory (or rollback_2, and similar values if it already exists) inside the checkpoints directory and moves each listed checkpoint into it. The checkpoint directories are not renamed; they keep their original names.
Parameters:
numbers(list[int]) – The checkpoint numbers to move.
Returns:
Path | None– The rollback directory path, or None if no checkpoints were moved.
register
register(name: str, file_type: type[CheckpointFile]) -> NoneRegister a checkpoint file type this manager can decode.
The type is registered under every extension it declares via file_type.extensions().
Parameters:
name(str) – The checkpoint file's base name; must not contain '.'.file_type(type[CheckpointFile]) – The concrete CheckpointFile subclass to decode files with this name into.
Raises:
ValueError– Ifnamecontains '.'.TypeError– Iffile_typeis not a concrete CheckpointFile subclass.
reserve_latest
reserve_latest(checkpoint: Checkpoint) -> PathAssign the next checkpoint number to checkpoint and create its directory without writing any files.
Reserving synchronously lets callers hand the (slow) file writes to a background thread while the checkpoint number is immediately visible via latest_checkpoint_number.
Parameters:
checkpoint(Checkpoint) – The checkpoint to reserve.
Returns:
Path– The reserved checkpoint directory.
save_latest
save_latest(checkpoint: Checkpoint) -> Pathtry_load_checkpoint
try_load_checkpoint(number: int) -> Checkpoint | Nonetry_load_checkpoint_file
try_load_checkpoint_file(number: int, name: str) -> CheckpointFile | NoneLoad a single checkpoint file from a checkpoint directory.
Parameters:
number(int) – The checkpoint number.name(str) – The registered checkpoint file base name.
Returns:
CheckpointFile | None– The loaded checkpoint file, orNonewhen not present.
try_load_latest
try_load_latest() -> Checkpoint | Nonewrite_reserved
write_reserved(checkpoint: Checkpoint) -> PathWrite the files of a checkpoint previously reserved with reserve_latest.
Parameters:
checkpoint(Checkpoint) – The reserved checkpoint to write.
Returns:
Path– The checkpoint directory.