`nemo_rl.algorithms.async_utils.replay_buffer`#

Module Contents#

Classes#

`ReplayBufferImpl`	Replay buffer storing per-prompt groups.
`ReplayBuffer`
`TQReplayBuffer`	Meta cache + TQ writer with reserve-then-commit slot semantics.

API#

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

Bases: nemo_rl.algorithms.async_utils.interfaces.ReplayBufferProtocol

Replay buffer storing per-prompt groups.

A single entry corresponds to 1 prompt repeated by grpo.num_generations_per_prompt (required to compute per-prompt advantages).

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.

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, int]#: Restore inside the actor and return only compact coordination metadata.

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)#: 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

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 – Per-group sample_id prefix; defaults to a fresh uuid4.

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 state_dict(*, saved_capacity: int) → dict[str, Any]#

Serialize ready groups (meta + DataPlane payloads) for checkpointing.

Snapshots the ready slots synchronously on the event loop first, then fetches each group’s rows from the DataPlane. Unready reservations are in-flight rollouts and are dropped, matching legacy semantics. The snapshot stays consistent during the async fetch: concurrent commits only append/flip other slots, and the train pump — the only remover — is the caller itself; groups committed mid-save land in the next checkpoint.

Parameters:: saved_capacity – max_buffered_rollouts at save time, recorded so load_state_dict can report capacity changes across restarts.
Returns:: {"partition_id": ..., "saved_capacity": ..., "groups": [{"meta", "start_weight", "end_weight", "target_step", "group_id", "fields_data"}, ...]}.
Return type:: Envelope

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

Validate and re-put checkpointed groups into the buffer.

The preflight runs entirely before any DataPlane write (legacy precedent: validate, then truncate):

Validate the envelope and raise ValueError on malformed state.
Truncate to max_groups, keeping the freshest groups, so the restored count can never exceed the buffer’s capacity. Groups carrying a target_step are never truncated — an over-capacity in-order checkpoint raises instead (see Raises).

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 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).

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), or if target-stamped groups exceed max_groups.

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

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

__len__() → int#

async _call_dp(method_name: str, **kwargs: Any) → Any#: Call a DataPlaneClient method, awaiting Ray remotes if needed.

nemo_rl.algorithms.async_utils.replay_buffer#

Module Contents#

Classes#

API#

`nemo_rl.algorithms.async_utils.replay_buffer`#