`nemo_rl.algorithms.rollout_checkpointing.snapshots`#

Rollout snapshots and the rollout state of trainer checkpoints.

The single controller calls RolloutSnapshotter to save rollout state; it never builds TQ cuts, sidecars or snapshot manifests itself. Trainer checkpoints and rollout snapshots share one capture and one sidecar writer, and one lock keeps them from publishing at the same time.

Module Contents#

Classes#

RolloutSnapshotter

Save rollout state for the single controller.

Functions#

`write_controller_sidecars`	Write the captured controller state beside TQ; return the files written.
`_warn_if_environment_paused_past_interval`	Warn when a snapshot held rollouts longer than the snapshot interval.

API#

async nemo_rl.algorithms.rollout_checkpointing.snapshots.write_controller_sidecars( checkpoint_path: nemo_rl.utils.checkpoint.PathLike, captured: nemo_rl.algorithms.rollout_checkpointing.state.CapturedRolloutState, ) → list[str]#: Write the captured controller state beside TQ; return the files written.

nemo_rl.algorithms.rollout_checkpointing.snapshots._warn_if_environment_paused_past_interval( paused_seconds: float, interval_s: Optional[float], ) → None#

Warn when a snapshot held rollouts longer than the snapshot interval.

Rollouts then spend more time paused than running between snapshots.

class nemo_rl.algorithms.rollout_checkpointing.snapshots.RolloutSnapshotter( parts: nemo_rl.algorithms.rollout_checkpointing.state.RolloutParts, *, checkpointer: nemo_rl.utils.checkpoint.CheckpointManager, timeout: nemo_rl.utils.timer.TimeoutChecker, bootstrap_identity: Optional[nemo_rl.algorithms.rollout_checkpointing.manifest.BootstrapCompatibilityIdentity], teacher_coordinator: Optional[Any], records_prompt_groups: Callable[[], bool], )#

Bases: nemo_rl.algorithms.rollout_checkpointing.state._UsesRolloutParts

Save rollout state for the single controller.

Owns the save lock shared by trainer checkpoints and rollout snapshots, the optimizer-commit window that rollout snapshots skip, and the periodic snapshot pump.

Initialization

Store the parts this saves and coordinates with.

Parameters:

parts – The controller’s parts shared with restore.
checkpointer – Owns the trainer checkpoint directory.
timeout – The run’s checkpoint deadline, shared with the train pump.
bootstrap_identity – Identifies the run before its first trainer checkpoint; needed for snapshots taken before step 1.
teacher_coordinator – Records which teacher tagged the rows, if any.
records_prompt_groups – Whether dispatch records prompt groups in the recovery ledger, which decides what a trainer checkpoint saves.

property stop_requested: bool#: Whether the pump stopped the run after the pre-step deadline snapshot.

async begin_optimizer_commit(*, drain_snapshot: bool) → None#

Open the window rollout snapshots skip, until end_optimizer_commit.

Parameters:: drain_snapshot – First wait for a snapshot that is already saving the old trainer identity, as a step that publishes a new trainer version before its cleanup must.

end_optimizer_commit() → None#: Close the window once the step’s TQ cleanup and counters are published.

async trainer_checkpoint() → collections.abc.AsyncIterator[None]#: Hold off rollout snapshots while a trainer checkpoint saves or publishes.

trainer_checkpoint_published(step: int) → None#

Wake the pump once the trainer checkpoint for step is published.

A trainer checkpoint holds no environment state, so the pump waits for it to be durable, off the training loop, and then captures the environment anchored to it.

async _save_data_plane_checkpoint( checkpoint_path: nemo_rl.utils.checkpoint.PathLike, *, train_steps: int, trainer_version: int, current_epoch: int, replay_metadata: Optional[nemo_rl.algorithms.async_utils.replay_buffer.TQReplayMetadataState] = None, rollout_recovery_payload_sha256: Optional[str] = None, rollout_recovery_group_count: Optional[int] = None, ) → None#

Save a required TQ snapshot inside an SC checkpoint bundle.

A sampler with replay-buffer recovery writes an authoritative native TQ snapshot bound to its metadata-only replay index by a digest. Other samplers retain shadow-mode snapshots until their recovery contract is defined. Failures propagate so a finalized bundle never silently omits the advertised data-plane component.

async capture( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, checkpoint_path: nemo_rl.utils.checkpoint.PathLike, kind: nemo_rl.algorithms.rollout_checkpointing.state.CheckpointKind, ) → nemo_rl.algorithms.rollout_checkpointing.state.CapturedRolloutState#

Capture restart-facing controller state and save TQ under the barrier.

A trainer checkpoint saves TQ only with checkpointing.save_data_plane, and refuses rows an unfinished streamed step still owns. A rollout snapshot always saves TQ and re-indexes those rows in the persisted cut only; the live trainer keeps accumulating gradients without modification.

_begin_optimizer_commit() → None#: Mark the step’s optimizer update as applied until its counters publish.

async _enter_environment_phase_unless_a_step_commits( stack: contextlib.AsyncExitStack, environment_phase: contextlib.AbstractAsyncContextManager[Optional[nemo_rl.environments.checkpointing.interfaces.EnvironmentSnapshot]], ) → tuple[bool, Optional[nemo_rl.environments.checkpointing.interfaces.EnvironmentSnapshot]]#

Enter the environment phase on stack, or stop it once a step commits.

The snapshot would be thrown away at the trainer-state check anyway, so the environment is not kept paused while it saves for it. Cancelling the entry leaves no cleanup: the phase resumes the environment and reopens dispatch on any exception. Returns whether the phase was entered; a phase that finished entering is on stack, and the trainer-state check decides whether its snapshot is kept.

async _resolve_anchor() → Optional[tuple[pathlib.Path, Optional[str]]]#

Choose the durable trainer checkpoint this snapshot attaches to.

Before the first train step that is the bootstrap anchor, identified by this run’s fingerprint; after it, step_<N> for the current step.

Returns:: The anchor directory and the bootstrap fingerprint (None past step 0), or None when the matching trainer checkpoint is not durable yet and the attempt is skipped.

async _publish( tmp_path: pathlib.Path, final_path: pathlib.Path, captured: nemo_rl.algorithms.rollout_checkpointing.state.CapturedRolloutState, manifest: nemo_rl.algorithms.rollout_checkpointing.manifest.RolloutSnapshotManifest, ) → tuple[int, float, float]#

Write the snapshot’s files, then commit it atomically under retention.

Returns:: The controller files’ on-disk bytes, and the seconds spent writing them and committing the snapshot.

async save_snapshot( *, force: bool = False, ) → nemo_rl.algorithms.rollout_checkpointing.state.RolloutCheckpointSaveResult#: Publish one rollout-only snapshot anchored to durable trainer state.

_warn_if_snapshots_keep_losing_to_train_steps( consecutive_discards: int, ) → None#: Warn once per streak when train steps keep discarding turn snapshots.

_log_outcome( *, outcome: nemo_rl.algorithms.metric_utils.RolloutCheckpointAttemptOutcome, reason: nemo_rl.algorithms.metric_utils.RolloutCheckpointAttemptReason, attempt_duration_seconds: float, ) → None#: Record every scheduled checkpoint attempt, including no-op cuts.

async run() → None#: Persist rollout state periodically, including during streamed train.

nemo_rl.algorithms.rollout_checkpointing.snapshots#

Module Contents#

Classes#

Functions#

API#

`nemo_rl.algorithms.rollout_checkpointing.snapshots`#