`nemo_rl.algorithms.async_utils.trajectory_collector`#

Module Contents#

Classes#

AsyncTrajectoryCollector

Collects trajectories asynchronously and adds them to replay buffer.

Data#

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

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

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

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 max_trajectory_age_steps = 4

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#

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

start_collection( dataloader: torchdata.stateful_dataloader.StatefulDataLoader, ) → 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.

_collection_loop()#: Run the collection loop in background thread.

_stamp_nemo_gym_task_indices( batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], ) → None#: Assign one stable, monotonic task index to every prompt in a batch.

_process_batch( batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], ) → None#: Process a single batch and generate for one target weight.

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.

For backends with an async engine in-flight weight updates allows ongoing generations to continue with their current KV caches while weights are updated. This significantly improves async performance.

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.

get_rollouts_state() → dict[str, int]#: Get collector-side rollout state for checkpointing.

_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, ) → 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)

Returns:

([B, S] teacher logprobs tensor, total_time_seconds)

async _iter_rollout_groups( repeated_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], num_generations: int, use_nemo_gym: bool, task_index_to_group_index: dict[int, int], ) → collections.abc.AsyncGenerator[nemo_rl.experience.rollouts.RolloutGroupResult, None]#: Yield prompt groups from either backend through one result type.

async _run_rollout_batch_worker( repeated_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], generation_weight_version: int, target_weight_version: int, num_generations: int, use_nemo_gym: bool, ) → None#: Own one target reservation while collecting its rollout batch.

static _build_task_index_map( repeated_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], num_generations: int, ) → dict[int, int]#: Map each Gym task index to its repeated prompt-group position.

async _enqueue_rollout_group( rollout_result: nemo_rl.experience.rollouts.RolloutGroupResult, generation_weight_version: int, target_weight_version: int, expected_prompt_groups: int, buffered_group_indices: set[int], collection_started_at: float, ) → None#: Push one prompt group to the replay buffer with bounded backoff.

async _collect_rollout_batch( repeated_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], generation_weight_version: int, target_weight_version: int, num_generations: int, use_nemo_gym: bool, ) → None#: Run one backend batch and enqueue every completed prompt group.

nemo_rl.algorithms.async_utils.trajectory_collector#

Module Contents#

Classes#

Data#

API#

`nemo_rl.algorithms.async_utils.trajectory_collector`#