`nemo_rl.experience.route_assembly`#

Shared route-plan executor for direct and deferred router replay.

One implementation turns a verified :class:RouteAssemblyPlan plus fetched

class:: RouteFragment payloads into a training route tensor. The finalizer runs it eagerly in direct mode (a failure rejects the rollout before publication as route_assembly:<reason>); the policy worker runs it after canonical publication in deferred mode (the same reason maps to the existing counted sentinel fallback). The executor owns per-span extras-digest verification against Gym’s receipt-bound commitments, span classification via Gym’s decision table, full/tail slicing, sentinel fill, and shape checks — it returns a tensor or a failure reason and never decides policy.

Module Contents#

Classes#

`RouteLayout`	Which columns of a staged route fragment belong to the trained MoE layers.
`RouteFragment`	One staged route payload plus the metadata its extras digest binds.

Functions#

`boundary_route_rows`	Include the last token of the previous message.
`verify_route_fragment_integrity`	Rebuild the staged extras envelope and verify the receipt-bound digest.
`execute_route_plan`	Assemble one canonical route tensor from staged fragments.

Data#

`ROUTE_MISSING_SENTINEL`
`ROUTE_FAILURE_CANONICAL_LENGTH`
`ROUTE_FAILURE_MISSING_FRAGMENT`
`ROUTE_FAILURE_INTEGRITY`
`ROUTE_FAILURE_RANK`
`ROUTE_FAILURE_LENGTH`
`ROUTE_FAILURE_MODEL_SHAPE`
`ROUTE_FAILURE_ASSEMBLED_LENGTH`

API#

nemo_rl.experience.route_assembly.ROUTE_MISSING_SENTINEL#: None

nemo_rl.experience.route_assembly.ROUTE_FAILURE_CANONICAL_LENGTH#: ‘canonical_length_mismatch’

nemo_rl.experience.route_assembly.ROUTE_FAILURE_MISSING_FRAGMENT#: ‘missing_fragment’

nemo_rl.experience.route_assembly.ROUTE_FAILURE_INTEGRITY#: ‘fragment_integrity’

nemo_rl.experience.route_assembly.ROUTE_FAILURE_RANK#: ‘fragment_rank’

nemo_rl.experience.route_assembly.ROUTE_FAILURE_LENGTH#: ‘fragment_length’

nemo_rl.experience.route_assembly.ROUTE_FAILURE_MODEL_SHAPE#: ‘fragment_model_shape’

nemo_rl.experience.route_assembly.ROUTE_FAILURE_ASSEMBLED_LENGTH#: ‘assembled_length_mismatch’

nemo_rl.experience.route_assembly.boundary_route_rows(prev_len: int) → int#

Include the last token of the previous message.

Chosen experts are only recorded for input tokens. Choice of the last generated token in the previous message was a placeholder. Rewrite it with the correct value, known from the new prefill.

class nemo_rl.experience.route_assembly.RouteLayout#

Which columns of a staged route fragment belong to the trained MoE layers.

A fragment’s layer axis comes in one of two layouts. MInf stages one column per MoE layer (num_moe_layers columns). vLLM stages one column per transformer layer (total_num_layers columns), so for a hybrid Mamba/attention/MoE stack the executor must pick out the MoE columns by moe_layer_numbers. Both counts coincide for a stack where every layer is MoE, in which case no selection happens.

moe_layer_numbers: tuple[int, ...]#

None

1-based global layer numbers of the MoE layers, ascending.

total_num_layers: int#: None

top_k: int#: None

__post_init__() → None#

classmethod compressed( *, num_moe_layers: int, top_k: int, ) → nemo_rl.experience.route_assembly.RouteLayout#

A MoE-only layout for callers that know only (num_moe_layers, top_k).

Direct mode learns these from the fetched fragments themselves and has no model config, so it can never need column selection.

property num_moe_layers: int#

property dims: tuple[int, int]#: The published (num_moe_layers, top_k) route tensor dims.

moe_column_indices( num_payload_layers: int, ) → Optional[tuple[int, ...]]#

Payload column for each MoE layer, in moe_layer_numbers order.

This is the one rule both replay paths share. A MoE-only payload (MInf) maps the i-th MoE layer to column i; a full-layer payload (vLLM) maps layer number n to column n - 1. Any other width is unsupported and returns None.

select_moe_columns( routes: torch.Tensor, ) → Optional[torch.Tensor]#: Return routes as [T, num_moe_layers, top_k] or None if it fits neither layout.

class nemo_rl.experience.route_assembly.RouteFragment#

One staged route payload plus the metadata its extras digest binds.

routes is the staged [staged_len, num_moe_layers, topk] tensor, encoding the ROUTE_ENCODING_* wire code the digest was committed over, and extras_metadata_json the canonical non-route extras JSON staged beside it.

routes: torch.Tensor#: None

encoding: int#: None

extras_metadata_json: bytes#: None

nemo_rl.experience.route_assembly.verify_route_fragment_integrity( fragment: nemo_rl.experience.route_assembly.RouteFragment, *, extras_digest_version: int, expected_extras_digest: str, ) → bool#: Rebuild the staged extras envelope and verify the receipt-bound digest.

nemo_rl.experience.route_assembly.execute_route_plan( plan: nemo_rl.experience.route_plan.RouteAssemblyPlan, fragments: Mapping[str, nemo_rl.experience.route_assembly.RouteFragment], *, layout: nemo_rl.experience.route_assembly.RouteLayout, canonical_len: int, ) → tuple[Optional[torch.Tensor], Optional[str]]#

Assemble one canonical route tensor from staged fragments.

Parameters:

plan – The verified assembly plan built by the finalizer.
fragments – Fetched fragments keyed by staging key. Sentinel spans need no entry.
layout – Which fragment columns are the trained MoE layers. The policy worker supplies the real model layout (the authoritative shape check, accepting both MoE-only and full-layer fragments); the direct-mode finalizer supplies a compressed layout learned from the fetched fragments.
canonical_len – The published row’s token length.

Returns:

(tensor, None) on success — [canonical_len, num_moe_layers, topk] int16, sentinel-filled wherever no fragment contributed — or (None, reason) with a reason from the shared vocabulary.

nemo_rl.experience.route_assembly#