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

skyrl/backends/skyrl_train/distributed/megatron/token_metadata.py

@@ -2,6 +2,7 @@
 
 from dataclasses import dataclass
 
+import numpy as np
 import torch
 
 from skyrl.backends.skyrl_train.distributed.megatron.packing_utils import (
@@ -205,3 +206,51 @@ def scatter_packed_token_values_to_batch(
     )
     batch_values[output_mask] = values[packed_mask]
     return batch_values
+
+
+class TokenMetadataTrace:
+    """Accumulate arrays whose first dimension is aligned to tokens."""
+
+    def __init__(self) -> None:
+        self._chunks: list[np.ndarray] = []
+        self._schema: tuple[tuple[int, ...], np.dtype] | None = None
+        self._num_rows = 0
+        self._finalized = False
+
+    @property
+    def num_rows(self) -> int:
+        return self._num_rows
+
+    def append(self, rows: np.ndarray, *, expected_rows: int) -> None:
+        if self._finalized:
+            raise RuntimeError("token metadata trace is already finalized")
+        if isinstance(expected_rows, bool) or not isinstance(expected_rows, int) or expected_rows < 0:
+            raise ValueError(f"expected_rows must be a non-negative integer, got {expected_rows!r}")
+        if not isinstance(rows, np.ndarray):
+            raise TypeError("token metadata rows must be a NumPy array")
+        if rows.ndim < 1:
+            raise ValueError("token metadata must have a token-row dimension")
+        if rows.shape[0] != expected_rows:
+            raise ValueError(f"token metadata has {rows.shape[0]} rows, expected {expected_rows}")
+        if not rows.flags.c_contiguous:
+            raise ValueError("token metadata rows must be contiguous")
+
+        schema = (rows.shape[1:], rows.dtype)
+        if self._schema is None:
+            self._schema = schema
+        elif schema != self._schema:
+            raise ValueError(f"token metadata schema changed from {self._schema} to {schema}")
+
+        self._chunks.append(rows)
+        self._num_rows += expected_rows
+
+    def finalize(self, *, expected_rows: int) -> np.ndarray:
+        if self._finalized:
+            raise RuntimeError("token metadata trace is already finalized")
+        if self._num_rows != expected_rows:
+            raise ValueError(f"token metadata trace has {self._num_rows} rows, expected {expected_rows}")
+        if not self._chunks:
+            raise ValueError("token metadata trace has no chunks")
+
+        self._finalized = True
+        return self._chunks[0] if len(self._chunks) == 1 else np.concatenate(self._chunks, axis=0)

skyrl/backends/skyrl_train/inference_servers/base.py

@@ -34,6 +34,7 @@ class InferenceEngineInput(TypedDict):
     # Optional prefix-cache salt forwarded to vLLM as the request ``cache_salt`` so cache blocks are
     # only shared between requests carrying the same salt. See ``GeneratorConfig.use_cache_salt``.
     cache_salt: Optional[str]
+    routed_experts_prompt_starts: Optional[List[int]]
 
 
 class InferenceEngineOutput(TypedDict):

skyrl/backends/skyrl_train/inference_servers/generate_wire.py

@@ -3,14 +3,21 @@
 ``VLLMServerActor`` writes these payloads and ``RemoteInferenceClient`` reads
 them; nothing else depends on the encoding. Both sides serialize with orjson,
 which rejects non-finite floats and has no notion of NumPy arrays, so the
-helpers here exist to get sampled logprobs and routed-expert IDs across that
+helpers here exist to get sampled logprobs and NumPy side channels across that
 boundary intact.
+
+Side-channel arrays use ``{data: <base64>, shape: [...], dtype: <name>}``
+envelopes. Keeping ``data`` first lets ``load_packed_body`` decode it from the
+raw response without materializing a large Python ``str``.
 """
 
 import math
-from typing import Any, Iterable, Mapping, Optional, Tuple
+from collections import deque
+from enum import StrEnum
+from typing import Any, Collection, Iterable, Mapping, Optional, Tuple
 
 import numpy as np
+import orjson
 import pybase64
 
 from skyrl.backends.skyrl_train.utils.routed_experts import (
@@ -22,7 +29,32 @@
 # Matches the floor vLLM applies at its own serving boundaries.
 CLAMPED_LOGPROB = -9999.0
 
-_DTYPES = {dtype.name: dtype for dtype in ROUTED_EXPERT_DTYPES}
+
+class PackedArrayKey(StrEnum):
+    """Envelope keys, with ``DATA`` first for ``load_packed_body``."""
+
+    DATA = "data"
+    SHAPE = "shape"
+    DTYPE = "dtype"
+
+
+class PackedField(StrEnum):
+    """Response-body fields whose value is a packed-array envelope."""
+
+    ROUTED_EXPERTS = "routed_experts"
+    ROLLOUT_SAMPLE_SUPPORT = "rollout_sample_support"
+
+
+PACKED_SIDE_CHANNEL_FIELDS: tuple[str, ...] = tuple(PackedField)
+
+_ENVELOPE_KEYS = frozenset(PackedArrayKey)
+
+_ROUTED_EXPERTS_NDIM = 3
+
+_QUOTE = b'"'
+
+# Base64 cannot contain this scan anchor.
+_PACKED_DATA_ANCHOR = f':{{"{PackedArrayKey.DATA}":"'.encode()
 
 
 def build_logprobs_content(
@@ -72,35 +104,151 @@ def _to_host_array(routed_experts: Any) -> Any:
     return routed_experts
 
 
-def pack_routed_experts(routed_experts: RoutedExpertIndices) -> dict[str, Any]:
-    compact = compact_routed_expert_indices(_to_host_array(routed_experts))
-    return {
-        "data": pybase64.b64encode(memoryview(compact)).decode("ascii"),
-        "shape": list(compact.shape),
-        "dtype": compact.dtype.name,
+def pack_ndarray(
+    arr: np.ndarray,
+    *,
+    allowed_dtypes: Collection[np.dtype],
+    extra: Optional[Mapping[str, Any]] = None,
+) -> dict[str, Any]:
+    """Encode ``arr`` as a base64 envelope carrying ``extra`` as sidecar fields."""
+    if not isinstance(arr, np.ndarray):
+        raise TypeError("packed array must be a NumPy array")
+    if arr.dtype not in allowed_dtypes:
+        allowed = sorted(dtype.name for dtype in allowed_dtypes)
+        raise ValueError(f"packed array {PackedArrayKey.DTYPE} {arr.dtype.name!r} is not one of {allowed}")
+    if extra is not None:
+        collisions = sorted(set(extra) & _ENVELOPE_KEYS)
+        if collisions:
+            raise ValueError(f"sidecar fields collide with envelope keys: {collisions}")
+
+    contiguous = np.ascontiguousarray(arr)
+    # `.value` keys: orjson rejects str subclasses as dict keys.
+    payload = {
+        PackedArrayKey.DATA.value: pybase64.b64encode(memoryview(contiguous)).decode("ascii"),
+        PackedArrayKey.SHAPE.value: list(contiguous.shape),
+        PackedArrayKey.DTYPE.value: contiguous.dtype.name,
     }
-
-
-def decode_packed_routed_experts(payload: dict[str, Any]) -> RoutedExpertIndices:
-    if not isinstance(payload, dict):
-        raise TypeError("packed routed expert indices must be an object")
+    if extra is not None:
+        payload.update(extra)
+    return payload
+
+
+def unpack_ndarray(
+    payload: Mapping[str, Any],
+    *,
+    allowed_dtypes: Collection[np.dtype],
+    ndim: int,
+) -> Tuple[np.ndarray, dict[str, Any]]:
+    """Decode an envelope whose base64 ``data`` may be a string or buffer."""
+    if not isinstance(payload, Mapping):
+        raise TypeError("packed array payload must be an object")
     try:
-        dtype = _DTYPES[payload["dtype"]]
-        shape = tuple(payload["shape"])
-        data = pybase64.b64decode_as_bytearray(payload["data"], validate=True)
+        dtype_name = payload[PackedArrayKey.DTYPE]
+        shape = tuple(payload[PackedArrayKey.SHAPE])
+        data = pybase64.b64decode_as_bytearray(payload[PackedArrayKey.DATA], validate=True)
     except (KeyError, TypeError, ValueError) as exc:
-        raise ValueError("invalid packed routed_experts payload") from exc
-    # bool is a subclass of int, so it needs an explicit rejection; np.integer is
-    # accepted for in-process callers, since orjson only ever yields plain ints.
-    if len(shape) != 3 or any(
+        raise ValueError(f"invalid packed array envelope: {exc}") from exc
+
+    dtypes = {dtype.name: dtype for dtype in allowed_dtypes}
+    if not isinstance(dtype_name, str) or dtype_name not in dtypes:
+        raise ValueError(f"packed array {PackedArrayKey.DTYPE} {dtype_name!r} is not one of {sorted(dtypes)}")
+    dtype = dtypes[dtype_name]
+    # Reject bool, an int subclass; accept np.integer for in-process callers.
+    if len(shape) != ndim or any(
         not isinstance(dim, (int, np.integer)) or isinstance(dim, bool) or dim < 0 for dim in shape
     ):
-        raise ValueError(f"invalid packed routed_experts shape: {shape}")
+        raise ValueError(f"packed array {PackedArrayKey.SHAPE} {shape} is not {ndim} non-negative dimensions")
     expected_size = math.prod(shape) * dtype.itemsize
     if len(data) != expected_size:
-        raise ValueError(f"packed routed_experts has {len(data)} bytes, expected {expected_size}")
-    decoded = np.frombuffer(data, dtype=dtype).reshape(shape)
+        raise ValueError(
+            f"packed array {PackedArrayKey.DATA} has {len(data)} bytes, "
+            f"expected {expected_size} for {dtype_name}{list(shape)}"
+        )
+
+    array = np.frombuffer(data, dtype=dtype).reshape(shape)
+    sidecar = {key: value for key, value in payload.items() if key not in _ENVELOPE_KEYS}
+    return array, sidecar
+
+
+def pack_routed_experts(routed_experts: RoutedExpertIndices) -> dict[str, Any]:
+    compact = compact_routed_expert_indices(_to_host_array(routed_experts))
+    return pack_ndarray(compact, allowed_dtypes=ROUTED_EXPERT_DTYPES)
+
+
+def decode_packed_routed_experts(payload: dict[str, Any]) -> RoutedExpertIndices:
+    decoded, _ = unpack_ndarray(payload, allowed_dtypes=ROUTED_EXPERT_DTYPES, ndim=_ROUTED_EXPERTS_NDIM)
     compact = compact_routed_expert_indices(decoded)
-    if compact.dtype != dtype:
-        raise ValueError(f"packed routed_experts uses non-canonical dtype {dtype.name}; expected {compact.dtype.name}")
+    if compact.dtype != decoded.dtype:
+        raise ValueError(
+            f"packed routed_experts uses non-canonical dtype {decoded.dtype.name}; expected {compact.dtype.name}"
+        )
     return compact
+
+
+def _data_prefix(field: str) -> bytes:
+    """The bytes an orjson-serialized packed ``field`` opens with."""
+    return f'"{field}"'.encode() + _PACKED_DATA_ANCHOR
+
+
+def load_packed_body(raw: bytes, *, fields: tuple[str, ...] = PACKED_SIDE_CHANNEL_FIELDS) -> dict[str, Any]:
+    """Parse a response after replacing registered base64 blobs with views.
+
+    Null fields pass through. An envelope layout the scan cannot splice raises
+    instead of falling back to materializing the base64 as a Python string.
+    """
+    prefixes = {field: _data_prefix(field) for field in fields}
+    blobs: dict[str, deque[memoryview]] = {field: deque() for field in fields}
+    view = memoryview(raw)
+    pieces: list[memoryview] = []
+    copied = 0
+    scan = 0
+    while (anchor := raw.find(_PACKED_DATA_ANCHOR, scan)) >= 0:
+        field = _match_packed_field(raw, anchor, prefixes)
+        if field is None:
+            scan = anchor + len(_PACKED_DATA_ANCHOR)
+            continue
+        start = anchor + len(_PACKED_DATA_ANCHOR)
+        end = raw.find(_QUOTE, start)
+        if end < 0:
+            raise ValueError(f"unterminated base64 {PackedArrayKey.DATA} for {field} in the response body")
+        pieces.append(view[copied:start])
+        blobs[field].append(view[start:end])
+        copied = scan = end
+
+    if pieces:
+        pieces.append(view[copied:])
+        body = orjson.loads(b"".join(pieces))
+    else:
+        body = orjson.loads(raw)
+    _restore_packed_data(body, blobs)
+
+    unplaced = {field: len(queue) for field, queue in blobs.items() if queue}
+    if unplaced:
+        raise ValueError(f"spliced packed blobs found no envelope in the response body: {unplaced}")
+    return body
+
+
+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 _restore_packed_data(node: Any, blobs: Mapping[str, deque[memoryview]]) -> None:
+    """Put each blob back on its envelope's ``data`` key, in document order."""
+    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:
+                if not queue:
+                    raise ValueError(f"packed {key} survived the scan unspliced; the response-body layout drifted")
+                value[PackedArrayKey.DATA.value] = queue.popleft()
+            elif isinstance(value, (dict, list)):
+                _restore_packed_data(value, blobs)
+    elif isinstance(node, list):
+        for item in node:
+            if isinstance(item, (dict, list)):
+                _restore_packed_data(item, blobs)