`nemo_rl.algorithms.rollout_checkpointing.state`#

Rollout state the single controller saves and restores, and the parts that hold it.

Module Contents#

Classes#

`RolloutCheckpointSaveResult`	Outcome returned by one rollout checkpoint save attempt.
`RolloutRuntimeState`	Rollout state that the pumps change and every checkpoint saves.
`CapturedRolloutState`	Controller state captured under one exclusive cut, with its TQ save.
`TrainingProgress`	The controller’s training counters, read when a checkpoint needs them.
`ResumeSource`	The checkpoint setup restored from; the defaults describe a fresh run.
`RolloutParts`	The controller’s parts that both saving and restoring rollout state use.
`_UsesRolloutParts`	Unpack the shared parts into the attributes both sides read.

Data#

CheckpointKind

API#

class nemo_rl.algorithms.rollout_checkpointing.state.RolloutCheckpointSaveResult#

Outcome returned by one rollout checkpoint save attempt.

saved: bool#: None

reason: nemo_rl.algorithms.metric_utils.RolloutCheckpointAttemptReason#: None

class nemo_rl.algorithms.rollout_checkpointing.state.RolloutRuntimeState#

Rollout state that the pumps change and every checkpoint saves.

The controller owns the one instance; restore fills it before the pumps start, and checkpoints read it under the exclusive data-plane cut.

replacement_reserve: collections.deque[nemo_rl.data.interfaces.DatumSpec]#: ‘field(…)’

batch_shortfall: dict[int, int]#: ‘field(…)’

sampler_stamps_target_steps: bool#: False

finalizer_metrics_by_group: dict[str, dict[str, float]]#: ‘field(…)’

finalizer_rewards_by_group: dict[str, list[nemo_rl.experience.reward_penalties.FinalizedReward]]#: ‘field(…)’

nemo_rl.algorithms.rollout_checkpointing.state.CheckpointKind#: None

class nemo_rl.algorithms.rollout_checkpointing.state.CapturedRolloutState#

Controller state captured under one exclusive cut, with its TQ save.

dataloader_state: dict[str, Any]#: None

sampler_dispatch_index: int#: None

next_nemo_gym_task_index: int#: None

replacement_reserve: list[nemo_rl.data.interfaces.DatumSpec]#: None

replay_metadata: Optional[nemo_rl.algorithms.async_utils.replay_buffer.TQReplayMetadataState]#: None

rollout_recovery_payload: Optional[bytes]#: None

rollout_recovery_group_count: Optional[int]#: None

replay_row_count: int#: None

staging_row_count: int#: None

rolled_back_train_group_count: int#: None

mutation_version: int#: None

tq_save_seconds: float#: None

class nemo_rl.algorithms.rollout_checkpointing.state.TrainingProgress#

The controller’s training counters, read when a checkpoint needs them.

train_steps: int#: None

trainer_version: int#: None

current_epoch: int#: None

class nemo_rl.algorithms.rollout_checkpointing.state.ResumeSource#

The checkpoint setup restored from; the defaults describe a fresh run.

path: Optional[str]#: None

from_rollout_snapshot: bool#: False

data_plane_metadata: Optional[nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneCheckpointMetadata]#: None

rollout_checkpoint_load_metrics: Optional[dict[str, float]]#: None

class nemo_rl.algorithms.rollout_checkpointing.state.RolloutParts#

The controller’s parts that both saving and restoring rollout state use.

Neither side changes which parts these are; runtime_state is the one mutable object, shared with the controller’s pumps.

master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig#: None

dp_client: Any#: None

partition_id: str#: None

buffer: nemo_rl.algorithms.async_utils.replay_buffer.TQReplayBuffer#: None

dataloader: Any#: None

sampler: nemo_rl.algorithms.async_utils.staleness_sampler.PromptGroupSampler#: None

rollout_manager: nemo_rl.experience.rollout_manager.RolloutManager#: None

data_plane_barrier: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneCheckpointBarrier#: None

environment_checkpoint: nemo_rl.environments.checkpointing.interfaces.EnvironmentCheckpoint#: None

runtime_state: nemo_rl.algorithms.rollout_checkpointing.state.RolloutRuntimeState#: None

read_progress: Callable[[], nemo_rl.algorithms.rollout_checkpointing.state.TrainingProgress]#: None

log_metrics: Callable[..., None]#: None

trace: nemo_rl.algorithms.rollout_checkpointing.trace.RolloutTrace#: None

class nemo_rl.algorithms.rollout_checkpointing.state._UsesRolloutParts( parts: nemo_rl.algorithms.rollout_checkpointing.state.RolloutParts, )#

Unpack the shared parts into the attributes both sides read.

Initialization