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

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

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

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

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.

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

API#

`nemo_rl.experience.rollout_reassembler`#