`nemo_rl.experience.rollout_reassembler`#

Blackbox finalization: token-free receipts + staged deltas -> canonical rows.

Orchestration only: per rollout, apply the rollout-level receipt guards, fetch the staged base rows the receipt manifest names through the TokenSource (normally validated StagedCallBaseSnapshot values), and delegate all token, digest, lineage, and terminal-chain semantics to Gym’s verify_and_linearize. Any rejection becomes a masked placeholder row — the group always publishes exactly N rows so GRPO group shape survives; validity folds into sample_mask (no new train field) and placeholders copy prompt_ids_for_adv from a valid sibling so per-prompt baselines stay well-formed.

Router replay runs one unified flow: both modes construct the same RouteAssemblyPlan from Gym’s link spans and extras commitments. Deferred mode publishes the encoded plan beside the canonical row and leaves staged route fragments live until policy consumption; direct mode executes the plan eagerly with fragments fetched in the same batch — any executor failure is a pre-publication route_assembly:<reason> rejection.

Module Contents#

Classes#

`FinalizedRollout`	One rollout’s canonical row, or its rejection.
`ActionOutputFlags`	RL’s existing detector result for one selected model action.
`RolloutSelection`	Eligible calls and flags within one attempt, independent of commit order.
`SegmentReceipt`	Metadata-only selection for one ordinary captured physical segment.
`_LogicalRows`
`FinalizedGroup`	What `finalize_group` hands back for `commit_finalized`.
`RolloutReassembler`	Receipts -> verified rows -> N-row publish, off the generation hot path.

Functions#

`_capture_token_metrics`	Count verified output per logical owner, before sample-level filtering.
`_concat_media`	Concatenate per-call media deltas along the terminal chain, in order.
`_trainer_media`	Wrap the engine’s media tensors as the trainer’s one-row PackedTensors.
`_empty_trainer_media`	Media fields for a group in which no rollout carried media.
`_media_fields_for_group`	Stack per-rollout media into group-level PackedTensors, one logical row each.

API#

class nemo_rl.experience.rollout_reassembler.FinalizedRollout#

One rollout’s canonical row, or its rejection.

rollout_id: str#: None

valid: bool#: None

rejection_reason: Optional[str]#: None

token_ids: list[int]#: None

token_mask: list[float]#: None

logprobs: list[float]#: None

prompt_len: int#: None

reward: float#: None

staging_keys: list[str]#: None

min_wv: Optional[int]#: None

max_wv: Optional[int]#: None

routed_experts: Optional[torch.Tensor]#: None

route_plan: Optional[nemo_rl.experience.route_plan.RouteAssemblyPlan]#: None

model_call_ids: tuple[str, ...]#: ()

media: Optional[dict[str, nemo_rl.data.multimodal_utils.PackedTensor]]#: None

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

reward_metrics: dict[str, float]#: ‘field(…)’

sampler_topk_token_ids: Optional[torch.Tensor]#: None

sampler_topk_logprobs: Optional[torch.Tensor]#: None

output_masks: dict[str, list[bool]]#: ‘field(…)’

generation_lengths: tuple[int, ...]#: ()

call_versions: tuple[tuple[str, int, int], ...]#: ()

class nemo_rl.experience.rollout_reassembler.ActionOutputFlags#

RL’s existing detector result for one selected model action.

invalid_tool_call: bool#: None

malformed_thinking: bool#: None

class nemo_rl.experience.rollout_reassembler.RolloutSelection#

Eligible calls and flags within one attempt, independent of commit order.

Selection order determines root/path priority, including the grouping prompt; parents determine reconstruction order. The receipt identifies the scored terminal separately. Providers need not use Gym’s accepted-history matcher.

response_ids: tuple[str, ...]#: None

action_flags: tuple[nemo_rl.experience.rollout_reassembler.ActionOutputFlags, ...]#: None

truncated: bool#: False

class nemo_rl.experience.rollout_reassembler.SegmentReceipt#

Metadata-only selection for one ordinary captured physical segment.

capture_rollout_id: str#: None

receipt: Optional[dict[str, Any]]#: None

selected_response_ids: tuple[str, ...]#: None

truncated: bool#: False

action_flags: tuple[nemo_rl.experience.rollout_reassembler.ActionOutputFlags, ...] | None#: None

class nemo_rl.experience.rollout_reassembler._LogicalRows#

rows: list[nemo_rl.experience.rollout_reassembler.FinalizedRollout]#: None

prompt_ids: list[list[int]]#: None

mask_sample: list[bool]#: None

tags: list[dict[str, Any]]#: None

staging_keys: list[str]#: None

truncated: list[bool]#: None

valid_owner_count: int#: None

class nemo_rl.experience.rollout_reassembler.FinalizedGroup#

What finalize_group hands back for commit_finalized.

meta: Optional[nemo_rl.data_plane.KVBatchMeta]#: None

group_min_wv: int#: None

group_max_wv: int#: None

staging_keys: list[str]#: None

canonical_output_tokens: int#: 0

metrics: dict[str, float]#: ‘field(…)’

reward_observations: list[nemo_rl.experience.reward_penalties.FinalizedReward]#: ‘field(…)’

dropped: bool#: False

drop_reason: Optional[str]#: None

valid_row_count: int#: 0

total_row_count: int#: 0

token_metrics: dict[str, Any]#: ‘field(…)’

call_versions: list[tuple[str, int, int]]#: ‘field(…)’

nemo_rl.experience.rollout_reassembler._capture_token_metrics( rows: list[nemo_rl.experience.rollout_reassembler.FinalizedRollout], slots: list[int], contexts: collections.abc.Sequence[nemo_rl.experience.reward_penalties.RewardLogContext | None], ) → dict[str, Any]#

Count verified output per logical owner, before sample-level filtering.

Called before execution padding. Invalid owners contribute zero; environment masks do not erase verified generations. Shared ancestors already have zero loss on later paths, and generation_lengths excludes those same owned calls.

nemo_rl.experience.rollout_reassembler._concat_media( parts: list[nemo_rl.data_plane.tq_token_sink.StagedMediaTensors], ) → nemo_rl.data_plane.tq_token_sink.StagedMediaTensors#

Concatenate per-call media deltas along the terminal chain, in order.

Packed patches ([1, total_patches, F]) join along the patch dim, the per-frame sizes and per-video frame counts along their only dim. Parts must agree on pixel dtype and patch feature width (no silent promotion) and either all or none may carry frame counts (no mixed image/video chains). Each part was already validated by validate_media_tensors.

nemo_rl.experience.rollout_reassembler._trainer_media( staged: nemo_rl.data_plane.tq_token_sink.StagedMediaTensors, ) → dict[str, nemo_rl.data.multimodal_utils.PackedTensor]#

Wrap the engine’s media tensors as the trainer’s one-row PackedTensors.

pixel_values keeps MInf’s packed-patch layout: [total_patches, C*P*P] per row, so rows concatenate along dim 0 with no padding. Bridge’s _patchify_dynamic_images accepts that 2-D layout directly. imgs_sizes and num_frames mirror what extract_multimodal_model_inputs emits on the token-echo path (stills get one frame per image).

nemo_rl.experience.rollout_reassembler._empty_trainer_media( spec: nemo_rl.data.captured_media.MediaColumnSpec, num_rows: int, ) → dict[str, nemo_rl.data.multimodal_utils.PackedTensor]#

Media fields for a group in which no rollout carried media.

Each logical row gets a zero-row segment in _trainer_media’s layout ([0, 3*P*P] pixels, [0, 2] sizes, [0] frames) rather than no segment at all: a segment-less row ships nothing, so the group would publish without the media columns and TransferQueue would narrow any train fetch that mixes its keys with media keys. A zero-row segment keeps the column present, in the pinned dtype, and rebuilds on the trainer as a [0, 3*P*P] tensor the Omni model treats as “no images”.

nemo_rl.experience.rollout_reassembler._media_fields_for_group( rows: list[nemo_rl.experience.rollout_reassembler.FinalizedRollout], ) → dict[str, nemo_rl.data.multimodal_utils.PackedTensor]#

Stack per-rollout media into group-level PackedTensors, one logical row each.

Rows without media (text rollouts, placeholders) contribute an empty logical row so every media field stays aligned with input_ids.

class nemo_rl.experience.rollout_reassembler.RolloutReassembler( dp_client: Any, *, partition_id: str, staging_partition: str, pad_token_id: int, max_seq_len: int, reward_penalty_config: nemo_rl.algorithms.grpo.RewardPenaltyConfig | None = None, effort_config: nemo_rl.experience.rollouts.EffortLevelsConfig | None = None, router_replay_enabled: bool = False, defer_routed_experts_to_policy: bool = False, capture_media: bool = False, media_columns: Optional[nemo_rl.data.captured_media.MediaColumnSpec] = None, tokenizer: Any = None, thinking_tags: collections.abc.Sequence[str] | None = None, top_logprobs: int = 0, )#

Receipts -> verified rows -> N-row publish, off the generation hot path.

Initialization

finalize_rollout( rollout_id: str, receipt: Optional[dict[str, Any]], *, reward: float, forest_training: bool = False, action_flags: tuple[nemo_rl.experience.rollout_reassembler.ActionOutputFlags, ...] | None = None, loss_excluded_call_ids: frozenset[str] = frozenset(), reward_checks: nemo_rl.experience.reward_penalties.RewardChecks | None = None, reward_call_id: str | None = None, ) → nemo_rl.experience.rollout_reassembler.FinalizedRollout#

Verify one receipt against its staged rows and linearize the main chain.

Never raises for rollout-level problems: every rejection returns an invalid row whose reason feeds the metrics; the group publisher substitutes a placeholder.

_decode(token_ids: collections.abc.Sequence[int]) → str#

_resolve_media( row: Any, fetched_by_call: dict[str, nemo_rl.data_plane.tq_token_sink.FetchedStagedCall], *, verify_occurrences: bool = False, ) → tuple[Optional[dict[str, nemo_rl.data.multimodal_utils.PackedTensor]], Optional[str]]#

Read and validate the media columns along the terminal chain.

Each call row holds only the media new to that call (the worker slices at the parent chain’s item count), so the chain’s parts are concatenated in order, like the token deltas. Returns (media, rejection_reason). Text-only partitions and text rollouts return (None, None) without touching storage; media-carrying rollouts cost exactly one batched tensor read. The media columns live on the call rows themselves, so cleanup needs no extra key.

_execute_direct_plan( plan: nemo_rl.experience.route_plan.RouteAssemblyPlan, fetched: list[Any], ) → tuple[Optional[torch.Tensor], Optional[str]]#

Run the shared executor eagerly with locally fetched fragments.

Returns (None, None) when the rollout staged no routes at all — the group tensor build fills those rows with sentinels, exactly like a deferred row whose plan is all-sentinel.

finalize_group( group_id: str, rollout_ids: list[str], receipts: list[Optional[dict[str, Any]]], rewards: list[float], *, mask_sample: list[bool], fallback_weight_version: int, prompt_idx: int, loss_multiplier: float = 1.0, canonical_sample_ids: Optional[list[str]] = None, logical_selections: Optional[list[nemo_rl.experience.rollout_reassembler.RolloutSelection]] = None, execution_row_multiple: int = 1, reward_checks: Optional[list[nemo_rl.experience.reward_penalties.RewardChecks | None]] = None, reward_log_contexts: Optional[list[nemo_rl.experience.reward_penalties.RewardLogContext | None]] = None, ) → nemo_rl.experience.rollout_reassembler.FinalizedGroup#

Publish one prompt group, with optional physical rows per logical owner.

Blocking (TQ round trips); run via asyncio.to_thread from the dispatch task. fallback_weight_version stamps a group none of whose rollouts produced a valid row (placeholder-only groups still need a staleness tag). mask_sample is the per-rollout advantage-stage flag the native pack_payload path emits from each Completion; it rides along unchanged so the train pump’s environment masking reads the same field on both paths (placeholder rows already train nothing through sample_mask 0). loss_multiplier supplies the dataset-level weight for every valid row, matching the ordinary record_to_train_batch path. truncated is not carried from the dispatcher – the receipt path has no real tokens to measure it from at dispatch time – so it is computed here instead, from each row’s rebuilt length against max_seq_len.

_truncated_flags( rows: collections.abc.Sequence[nemo_rl.experience.rollout_reassembler.FinalizedRollout], logical: Optional[nemo_rl.experience.rollout_reassembler._LogicalRows], ) → list[bool]#

Per-row truncation, shared by the TRUNCATED column and termination rates.

Forest rows carry their rollout’s Gym-reported truncation; ordinary rows count as truncated when they fill max_seq_len.

_publish_rows( group_id: str, rows: list[nemo_rl.experience.rollout_reassembler.FinalizedRollout], rollout_ids: list[str], *, mask_sample: list[bool], fallback_weight_version: int, prompt_idx: int, staging_keys: list[str], metrics: dict[str, float], rollouts_ms: float, reward_observations: list[nemo_rl.experience.reward_penalties.FinalizedReward], logical: Optional[nemo_rl.experience.rollout_reassembler._LogicalRows], execution_row_multiple: int, canonical_sample_ids: list[str], loss_multiplier: float, ) → nemo_rl.experience.rollout_reassembler.FinalizedGroup#: Shared ordinary/Forest tensorization, single publication, and cleanup.

static _plan_selected_calls( rollout_ids: list[str], receipts: list[Optional[dict[str, Any]]], selections: list[nemo_rl.experience.rollout_reassembler.RolloutSelection], ) → list[list[nemo_rl.experience.rollout_reassembler.SegmentReceipt]]#: Plan paths to selected leaves using committed capture parents.

_finalize_logical_rows( group_id: str, rollout_ids: list[str], logical_segments: list[list[nemo_rl.experience.rollout_reassembler.SegmentReceipt]], rewards: list[float], mask_sample: list[bool], *, canonical_sample_ids: list[str], receipts: list[Optional[dict[str, Any]]], reward_checks: tuple[nemo_rl.experience.reward_penalties.RewardChecks | None, ...], ) → nemo_rl.experience.rollout_reassembler._LogicalRows#

Verify ordinary segments, then keep all or none of each owner’s rows.

A missing initial prompt is a structural group failure: a sibling’s possibly different post-agent prompt is not evidence for a failed owner. No publication or cleanup happens before this method completes.

_build_routed_experts_tensor( rows: list[nemo_rl.experience.rollout_reassembler.FinalizedRollout], *, max_len: int, metrics: dict[str, float], ) → torch.Tensor#

[n, max_len, L, K] int16 routes for the group; sentinel elsewhere.

Padding, placeholder rows, and valid rows whose rebuild carried no routes are all-sentinel: Megatron’s replay falls back to its own router for exactly those positions. (L, K) is learned from the first rebuilt row that carries routes and cached for placeholder-only groups; a group arriving before any routed row has been seen cannot be shaped and fails loudly (unreachable once the first real rollout of the run finalizes).

_drop_group( group_id: str, *, reason: str, staging_keys: list[str], group_min_wv: int, group_max_wv: int, metrics: dict[str, float], ) → nemo_rl.experience.rollout_reassembler.FinalizedGroup#

Reject a whole group and hand the caller a dropped result.

Logs the drop, releases the group’s staging rows, marks the finalize/group_dropped metric, and carries reason through as drop_reason for the caller’s log line.

_clear_staging(staging_keys: list[str]) → None#

_call_dp(method_name: str, **kwargs: Any) → Any#

nemo_rl.experience.rollout_reassembler#

Module Contents#

Classes#

Functions#

API#

`nemo_rl.experience.rollout_reassembler`#