`nemo_rl.algorithms.async_utils.replay_buffer`#

Module Contents#

Classes#

`TQReplayGroupMetadata`	Controller-local index for one training-ready group stored in TQ.
`TQReplayMetadataState`	Versioned metadata-only replay index paired with a TQ snapshot.
`DataPlaneCheckpointMetadata`	SC metadata envelope stored with a native data-plane checkpoint.
`DataPlaneMutationCut`	Live capability proving code runs inside a data-plane barrier cut.
`DataPlaneCheckpointBarrier`	Allow concurrent mutations while giving live checkpoints exclusivity.
`ReplayBufferImpl`	Replay buffer storing per-prompt groups.
`ReplayBuffer`
`TQReplayBuffer`	Meta cache + TQ writer with reserve-then-commit slot semantics.

Functions#

`_canonical_manifest_value`	Return a deterministic JSON value or reject unsupported metadata.
`replay_manifest_digest`	Return a stable digest binding replay metadata to a TQ checkpoint.

Data#

`DATA_PLANE_CHECKPOINT_DIR`
`REPLAY_BUFFER_METADATA_FILENAME`
`LEGACY_REPLAY_BUFFER_FILENAME`
`REPLACEMENT_RESERVE_FILENAME`
`REPLAY_BUFFER_METADATA_SCHEMA_VERSION`
`REPLAY_BUFFER_METADATA_STORAGE`

API#

nemo_rl.algorithms.async_utils.replay_buffer.DATA_PLANE_CHECKPOINT_DIR#: ‘data_plane’

nemo_rl.algorithms.async_utils.replay_buffer.REPLAY_BUFFER_METADATA_FILENAME#: ‘replay_buffer_metadata.pt’

nemo_rl.algorithms.async_utils.replay_buffer.LEGACY_REPLAY_BUFFER_FILENAME#: ‘replay_buffer.pt’

nemo_rl.algorithms.async_utils.replay_buffer.REPLACEMENT_RESERVE_FILENAME#: ‘replacement_reserve.pt’

nemo_rl.algorithms.async_utils.replay_buffer.REPLAY_BUFFER_METADATA_SCHEMA_VERSION#: 1

nemo_rl.algorithms.async_utils.replay_buffer.REPLAY_BUFFER_METADATA_STORAGE: Literal[tq_checkpoint]#: ‘tq_checkpoint’

class nemo_rl.algorithms.async_utils.replay_buffer.TQReplayGroupMetadata#

Bases: typing.TypedDict

Controller-local index for one training-ready group stored in TQ.

Initialization

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

meta: nemo_rl.data_plane.KVBatchMeta#: None

start_weight: int#: None

end_weight: int#: None

target_step: Optional[int]#: None

group_id: str#: None

class nemo_rl.algorithms.async_utils.replay_buffer.TQReplayMetadataState#

Bases: typing.TypedDict

Versioned metadata-only replay index paired with a TQ snapshot.

Initialization

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

schema_version: int#: None

storage: Literal[tq_checkpoint]#: None

partition_id: str#: None

saved_capacity: int#: None

manifest_digest: str#: None

groups: list[nemo_rl.algorithms.async_utils.replay_buffer.TQReplayGroupMetadata]#: None

class nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneCheckpointMetadata#

Bases: typing.TypedDict

SC metadata envelope stored with a native data-plane checkpoint.

The replay fields are present together in authoritative mode and absent in shadow mode.

Initialization

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

data_plane_checkpoint_schema_version: int#: None

single_controller_train_steps: int#: None

single_controller_trainer_version: int#: None

single_controller_epoch: int#: None

partition_id: str#: None

sampler_name: str#: None

mode: Literal[authoritative, shadow]#: None

replay_metadata_schema_version: NotRequired[int]#: None

replay_manifest_digest: NotRequired[str]#: None

replay_group_count: NotRequired[int]#: None

rollout_recovery_schema_version: NotRequired[int]#: None

rollout_recovery_payload_sha256: NotRequired[str]#: None

rollout_recovery_group_count: NotRequired[int]#: None

nemo_rl.algorithms.async_utils.replay_buffer._canonical_manifest_value( value: Any, *, path: str, ) → Any#: Return a deterministic JSON value or reject unsupported metadata.

nemo_rl.algorithms.async_utils.replay_buffer.replay_manifest_digest( groups: list[nemo_rl.algorithms.async_utils.replay_buffer.TQReplayGroupMetadata], ) → str#: Return a stable digest binding replay metadata to a TQ checkpoint.

class nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut( barrier: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneCheckpointBarrier, )#

Live capability proving code runs inside a data-plane barrier cut.

Initialization

__slots__#: (‘_barrier’, ‘_live’)

require_live() → None#: Fail when a mutation tries to reuse an absent or expired cut.

_invalidate() → None#

class nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneCheckpointBarrier#

Allow concurrent mutations while giving live checkpoints exclusivity.

At most one checkpoint holder is active. New mutations queue behind it, and a checkpoint waits for all active mutations before yielding. Every live canonical TQ commit/clear and native save must use this barrier so the snapshot and controller replay index describe the same rows.

Initialization

async mutation() → collections.abc.AsyncIterator[nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut]#: Yield a live mutation capability after any active checkpoint exits.

async checkpoint() → collections.abc.AsyncIterator[nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut]#: Yield a live capability after blocking and draining all mutations.

exception nemo_rl.algorithms.async_utils.replay_buffer.PostWriteEnrichmentError#

Bases: RuntimeError

A rollout reached TQ but failed in required post-write processing.

Initialization

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

class nemo_rl.algorithms.async_utils.replay_buffer.ReplayBufferImpl( max_size: int, drop_incomplete_targets_on_restore: bool, )#

Bases: nemo_rl.algorithms.async_utils.interfaces.ReplayBufferProtocol

Replay buffer storing per-prompt groups.

A single entry corresponds to 1 prompt repeated by the algorithm’s num_generations_per_prompt setting.

Initialization

static _rollout_metrics_turn_count_for_diagnostics( rm: dict[str, Any], ) → Optional[float]#

One scalar turn-depth per buffered trajectory for starvation diagnostics.

Supports sync multi-turn rollouts (max_turns_per_sample / avg_turns_per_sample) and NeMo Gym (turns_per_sample/max / turns_per_sample/mean).

add( trajectory: dict[str, Any], weight_version: int, target_weight_version: int, ) → str#

Add a per-prompt trajectory group with metadata.

Parameters:

trajectory – data dict
weight_version – version of the model weights used for generation
target_weight_version – version of the model weights this trajectory is intended for training

get_debug_info() → dict#: Get debug information about buffer state.

get_last_target_weight_already_generated() → int#

get_existing_target_weights() → set[int]#: Get set of target weight versions that already have trajectories.

_remove_indices(indices: Iterable[int]) → None#: Remove trajectories at the given indices.

sample( num_prompt_groups: int, current_weight_version: int, max_age_steps: int, ) → Optional[dict[str, Any]]#

Sample per-prompt trajectory groups intended for the current training step.

Only returns trajectories with target_weight_version == current_weight_version. If insufficient trajectories are available, returns None to stall training until the remaining trajectories are generated. This ensures no trajectory loses its last chance to be used for its intended training step.

Returns:: Dictionary with ‘trajectories’ and ‘avg_trajectory_age’ keys, or None if insufficient data

size() → int#: Return current buffer size.

get_held_task_indices() → list[int]#

Ordinals of every prompt group currently held in the buffer.

All held groups are untrained (sampling removes trained ones). The checkpoint cut must not exceed any of these ordinals: with checkpointing.load_replay_buffer=false the buffer is discarded on resume, and ordinals below the cut are never re-yielded.

clear() → None#: Clear the buffer.

state_dict() → dict[str, Any]#: Return serializable state for checkpointing.

save_to_path(path: str) → int#: Serialize inside the actor without materializing the buffer on the driver.

load_from_path( path: str, num_prompts_per_step: int | None = None, current_training_step: int | None = None, max_age_steps: int | None = None, ) → dict[str, Any]#

Restore inside the actor and return only compact coordination metadata.

Returns:: Mapping with num_trajectories (pre-filter count), NEXT_NEMO_GYM_TASK_INDEX_KEY (one past the highest saved task index, computed before age/step filtering; on a legacy resume this keeps used indices from being re-issued, while a frontier-aligned resume deliberately rewinds the counter to the saved base ordinal so the covered window re-yields under its original indices), and RETAINED_TASK_INDICES_KEY (the sorted task indices of the groups that survived filtering — what a frontier-aligned resume must not regenerate).

load_state_dict( state: dict[str, Any], num_prompts_per_step: int | None = None, current_training_step: int | None = None, max_age_steps: int | None = None, ) → None#

Restore replay buffer state from a checkpoint.

Parameters:

state – State returned by state_dict.
num_prompts_per_step – Number of prompt groups required for one training step. When provided, incomplete target steps can be removed or prepared for gap filling.
current_training_step – Step being resumed. When provided with num_prompts_per_step, past target steps are dropped and incomplete current/future target steps are kept for gap filling.
max_age_steps – Maximum allowed age for restored trajectories. When provided, stale trajectories are removed during restore.

Raises:

ValueError – If the checkpoint is missing required fields or has inconsistent parallel list lengths.

_prepare_for_training_step( current_step: int, num_prompts_per_step: int, ) → None#: Prepare restored state so training can resume at current_step.

static _is_valid_for_target( trajectory_version: int, target_step: int, max_age_steps: int | None, ) → bool#

_remove_stale_trajectories(max_age_steps: int) → None#

Remove restored trajectories that are stale for their target step.

Must be called while holding self._lock.

_count_for_target( target_step: int, max_age_steps: int | None = None, ) → int#

Count trajectories usable for target_step.

Must be called while holding self._lock.

_truncate_to_max_size( current_training_step: int | None = None, ) → None#

Truncate restored state to max_size after resume cleanup.

Must be called while holding self._lock.

get_trajectories_needed( target_step: int, num_prompts_per_step: int, max_age_steps: int | None = None, ) → int#: Return additional trajectories needed for target_step.

has_complete_batch( target_step: int, num_prompts_per_step: int, max_age_steps: int | None = None, ) → bool#: Return whether target_step has enough trajectories to train.

_remove_incomplete_target_steps(num_prompts_per_step: int) → None#

Remove target steps without a complete batch.

Must be called while holding self._lock.

class nemo_rl.algorithms.async_utils.replay_buffer.ReplayBuffer(max_size: int, drop_incomplete_targets_on_restore: bool)#: Bases: nemo_rl.algorithms.async_utils.replay_buffer.ReplayBufferImpl

class nemo_rl.algorithms.async_utils.replay_buffer.TQReplayBuffer( dp_client: Any, partition_id: str, *, pad_value_dict: collections.abc.Mapping[str, int], require_routed_experts: bool = False, )#

Meta cache + TQ writer with reserve-then-commit slot semantics.

meta_list, weight_list, ready_list, _group_ids are parallel; a slot stays ready=False until commit fills it.

Initialization

set_data_plane_checkpoint_barrier( barrier: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneCheckpointBarrier, ) → None#

Bind the controller’s shared checkpoint/mutation barrier once.

A private fallback barrier would not coordinate with controller-owned saves and clears, so destructive operations fail loudly until the SC actor supplies its barrier.

property data_plane_checkpoint_barrier: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneCheckpointBarrier#: Return the shared barrier used by controller and post-commit ownership.

set_post_write_enricher( enricher: Callable[[nemo_rl.data_plane.KVBatchMeta, nemo_rl.experience.interfaces.PromptGroupRecord], Awaitable[nemo_rl.data_plane.KVBatchMeta]], ) → None#: Install the required enrichment stage run before slots become ready.

reserve( *, weight_version: int, target_step: Optional[int] = None, group_id: Optional[str] = None, ) → str#

Append an unready slot tagged with weight_version.

Parameters:

weight_version – Weight version stamped on the slot.
target_step – Training step this slot targets; only consulted by StalenessSampler.force_in_order.
group_id – Pre-minted logical group ID and sample-ID prefix. The checkpoint-enabled lineage path always supplies this. None creates a fresh UUID only for untracked callers.

Returns:

group_id used by the matching commit.

async commit( group_id: str, record: nemo_rl.experience.interfaces.PromptGroupRecord, start_weight_version: int, end_weight_version: int, ) → nemo_rl.data_plane.KVBatchMeta#

Tensorize record, write N rows to TQ, and mark the slot ready.

Parameters:

group_id – group_id returned by the matching reserve call.
record – PromptGroupRecord to tensorize.
start_weight_version – Weight version stamped on the slot before rollout. The same as the one from reserve, passed again to avoid race condition when lookup.
end_weight_version – Weight version stamped on the slot after rollout.

Returns:

KVBatchMeta for the committed group.

Raises:

ValueError – group_id has no live slot (removed or never reserved).
RuntimeError – router replay is enabled but the payload has no routes.

async remove_group(group_id: str, *, remove_in_dp: bool = False) → int#

Remove the live slot identified by group_id.

Parameters:

group_id – Group identifier returned by :meth:reserve.
remove_in_dp – Whether to clear rows referenced by a committed slot.

Returns:

Number of removed slots (always one on success).

Raises:

ValueError – group_id has no live slot.

async remove(idxs: list[int], remove_in_dp: bool) → int#

Drop entries at the given indices and optionally clear them from DataPlane.

Parameters:

idxs – Entry indices to drop. Must be within [0, size).
remove_in_dp – If True, also clear the dropped rows from DataPlane.

Returns:

Number of group entries removed from the buffer.

async _remove_unlocked( drop_idxs: list[int], *, clear_data_plane: bool, ) → int#: Remove validated indices while the caller owns any required lock.

metadata_state_dict( *, saved_capacity: int, ) → nemo_rl.algorithms.async_utils.replay_buffer.TQReplayMetadataState#

Capture the controller index for ready groups without tensor payloads.

The caller must hold the exclusive side of the shared data-plane checkpoint barrier through this capture and the matching TQ save. Commits and destructive clears use shared mutation slots, so the replay index and native snapshot describe one exact set of training-ready groups. Every operation that mutates the canonical rollout partition or its controller-local replay membership must either participate in that barrier across the complete publish/index or clear/remove transition, or run in the same asyncio task as the checkpoint save. The advantage stage relies on the latter: it and _save_checkpoint both live in _train_pump, so they cannot interleave. Any new writer outside _train_pump – including future finalizer paths – must take a mutation slot; canonical writes are not required to originate specifically from :meth:commit. In-flight reservations are intentionally omitted.

async load_state_dict( state: dict[str, Any], *, max_groups: int, expected_partition_id: str, expected_group_size: int, expected_manifest_digest: str, ) → int#

Restore the local replay index for an already-restored TQ snapshot.

The replay index never contains tensor payloads and this method never writes to the DataPlane. TQ must be restored first; the caller binds the two artifacts by passing the manifest digest returned by TQ checkpoint loading.

Staleness is intentionally NOT handled here — load only loads. The train pump’s first sampler.evict drops any restored group that is outside the staleness window and releases its capacity permit, keeping eviction in one place.

Parameters:

state – Envelope produced by metadata_state_dict.
max_groups – Current max_buffered_rollouts; the restored count never exceeds it.
expected_partition_id – Partition this buffer writes to; must match the envelope.
expected_group_size – num_generations_per_prompt; every group must hold exactly this many rows (a changed group size silently breaks the group-relative baseline).
expected_manifest_digest – Digest returned by the matching native TQ checkpoint load. It must match the replay metadata file.

Returns:

Number of groups restored into the buffer.

Raises:

ValueError – If the envelope is malformed (missing keys, partition mismatch, misaligned or wrongly sized groups, duplicate sample_ids), disagrees with the native TQ snapshot, or exceeds max_groups.

count_for_target_step(target_step: int) → int#: Return how many slots are stamped with target_step.

promote_ready_group(*, to_target_step: int) → Optional[int]#

Re-stamp a finished group from a later step so it lands in this one.

Fills a hole left by a dropped prompt with generation that is already done, which is the point: the step closes immediately instead of waiting out a fresh rollout. The step it was borrowed from is returned so the caller can repay it, and the caller must – an unrepaid loan is the same hole one step later, carried forward until it reaches the last step, which has nobody to borrow from.

The furthest future step is preferred because it is due last and so has the most slack to absorb the repayment. Only ready slots qualify: an unready one is a reservation whose rollout is still running, so moving its stamp would hand this step the same wait it is trying to avoid.

Promotion can only make a step fresher, never staler. Slots are appended in dispatch order and the trainer version never decreases, so a group stamped for a later step was generated at a weight version at least as new as the ones already in this step.

Synchronous on purpose. remove deletes its indices before its own first await, so as long as nothing here yields, the index picked below cannot be shifted out from under the write by a selection running concurrently.

Parameters:: to_target_step – Training step to re-stamp the borrowed group onto – the step that lost a prompt. Must be at or ahead of the trainer version: a group re-stamped onto a step already trained is never selectable again and would only be evicted.
Returns:: The target step the group was taken from, or None when no later step has a ready group to lend.

size() → int#: Return the number of prompt-group entries currently held.

__len__() → int#

async _clear_samples_unlocked(*, sample_ids: list[str]) → None#: Clear rows while the caller holds a barrier mutation slot.

nemo_rl.algorithms.async_utils.replay_buffer#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.algorithms.async_utils.replay_buffer`#