`nemo_rl.experience.rollout_recovery`#

Controller-owned lineage for recoverable token-capture prompt groups.

The ledger deliberately contains control-plane metadata only. Token tensors and router-replay payloads remain in TQ. The versioned state_dict boundary here is what the controller writes into rollout_recovery.pt at each checkpoint and reads back on restore.

Module Contents#

Classes#

`PromptGroupPhase`	Durable sampler-admission phase for an unfinished prompt group.
`RecoveryGranularity`	Unit of completed work reused after a live failure or process restart.
`RolloutAttemptStatus`	Lifecycle of one physical Gate execution attempt.
`PromptGroupStatus`	Ownership lifecycle of one logical prompt group.
`PromptRef`	Small durable locator for a prompt owned by the input dataset.
`RolloutAttemptRecord`	One physical attempt for a stable logical sibling.
`RolloutSiblingRecord`	One stable GRPO generation slot and its physical attempts.
`PromptGroupRecoveryRecord`	Lineage and ownership for one logical prompt group.
`ParsedRolloutRecoveryState`	Validated controller and ledger state loaded from one checkpoint sidecar.
`SiblingSealResult`	One terminal sibling result waiting for an atomic prompt-group seal.
`RolloutRecoveryLedger`	In-memory source of truth for token-capture rollout ownership.

Functions#

`_reject_unknown_fields`	Reject fields that are not part of the current versioned schema.
`_validate_prompt_identity`	Require a runtime prompt to resolve the ledger’s durable dataset key.
`_new_attempt`
`_receipt_staging_keys`	Validate a terminal Gate receipt and return its ordered staging keys.
`_validate_batch_shortfall`	Return a defensive copy of per-step permanent rollout losses.
`_validate_finalizer_metrics`	Copy checkpoint-owned per-group finalizer metrics.
`_validate_finalizer_rewards`	Validate per-group observations without retaining checkpoint aliases.
`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`
`SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS`
`ROLLOUT_RECOVERY_STATE_FILENAME`
`RolloutRecoveryState`
`_LEDGER_STATE_FIELDS`
`_SIDECAR_STATE_FIELDS`
`_GROUP_STATE_FIELDS`
`_PROMPT_REF_STATE_FIELDS`
`_SIBLING_STATE_FIELDS`
`_ATTEMPT_STATE_FIELDS`

API#

nemo_rl.experience.rollout_recovery.ROLLOUT_RECOVERY_SCHEMA_VERSION#: 3

nemo_rl.experience.rollout_recovery.SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS#: None

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

nemo_rl.experience.rollout_recovery.RolloutRecoveryState: TypeAlias#: None

nemo_rl.experience.rollout_recovery._LEDGER_STATE_FIELDS#: ‘frozenset(…)’

nemo_rl.experience.rollout_recovery._SIDECAR_STATE_FIELDS#: ‘frozenset(…)’

nemo_rl.experience.rollout_recovery._GROUP_STATE_FIELDS#: ‘frozenset(…)’

nemo_rl.experience.rollout_recovery._PROMPT_REF_STATE_FIELDS#: ‘frozenset(…)’

nemo_rl.experience.rollout_recovery._SIBLING_STATE_FIELDS#: ‘frozenset(…)’

nemo_rl.experience.rollout_recovery._ATTEMPT_STATE_FIELDS#: ‘frozenset(…)’

nemo_rl.experience.rollout_recovery._reject_unknown_fields( mapping: collections.abc.Mapping[Any, Any], *, expected: frozenset[str], context: str, ) → None[source]#: Reject fields that are not part of the current versioned schema.

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

Bases: enum.StrEnum

Durable sampler-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.RecoveryGranularity[source]#

Bases: enum.StrEnum

Unit of completed work reused after a live failure or process restart.

Initialization

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

SIBLING#: ‘sibling’

PROMPT_GROUP#: ‘prompt_group’

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

Bases: enum.StrEnum

Lifecycle of one physical Gate execution attempt.

Initialization

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

RESERVED#: ‘reserved’

DISPATCHED#: ‘dispatched’

SEALED#: ‘sealed’

FAILED#: ‘failed’

ABANDONED#: ‘abandoned’

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

Bases: enum.StrEnum

Ownership lifecycle of one logical prompt group.

Initialization

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

GENERATING#: ‘generating’

READY_TO_FINALIZE#: ‘ready_to_finalize’

FINALIZING#: ‘finalizing’

FINALIZATION_UNKNOWN#: ‘finalization_unknown’

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

Small durable locator for a prompt owned by the input dataset.

The persistence layer will resolve this reference and validate the dataset identity before redispatch. The full DatumSpec is runtime-only state and is deliberately excluded from the serialized ledger.

sample_id: str#: None

task_name: Optional[str]#: None

__post_init__() → 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]#: Require a runtime prompt to resolve the ledger’s durable dataset key.

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

One physical attempt for a stable logical sibling.

attempt_uuid: uuid.UUID#: None

status: nemo_rl.experience.rollout_recovery.RolloutAttemptStatus#: None

receipt: Optional[dict[str, Any]]#: None

reward: Optional[float]#: None

mask_sample: Optional[bool]#: None

reward_checks: nemo_rl.experience.reward_penalties.RewardChecks | None#: None

reward_log_context: nemo_rl.experience.reward_penalties.RewardLogContext | None#: None

staging_keys: list[str]#: ‘field(…)’

property attempt_id: str#: Return the compact external representation of this attempt UUID.

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

One stable GRPO generation slot and its physical attempts.

generation_index: int#: None

attempts: list[nemo_rl.experience.rollout_recovery.RolloutAttemptRecord]#: None

property current_attempt: nemo_rl.experience.rollout_recovery.RolloutAttemptRecord#

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

Lineage and ownership for one logical prompt group.

group_id: str#: None

admission_id: str#: None

prompt_id: str#: None

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

task_source: Optional[str]#: None

recovery_granularity: nemo_rl.experience.rollout_recovery.RecoveryGranularity#: None

runtime_prompt_payload: Optional[nemo_rl.data.interfaces.DatumSpec]#: None

expected_generations: int#: None

target_step: Optional[int]#: None

start_weight_version: int#: None

siblings: list[nemo_rl.experience.rollout_recovery.RolloutSiblingRecord]#: None

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

status: nemo_rl.experience.rollout_recovery.PromptGroupStatus#: None

property prompt_payload: nemo_rl.data.interfaces.DatumSpec#: Return the runtime prompt required to redispatch unfinished work.

property logical_rollout_ids: list[str]#

property gate_rollout_ids: list[str]#

logical_rollout_id(generation_index: int) → str[source]#: Derive the stable sibling ID instead of storing another UUID string.

gate_rollout_id(generation_index: int) → str[source]#: Derive the physical Gate ID from group, sibling, and attempt UUID.

property sealed_generation_indices: list[int]#

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.RolloutRecoveryState#: None

batch_shortfall: dict[int, int]#: None

sampler_stamps_target_steps: Optional[bool]#: None

finalizer_metrics_by_group: dict[str, dict[str, float]]#: None

finalizer_rewards_by_group: dict[str, list[nemo_rl.experience.reward_penalties.FinalizedReward]]#: None

reward_settings: nemo_rl.experience.reward_penalties.CaptureRewardSettings | None#: None

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

One terminal sibling result waiting for an atomic prompt-group seal.

gate_rollout_id: str#: None

receipt: Optional[dict[str, Any]]#: None

reward: float#: None

mask_sample: bool#: None

reward_checks: nemo_rl.experience.reward_penalties.RewardChecks | None#: None

reward_log_context: nemo_rl.experience.reward_penalties.RewardLogContext | None#: None

nemo_rl.experience.rollout_recovery._new_attempt() → nemo_rl.experience.rollout_recovery.RolloutAttemptRecord[source]#

nemo_rl.experience.rollout_recovery._receipt_staging_keys( receipt: Optional[dict[str, Any]], ) → list[str][source]#: Validate a terminal Gate receipt and return its ordered staging keys.

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

In-memory source of truth for token-capture rollout ownership.

Controller-owned mutations require a live data-plane cut so lineage cannot change outside the checkpoint barrier’s consistent snapshot boundary.

Initialization

groups() → list[nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord][source]#

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: Optional[int], start_weight_version: int, task_source: Optional[str] = None, recovery_granularity: nemo_rl.experience.rollout_recovery.RecoveryGranularity = RecoveryGranularity.SIBLING, admitted: bool = True, group_id: Optional[str] = None, admission_id: Optional[str] = None, prompt_ref: Optional[nemo_rl.experience.rollout_recovery.PromptRef] = None, ) → nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord[source]#: Create one logical group and its first physical sibling attempts.

mark_group_admitted( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, *, target_step: Optional[int], start_weight_version: int, ) → None[source]#: Commit sampler admission without replacing sibling lineage.

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-reconstructed prompt after identity validation.

prepare_for_restart( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, ) → None[source]#: Apply each group’s persisted restore policy to interrupted attempts.

_abandon_entire_group( record: nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord, ) → None[source]#

Discard every current sibling when an incomplete group is atomic.

Sealed staging rows become unreferenced here. The controller’s restore inventory pass removes those rows from TQ before redispatch.

assert_checkpoint_safe() → None[source]#: Reject states whose canonical publication outcome is ambiguous.

expected_staging_keys() → set[str][source]#: Return staged token rows still owned by sealed sibling attempts.

prepare_incomplete_retry( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, ) → nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord[source]#: Mint fresh physical attempts according to the persisted granularity.

mark_group_dispatched( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, *, generation_indices: Optional[list[int]] = None, ) → None[source]#: Move the selected current sibling attempts to dispatched.

mark_sibling_sealed( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, *, generation_index: int, gate_rollout_id: str, receipt: Optional[dict[str, Any]], reward: float, mask_sample: bool, reward_checks: nemo_rl.experience.reward_penalties.RewardChecks | None = None, reward_log_context: nemo_rl.experience.reward_penalties.RewardLogContext | None = None, ) → None[source]#: Record one streamed sibling receipt as soon as the row arrives.

mark_group_sealed( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, results: dict[int, nemo_rl.experience.rollout_recovery.SiblingSealResult], ) → None[source]#: Atomically seal one complete prompt-group-scoped physical cohort.

abandon_unsealed( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, ) → None[source]#: Abandon failed work at the group’s persisted recovery granularity.

finalization_inputs( group_id: str, ) → tuple[list[str], list[str], list[Optional[dict[str, Any]]], list[float], list[bool], list[nemo_rl.experience.reward_penalties.RewardChecks | None], list[nemo_rl.experience.reward_penalties.RewardLogContext | None]][source]#: Return sealed finalization inputs in stable sibling order.

mark_finalization_started( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, ) → None[source]#

mark_finalization_unknown( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, ) → None[source]#

discard_group( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, ) → None[source]#: Drop a group only after its external TQ/Gate ownership is cleaned.

discard_canonical_groups( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_ids: set[str], ) → int[source]#: Prefer canonical TQ ownership over a stale unfinished sidecar row.

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

__len__() → int[source]#

__contains__(group_id: object) → bool[source]#

state_dict() → dict[str, Any][source]#: Return the versioned metadata persisted in rollout_recovery.pt.

classmethod from_state_dict(state: dict[str, Any]) → Self[source]#: Restore and validate a ledger metadata envelope.

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

classmethod _group_from_state( raw_group: Any, *, seen_attempt_uuids: set[uuid.UUID], ) → nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord[source]#

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

static _copy_group( record: nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord, ) → nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord[source]#: Copy mutable lineage metadata without duplicating the prompt payload.

static _require_sibling( record: nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord, generation_index: int, ) → nemo_rl.experience.rollout_recovery.RolloutSiblingRecord[source]#

static _require_group_status( record: nemo_rl.experience.rollout_recovery.PromptGroupRecoveryRecord, *, allowed: set[nemo_rl.experience.rollout_recovery.PromptGroupStatus], transition: str, ) → None[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._validate_finalizer_metrics( value: object, ) → dict[str, dict[str, float]][source]#: Copy checkpoint-owned per-group finalizer metrics.

nemo_rl.experience.rollout_recovery._validate_finalizer_rewards( value: object, ) → dict[str, list[nemo_rl.experience.reward_penalties.FinalizedReward]][source]#: Validate per-group observations without retaining checkpoint aliases.

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, finalizer_metrics_by_group: dict[str, dict[str, float]] | None = None, finalizer_rewards_by_group: dict[str, list[nemo_rl.experience.reward_penalties.FinalizedReward]] | None = None, reward_settings: nemo_rl.experience.reward_penalties.CaptureRewardSettings | None = None, canonical_group_ids: set[str] | None = None, ) → 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`#