`nemo_rl.experience.rollout_manager`#

Module Contents#

Classes#

`RolloutOutcome`	How :meth:`RolloutManager.generate_and_push` finished for one prompt.
`RolloutRetryPolicy`	Retry budgets for one prompt, resolved from `async_rl.rollout_failure`.
`RolloutStats`	Counters describing what the retry policy has been doing.
`RolloutTimeouts`	Deadlines for the blocking waits inside one rollout.
`_Deadline`	`asyncio.timeout` that reports expiry as a typed :class:`RolloutTimeout`.
`AsyncRolloutImpl`	Manages per-prompt multi-turn rollouts, producing a PromptGroupRecord per call.
`AsyncNemoGymRolloutImpl`	Manages per-prompt NeMo-Gym rollouts, producing a PromptGroupRecord per call.
`RolloutManager`	Routes to AsyncRolloutImpl (native async) or AsyncNemoGymRolloutImpl (NeMo-Gym), and pushes results to a TQReplayBuffer.

Functions#

`_classify_generation_failure`	Wrap a generation error in the typed failure its class implies.
`_gather_cancelling_siblings`	Gather coroutines, cancelling the remainder as soon as one fails.

Data#

TokenizerType

API#

nemo_rl.experience.rollout_manager.TokenizerType#: None

class nemo_rl.experience.rollout_manager.RolloutOutcome#

Bases: str, enum.Enum

How :meth:RolloutManager.generate_and_push finished for one prompt.

Initialization

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

COMMITTED#: ‘committed’

SKIPPED#: ‘skipped’

class nemo_rl.experience.rollout_manager.RolloutRetryPolicy#

Retry budgets for one prompt, resolved from async_rl.rollout_failure.

The attempt budgets are required. They previously defaulted to 1/1/1, which contradicted RolloutFailureConfig’s 5/2/3 and put a second set of defaults in the codebase – a reader here came away believing the shipped budget was 1. The only place a retry default lives is RolloutFailureConfig; callers that need the historical no-retry behaviour ask for it by name via :meth:single_attempt.

max_infra_attempts: int#: None

max_data_attempts: int#: None

max_gym_row_attempts: int#: None

backoff_base_s: float#: 1.0

max_backoff_s: float#: 30.0

max_skipped_prompts: int#: 0

max_consecutive_dropped_prompts: int#: 0

classmethod single_attempt(

**overrides: Any,

) → nemo_rl.experience.rollout_manager.RolloutRetryPolicy#

The historical no-retry policy, with optional overrides.

An explicit choice for callers constructing a RolloutManager directly, who must not silently gain retries – not a second set of defaults.

__post_init__() → None#

backoff_for(attempt: int) → float#: Return the delay before infra attempt attempt + 1 (1-based attempts).

class nemo_rl.experience.rollout_manager.RolloutStats#

Counters describing what the retry policy has been doing.

Read by the SingleController for logging. A stall or a rising redispatch count is the only externally visible sign that the fleet is degrading, so these are not optional bookkeeping.

committed: int#: 0

skipped: int#: 0

redispatches_by_reason: dict[str, int]#: ‘field(…)’

data_retries_by_reason: dict[str, int]#: ‘field(…)’

data_failures_by_reason: dict[str, int]#: ‘field(…)’

infra_drops_by_reason: dict[str, int]#: ‘field(…)’

max_consecutive_infra_drops: int#: 0

gym_row_redispatches: int#: 0

record_redispatch(reason: str) → None#

record_data_retry(reason: str) → None#

record_data_failure(reason: str) → None#

record_infra_drop(reason: str, consecutive: int) → None#

record_gym_row_redispatch(rows: int = 1) → None#

as_metrics() → dict[str, float]#: Flatten into a metric dict for the SingleController logger.

class nemo_rl.experience.rollout_manager.RolloutTimeouts#

Deadlines for the blocking waits inside one rollout.

Resolved from async_rl.rollout_failure.nemo_gym.rollout_timeout_s and async_rl.rollout_failure.native.{generation,env}_timeout_s, which own the user-facing defaults. None means no deadline, reproducing the historical behaviour of waiting indefinitely.

rollout_s: Optional[float]#: None

generation_s: Optional[float]#: None

env_s: Optional[float]#: None

nemo_rl.experience.rollout_manager._classify_generation_failure( exc: Exception, *, prompt_idx: Any, traj_idx: int, ) → nemo_rl.experience.failures.RolloutFailure#

Wrap a generation error in the typed failure its class implies.

The original exception is preserved as __cause__; the prompt and trajectory coordinates are attached because a raw generation traceback does not say which rollout it belonged to.

Any Ray-boundary error is infrastructure here, by context. An exception raised inside a still-living generation worker – vLLM EngineDeadError, a CUDA OOM – arrives as a bare RayTaskError whose cause the boundary degraded, so classify_rollout_failure can only fall through to DATA and the prompt gets two attempts instead of the five the fleet-failure path is built for.

Scoped to this call site rather than widened in classify_rollout_failure, deliberately. Globally, “unrecognized means DATA” is the right default: it is about exceptions we can inspect and do not recognize, and flipping it would retry genuine bugs everywhere. Here we have information the classifier does not – this exception came from a generation RPC, so “that shard could not serve it” is the correct reading whatever the destroyed cause was, and re-dispatching to another shard is exactly the right response. A real bug in the worker still surfaces, chained, once the bounded infra budget runs out.

This also makes the two paths agree: part 2/4’s _generate_on_shard already maps ray.exceptions.RayError to GenerationUnavailable, so without this the same exception classified INFRA there and DATA here.

Parameters:

exc – The exception raised while generating a turn.
prompt_idx – Index of the prompt whose rollout failed.
traj_idx – Index of the failing generation within the prompt group.

Returns:

GenerationUnavailable for infrastructure failures (retriable on another shard), RolloutDataFailure otherwise.

async nemo_rl.experience.rollout_manager._gather_cancelling_siblings( coros: list[Any], ) → list[Any]#

Gather coroutines, cancelling the remainder as soon as one fails.

asyncio.gather propagates the first exception but leaves the other awaitables running detached. On the rollout path those keep occupying generation capacity for a prompt group whose result is already being discarded, so they are cancelled and drained before unwinding.

Parameters:: coros – Coroutines to run concurrently.
Returns:: Their results, in input order.

class nemo_rl.experience.rollout_manager._Deadline(seconds: Optional[float], description: str)#

asyncio.timeout that reports expiry as a typed :class:RolloutTimeout.

A bare asyncio.timeout surfaces expiry as TimeoutError, which is indistinguishable from a TimeoutError raised by the wrapped code itself. This consults expired() so only a real deadline breach is relabelled, and anything else propagates untouched.

seconds=None disables the deadline, matching asyncio.timeout semantics.

Initialization

async __aenter__() → nemo_rl.experience.rollout_manager._Deadline#

async __aexit__(exc_type, exc, tb) → Optional[bool]#

class nemo_rl.experience.rollout_manager.AsyncRolloutImpl(

tokenizer: nemo_rl.experience.rollout_manager.TokenizerType,

task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface],

num_generations_per_prompt: int,

max_seq_len: int,

max_rollout_turns: int,

policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface,

timeouts: nemo_rl.experience.rollout_manager.RolloutTimeouts = RolloutTimeouts(),

**kwargs: Any,

)#

Manages per-prompt multi-turn rollouts, producing a PromptGroupRecord per call.

Each run_rollout takes one prompt and returns num_generations_per_prompt completions generated concurrently via asyncio.gather.

Initialization

async run_rollout( input_sample: nemo_rl.data.interfaces.DatumSpec, ) → nemo_rl.experience.interfaces.PromptGroupRecord#

Run num_generations_per_prompt rollouts for one prompt.

Parameters:: input_sample – A single prompt (one DatumSpec entry).
Returns:: PromptGroupRecord with num_generations_per_prompt completions.

async _run_single_rollout( input_sample: nemo_rl.data.interfaces.DatumSpec, traj_idx: int, ) → tuple[nemo_rl.experience.interfaces.Completion, dict]#: Run one multi-turn rollout for a single generation index.

async _generate_response( message_log: list[dict], stop_strings: list[str] | None, ) → tuple[dict, torch.Tensor, dict[str, Any]]#

Generate a single-turn response for one sample.

Returns:: Tuple of (assistant_message, input_lengths, gen_metrics)

_aggregate_rollout_metrics( completions: list[nemo_rl.experience.interfaces.Completion], all_sample_metrics: list[dict], ) → dict[str, Any]#: Aggregate per-sample metrics across all completions.

class nemo_rl.experience.rollout_manager.AsyncNemoGymRolloutImpl(

tokenizer: nemo_rl.experience.rollout_manager.TokenizerType,

task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface],

num_generations_per_prompt: int,

max_seq_len: int,

max_rollout_turns: int,

generation_config: nemo_rl.models.generation.interfaces.GenerationConfig,

mask_env_flagged_samples: bool = True,

timeouts: Optional[nemo_rl.experience.rollout_manager.RolloutTimeouts] = None,

retry_policy: Optional[nemo_rl.experience.rollout_manager.RolloutRetryPolicy] = None,

stats: Optional[nemo_rl.experience.rollout_manager.RolloutStats] = None,

**kwargs: Any,

)#

Manages per-prompt NeMo-Gym rollouts, producing a PromptGroupRecord per call.

Each run_rollout takes one prompt and returns num_generations_per_prompt completions batched through a single NeMo-Gym run_rollouts call.

Initialization

async run_rollout( input_sample: nemo_rl.data.interfaces.DatumSpec, ) → nemo_rl.experience.interfaces.PromptGroupRecord#

Run num_generations_per_prompt rollouts for one prompt.

Parameters:: input_sample – A single prompt (one DatumSpec entry).
Returns:: PromptGroupRecord with num_generations_per_prompt completions.

_validate_init_params() → None#: Validate initialization parameters.

_build_inputs( input_sample: nemo_rl.data.interfaces.DatumSpec, ) → list[dict]#: Build N row dicts from input_sample, applying generation config params.

async _stream_rows( nemo_gym_env: Any, pending: list[dict], results: list[Optional[dict]], total_rows: int, timer_prefix: str, ) → Optional[dict[str, Any]]#

Dispatch pending rows and fill their slots in results as they land.

Parameters:

nemo_gym_env – The NeMo-Gym environment actor handle.
pending – Rows still awaiting a result; each carries its original _rowidx.
results – Full-length result list, mutated in place.
total_rows – Size of the original prompt group, used to validate row indices.
timer_prefix – Timer namespace forwarded to the environment.

Returns:

The environment’s timing metrics, or None if the stream ended without them.

async _run_rollouts( inputs: list[dict], timer: nemo_rl.utils.timer.Timer, timer_prefix: str, ) → tuple[list[nemo_rl.experience.interfaces.Completion], nemo_rl.data.interfaces.LLMMessageLogType, dict[str, Any]]#

Dispatch rows to NeMo-Gym; return completions, prompt, and metrics.

Rows that never arrive are re-dispatched on their own rather than by redoing the whole group. NeMo-Gym’s stream dies on the first failing row, so one bad row takes every later row with it; at num_generations_per_prompt=16 a naive whole group retry pays 16 generations to recover one. Completed rows are kept across attempts, which is the same shape as the legacy collector’s pending-group retry.

_result_to_completion( result: dict, ) → nemo_rl.experience.interfaces.Completion#: Convert one run_rollouts result dict into a Completion.

_compute_rollout_metrics( completions: list[nemo_rl.experience.interfaces.Completion], agent_name: str, ) → dict[str, Any]#: Aggregate per-sample and per-agent metrics.

class nemo_rl.experience.rollout_manager.RolloutManager( tokenizer: nemo_rl.experience.rollout_manager.TokenizerType, task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], num_generations_per_prompt: int, max_seq_len: int, max_rollout_turns: int = 1, policy_generation: Optional[nemo_rl.models.generation.interfaces.GenerationInterface] = None, generation_config: Optional[nemo_rl.models.generation.interfaces.GenerationConfig] = None, use_nemo_gym: bool = False, mask_env_flagged_samples: bool = True, tq_buffer: Optional[nemo_rl.algorithms.async_utils.replay_buffer.TQReplayBuffer] = None, timeouts: Optional[nemo_rl.experience.rollout_manager.RolloutTimeouts] = None, retry_policy: Optional[nemo_rl.experience.rollout_manager.RolloutRetryPolicy] = None, )#

Routes to AsyncRolloutImpl (native async) or AsyncNemoGymRolloutImpl (NeMo-Gym), and pushes results to a TQReplayBuffer.

Initialization

property stats: nemo_rl.experience.rollout_manager.RolloutStats#: Counters describing retry/skip activity so far.

set_weight_version(version: int) → None#

Set the weight_version used for rollout tags.

Parameters:: version – Trainer weight version to stamp on future rollout tags.

async run_rollout( input_sample: nemo_rl.data.interfaces.DatumSpec, ) → nemo_rl.experience.interfaces.PromptGroupRecord#

async generate_and_push( input_sample: nemo_rl.data.interfaces.DatumSpec, *, target_step: Optional[int] = None, inflight_registry: Optional[dict[str, tuple[asyncio.Task[None], int]]] = None, ) → nemo_rl.experience.rollout_manager.RolloutOutcome#

Roll out one prompt and commit it, re-dispatching on infrastructure failure.

No prompt is discarded for infrastructure reasons. An infra failure means the fleet is unwell, not the prompt, so the attempt is retried – and because each retry re-enters generation-shard selection, it naturally lands somewhere else without this method needing to know anything about shard health. Exhausting the infra budget therefore means the failure follows the prompt across the whole fleet, which is reported as fleet-wide failure rather than absorbed.

Deterministic failures get their own, much smaller budget: another shard would reject the prompt identically, so retrying mostly burns time. One retry is still worth taking because a shard under memory pressure can return an empty generation that looks deterministic and is not.

Parameters:

input_sample – A single prompt (one DatumSpec entry).
target_step – Training step this rollout targets; stamped on the buffer slot for StalenessSampler.force_in_order.
inflight_registry – Optional controller-owned mapping from group ID to its dispatch task and start weight version.

Returns:

COMMITTED when the group reached the buffer, or SKIPPED when the prompt was given up on within a budget: its data budget within max_skipped_prompts, or its infra budget within max_consecutive_dropped_prompts. A SKIPPED prompt committed nothing, so the caller owns both its backpressure permit and the shortfall for the training step it was stamped for.

Raises:

RolloutRedispatchExhausted – The infra budget ran out and the fleet has not committed anything since max_consecutive_dropped_prompts drops ago.
RolloutDataFailure – The data budget ran out beyond max_skipped_prompts.

nemo_rl.experience.rollout_manager#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.experience.rollout_manager`#