`nemo_rl.experience.rollout_recovery`#

Versioned ownership state for unfinished SingleController prompt groups.

Module Contents#

Classes#

`PromptGroupPhase`	Durable admission phase for an unfinished prompt group.
`PromptRefState`	Serializable locator for rebuilding one prompt from the dataset.
`PromptGroupRecoveryState`	Serializable ownership state for one unfinished prompt group.
`RolloutRecoveryLedgerState`	Versioned prompt-group ownership state managed by the ledger.
`RolloutRecoveryState`	Complete checkpoint sidecar for unfinished rollout scheduling state.
`PromptRef`	Stable dataset identity for rebuilding one prompt.
`PromptGroupRecoveryRecord`	In-memory ownership record for one prompt group.
`ParsedRolloutRecoveryState`	Validated controller and ledger state loaded from one checkpoint sidecar.
`RolloutRecoveryLedger`	Own prompts after dataloader advance and before canonical TQ commit.

Functions#

`_require_int`	Validate one integer field without accepting booleans.
`_prompt_task_name`
`_validate_prompt_identity`
`_validate_batch_shortfall`	Return a defensive copy of per-step permanent rollout losses.
`build_rollout_recovery_state`	Build the complete versioned sidecar from ledger and controller state.
`parse_rollout_recovery_state`	Validate and split a complete checkpoint sidecar by runtime owner.

Data#

`ROLLOUT_RECOVERY_SCHEMA_VERSION`
`ROLLOUT_RECOVERY_STATE_FILENAME`

API#

nemo_rl.experience.rollout_recovery.ROLLOUT_RECOVERY_SCHEMA_VERSION#: 1

nemo_rl.experience.rollout_recovery.ROLLOUT_RECOVERY_STATE_FILENAME#: ‘rollout_recovery.pt’

class nemo_rl.experience.rollout_recovery.PromptGroupPhase[source]#

Bases: enum.StrEnum

Durable admission phase for an unfinished prompt group.

Initialization

Initialize self. See help(type(self)) for accurate signature.

RESERVED#: ‘reserved’

ADMITTED#: ‘admitted’

class nemo_rl.experience.rollout_recovery.PromptRefState[source]#

Bases: typing.TypedDict

Serializable locator for rebuilding one prompt from the dataset.

Initialization

Initialize self. See help(type(self)) for accurate signature.

sample_id: str#: None

task_name: str | None#: None

class nemo_rl.experience.rollout_recovery.PromptGroupRecoveryState[source]#

Bases: typing.TypedDict

Serializable ownership state for one unfinished prompt group.

Initialization

Initialize self. See help(type(self)) for accurate signature.

group_id: str#: None

admission_id: str#: None

prompt_id: str#: None

prompt_ref: nemo_rl.experience.rollout_recovery.PromptRefState#: None

expected_generations: int#: None

target_step: int | None#: None

start_weight_version: int#: None

phase: str#: None

class nemo_rl.experience.rollout_recovery.RolloutRecoveryLedgerState[source]#

Bases: typing.TypedDict

Versioned prompt-group ownership state managed by the ledger.

Initialization

Initialize self. See help(type(self)) for accurate signature.

schema_version: int#: None

groups: list[nemo_rl.experience.rollout_recovery.PromptGroupRecoveryState]#: None

class nemo_rl.experience.rollout_recovery.RolloutRecoveryState[source]#

Bases: nemo_rl.experience.rollout_recovery.RolloutRecoveryLedgerState

Complete checkpoint sidecar for unfinished rollout scheduling state.

Initialization

Initialize self. See help(type(self)) for accurate signature.

batch_shortfall: NotRequired[dict[int, int]]#: None

sampler_stamps_target_steps: NotRequired[bool]#: None

class nemo_rl.experience.rollout_recovery.PromptRef[source]#

Stable dataset identity for rebuilding one prompt.

sample_id: str#: None

task_name: str | None#: None

class nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord[source]#

In-memory ownership record for one prompt group.

group_id: str#: None

admission_id: str#: None

prompt_id: str#: None

prompt_ref: nemo_rl.experience.rollout_recovery.PromptRef#: None

runtime_prompt_payload: nemo_rl.data.interfaces.DatumSpec | None#: None

expected_generations: int#: None

target_step: int | None#: None

start_weight_version: int#: None

phase: nemo_rl.experience.rollout_recovery.PromptGroupPhase#: None

property prompt_payload: nemo_rl.data.interfaces.DatumSpec#: Return the rehydrated prompt required for rollout redispatch.

class nemo_rl.experience.rollout_recovery.ParsedRolloutRecoveryState[source]#

Validated controller and ledger state loaded from one checkpoint sidecar.

ledger_state: nemo_rl.experience.rollout_recovery.RolloutRecoveryLedgerState#: None

batch_shortfall: dict[int, int]#: None

sampler_stamps_target_steps: bool | None#: None

nemo_rl.experience.rollout_recovery._require_int(value: Any, *, field: str, minimum: int) → int[source]#: Validate one integer field without accepting booleans.

nemo_rl.experience.rollout_recovery._prompt_task_name( prompt_payload: nemo_rl.data.interfaces.DatumSpec, ) → str | None[source]#

nemo_rl.experience.rollout_recovery._validate_prompt_identity( prompt_ref: nemo_rl.experience.rollout_recovery.PromptRef, prompt_payload: nemo_rl.data.interfaces.DatumSpec, *, group_id: str, ) → None[source]#

class nemo_rl.experience.rollout_recovery.RolloutRecoveryLedger[source]#

Own prompts after dataloader advance and before canonical TQ commit.

Every mutating operation requires a live data-plane cut so ownership cannot change outside the checkpoint barrier’s consistent snapshot boundary.

Initialization

reserve_group( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, *, prompt_id: str, prompt_payload: nemo_rl.data.interfaces.DatumSpec, expected_generations: int, target_step: int | None, start_weight_version: int, admitted: bool, group_id: str | None = None, admission_id: str | None = None, ) → nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord[source]#

Record ownership before the prompt can disappear from the dataloader.

Parameters:

cut – Live capability yielded by the shared data-plane barrier.
prompt_id – Dataset-level prompt identity used for diagnostics.
prompt_payload – Runtime prompt used for whole-group regeneration. Only its stable dataset reference is checkpointed.
expected_generations – Number of GRPO siblings in the prompt group.
target_step – Original gated training step, when the sampler stamps one.
start_weight_version – Policy version visible at reservation time.
admitted – Whether sampler admission already completed. This is explicit because target_step=None is also valid for admitted ungated groups.
group_id – Stable logical and canonical TQ group ID. Generated when absent.
admission_id – Stable identity shared by every prompt in one sampler admission. Defaults to group_id for single-prompt direct callers.

Returns:

A defensive copy of the new record.

mark_group_admitted( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, *, target_step: int | None, start_weight_version: int, ) → None[source]#: Attach the sampler result to a previously reserved prompt group.

bind_runtime_prompt( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, prompt_payload: nemo_rl.data.interfaces.DatumSpec, ) → None[source]#

Attach a dataset-rehydrated prompt after identity validation.

The current reference is a positional index into a map-style dataset. Recovery therefore requires dataset ordering to remain unchanged between checkpoint and restart.

get_group( group_id: str, ) → nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord[source]#: Return a record copy while sharing its immutable runtime prompt.

groups() → list[nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord][source]#: Return record copies in reservation order without cloning prompts.

discard_group( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, ) → None[source]#: Release ownership after canonical commit or intentional discard.

discard_canonical_groups( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_ids: set[str], ) → int[source]#: Drop ledger copies already owned by canonical replay metadata.

state_dict() → nemo_rl.experience.rollout_recovery.RolloutRecoveryLedgerState[source]#: Return versioned references without serializing full prompt payloads.

load_state_dict( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, state: nemo_rl.experience.rollout_recovery.RolloutRecoveryLedgerState, ) → None[source]#: Replace this empty ledger from a validated checkpoint payload.

_require_group( group_id: str, ) → nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord[source]#

__len__() → int[source]#

nemo_rl.experience.rollout_recovery._validate_batch_shortfall(value: object) → dict[int, int][source]#: Return a defensive copy of per-step permanent rollout losses.

nemo_rl.experience.rollout_recovery.build_rollout_recovery_state( ledger: nemo_rl.experience.rollout_recovery.RolloutRecoveryLedger, *, batch_shortfall: dict[int, int], sampler_stamps_target_steps: bool, ) → nemo_rl.experience.rollout_recovery.RolloutRecoveryState[source]#: Build the complete versioned sidecar from ledger and controller state.

nemo_rl.experience.rollout_recovery.parse_rollout_recovery_state( state: object, ) → nemo_rl.experience.rollout_recovery.ParsedRolloutRecoveryState[source]#: Validate and split a complete checkpoint sidecar by runtime owner.

nemo_rl.experience.rollout_recovery#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.experience.rollout_recovery`#