`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.
`RequestDeadlineRegistry`	Live request deadlines, pausable while a colocated engine has switched to training.
`_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#

`_contains_post_write_enrichment_error`	Whether an error, including a rollback ExceptionGroup, is post-write.
`_nemo_gym_metric_namespace`	Return the best available namespace for NeMo-Gym rollout metrics.
`_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`
`RolloutCompletionCallback`

API#

nemo_rl.experience.rollout_manager.TokenizerType#: None

nemo_rl.experience.rollout_manager.RolloutCompletionCallback#: None

nemo_rl.experience.rollout_manager._contains_post_write_enrichment_error(error: BaseException) → bool#: Whether an error, including a rollback ExceptionGroup, is post-write.

nemo_rl.experience.rollout_manager._nemo_gym_metric_namespace( row: collections.abc.Mapping[str, Any], ) → str#: Return the best available namespace for NeMo-Gym rollout metrics.

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.RequestDeadlineRegistry#

Live request deadlines, pausable while a colocated engine has switched to training.

Initialization

add(deadline: nemo_rl.experience.rollout_manager._Deadline) → None#

discard( deadline: nemo_rl.experience.rollout_manager._Deadline, ) → None#

suspend() → None#

resume() → None#

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

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. registry puts the deadline clock in units of inference clock-time, not wall clock-time: When a colocated engine is suspended for training, inference deadlines should not tick down.

Initialization

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

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

suspend() → None#: Disarm the clock, banking whatever budget is left.

resume() → None#: Re-arm the clock with the banked budget.

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

deadline_registry: Optional[nemo_rl.experience.rollout_manager.RequestDeadlineRegistry] = None,

**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, *, rollout_ids: Optional[list[str]] = None, generation_indices: Optional[list[int]] = None, on_completion: Optional[nemo_rl.experience.rollout_manager.RolloutCompletionCallback] = None, recovery_granularity: nemo_rl.experience.rollout_recovery.RecoveryGranularity = RecoveryGranularity.SIBLING, ) → nemo_rl.experience.interfaces.PromptGroupRecord#

Run num_generations_per_prompt rollouts for one prompt.

Parameters:

input_sample – A single prompt (one DatumSpec entry).
rollout_ids – Unsupported here — token capture is NeMo-Gym only.

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, *, native_generation_data: dict[str, Any] | None = 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,

reward_penalty_config: Optional[dict[str, Any]] = None,

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

deadline_registry: Optional[nemo_rl.experience.rollout_manager.RequestDeadlineRegistry] = None,

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

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

effort_config: Optional[nemo_rl.experience.rollouts.EffortLevelsConfig] = None,

log_full_result_tables: bool = False,

**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, *, rollout_ids: Optional[list[str]] = None, generation_indices: Optional[list[int]] = None, on_completion: Optional[nemo_rl.experience.rollout_manager.RolloutCompletionCallback] = None, recovery_granularity: nemo_rl.experience.rollout_recovery.RecoveryGranularity = RecoveryGranularity.SIBLING, ) → nemo_rl.experience.interfaces.PromptGroupRecord#

Run num_generations_per_prompt rollouts for one prompt.

Parameters:

input_sample – A single prompt (one DatumSpec entry).
rollout_ids – Token-capture mode: gate-registered rollout ids, one per generation, riding each row’s run body as the opaque _ng_rollout_id key (agents stamp /ng-rollout/ from it; zero agent changes).

Returns:

PromptGroupRecord with num_generations_per_prompt completions.

_validate_init_params() → None#: Validate initialization parameters.

_build_inputs( input_sample: nemo_rl.data.interfaces.DatumSpec, *, rollout_ids: Optional[list[str]] = None, generation_indices: Optional[list[int]] = None, ) → 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]], shaping_by_rowidx: list[Optional[nemo_rl.experience.rollouts._EffortShapingMetrics]], total_rows: int, timer_prefix: str, on_completion: Optional[nemo_rl.experience.rollout_manager.RolloutCompletionCallback] = None, ) → 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.
shaping_by_rowidx – Per-row shaping metrics, populated before a completion can be published to the recovery ledger.
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, *, on_completion: Optional[nemo_rl.experience.rollout_manager.RolloutCompletionCallback] = None, recovery_granularity: nemo_rl.experience.rollout_recovery.RecoveryGranularity = RecoveryGranularity.SIBLING, ) → 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.

Sibling recovery re-dispatches only rows that never arrive. Prompt-group recovery performs one physical Gym dispatch here and delegates a complete cohort replacement to the outer recovery loop.

_results_to_completions( results: list[dict], ) → tuple[list[nemo_rl.experience.interfaces.Completion], dict[str, int]]#

Apply configured penalties and convert a Gym result batch.

Receipt-mode (token-capture) results are token-free — the message_log is empty and the canonical row is rebuilt by the finalizer from staged deltas — so they skip tensorization, truncation, and the token/text reward penalties; the receipt and rollout id ride env_extras for the finalize step.

_compute_reward_penalty_metrics( penalty_counts: dict[str, int], num_results: int, ) → dict[str, float]#: Return enabled penalty rates using the legacy Gym metric names.

_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, rollout_recovery_config: nemo_rl.algorithms.single_controller_utils.config.RolloutRecoveryConfig, 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, reward_penalty_config: Optional[dict[str, Any]] = None, 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, effort_config: Optional[nemo_rl.experience.rollouts.EffortLevelsConfig] = None, log_full_result_tables: bool = False, )#

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.

suspend_request_deadlines() → None#: Pause live request-deadline clocks while a colocated engine is in training mode.

resume_request_deadlines() → None#: Resume live request-deadline clocks when a colocated engine exits training mode.

property recovery_ledger: nemo_rl.experience.rollout_recovery.RolloutRecoveryLedger#: Return the prompt-group ownership ledger shared with the controller.

record_finalizer_dropped_prompt() → None#

Count a controller-side drop after generation and finalization succeeded.

A group whose valid-row fraction fell below token_capture.min_valid_fraction_per_group is not an infra failure of the kind the retry loop above tracks, but it is the same signal for an operator watching max_consecutive_dropped_prompts – no rollout got committed for this prompt – so it shares that budget’s counters rather than going uncounted.

set_data_plane_checkpoint_barrier( barrier: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneCheckpointBarrier, ) → None#: Join streamed sibling transitions to the SC snapshot barrier.

async _recovery_mutation( kind: nemo_rl.algorithms.async_utils.replay_buffer.CheckpointMutationKind = 'recovery_retries', ) → collections.abc.AsyncIterator[nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut]#: Serialize short lineage transitions with native TQ snapshots.

telemetry_snapshot() → dict[str, int]#: Return cumulative committed-publication and recovery counters.

record_canonical_publication(output_tokens: int) → None#: Count one prompt group after its canonical TQ commit succeeds.

record_recovery_siblings(*, reused: int, redispatched: int) → None#: Count sibling work avoided and repeated after a process restart.

reserve_prompt_group( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, input_sample: nemo_rl.data.interfaces.DatumSpec, *, target_step: Optional[int], admitted: bool = True, admission_id: Optional[str] = None, ) → str#: Own a prompt before controller dispatch can yield or checkpoint.

mark_prompt_group_admitted( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, *, target_step: Optional[int], ) → None#: Attach sampler admission state to a pre-admission reservation.

discard_prompt_group( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, ) → None#: Release a reservation that will intentionally never be dispatched.

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, *, rollout_ids: Optional[list[str]] = None, generation_indices: Optional[list[int]] = None, on_completion: Optional[nemo_rl.experience.rollout_manager.RolloutCompletionCallback] = None, recovery_granularity: nemo_rl.experience.rollout_recovery.RecoveryGranularity = RecoveryGranularity.SIBLING, ) → 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, lineage_group_id: Optional[str] = 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.
lineage_group_id – Stable group minted by the rollout ledger before dataloader dispatch. TQ records this same ID rather than minting one. None preserves the ordinary non-checkpointed fresh-ID retry path.

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 checkpoint-atomic transition from its retained ledger record to either a replacement prompt or 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.

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

Capture siblings with stable lineage and configured retry granularity.

Returns None when infrastructure retries are exhausted within the configured drop budget. The caller then owns the backpressure permit and target-step shortfall.

async _generate_for_finalization_attempt( input_sample: nemo_rl.data.interfaces.DatumSpec, *, recovery_group_id: str, inflight_registry: Optional[dict[str, tuple[asyncio.Task[None], int]]], ) → nemo_rl.experience.rollout_reassembler_actor.ReassemblyRequest#: Dispatch the current sibling cohort and leave one slot unready.

async discard_recovery_group( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, group_id: str, ) → None#: Clean known staged rows before intentionally dropping lineage.

nemo_rl.experience.rollout_manager#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.experience.rollout_manager`#