`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.
`FinalizedGroup`	What `finalize_group` hands back for `commit_finalized`.
`RolloutReassembler`	Receipts -> verified rows -> N-row publish, off the generation hot path.

Functions#

`_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.
`_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

media: Optional[dict[str, nemo_rl.data.multimodal_utils.PackedTensor]]#: 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(…)’

dropped: bool#: False

drop_reason: Optional[str]#: None

valid_row_count: int#: 0

total_row_count: int#: 0

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._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, router_replay_enabled: bool = False, defer_routed_experts_to_policy: bool = False, capture_media: bool = False, )#

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

Initialization

finalize_rollout( rollout_id: str, receipt: Optional[dict[str, Any]], *, reward: float, ) → 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.

_resolve_media( row: Any, fetched_by_call: dict[str, nemo_rl.data_plane.tq_token_sink.FetchedStagedCall], ) → 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, ) → nemo_rl.experience.rollout_reassembler.FinalizedGroup#

Publish exactly N canonical rows for one prompt group.

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.

_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).

_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`#