refactor(wire): generic packed-ndarray codec and an N-field response-body splice - #2079
refactor(wire): generic packed-ndarray codec and an N-field response-body splice#2079dyurk-lila wants to merge 8 commits into
Conversation
Extract the single-request HTTP generation path out of RemoteInferenceClient into a standalone RemoteGenerateClient plus a RemoteGenerateResult dataclass, so routed-expert results can be obtained without constructing the full inference/control-plane client. RemoteInferenceClient now owns an internal RemoteGenerateClient and delegates session management, _post, and _generate_single to it. Endpoint routing, retry/backoff, cache_salt handling, serialization, and lifecycle behavior are unchanged.
For multi-turn rollouts, build routed-expert (R3) trace data incrementally as the conversation grows instead of re-gathering the whole conversation's routes on every turn. - Add `routed_experts_prompt_starts` to `InferenceEngineInput` and thread a per-request `routed_experts_prompt_start` through `RemoteInferenceClient` and `RemoteInferenceGenerator.generate()`, forwarded as a sampling param so the engine only returns routes for the newly generated suffix. - Introduce `TokenMetadataTrace` (token-aligned array accumulator) and `RoutedExpertTrace`, which records each generation's routes and finalizes a full per-token routed-expert array with loss-mask-aware terminal padding. - Track a `RoutedExpertTrace` on `AgentLoopState` and record routes per turn, replacing the previous whole-conversation re-gather in `SkyRLGymGenerator.agent_loop`. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for packed side-channel arrays (such as routed experts) in the inference server wire format using base64 envelopes, along with incremental routing trace tracking and dynamic thread pool sizing based on cgroup quotas and process affinity. The review feedback highlights several key improvements for robustness and type safety, including catching ValueError during response parsing to ensure transient failures are retried, preventing false-positive key matches in the JSON stream scanner, consistently using .value on StrEnum keys, and catching OSError on os.sched_getaffinity to support restricted environments.
| try: | ||
| raw = await resp.read() | ||
| body = load_packed_body(raw) if packed_side_channels else orjson.loads(raw) | ||
| except orjson.JSONDecodeError as exc: |
There was a problem hiding this comment.
If load_packed_body(raw) raises a ValueError (e.g., due to a truncated/incomplete response or layout drift), it will not be caught by except orjson.JSONDecodeError. This causes transient network truncation errors to immediately fail the request instead of being retried. Catching ValueError instead of orjson.JSONDecodeError ensures both JSON decoding and parsing/splicing errors are safely retried.
| except orjson.JSONDecodeError as exc: | |
| except ValueError as exc: |
|
|
||
| PACKED_SIDE_CHANNEL_FIELDS: tuple[str, ...] = tuple(PackedField) | ||
|
|
||
| _ENVELOPE_KEYS = frozenset(PackedArrayKey) |
There was a problem hiding this comment.
Using PackedArrayKey directly in frozenset creates a set of StrEnum instances. For cleaner and safer operations (especially when intersecting with sets of plain str keys), it is better to explicitly use their .value strings.
| _ENVELOPE_KEYS = frozenset(PackedArrayKey) | |
| _ENVELOPE_KEYS = frozenset(key.value for key in PackedArrayKey) |
| def _match_packed_field(raw: bytes, anchor: int, prefixes: Mapping[str, bytes]) -> Optional[str]: | ||
| """Name the registered field whose prefix ends at ``anchor``, if any.""" | ||
| for field, prefix in prefixes.items(): | ||
| begin = anchor + len(_PACKED_DATA_ANCHOR) - len(prefix) | ||
| if begin >= 0 and raw.startswith(prefix, begin): | ||
| return field | ||
| return None |
There was a problem hiding this comment.
The current suffix matching logic in _match_packed_field can lead to false positives if another key in the JSON response ends with a registered field name (e.g., "not_routed_experts" matching "routed_experts"). This would cause a ValueError during restoration and crash the request. To prevent this, ensure the matched prefix is preceded by a valid JSON delimiter like {, ,, or whitespace.
| def _match_packed_field(raw: bytes, anchor: int, prefixes: Mapping[str, bytes]) -> Optional[str]: | |
| """Name the registered field whose prefix ends at ``anchor``, if any.""" | |
| for field, prefix in prefixes.items(): | |
| begin = anchor + len(_PACKED_DATA_ANCHOR) - len(prefix) | |
| if begin >= 0 and raw.startswith(prefix, begin): | |
| return field | |
| return None | |
| def _match_packed_field(raw: bytes, anchor: int, prefixes: Mapping[str, bytes]) -> Optional[str]: | |
| """Name the registered field whose prefix ends at ``anchor``, if any.""" | |
| for field, prefix in prefixes.items(): | |
| begin = anchor + len(_PACKED_DATA_ANCHOR) - len(prefix) | |
| if begin >= 0 and raw.startswith(prefix, begin): | |
| # Ensure the match is not a suffix of a longer key (e.g., "not_routed_experts") | |
| if begin == 0 or raw[begin - 1] in (123, 44, 32, 9, 10, 13): | |
| return field | |
| return None |
| if isinstance(node, dict): | ||
| for key, value in node.items(): | ||
| queue = blobs.get(key) | ||
| if queue is not None and isinstance(value, dict) and PackedArrayKey.DATA in value: |
There was a problem hiding this comment.
For consistency with line 248 and to ensure compatibility with strict type checkers, use PackedArrayKey.DATA.value instead of PackedArrayKey.DATA when checking for key existence in the dictionary.
| if queue is not None and isinstance(value, dict) and PackedArrayKey.DATA in value: | |
| if queue is not None and isinstance(value, dict) and PackedArrayKey.DATA.value in value: |
|
|
||
| routed_experts = None | ||
| if return_routed_experts: | ||
| packed_routed_experts = choice.get(PackedField.ROUTED_EXPERTS) |
There was a problem hiding this comment.
For consistency with other parts of the codebase (such as vllm_server_actor.py and test_remote_inference_client.py) and to prevent potential type-checking warnings, use .value on the PackedField enum member.
| packed_routed_experts = choice.get(PackedField.ROUTED_EXPERTS) | |
| packed_routed_experts = choice.get(PackedField.ROUTED_EXPERTS.value) |
| try: | ||
| affinity = len(os.sched_getaffinity(0)) | ||
| except AttributeError: | ||
| affinity = os.cpu_count() or 1 |
There was a problem hiding this comment.
In restricted, sandboxed, or containerized environments, os.sched_getaffinity can raise an OSError rather than AttributeError. Catching OSError as well ensures the function robustly falls back to os.cpu_count() instead of crashing.
| try: | |
| affinity = len(os.sched_getaffinity(0)) | |
| except AttributeError: | |
| affinity = os.cpu_count() or 1 | |
| try: | |
| affinity = len(os.sched_getaffinity(0)) | |
| except (AttributeError, OSError): | |
| affinity = os.cpu_count() or 1 |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.
Reviewed by Cursor Bugbot for commit a604e5c. Configure here.
| rollout_expert_indices_out = agent_loop_state.routed_expert_trace.finalize( | ||
| token_count=len(prompt_ids) + len(response_ids), | ||
| loss_mask=[0] * len(prompt_ids) + loss_mask, | ||
| ) |
There was a problem hiding this comment.
Synthetic EOS breaks R3 finalize
High Severity
In single-turn mode the loop appends a synthetic EOS with loss_mask 1, then RoutedExpertTrace.finalize rejects any loss-active target past the captured route prefix. Trajectories that previously padded that terminal gap now raise during finalize, so single-turn R3 generation fails whenever EOS is appended.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit a604e5c. Configure here.
| f"{worker_type}.megatron_config: moe_enable_routing_replay is incompatible with " | ||
| "virtual_pipeline_model_parallel_size -- interleaved chunks desync the replay FIFO. " | ||
| "Unset virtual_pipeline_model_parallel_size." | ||
| ) |
There was a problem hiding this comment.
VPP=1 rejected inconsistently
Medium Severity
Config validation treats any truthy virtual_pipeline_model_parallel_size as incompatible with routing replay, including 1. The Megatron worker only errors when the resolved size is greater than 1, since size 1 is not interleaved. Valid configs with VPP explicitly set to 1 are refused at startup despite being allowed later.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit a604e5c. Configure here.
2 participants
Problem
The generate endpoint returns large NumPy side channels inside a JSON response. Parsing the entire body normally first materializes each base64 blob as a Python string, adding a large allocation and copy before decoding. The existing optimization recognized only the routed-expert field through one hard-coded byte prefix; adding another packed field would silently leave that second blob on the expensive path.
Wire contract
Packed arrays use one typed envelope:
{"data": "<base64>", "shape": [12, 40, 8], "dtype": "int16"}datais deliberately the first key. The client can then replace registered blobs withmemoryviewobjects before handing the remaining JSON toorjson. Shape, dimensionality, dtype, and decoded byte count are still validated by the field-specific decoder.Implementation
pack_ndarrayandunpack_ndarrayhelpers with explicit allowed-dtype and dimensionality contracts.nullfields, sidecar metadata, multiple choices, and all ordinary JSON content.orjsonparsing.Transport errors and undecodable gateway responses retain the existing retry behavior. A valid JSON response with a malformed packed contract is treated as a deterministic protocol error and is not retried.
Performance
For a representative 121 MiB response containing 90 MiB of routes, parsing and decoding improved from 268 ms to 83 ms. Registering a second packed field did not add another full-body scan. The optimization is therefore tied to the response as a whole rather than to one particular side channel.
Testing
nullvalues, JSON lookalikes, reordered envelopes, unregistered fields, and unterminated data.Note
High Risk
Changes the generate HTTP wire contract, client parsing of large binary side channels, and MoE routed-expert bookkeeping used in training. Bugs here can silently drop or misalign replay routes.
Overview
Makes generate-side NumPy payloads a generic packed-array envelope (
data/shape/dtype,datafirst) so the client can splice all registered blobs intomemoryviews beforeorjsonparses the rest.pack_ndarray/unpack_ndarrayplusload_packed_bodyreplace the routed-experts-only codec;routed_expertspacking is a thin validator on top. Layout drift fails instead of materializing huge base64 strings. Ordinary generate still uses direct JSON parse.Extracts
RemoteGenerateClientfor one-shot token generation and retries. Generation can sendrouted_experts_prompt_starts. Multi-turn R3 no longer overwrites a full-sequence route array each turn:RoutedExpertTrace(onTokenMetadataTrace) records suffixes, pads only loss-masked tails, and rejects missing routes on loss-active tokens.Controller collation writes left-padded route tensors in parallel via
fill_batch_rows, sized with newpool_workers(affinity + cgroup quota). Weight-sync apply/publish uses the same helper. Canonical dtypes are validated, not re-compacted. Routing replay is rejected with virtual pipeline parallelism (FIFO desync).Reviewed by Cursor Bugbot for commit a604e5c. Bugbot is set up for automated code reviews on this repo. Configure here.