refactor(wire): generic packed-ndarray codec and an N-field response-body splice

[erictang000]

Note

Standalone-reviewable reopen of #2079 — original work by @dyurk-lila.
Commits are pushed unchanged, so author and committer metadata are preserved.
The branch lives on NovaSky-AI/SkyRL rather than a fork so CI runs.

Stack position 4/11 — based on #2233 (branch dyurk/r3-pooled-route-collation).
Because the base is the previous branch in the stack, the Files changed tab above shows only this PR's own changes, not the cumulative stack. Merge bottom-up.

Supersedes #2079.

---

Note on the diff: This PR is part of a routed-expert-replay / sampler-support series and builds on the PRs below. GitHub cannot show the intermediate branches here, so the diff is cumulative on top of main — the changes new to this PR sit on top of:

• refactor(inference): extract reusable RemoteInferenceGenerator #1911 — refactor extract RemoteInferenceGenerator
• refactor(generators): assemble routed-expert traces incrementally #1912 — refactor assemble routed-expert traces incrementally
• perf(r3): collate routes through a container-aware pool, and refuse replay under VPP #2078 — perf(r3) collate routes through a container-aware pool

Reviewing in PR order (lowest number first) shows each incremental change cleanly.

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"}

data is deliberately the first key. The client can then replace registered blobs with memoryview objects before handing the remaining JSON to orjson. Shape, dimensionality, dtype, and decoded byte count are still validated by the field-specific decoder.

Implementation

• Introduce shared pack_ndarray and unpack_ndarray helpers with explicit allowed-dtype and dimensionality contracts.
• Keep routed-expert packing as a thin field-specific validator over the generic codec.
• Register packed response fields centrally so future side channels do not add another parser.
• Scan the response once, splice every registered blob, and restore blobs to envelopes in document order.
• Preserve absent and null fields, sidecar metadata, multiple choices, and all ordinary JSON content.
• Fail if the serialized envelope layout drifts instead of silently materializing a large base64 string.
• Use packed parsing only for generate requests that can return side channels; ordinary generation retains direct orjson parsing.

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

• Codec tests cover contiguous and non-contiguous arrays, dtype and shape validation, truncated buffers, sidecar fields, and canonical routed-expert dtypes.
• Body-splice tests cover one and multiple fields, multiple choices, absent and null values, JSON lookalikes, reordered envelopes, unregistered fields, and unterminated data.
• Remote-client integration tests cover successful two-field decoding, transient-response retries, deterministic layout failures, text error bodies, and bypassing the scanner for ordinary generation.

---

Note

Medium Risk
Changes the generate response parsing contract and multi-turn routed-expert bookkeeping used for training replay; mistakes could corrupt routes or silently mis-parse payloads, though layout drift is explicitly detected.

Overview
Generate wire format moves from routed-expert-only packing to shared pack_ndarray / unpack_ndarray envelopes (data first for scanning) and a central PackedField registry (including rollout_sample_support). load_packed_body splices registered base64 blobs into memoryview before orjson parses the rest, so multiple large side channels avoid giant Python strings; layout drift fails fast instead of silently falling back.

Remote inference pulls raw /skyrl/v1/generate HTTP into RemoteGenerateClient, enables packed parsing only when returning routed experts, and forwards routed_experts_prompt_starts as routed_experts_prompt_start in sampling params.

Multi-turn R3 adds TokenMetadataTrace and RoutedExpertTrace so the gym generator records incremental expert rows per turn and finalizes against token count and loss mask; batch preprocessing collates routes in parallel via fill_batch_rows and cpu_topology.pool_workers (also used for delta weight-sync thread counts).

Megatron routing replay now rejects virtual pipeline parallelism (config validation and worker setup) because it desyncs replay FIFOs.

^{Reviewed by Cursor Bugbot for commit 7dd3ff3. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request generalizes the serialization of side-channel data by introducing helpers to pack and unpack arbitrary NumPy arrays as base64 envelopes with sidecar fields. It also adds an optimized response-parsing utility (load_packed_body) that splices base64 blobs directly into memoryviews to avoid large string allocations. The review feedback suggests consistently using the .value attribute of StrEnum members (such as PackedArrayKey and PackedField) during dictionary lookups and membership checks to prevent potential issues with strict key matching and third-party serialization libraries.

[gemini-code-assist]

For consistency and to prevent any potential issues with custom mapping implementations or strict key type checks, it is safer to use .value when accessing dictionary keys with StrEnum members, especially since orjson and other libraries can be sensitive to str subclasses.

Suggested change

	dtype_name = payload[PackedArrayKey.DTYPE]
	shape = tuple(payload[PackedArrayKey.SHAPE])
	data = pybase64.b64decode_as_bytearray(payload[PackedArrayKey.DATA], validate=True)
	dtype_name = payload[PackedArrayKey.DTYPE.value]
	shape = tuple(payload[PackedArrayKey.SHAPE.value])
	data = pybase64.b64decode_as_bytearray(payload[PackedArrayKey.DATA.value], validate=True)

[gemini-code-assist]

Use PackedArrayKey.DATA.value instead of PackedArrayKey.DATA for consistency with line 248 where .value is explicitly used. This ensures uniform dictionary key lookups and avoids mixing StrEnum members with their underlying string values.

Suggested change

	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:

[gemini-code-assist]

Use PackedField.ROUTED_EXPERTS.value instead of the StrEnum member PackedField.ROUTED_EXPERTS directly when performing dictionary lookups. This maintains consistency with how dictionary keys are handled elsewhere in the codebase (e.g., in vllm_server_actor.py and generate_wire.py) and avoids potential issues with strict key matching.

Suggested change

	packed_routed_experts = choice.get(PackedField.ROUTED_EXPERTS)
	packed_routed_experts = choice.get(PackedField.ROUTED_EXPERTS.value)

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; no concrete correctness, security, or repository-rule violations remain.

Diagram

sequenceDiagram
    participant Server as VLLMServerActor
    participant Wire as pack_ndarray
    participant Client as RemoteGenerateClient
    participant Scanner as load_packed_body
    participant Decoder as unpack_ndarray

    Server->>Wire: Pack routed-expert ndarray
    Wire-->>Server: data-first base64 envelope
    Server-->>Client: orjson response bytes
    alt Packed side channels requested
        Client->>Scanner: Parse raw response bytes
        Scanner->>Scanner: Scan once and splice registered blobs
        Scanner->>Scanner: Parse remaining JSON
        Scanner-->>Client: Envelopes with memoryview data
        Client->>Decoder: Validate and decode routed_experts
        Decoder-->>Client: Canonical ndarray
    else Ordinary generation
        Client->>Client: orjson.loads(raw)
    end

Loading

_{Reviews (1) · Last reviewed commit: "style: tighten wire codec comments and t..."}