`nemo_rl.algorithms.async_utils.trajectory_collector`#

Module Contents#

Classes#

AsyncTrajectoryCollector

Collects trajectories asynchronously and adds them to replay buffer.

Functions#

`_stamped_task_indices`	Stamped per-group ordinals of a pre-repeat slice ([] when unstamped).
`_unanimous_task_index`	Return the task index every row agrees on, or `None`.

Data#

`TokenizerType`
`_MAX_NEMO_GYM_STREAM_RETRIES`
`_NEMO_GYM_RETRY_DELAY_BASE_SECONDS`
`_REPLAY_BUFFER_MAX_BACKOFF_SECONDS`
`_WAKE_RETRY_INTERVAL_S`

API#

nemo_rl.algorithms.async_utils.trajectory_collector.TokenizerType#: None

nemo_rl.algorithms.async_utils.trajectory_collector._MAX_NEMO_GYM_STREAM_RETRIES#: 3

nemo_rl.algorithms.async_utils.trajectory_collector._NEMO_GYM_RETRY_DELAY_BASE_SECONDS#: 1.0

nemo_rl.algorithms.async_utils.trajectory_collector._REPLAY_BUFFER_MAX_BACKOFF_SECONDS#: 0.5

nemo_rl.algorithms.async_utils.trajectory_collector._WAKE_RETRY_INTERVAL_S#: 0.1

nemo_rl.algorithms.async_utils.trajectory_collector._stamped_task_indices( batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], ) → list[int]#: Stamped per-group ordinals of a pre-repeat slice ([] when unstamped).

nemo_rl.algorithms.async_utils.trajectory_collector._unanimous_task_index(rows: list[Any]) → Optional[int]#

Return the task index every row agrees on, or None.

A prompt group is one prompt repeated num_generations times, so its stamped rows all carry the same ordinal. Returns None for unstamped rows (legacy runs) rather than guessing.

class nemo_rl.algorithms.async_utils.trajectory_collector.AsyncTrajectoryCollector( policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, tokenizer: nemo_rl.algorithms.async_utils.trajectory_collector.TokenizerType, task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], master_config: nemo_rl.algorithms.grpo.MasterConfig | nemo_rl.algorithms.ppo.MasterConfig, replay_buffer: Any, start_step: int = 0, teacher_worker_groups: Optional[dict[str, Any]] = None, alias_to_group_alias: Optional[dict[str, str]] = None, on_policy_distillation_cfg: Optional[dict[str, Any]] = None, next_nemo_gym_task_index: int = 0, processor: Any = None, pending_batch: Optional[nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec]] = None, ordinals_frontier_aligned: bool = True, resume_frontier_ordinal: Optional[int] = None, resume_covered_task_indices: Optional[list[int]] = None, trace_carrier: Optional[dict[str, str]] = None, )#

Collects trajectories asynchronously and adds them to replay buffer.

Initialization

_calculate_target_weights(generation_weight_version: int) → list[int]#

Calculate target weight versions for given generation weight version.

The list of versions returned enumerate the possible version a generation server can target. These versions are looped over to see what training step they can target. If all target versions are exhausted, this generation server will remain idle until the next weight update.

Example: generation_weight_version = 10 generation_lead_steps = 4

Generation lead usually equals maximum trajectory age, but PPO critic warmup can temporarily configure them independently.

Returns:: [11, 12, 13, 14] # Meaning this generation server can create trajectories for training step 11, 12, 13, 14

_get_next_target_for_generation( generation_weight_version: int, ) → Optional[int]#: Get the next target weight that needs generation (if any).

set_weight_version(version: int) → None#

set_generation_window( *, weight_version: int, generation_lead_steps: int, max_trajectory_age_steps: int, ) → None#: Update the PPO generation version, lead, and buffer-validity age.

_should_pause_for_generation_limits() → bool#: Check if collection should be paused due to generation limits.

start_collection( dataloader: torchdata.stateful_dataloader.StatefulDataLoader | nemo_rl.data.dataloader.CyclingDataLoader, ) → None#: Start collecting trajectories from dataloader.

is_data_exhausted() → bool#: Check if collection stopped because the dataloader ran out of data.

get_status() → dict#: Return a snapshot of the collector’s internal state for driver-side diagnostics.

_mark_collection_failed(error: Exception) → None#: Record the first collection-loop failure.

_collection_loop()#

Run the collection loop in background thread.

Prompts left over from a gap-fill slice (_pending_batch) are consumed before the next dataloader pull, so a partially used batch is never discarded. The dataloader counts as exhausted only when the iterator drains with no pending prompts remaining.

_run_collection_loop( dataloader: torchdata.stateful_dataloader.StatefulDataLoader | nemo_rl.data.dataloader.CyclingDataLoader, ) → None#: Body of :meth:_collection_loop, inside the driver’s trace context.

_stamp_task_indices( batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], ) → bool#

Assign one stable, monotonic task index to every prompt in a batch.

The ordinal equals the prompt’s global position in the dataloader stream, which is what frontier-aligned checkpointing keys on. Batches whose extra_env_info rows are not dicts cannot carry a stamp; the collector then permanently falls back to live-cursor checkpoints.

Returns:: Whether the batch was stamped.

_capture_dataloader_state() → Optional[dict]#: Snapshot the dataloader state, or None when it has no state_dict.

_filter_covered_rows( batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], ) → Optional[nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec]]#

Drop re-yielded rows a frontier restore already accounts for.

After a frontier-aligned restore the dataloader re-yields the window between the cut and the old cursor. Rows already covered (ordinal below the cut, or in the covered set) are dropped; the normal gap-fill path regenerates the rest.

Returns:: The (possibly row-filtered) batch, or None when every row was already covered.

get_checkpoint_dataloader_state( frontier_ordinal: int, ) → dict[str, Any]#

Return the dataloader state a checkpoint should persist.

Returns the newest ring snapshot at or below frontier_ordinal, so a resume re-yields everything past it. Falls back to the live cursor when ordinals are not frontier-aligned or the snapshot ring no longer covers the frontier.

Parameters:: frontier_ordinal – The trained-prompt count (consumed_samples).
Returns:: Mapping with dataloader_state, base_ordinal (the ordinal the saved state resumes yielding from; None on fallback), and frontier_aligned.

get_checkpoint_state(frontier_ordinal: int) → dict[str, Any]#

Return the dataloader snapshot and rollout state as one consistent pair.

Both are read while holding the pending lock so the collection loop cannot consume or replace the pending batch in between; a torn pair (cursor newer than its pending suffix) would duplicate prompts on a fallback resume.

The pending remainder is only included on fallback snapshots. On a frontier-aligned snapshot the rewound dataloader re-yields it (the restore path forces pending_batch=None), so persisting it would be dead weight in every frontier checkpoint.

The snapshot is taken at the conservative cut — the minimum of the trained frontier, the lowest ordinal still out with a rollout worker, and the lowest ordinal held in the replay buffer (buffered- but-untrained groups are covered only by the buffer, which a load_replay_buffer=false resume discards). Normally everything in flight or buffered sits at or above the frontier and the cut equals it; when a target is refilled from later prompts (after a tolerated generation failure, or when gap-filling an incomplete target restored from a checkpoint), training it can advance past another target’s in-flight groups, and cutting at the frontier would strand those prompts. Cutting below the frontier instead re-yields a window that includes already-trained prompts; the driver persists those trained ordinals in the checkpoint (TRAINED_TASK_INDICES_KEY) so the resume covers them like retained groups — nothing is skipped and nothing is re-trained.

Returns:

Mapping with dataloader (shape of

meth:: get_checkpoint_dataloader_state, plus frontier_ordinal = the cut, which the checkpoint must persist as its filter threshold) and rollouts (shape of :meth:get_rollouts_state).

_process_batch( batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], ) → Optional[nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec]]#

Process a batch, generating for one target weight.

Parameters:: batch – Prompt batch pulled from the dataloader (or carried over from a previous gap-fill remainder).
Returns:: The unconsumed remainder of batch — the whole batch when no target currently needs generation, the sliced-off suffix when this target needed fewer prompts than the batch holds, or None when every prompt was consumed. The caller re-queues it so no yielded prompt is ever discarded.

get_weight_version() → int#

check_health() → None#

Raise the stored fatal worker error, if any.

Called by the trainer between sampling iterations. When a generation worker has recorded a fatal failure (consecutive count exceeded max_generation_failures), this raises it so the training job dies instead of stalling on an empty replay buffer. Safe to call repeatedly: returns silently when no fatal error is set, and raises every time once one is.

pause() → None#: Pause trajectory collection.

resume() → None#: Resume trajectory collection.

prepare_for_refit() → None#

Pause new generation starts and optionally wait for pending generations.

Every async backend configured for in-flight weight updates, except managed Dynamo, is asked to pause generation. vLLM preserves in-flight request state with its native keep-mode pause. Backends without pause support warn and retain their existing behavior. Managed Dynamo drains active trajectories.

For non-async engines, waits for all pending generations to complete before refit.

resume_after_refit() → None#: Resume new generation starts after refit is complete.

wait_for_pending_generations() → None#: Wait for all in-flight generation threads to complete.

get_dataloader_state() → dict#: Get the current dataloader state for checkpointing.

get_efficiency_metrics() → dict[str, float]#

Return accumulated efficiency metrics (sum of durations per category).

Called by the driver process each step to merge collector-side metrics.

flush_telemetry(quiesce_timeout_s: float = 5.0) → None#

Stop collecting, then export whatever spans are still buffered here.

The driver reaps this actor with ray.kill, which runs no atexit handler, so the span processor’s pending batch would otherwise be dropped – including the last rollout batches of the run. Call this before the kill.

Quiesces first because the shutdown is terminal, not a checkpoint: once the provider is gone, a still-running loop thread or batch worker keeps opening spans against a dead processor, which drops them and logs a line per span. Clearing running and waiting for the in-flight batches is what makes the flush cover the batches it exists to save.

Joins _live_threads rather than _inflight_threads, and joins rather than polling is_alive: a batch worker leaves the latter from inside its own finally, which still sits inside the rl.grpo.generation span, so neither an empty set nor a set snapshot means the spans are closed. Thread death does.

The loop is woken before it is joined, since the refit and generation-limit waits each hold an open span and each is followed by a running re-check, so releasing them makes the loop exit rather than pick up another batch. It is joined before the batch workers because while it is alive it can still spawn one.

The loop gets at most half the budget. The two joins are sequential, so a loop wedged in a slow batch dispatch would otherwise spend all of it and leave the batch workers – whose spans are the ones this exists to save – with nothing.

Bounded rather than reusing :meth:wait_for_pending_generations, which waits indefinitely: that is right mid-run before a refit, but here the caller is on its way to ray.kill, so a wedged rollout must not be able to hold the run’s teardown open. The caller’s RPC timeout has to leave room for this budget plus the export that follows it.

Parameters:: quiesce_timeout_s – Total time to wait for the collection loop and the in-flight batch workers before flushing anyway.

_wake_waits() → None#: Release every event the collection loop can be parked on.

_drain_thread(thread: threading.Thread, deadline: float) → None#

Join thread, re-arming the pause events until it exits or time is up.

Re-armed every pass rather than set once, because the loop checks running and then clears an event: a single set landing in that window is swallowed by the clear, and the wait that follows it has nothing left to release it – the driver is on its way to ray.kill, so no refit or weight update is coming.

static _may_still_run(thread: threading.Thread) → bool#

Whether thread could still be opening spans.

is_alive() alone is false for a worker that has been registered but has not reached start() yet, which would drop it from both the join set and the warning.

static _join_until(thread: threading.Thread, deadline: float) → None#: Join thread with whatever is left of the budget, if anything.

async drain_payload_metrics() → dict[str, int | float]#

Close one drain-to-drain collector/Gym telemetry interval.

Rollout collection is concurrent with training, so the interval is not claimed to own the sampled training batch. Call-normalized metrics make intervals comparable even when their background transfer counts differ.

_build_rollouts_state( *, include_pending: bool, ) → dict[str, Any]#

Build the rollout-state mapping. Caller must hold _pending_lock.

The mapping carries the next task index and, when requested and present, the pending prompt batch under PENDING_PROMPTS_KEY. Serializing the remainder keeps yielded prompts recoverable on live-cursor checkpoints: the dataloader cursor has already advanced past them, so a checkpoint that dropped them would skip those prompts for the rest of the run.

get_rollouts_state() → dict[str, Any]#

Get collector-side rollout state (always including any pending batch).

The driver’s save path reads this through

Meth:: get_checkpoint_state, which pairs it with the dataloader snapshot under one lock; this standalone accessor serves tests and diagnostics.

_cleanup_finished_threads() → None#

_release_target(target_weight_version: int) → None#: Release the reservation owned by a completed batch worker.

_compute_teacher_logprobs( input_ids: torch.Tensor, agent_refs: list[dict[str, Any]], input_lengths: Optional[torch.Tensor] = None, multimodal_data: Optional[dict[str, Any]] = None, ) → tuple[torch.Tensor, float]#

Compute teacher logprobs for non-colocated teachers.

Groups samples by teacher, fans out in parallel, stitches results.

Parameters:

input_ids – [B, S] tokenized input tensor
agent_refs – list of B agent reference dicts
input_lengths – [B] per-sample lengths (required for sequence packing)
multimodal_data – batch-level multimodal inputs, row-aligned with input_ids and sliced per teacher