`nemo_rl.algorithms.async_utils.replay_buffer`#

Module Contents#

Classes#

`DataPlaneCheckpointBarrierTelemetry`	Bounded interval telemetry for checkpoint-induced mutation waits.
`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.
`_canonical_manifest_extra_info`	Canonicalize identity metadata without hashing advisory rollout metrics.
`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`
`CheckpointMutationKind`
`CHECKPOINT_MUTATION_KINDS`

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’

nemo_rl.algorithms.async_utils.replay_buffer.CheckpointMutationKind#: None

nemo_rl.algorithms.async_utils.replay_buffer.CHECKPOINT_MUTATION_KINDS#: ‘cast(…)’

class nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneCheckpointBarrierTelemetry#

Bounded interval telemetry for checkpoint-induced mutation waits.

blocked_by_kind: dict[nemo_rl.algorithms.async_utils.replay_buffer.CheckpointMutationKind, int]#: None

wait_durations_s: tuple[float, ...]#: None

active_mutations: int#: None

waiting_mutations: int#: None

max_waiting_mutations: int#: None

checkpoint_active: bool#: None

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

opd_full_teacher_checkpoints: NotRequired[list[str]]#: 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._canonical_manifest_extra_info( value: Any, *, path: str, ) → Any#

Canonicalize identity metadata without hashing advisory rollout metrics.

Rollout metrics may contain logger payloads that are not JSON-compatible. They remain in the serialized KVBatchMeta for post-restore logging, but do not identify the replay rows bound to the native TQ checkpoint.

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

_current_task() → asyncio.Task[Any]#: Return the task entering a barrier section and reject reentrancy.

property mutation_version: int#: Return a monotonic marker for completed outer mutation sections.

async mutation( kind: nemo_rl.algorithms.async_utils.replay_buffer.CheckpointMutationKind = 'other', ) → collections.abc.AsyncIterator[nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut]#: Yield one task-local live cut after any active checkpoint exits.

async drain_telemetry() → nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneCheckpointBarrierTelemetry#: Return and reset interval waits while preserving current state.

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 an unconsumed 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], include_message_violation_fields: bool, staging_partition_id: Optional[str] = None, 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.

property group_ids: tuple[str, ...]#: Return a stable snapshot of controller-local replay ownership.

reserve( *, weight_version: int, target_step: Optional[int] = None, group_id: Optional[str] = None, rollout_ids: Optional[list[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.
rollout_ids – Token-capture mode: the ledger-registered rollout ids this slot dispatched, recorded so cleanup can name what it owns even before a receipt exists.

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:

One when this call removes the slot, or zero if another concurrent mutation removed it while DataPlane cleanup was awaiting.

Raises:

ValueError – group_id has no live slot.

async clear_staging_keys( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, staging_keys: list[str], ) → None#: Clear known token-capture staging rows under a caller-owned cut.

async commit_finalized( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, meta: nemo_rl.data_plane.KVBatchMeta, group_min_wv: int, group_max_wv: int, *, staging_keys: Optional[list[str]] = None, ) → nemo_rl.data_plane.KVBatchMeta#

Mark a slot ready from finalizer output (token-capture mode).

Unlike :meth:commit, the canonical rows are already in TQ — the finalizer tensorized and put them — so this only fills the slot. The slot’s effective version is the group’s OLDEST call version (group_min_wv): staleness accounting stays conservative when a rollout straddles a refit.

Parameters:

cut – Live cut acquired by the owner coordinating finalization.
group_id – group_id returned by the matching reserve call.
meta – KVBatchMeta the finalizer built over its published rows.
group_min_wv – Oldest weight version any call in the group used.
group_max_wv – Newest weight version any call in the group used.
staging_keys –
The group’s staged delta keys, recorded so

meth:

remove can clear the staging partition too.

Raises:

ValueError – group_id has no live slot (removed or never reserved).

abort(group_id: str) → bool#

Drop an unready slot whose dispatch failed or was cancelled.

Token-capture mode; called from the failed dispatch path. No DataPlane rows are cleared here. Callers with sealed receipts must clear their deterministic canonical IDs and full staging manifests before dropping this ownership record. Before receipt sealing, orphan cleanup remains an explicit controlled-validation limitation.

Returns:: True when a slot was dropped; False when the group_id has no live slot (already committed+consumed or never reserved).

_delete_slot(idx: int) → None#

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

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

In token-capture mode (staging_partition_id set), clearing a group also clears its recorded staged delta rows, so eviction leaves neither canonical nor staging bytes behind.

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 claim_for_training(idxs: list[int]) → int#

Transfer ready groups from sampler ownership to an open train step.

The canonical rows remain in TQ. Their metadata stays checkpoint-visible until :meth:release_training_claims runs after optimizer success and data-plane cleanup.

training_owned_replay_groups() → list[nemo_rl.algorithms.async_utils.replay_buffer.TQReplayGroupMetadata]#: Return metadata for canonical rows owned by the open train step.

training_owned_group_ids() → set[str]#: Return stable IDs currently owned by the open train step.

release_training_claims(group_ids: list[str]) → None#: Release checkpoint ownership after consumed TQ rows are cleared.

async _remove_groups_unlocked( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_ids: list[str], *, clear_data_plane: bool, retain_training_claims: bool = False, ) → int#: Remove stable groups while the caller owns a live mutation cut.

metadata_state_dict( *, saved_capacity: int, additional_groups: Optional[list[nemo_rl.algorithms.async_utils.replay_buffer.TQReplayGroupMetadata]] = None, ) → 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 or staging partitions or the controller-local replay membership must hold a mutation cut across the complete publish/index or clear/remove transition. No writer is exempt, including post-train cleanup in _train_pump; canonical writes are not required to originate specifically from :meth:commit. The advantage stage also takes a mutation slot because the periodic checkpoint pump runs concurrently with _train_pump. In-flight reservations are intentionally omitted. additional_groups is used by periodic snapshots to re-index rows claimed by an unfinished streamed optimizer step.

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( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, *, sample_ids: list[str], ) → None#: Clear rows while the caller owns the provided live mutation cut.

nemo_rl.algorithms.async_utils.replay_buffer#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.algorithms.async_utils.replay_buffer`#