t5x.checkpoint_utils package

Checkpoint helper functions for managing checkpoints.

Supports marking checkpoints as pinned to exclude them from the checkpointer removal process.

class t5x.checkpoint_utils.CheckpointTypes(value)

An enumeration.

t5x.checkpoint_utils.detect_checkpoint_type(checkpoint_path, expected=None)

Returns the checkpoint type by reading the .checkpoint metadata file.

Parameters:

• checkpoint_path – The path of the .checkpoint file.

• expected – The expected checkpoint type. If the checkpoint type is not as expected, this function will log a warning but will not raise an error.

Returns:

The checkpoint type.

t5x.checkpoint_utils.is_pinned_checkpoint(ckpt_dir)

Returns whether the checkpoint is pinned, and should NOT be removed.

t5x.checkpoint_utils.pin_checkpoint(ckpt_dir, txt='1')

Pin a checkpoint so it does not get deleted by the normal pruning process.

Creates a PINNED file in the checkpoint directory to indicate the checkpoint should be excluded from the deletion of old checkpoints.

Parameters:

• ckpt_dir – The checkpoint step dir that is to be always kept.

• txt – Text to be written into the checkpoints ALWAYS_KEEP me file.

t5x.checkpoint_utils.pinned_checkpoint_filepath(ckpt_dir)

Full path of the pinned checkpoint file.

t5x.checkpoint_utils.remove_checkpoint_dir(ckpt_dir)

Removes the checkpoint dir if it is not pinned.

t5x.checkpoint_utils.remove_dataset_checkpoint(ckpt_dir, train_ds_prefix)

Removes dataset checkpoints if the checkpoint is not pinned.

t5x.checkpoint_utils.unpin_checkpoint(ckpt_dir)

Removes the pinned status of the checkpoint so it is open for deletion.