feat: let the host keep finished episodes' training arrays

tests/unit/serve/test_delta.py

@@ -172,8 +172,7 @@ async def send(data: dict) -> None:
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("routing_dtype", [np.uint8, np.uint16])
-async def test_deltas_stream_once_and_reassemble_the_episode(routing_dtype):
+async def test_deltas_stream_once_and_reassemble_the_episode():
     traces: list[vf.Trace] = []
     frames: list[bytes] = []
 
@@ -192,12 +191,6 @@ async def send(data: dict) -> None:
         trace.nodes[1].routed_experts = np.arange(6, dtype=np.uint8).reshape(3, 2, 1)
         await settle()
         add_turn(trace, "a2")
-        # The next prefill corrects the previous turn's last routing row, potentially
-        # widening its dtype. The previous node has already crossed the wire.
-        trace.nodes[1].routed_experts = np.array(
-            [[[0], [1]], [[2], [3]], [[100], [np.iinfo(routing_dtype).max]]],
-            dtype=routing_dtype,
-        )
         trace.nodes[0].semantic_parents.append(ParentLink(node=1, type="reply"))
         await settle()
         trace.record_reward("match", 1.0)
@@ -228,10 +221,8 @@ async def send(data: dict) -> None:
     rebuilt = WireEpisode.model_validate(unpack(pack(assembly.finish(head, summaries))))
     assert rebuilt.ok and rebuilt.traces[0].id == trace.id
     assert rebuilt.traces[0].messages == trace.messages
-    np.testing.assert_array_equal(
-        rebuilt.traces[0].nodes[1].routed_experts, trace.nodes[1].routed_experts
-    )
-    assert rebuilt.traces[0].nodes[1].routed_experts.dtype == routing_dtype
+    # Training arrays stay in the worker.
+    assert rebuilt.traces[0].nodes[1].routed_experts is None
     assert (
         rebuilt.traces[0].nodes[0].semantic_parents == trace.nodes[0].semantic_parents
     )

verifiers/v1/serve/__init__.py

@@ -1,7 +1,7 @@
 from verifiers.v1.serve.client import EnvClient
 from verifiers.v1.serve.delta import EpisodeAssembly
 from verifiers.v1.serve.pool import EnvServerPool, env_config_data, serve_env
-from verifiers.v1.serve.server import EnvServer
+from verifiers.v1.serve.server import EnvServer, set_episode_hook
 from verifiers.v1.serve.types import (
     HealthRequest,
     HealthResponse,
@@ -20,4 +20,5 @@
     "RunResponse",
     "env_config_data",
     "serve_env",
+    "set_episode_hook",
 ]

verifiers/v1/serve/client.py

@@ -169,12 +169,13 @@ async def _request(
         response_type: type[ResponseT],
         timeout: float | None = None,
         on_delta: Callable[[dict], None] | None = None,
+        request_id: str | None = None,
     ) -> ResponseT:
         """Send a typed request and validate the reply into `response_type`. A
         `timeout` is only used for health polling — rollouts run untimed. `on_delta`
         receives each decoded `delta` the request streams before its reply."""
         self._ensure_receiver()
-        request_id = uuid.uuid4().hex
+        request_id = request_id or uuid.uuid4().hex
         future: asyncio.Future[bytes] = asyncio.get_running_loop().create_future()
         self._pending[request_id] = future
         if on_delta is not None:
@@ -191,14 +192,7 @@ async def _request(
         except asyncio.CancelledError:
             self._pending.pop(request_id, None)
             if request.method == RunRequest.method:
-                # Fire-and-forget: tell the server to abort the rollout so the
-                # episode stops consuming inference and env-runtime resources.
-                # A task on the loop outlives this (cancelled) caller.
-                task = asyncio.get_running_loop().create_task(
-                    self._send_cancel(request_id)
-                )
-                self._cancel_tasks.add(task)
-                task.add_done_callback(self._cancel_tasks.discard)
+                self._cancel_run(request_id)
             raise
         finally:
             self._deltas.pop(request_id, None)
@@ -223,6 +217,14 @@ async def _decode(self, decode: Callable[..., DecodedT], *args: object) -> Decod
             )
             raise
 
+    def _cancel_run(self, run_request_id: str) -> None:
+        """Fire-and-forget: tell the server the caller abandoned the run, so it aborts
+        the rollout (or frees what it kept for an episode already sent). A task on the
+        loop outlives the cancelled caller."""
+        task = asyncio.get_running_loop().create_task(self._send_cancel(run_request_id))
+        self._cancel_tasks.add(task)
+        task.add_done_callback(self._cancel_tasks.discard)
+
     async def _send_cancel(self, run_request_id: str) -> None:
         """Best-effort server-side abort of an abandoned run."""
         with contextlib.suppress(Exception):
@@ -276,6 +278,7 @@ def apply(delta: dict) -> None:
             if on_delta is not None:
                 on_delta(delta)
 
+        request_id = uuid.uuid4().hex
         response = await self._request(
             RunRequest(
                 task_data=task_data,
@@ -285,11 +288,19 @@ def apply(delta: dict) -> None:
             ),
             RunResponse,
             on_delta=apply,
+            request_id=request_id,
         )
         assert response.head is not None
-        return await self._decode(
-            WireEpisode.model_validate, assembly.finish(response.head, response.traces)
-        )
+        try:
+            return await self._decode(
+                WireEpisode.model_validate,
+                assembly.finish(response.head, response.traces),
+            )
+        except BaseException:
+            # The reply already arrived, so the server counts the episode delivered;
+            # tell it this caller never got the episode.
+            self._cancel_run(request_id)
+            raise
 
     async def close(self) -> None:
         if self._receiver is not None:

verifiers/v1/serve/delta.py

@@ -4,12 +4,13 @@
 (`Trace.notify`, fired by the rollout at every phase change and by the interception proxy
 after every recorded turn); the `DeltaStreamer` then diffs the run's live traces against
 what it has already sent and ships only the new part — the trace header once, then
-appended nodes / calls / errors, semantic links and routing-row repairs on earlier nodes, the scalar
+appended nodes / calls / errors, semantic links on earlier nodes, the scalar
 fields whose value changed (timing spans, stop condition, rewards, ...), and the
 `pending` preview — the messages of the request in flight that no node holds yet, so a
 watcher sees a tool result before the model has answered it. The `Trace` is
-append-only at turn granularity except for links and the last routing row of a node,
-which the next prefill can repair. Apart from these and the preview, every byte of the
+append-only at turn granularity except for links. Nodes stream without their per-token
+training arrays (`TRAINING_ARRAY_FIELDS`): a host that trains on them keeps them itself
+(`set_episode_hook`). Apart from links and the preview, every byte of the
 episode crosses the wire once and the stream costs about what a single reply would; the
 reply that ends the run carries only the episode head and per-trace counts the client
 checks its assembly against. A cursor advances only once its delta is on the wire, so a
@@ -29,10 +30,8 @@
 from typing import TYPE_CHECKING, Any, Self
 
 import msgpack
-import numpy as np
 from pydantic import BaseModel
 
-from verifiers.v1.graph import decode_ndarray, encode_ndarray
 from verifiers.v1.serve.encoding import msgpack_encoder
 
 if TYPE_CHECKING:
@@ -74,6 +73,10 @@
 """Re-sent whole whenever their dumped value changes."""
 
 
+TRAINING_ARRAY_FIELDS = frozenset({"routed_experts", "sampling_mask", "top_logprobs"})
+"""Node fields holding per-token training arrays; never streamed."""
+
+
 def pack(payload: Any) -> bytes:
     return msgpack.packb(payload, default=msgpack_encoder, use_bin_type=True)
 
@@ -96,8 +99,6 @@ def __init__(self) -> None:
         self.sent = dict.fromkeys(LIST_FIELDS, 0)
         self.links: list[int] = []
         """Per sent node, how many of its semantic links went out with or after it."""
-        self.final_rows: dict[int, bytes] = {}
-        """Per node that carries routing, its final row as packed when last sent."""
         self.scalars: dict[str, bytes] = {}
         self.pending: bytes = pack([])
 
@@ -169,31 +170,6 @@ async def flush(self) -> None:
                 else:
                     self.cursors[trace_id] = cursor
 
-    def _maybe_add_routing_repairs(
-        self,
-        delta: dict[str, Any],
-        trace: Trace,
-        cursor: TraceCursor,
-        sent_nodes: int,
-    ) -> None:
-        """Add final rows repaired since they were sent, keyed by node index.
-
-        Record every node's current row on the cursor. `sent_nodes` is the pre-flush
-        count, so a node first sent in this delta is never reported as a repair.
-        """
-        repairs: dict[int, dict] = {}
-        for index, node in enumerate(trace.nodes):
-            if node.routed_experts is None:
-                continue
-            row = encode_ndarray(node.routed_experts[-1:])
-            packed = pack(row)
-            if cursor.final_rows.get(index) != packed:
-                cursor.final_rows[index] = packed
-                if index < sent_nodes:
-                    repairs[index] = row
-        if repairs:
-            delta["routing_repairs"] = repairs
-
     def diff(self) -> list[tuple[str, dict, TraceCursor | None]]:
         """Each trace's delta against its sent cursor, with the cursor as it stands once
         that delta is sent (None for a discard). Nothing here is committed: `flush`
@@ -220,15 +196,19 @@ def diff(self) -> list[tuple[str, dict, TraceCursor | None]]:
                     cursor.links[index] = len(node_links)
             if links:
                 delta["links"] = links
-            self._maybe_add_routing_repairs(delta, trace, cursor, cursor.sent["nodes"])
             for field in LIST_FIELDS:
                 items = getattr(trace, field)
                 sent = cursor.sent[field]
                 if len(items) > sent:
                     # Nodes' per-token streams ride as raw-bytes arrays.
-                    context = {"token_arrays": True} if field == "nodes" else None
+                    is_nodes = field == "nodes"
                     delta[field] = [
-                        dump(item, context=context) for item in items[sent:]
+                        dump(
+                            item,
+                            context={"token_arrays": True} if is_nodes else None,
+                            exclude=TRAINING_ARRAY_FIELDS if is_nodes else None,
+                        )
+                        for item in items[sent:]
                     ]
                     cursor.sent[field] = len(items)
                     if field == "nodes":
@@ -262,20 +242,6 @@ class EpisodeAssembly:
     def __init__(self) -> None:
         self.traces: dict[str, dict] = {}
 
-    def _maybe_apply_routing_repairs(self, delta: dict[str, Any], trace: dict) -> None:
-        """Replace repaired final rows in nodes the client already holds.
-
-        Rebuild each array: a repair can widen its dtype, and decoded rows alias
-        the original delta's bytes, which must remain unchanged for consumers.
-        """
-        for index, row in (delta.get("routing_repairs") or {}).items():
-            node = trace["nodes"][int(index)]
-            node["routed_experts"] = encode_ndarray(
-                np.concatenate(
-                    [decode_ndarray(node["routed_experts"])[:-1], decode_ndarray(row)]
-                )
-            )
-
     def apply(self, delta: dict) -> None:
         trace_id = delta["trace"]
         if delta.get("discard"):
@@ -290,7 +256,6 @@ def apply(self, delta: dict) -> None:
             }
         for index, links in (delta.get("links") or {}).items():
             trace["nodes"][int(index)]["semantic_parents"].extend(links)
-        self._maybe_apply_routing_repairs(delta, trace)
         # a later `links` delta grows a node's semantic_parents in place, so the node
         # is copied: the delta stays as it was when the caller received it
         if "nodes" in delta:

verifiers/v1/serve/pool.py

@@ -28,14 +28,15 @@
 import signal
 import tempfile
 import threading
+import time
 from collections.abc import Callable
 
 import msgpack
 import zmq
 import zmq.asyncio
 
 from verifiers.v1.configs.env import EnvConfig
-from verifiers.v1.serve.server import EnvServer
+from verifiers.v1.serve.server import LATE_CANCEL_SECONDS, EnvServer
 from verifiers.v1.serve.types import CancelResponse, HealthResponse
 
 logger = logging.getLogger(__name__)
@@ -172,6 +173,9 @@ async def run(self) -> None:
             self._spawn_worker()
         # request_id -> {client_id, worker}
         pending: dict[bytes, dict] = {}
+        # request_id -> (replied_at, worker), oldest first: a late cancel still reaches
+        # the worker that may hold the run's kept arrays (`LATE_CANCEL_SECONDS`).
+        replied: dict[bytes, tuple[float, dict]] = {}
         logger.info(
             "EnvServerPool up: address=%s workers=%d/%s multiplex=%d elastic=%s",
             self.address,
@@ -199,8 +203,10 @@ async def run(self) -> None:
                                 [client_id, request_id, b"reply", _HEALTH]
                             )
                         elif method == b"cancel":
-                            # Route the cancel to the worker holding the target run;
-                            # an unknown/finished target (or an unparseable payload
+                            # Route the cancel to the worker holding the target run,
+                            # or to the one that replied to it within
+                            # `LATE_CANCEL_SECONDS` (it may hold the run's kept
+                            # arrays); any other target (or an unparseable payload
                             # — one bad frame must not tear down the broker) is
                             # answered inline
                             try:
@@ -209,12 +215,21 @@ async def run(self) -> None:
                             except Exception:  # noqa: BLE001 - one bad frame must not kill the broker
                                 target = b""
                             target_entry = pending.get(target)
-                            if target_entry is None:
+                            late = replied.get(target)
+                            if target_entry is not None:
+                                worker = target_entry["worker"]
+                            elif (
+                                late is not None
+                                and time.monotonic() - late[0] < LATE_CANCEL_SECONDS
+                            ):
+                                worker = late[1]
+                            else:
+                                worker = None
+                            if worker is None:
                                 await self.frontend.send_multipart(
                                     [client_id, request_id, b"reply", _CANCEL_MISS]
                                 )
                             else:
-                                worker = target_entry["worker"]
                                 worker["active"] += 1
                                 pending[request_id] = {
                                     "client_id": client_id,
@@ -253,6 +268,12 @@ async def run(self) -> None:
                                 continue
                             entry["worker"]["active"] -= 1
                             in_flight -= 1
+                            now = time.monotonic()
+                            replied[request_id.bytes] = (now, entry["worker"])
+                            for old, (replied_at, _) in list(replied.items()):
+                                if now - replied_at < LATE_CANCEL_SECONDS:
+                                    break
+                                del replied[old]
                         else:
                             entry = pending.get(request_id.bytes)
                             if entry is None: