nemo_rl.algorithms.rollout_checkpointing.restore#

Restore the rollout parts of a resumed checkpoint.

Restores the replay index (replay_buffer_metadata.pt), the recovery ledger (rollout_recovery.pt), the spare-prompt pool (replacement_reserve.pt) and, for turn-level groups from a rollout snapshot, the environment’s rollouts. Setup restores the TQ rows themselves before this runs. Runs once before the pumps start; the rollout pump then redispatches what it loaded.

Module Contents#

Classes#

RolloutRestorer

Restore the parts listed in the module docstring, then redispatch them.

Functions#

read_recovery_sidecar

Read the recovery-ledger sidecar and check it against its TQ checkpoint.

API#

async nemo_rl.algorithms.rollout_checkpointing.restore.read_recovery_sidecar(
recovery_path: pathlib.Path,
metadata: collections.abc.Mapping[str, Any],
*,
master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig,
) → nemo_rl.experience.rollout_recovery.ParsedRolloutRecoveryState#

Read the recovery-ledger sidecar and check it against its TQ checkpoint.

The native TQ metadata advertises the sidecar’s checksum, ledger schema and group count; the sidecar must match all three, and its reward settings must match this run’s when token capture is on.

class nemo_rl.algorithms.rollout_checkpointing.restore.RolloutRestorer(
parts: nemo_rl.algorithms.rollout_checkpointing.state.RolloutParts,
*,
algo_cfg: Any,
buffer_capacity: asyncio.Semaphore,
resume: nemo_rl.algorithms.rollout_checkpointing.state.ResumeSource,
)#

Bases: nemo_rl.algorithms.rollout_checkpointing.state._UsesRolloutParts

Restore the parts listed in the module docstring, then redispatch them.

Initialization

Store the parts restore reads and rebuilds.

Parameters:

  • parts – The controller’s parts shared with saving.

  • algo_cfg – The algorithm config; restore reads its prompt-group shape.

  • buffer_capacity – Replay-capacity permits; restore takes one per restored group.

  • resume – What setup restored from, which restore continues.

async restore() → bool#

Restore the parts listed in the module docstring.

Call it once, after the weights are in sync and before the pumps start. Restores the replay buffer, the recovery ledger, the spare prompts and the environment’s rollouts (turn-level only), checks the sampler cursor against them and logs restore timings.

Returns:

Whether restore queued groups only ledger-recorded dispatch can launch, so the rollout pump must record prompt groups.

_log_restore_metrics(
*,
replay_metadata_load_seconds: float,
recovery_prepare_seconds: float,
restored_replay_groups: int,
) → None#

Log rollout restore phases after the controller is ready to dispatch.

_validate_forest_checkpoint_metadata() → None#

Validate checkpoint row semantics before replay or custody changes.

async _restore_replay_buffer() → int#

Restore the local replay index for the native TQ checkpoint.

Recovery is authoritative only for samplers that explicitly support buffered-group restoration. The native snapshot and replay metadata file must both be present and agree on their manifest and group count.

async _restore_recovery_ledger(*, restored_replay_groups: int) → None#

Restore unfinished ownership for prioritized rollout-pump redispatch.

async _queue_replay_groups_for_regeneration(
cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut,
) → None#

Convert discarded canonical replay groups into fresh prompt work.

async _load_recovery_prompt(
*,
group_id: str,
sample_id: str,
) → nemo_rl.data.interfaces.DatumSpec#

Resolve and collate one stable dataset prompt reference.

_validate_sampler_cursor() → None#

Require the sampler cursor to cover every restored target step.

async _rehydrate_prompts(
cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut,
) → None#

Resolve positional prompt references against a stable map-style dataset.

This assumes the dataset exposes integer __getitem__ and retains the same ordering across checkpoint and restart.

async redispatch_restored(
launch: Callable[[nemo_rl.data.interfaces.DatumSpec, Optional[int], str], Awaitable[None]],
admit: Callable[[list[str]], Awaitable[Any]],
) → None#

Prioritize durable unfinished groups while the train pump drains TQ.

Launching happens inside the ordinary rollout pump so restored groups use the same in-flight and replay-capacity semaphores as new work. The train pump runs concurrently and releases replay capacity as it consumes canonical or newly recovered groups; therefore recovery cannot deadlock merely because the checkpoint contained more unfinished ownership records than free replay slots.

async _restore_spare_prompts() → None#

Restore spare prompts diverted before the previous run’s checkpoint.

These were pulled from the dataloader and never dispatched, so the restored dataloader resumes past them. Nothing else in the checkpoint holds them, and without this they are simply gone: one batch of the dataset per divert, plus the training step _clamp_max_num_steps had budgeted for it.

No sampler-name guard, unlike the buffer restore. Spares carry no stamp – they are prompts that never reached admit – so nothing about them depends on which sampler wrote the checkpoint. They are restored even into a run that has since switched to on_dropped_prompt="shrink", where the pool is never drawn on but is still drained back into training at the end of the dataloader.