nemo_automodel.components.checkpoint.lifecycle

Checkpoint directory publication and retention lifecycle.

Module Contents

Classes

Name	Description
`CheckpointLifecycle`	Own checkpoint reservation, publication, pointers, and retention.
`_PendingBestCheckpoint`	Best-checkpoint pointer update deferred until an async save completes.

Data

_RESERVATION_STAGING_SUFFIX

__all__

logger

API

class nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle(
    config: nemo_automodel.components.checkpoint.config.CheckpointingConfig,
    process_group: torch.distributed.ProcessGroup | None = None
)

Own checkpoint reservation, publication, pointers, and retention.

The lifecycle is built with the same process group as its Checkpointer. Filesystem work runs on that group’s rank 0 and failures are reduced over the same participants, so every participating rank either continues or raises.

Parameters:

config

CheckpointingConfig

Declarative checkpoint configuration.

process_group

ProcessGroup | NoneDefaults to None

Optional model-local group participating in checkpointing.

_best_val_loss

= float('inf')

_pending_best_checkpoint

_PendingBestCheckpoint | None = None

_pending_checkpoint_dir

str | None = None

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle._any_participant_reported(
    reported: bool
) -> bool

Return whether any checkpoint process-group participant reported a condition.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle._barrier() -> None

Synchronize checkpoint participants when distributed is initialized.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle._initialize_best_val_loss_from_pointer(
    metric_key: str | None
) -> None

Initialize best-metric state from the current complete LOWEST_VAL target.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle._is_coordinator() -> bool

Return whether this rank coordinates lifecycle filesystem work.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle._prune_old_checkpoints() -> None

Prune checkpoint directories according to the configured recent window.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle._publish_checkpoint(
    path: str
) -> None

Advance LATEST and commit a filesystem checkpoint’s marker.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle._remove_checkpoint_pointer(
    link_name: str
) -> None

Remove a checkpoint pointer symlink and fallback text file.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle._remove_stale_checkpoint_pointer(
    link_name: str
) -> None

Remove a checkpoint pointer whose target no longer exists.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle._update_best_checkpoint(
    target_dir: str,
    val_loss: float,
    metric_key: str | None
) -> None

Update LOWEST_VAL when a finite validation metric improves.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle._update_checkpoint_pointer(
    link_name: str,
    target_dir: str
) -> None

Atomically update a checkpoint-root pointer.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle.complete_pending() -> None

Publish a completed async checkpoint and apply pointer retention.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle.defer_publication(
    path: str,
    best_val_metric: float | None,
    metric_key: str | None
) -> None

Record an async checkpoint for publication after its writes complete.

Parameters:

path

str

Checkpoint directory being written asynchronously.

best_val_metric

float | None

Optional validation value eligible for LOWEST_VAL.

metric_key

str | None

Validation metric key stored in checkpoint loss metadata.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle.publish(
    path: str,
    best_val_metric: float | None,
    metric_key: str | None
) -> None

Publish a synchronous checkpoint, update pointers, and apply retention.

Parameters:

path

str

Fully written checkpoint directory.

best_val_metric

float | None

Optional validation value eligible for LOWEST_VAL.

metric_key

str | None

Validation metric key stored in checkpoint loss metadata.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle.reserve(
    path: str
) -> None

Reserve path for a save, replacing only an interrupted checkpoint.

Local checkpoint directories are created under a marked staging name and atomically renamed into place. An existing marked directory can be replaced, while an unmarked published checkpoint is never overwritten. msc:// roots retain their existing local recipe-metadata shadow but do not use filesystem lifecycle markers.

Parameters:

path

str

Checkpoint directory for the current step.

Raises:

FileExistsError: If path already holds a published checkpoint.
RuntimeError: If the directory could not be prepared.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle.run_coordinator_step(
    operation: collections.abc.Callable[[], None],
    description: str
) -> None

Run one coordinator operation and report any failure to all participants.

Parameters:

operation

Callable[[], None]

Work to run on process-group rank 0.

description

str

Step name used for logging and collective error reporting.

Raises:

RuntimeError: If the coordinator could not complete the operation.

nemo_automodel.components.checkpoint.lifecycle.CheckpointLifecycle.validate_checkpoint_dir_exists(
    ckpt_dir: str,
    restore_from: str
) -> None

Validate a resolved restore directory collectively.

Parameters:

ckpt_dir

str

Resolved checkpoint directory.

restore_from

str

User-provided restore selector.

Raises:

FileNotFoundError: If ckpt_dir does not exist.

class nemo_automodel.components.checkpoint.lifecycle._PendingBestCheckpoint(
    path: str,
    value: float | None,
    metric_key: str | None
)

Dataclass

Best-checkpoint pointer update deferred until an async save completes.

metric_key

str | None

path

str

value

float | None

nemo_automodel.components.checkpoint.lifecycle._RESERVATION_STAGING_SUFFIX = '.reserving-'

nemo_automodel.components.checkpoint.lifecycle.__all__ = ['CheckpointLifecycle']

nemo_automodel.components.checkpoint.lifecycle.logger = logging.getLogger(__name__)