[Router Replay]: Improve performance by removing Pydantic validation (#1394)

Author: S1ro1

docs/reference.md

@@ -181,6 +181,14 @@ class TrajectoryStep(TypedDict):
 
 A single turn in a multi-turn rollout.
 
+### RoutedExpertsPayload
+
+```python
+class RoutedExpertsPayload(TypedDict):
+    data: Any  # actually memoryview; kept opaque so Pydantic skips schema validation
+    shape: list[int]
+```
+
 ### TrajectoryStepTokens
 
 ```python
@@ -192,7 +200,7 @@ class TrajectoryStepTokens(TypedDict):
     completion_logprobs: list[float]
     overlong_prompt: bool
     is_truncated: bool
-    routed_experts: list[list[list[int]]] | None  # [seq_len, layers, topk] to enable router replay
+    routed_experts: RoutedExpertsPayload | None
     multi_modal_data: NotRequired[Any]  # renderers.MultiModalData sidecar (pixel_values, placeholder ranges) — set only on multimodal rollouts
 ```
 

tests/test_client_auth_errors.py

@@ -130,6 +130,9 @@ def __init__(self, message: str) -> None:
     def __init__(self, message: str) -> None:
         self.chat = self._Chat(message)
 
+    async def post(self, *args, **kwargs):  # noqa: ANN002, ANN003
+        return await self.chat.completions.create(*args, **kwargs)
+
 
 @pytest.mark.parametrize(
     "error_message",

tests/test_openai_chat_completions_token_client.py

@@ -1,5 +1,6 @@
 from typing import Any, cast
 
+import httpx
 import pytest
 
 from verifiers.clients.openai_chat_completions_client import OpenAIChatCompletionsClient
@@ -24,7 +25,25 @@ async def post(
         self, path: str, body: dict[str, Any], cast_to: type, **kwargs: Any
     ) -> Any:
         self.calls.append({"path": path, "body": body, "cast_to": cast_to})
-        return {"ok": True, "path": path, "body": body}
+        return httpx.Response(
+            200,
+            json={
+                "id": path,
+                "object": "chat.completion",
+                "created": 1,
+                "model": body["model"],
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": "ok"},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "ok": True,
+                "path": path,
+                "body": body,
+            },
+        )
 
 
 class _PromptIdTestClient(OpenAIChatCompletionsTokenClient):
@@ -270,7 +289,7 @@ async def fake_get_prompt_ids(  # noqa: ANN001
         state=state,
     )
 
-    assert response["ok"] is True
+    assert response.model_extra["ok"] is True
     assert len(recording_client.calls) == 1
     assert recording_client.calls[0]["path"] == "/chat/completions/tokens"
     assert recording_client.calls[0]["body"]["tokens"] == [10, 20]

verifiers/clients/openai_chat_completions_client.py

@@ -56,8 +56,10 @@
     Usage,
     UserMessage,
 )
-from verifiers.utils.client_utils import setup_openai_client
-from verifiers.utils.response_utils import parse_routed_experts
+from verifiers.utils.client_utils import (
+    post_chat_completion_with_routed_experts_sidecar,
+    setup_openai_client,
+)
 
 
 def handle_openai_overlong_prompt(func):
@@ -300,23 +302,24 @@ def normalize_sampling_args(sampling_args: SamplingArgs):
             sampling_args = {**sampling_args, "modalities": ["text"]}
 
         extra_headers = kwargs.pop("extra_headers", None)
-
+        request_args = normalize_sampling_args(sampling_args)
+        extra_body = request_args.pop("extra_body", {})
+
+        body: dict[str, Any] = {
+            "model": model,
+            "messages": prompt,
+            **request_args,
+            **extra_body,
+        }
         if tools:
-            response = await self.client.chat.completions.create(
-                model=model,
-                messages=prompt,
-                tools=tools,
-                extra_headers=extra_headers,
-                **normalize_sampling_args(sampling_args),
-            )
-        else:
-            response = await self.client.chat.completions.create(
-                model=model,
-                messages=prompt,
-                extra_headers=extra_headers,
-                **normalize_sampling_args(sampling_args),
-            )
-        return response
+            body["tools"] = tools
+
+        return await post_chat_completion_with_routed_experts_sidecar(
+            self.client,
+            "/chat/completions",
+            body=body,
+            extra_headers=extra_headers,
+        )
 
     async def raise_from_native_response(self, response: OpenAIChatResponse) -> None:
         if response is None:
@@ -483,14 +486,13 @@ def parse_tokens(response: OpenAIChatResponse) -> ResponseTokens | None:
                 completion_logprobs = [token["logprob"] for token in logprobs_content]
 
             choice_extra = choice.model_extra or {}
-            routed_experts = parse_routed_experts(choice_extra.get("routed_experts"))
             return ResponseTokens(
                 prompt_ids=prompt_ids,
                 prompt_mask=prompt_mask,
                 completion_ids=completion_ids,
                 completion_mask=completion_mask,
                 completion_logprobs=completion_logprobs,
-                routed_experts=routed_experts,
+                routed_experts=choice_extra.get("routed_experts"),
             )
 
         def parse_reasoning_content_from_response(

verifiers/clients/openai_chat_completions_token_client.py

@@ -3,7 +3,6 @@
 
 from openai import AsyncOpenAI, BaseModel
 from openai.types.chat import (
-    ChatCompletion,
     ChatCompletionAssistantMessageParam,
 )
 from openai.types.chat.chat_completion_message_function_tool_call_param import (
@@ -20,6 +19,9 @@
     handle_openai_overlong_prompt,
 )
 from verifiers.types import SamplingArgs, State
+from verifiers.utils.client_utils import (
+    post_chat_completion_with_routed_experts_sidecar,
+)
 
 
 def _has_multimodal_content(messages) -> bool:
@@ -140,20 +142,20 @@ def normalize_sampling_args(sampling_args: SamplingArgs):
             )
 
         extra_body = sampling_args.pop("extra_body", {})
-        body = dict(
-            model=model,
-            messages=prompt,
-            tools=tools,
-            tokens=prompt_ids,
+        body = {
+            "model": model,
+            "messages": prompt,
+            "tools": tools,
+            "tokens": prompt_ids,
             **sampling_args,
             **extra_body,
-        )
+        }
 
-        return await self.client.post(
+        return await post_chat_completion_with_routed_experts_sidecar(
+            self.client,
             "/chat/completions/tokens",
             body=body,
-            cast_to=ChatCompletion,
-            options={"headers": extra_headers} if extra_headers else {},
+            extra_headers=extra_headers,
         )
 
     async def get_prompt_ids(

verifiers/clients/openai_completions_client.py

@@ -25,7 +25,6 @@
     Usage,
 )
 from verifiers.utils.client_utils import setup_openai_client
-from verifiers.utils.response_utils import parse_routed_experts
 
 OpenAITextMessages = str
 OpenAITextResponse = Completion
@@ -170,15 +169,12 @@ def parse_tokens(response: OpenAITextResponse) -> ResponseTokens | None:
             )
             if completion_logprobs is None:
                 return None
-            choice_extra = response.choices[0].model_extra or {}
-            routed_experts = parse_routed_experts(choice_extra.get("routed_experts"))
             return ResponseTokens(
                 prompt_ids=prompt_ids,
                 prompt_mask=prompt_mask,
                 completion_ids=completion_ids,
                 completion_mask=completion_mask,
                 completion_logprobs=completion_logprobs,
-                routed_experts=routed_experts,
             )
 
         return Response(

verifiers/types.py

@@ -183,13 +183,19 @@ class Usage(CustomBaseModel):
     total_tokens: int
 
 
+class RoutedExpertsPayload(TypedDict):
+    # Keep the raw response sidecar opaque so Pydantic does not validate memoryview.
+    data: Any
+    shape: list[int]
+
+
 class ResponseTokens(CustomBaseModel):
     prompt_ids: list[int]
     prompt_mask: list[int]
     completion_ids: list[int]
     completion_mask: list[int]
     completion_logprobs: list[float]
-    routed_experts: str | None = None  # base64 NumPy [seq_len, layers, topk]
+    routed_experts: RoutedExpertsPayload | None = None
     # Renderer-emitted multimodal sidecar (renderers.base.MultiModalData)
     # carrying processed pixel_values / placeholder ranges per modality.
     # Populated by the renderer client when the rollout went through a
@@ -232,7 +238,7 @@ class TrajectoryStepTokens(TypedDict):
     completion_logprobs: list[float]
     overlong_prompt: bool
     is_truncated: bool
-    routed_experts: str | None  # base64 NumPy [seq_len, layers, topk]
+    routed_experts: RoutedExpertsPayload | None
     # Renderer-emitted multimodal sidecar (renderers.base.MultiModalData)
     # carrying processed pixel_values / placeholder ranges per modality.
     # ``NotRequired`` because text-only rollouts (and non-renderer client

verifiers/utils/client_utils.py

@@ -1,16 +1,20 @@
 import json
 import logging
 import os
+from collections.abc import Mapping
+from typing import Any
 from pathlib import Path
 
 import httpx
 from anthropic import AsyncAnthropic
 from openai import AsyncOpenAI
+from openai.types.chat import ChatCompletion
 
 from verifiers.types import (
     ClientConfig,
     EndpointClientConfig,
 )
+from verifiers.utils.response_utils import strip_routed_experts_data
 
 logger = logging.getLogger(__name__)
 
@@ -97,6 +101,32 @@ def setup_http_client(config: ClientConfig) -> httpx.AsyncClient:
     return _build_http_client(resolved_config, headers)
 
 
+def parse_chat_completion_with_routed_experts_sidecar(raw: bytes) -> ChatCompletion:
+    stripped, routed_data = strip_routed_experts_data(raw)
+    response = ChatCompletion.model_validate_json(stripped)
+    if routed_data is not None:
+        choice_extra = response.choices[0].model_extra
+        assert choice_extra is not None
+        choice_extra["routed_experts"]["data"] = routed_data
+    return response
+
+
+async def post_chat_completion_with_routed_experts_sidecar(
+    client: AsyncOpenAI,
+    path: str,
+    *,
+    body: dict[str, Any],
+    extra_headers: Mapping[str, str] | None = None,
+) -> ChatCompletion:
+    raw_response = await client.post(
+        path,
+        body=body,
+        cast_to=httpx.Response,
+        options={"headers": extra_headers} if extra_headers else {},
+    )
+    return parse_chat_completion_with_routed_experts_sidecar(raw_response.content)
+
+
 def _setup_openai_client_from_resolved(config: ClientConfig) -> AsyncOpenAI:
     headers, api_key = _build_headers_and_api_key(config)
     return AsyncOpenAI(

verifiers/utils/response_utils.py

@@ -1,34 +1,23 @@
-import base64
-from io import BytesIO
-from typing import Any, cast
-
-import numpy as np
-
 from verifiers.types import (
     AssistantMessage,
     Messages,
     Response,
     TrajectoryStepTokens,
 )
 
+ROUTED_EXPERTS_DATA_PREFIX = b'"routed_experts":{"data":"'
 
-def parse_routed_experts(raw: Any) -> str | None:
-    if raw is None:
-        return None
-    return cast(str, raw)
-
-
-def truncate_routed_experts(routed_experts: str | None, seq_len: int) -> str | None:
-    if routed_experts is None:
-        return None
 
-    array = np.load(BytesIO(base64.b64decode(routed_experts)), allow_pickle=False)
-    assert array.ndim == 3
-    assert 0 <= seq_len <= array.shape[0]
+def strip_routed_experts_data(raw: bytes) -> tuple[bytes, memoryview | None]:
+    data_start = raw.find(ROUTED_EXPERTS_DATA_PREFIX)
+    if data_start < 0:
+        return raw, None
 
-    buffer = BytesIO()
-    np.save(buffer, np.ascontiguousarray(array[:seq_len]), allow_pickle=False)
-    return base64.b64encode(buffer.getvalue()).decode("ascii")
+    data_start += len(ROUTED_EXPERTS_DATA_PREFIX)
+    data_end = raw.index(b'"', data_start)
+    routed_data = memoryview(raw)[data_start:data_end]
+    stripped = raw[:data_start] + raw[data_end:]
+    return stripped, routed_data
 
 
 async def parse_response_message(response: Response) -> Messages:
@@ -73,15 +62,11 @@ async def parse_response_tokens(
             completion_ids = []
             completion_mask = []
             completion_logprobs = []
-            routed_experts = truncate_routed_experts(routed_experts, len(prompt_ids))
         elif prompt_len + completion_len > max_seq_len:
             is_truncated = True
             completion_ids = tokens.completion_ids[: max_seq_len - prompt_len]
             completion_mask = tokens.completion_mask[: max_seq_len - prompt_len]
             completion_logprobs = tokens.completion_logprobs[: max_seq_len - prompt_len]
-            routed_experts = truncate_routed_experts(
-                routed_experts, prompt_len + len(completion_ids)
-            )
         else:
             is_truncated = False
     else:
@@ -104,4 +89,6 @@ async def parse_response_tokens(
         # step. Leaving it on ``response.message.tokens`` too means every
         # downstream pass (msgpack, save) has to dedupe the duplicate.
         tokens.multi_modal_data = None
+    if routed_experts is not None:
+        tokens.routed_experts = None
     return out