nemo_rl.algorithms.rollout_checkpointing.inventory

Checks that TQ holds exactly the rows a checkpoint’s controller state references.

Saving and restoring both run them, so they are plain functions.

Module Contents

Functions

capture_reward_settings	Save run-scoped settings once, alongside raw captured rewards.
validate_replay_inventory	Require the canonical TQ keys to match the SC replay index exactly.
validate_rollout_recovery_inventory	Validate staging ownership while the caller holds a stable cut.

API

nemo_rl.algorithms.rollout_checkpointing.inventory.capture_reward_settings(

master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig,

) → nemo_rl.experience.reward_penalties.CaptureRewardSettings | None

Save run-scoped settings once, alongside raw captured rewards.

async nemo_rl.algorithms.rollout_checkpointing.inventory.validate_replay_inventory(

replay_metadata: nemo_rl.algorithms.async_utils.replay_buffer.TQReplayMetadataState,

*,

dp_client: Any,

partition_id: str,

) → None

Require the canonical TQ keys to match the SC replay index exactly.

Live checkpoint callers must hold the exclusive data-plane barrier so commits and clears cannot race this inventory read. Restore calls are also safe before the rollout and train pumps start any live writers.

async nemo_rl.algorithms.rollout_checkpointing.inventory.validate_rollout_recovery_inventory(

cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut,

*,

dp_client: Any,

recovery_ledger: nemo_rl.experience.rollout_recovery.RolloutRecoveryLedger,

staging_partition: str,

replay_metadata: Optional[nemo_rl.algorithms.async_utils.replay_buffer.TQReplayMetadataState],

clear_unreferenced: bool,

additional_expected_staging_keys: Optional[set[str]] = None,

) → int

Validate staging ownership while the caller holds a stable cut.