From b42a26ff792d28dc12bc31b2f65646fc7c9dbbff Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 21 Oct 2026 14:37:23 +0000 Subject: [PATCH 01/10] feat(serve): let the host keep finished episodes' training arrays `set_episode_hook` registers a hook the env server runs, off the event loop, on each finished episode before its final flush. A trainer's env servers use it to keep per-token training arrays (routing rows, sampling masks, top-k heads) themselves and leave references in `trace.info`. While a hook is set, streamed nodes carry no training arrays and no routing repairs. Co-Authored-By: Claude Opus 5.5 --- tests/unit/serve/test_delta.py | 27 +++++++++++++++++++++++++++ verifiers/v1/serve/__init__.py | 3 ++- verifiers/v1/serve/delta.py | 18 +++++++++++++++++- verifiers/v1/serve/server.py | 25 ++++++++++++++++++++++++- 4 files changed, 70 insertions(+), 3 deletions(-) diff --git a/tests/unit/serve/test_delta.py b/tests/unit/serve/test_delta.py index cb2b9ef5a2..9d55fb7eec 100644 --- a/tests/unit/serve/test_delta.py +++ b/tests/unit/serve/test_delta.py @@ -272,3 +272,30 @@ def test_finish_refuses_a_gap(): assembly.apply({"trace": "t", "open": {"id": "t"}, "nodes": [{}]}) with pytest.raises(RuntimeError, match="assembled 1 nodes"): assembly.finish({}, [TraceSummary(id="t", nodes=2, calls=0)]) + + +@pytest.mark.asyncio +async def test_streamer_without_arrays_streams_nodes_without_them(): + """A host that keeps the training arrays itself (`set_episode_hook`) streams nodes + without them, and sends no routing repairs.""" + traces: list[vf.Trace] = [] + frames: list[bytes] = [] + + async def send(data: dict) -> None: + frames.append(pack(data)) + + async with DeltaStreamer(lambda: traces, send, keep_arrays=False) as streamer: + trace = make_trace() + traces.append(trace) + streamer.watch(trace) + add_turn(trace, "a1") + trace.nodes[1].routed_experts = np.arange(6, dtype=np.uint8).reshape(3, 2, 1) + await settle() + add_turn(trace, "a2") + trace.nodes[1].routed_experts = np.zeros((3, 2, 1), dtype=np.uint8) + await settle() + deltas = [unpack(frame) for frame in frames] + nodes = [node for delta in deltas for node in delta.get("nodes", [])] + assert len(nodes) == 3 + assert all("routed_experts" not in node for node in nodes) + assert not any("routing_repairs" in delta for delta in deltas) diff --git a/verifiers/v1/serve/__init__.py b/verifiers/v1/serve/__init__.py index c7e2433527..077ff5abbf 100644 --- a/verifiers/v1/serve/__init__.py +++ b/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", ] diff --git a/verifiers/v1/serve/delta.py b/verifiers/v1/serve/delta.py index abf4a2f308..fef596913b 100644 --- a/verifiers/v1/serve/delta.py +++ b/verifiers/v1/serve/delta.py @@ -74,6 +74,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; streamed only while `keep_arrays`.""" + + def pack(payload: Any) -> bytes: return msgpack.packb(payload, default=msgpack_encoder, use_bin_type=True) @@ -114,9 +118,12 @@ def __init__( self, traces: Callable[[], Iterable[Trace]], send: Callable[[dict], Awaitable[None]], + keep_arrays: bool = True, ) -> None: self.traces = traces self.send = send + self.keep_arrays = keep_arrays + """Stream nodes' training arrays; off when the host keeps them (`set_episode_hook`).""" self.cursors: dict[str, TraceCursor] = {} self._scheduled: asyncio.Handle | None = None self._flushes: set[asyncio.Task] = set() @@ -181,6 +188,8 @@ def _maybe_add_routing_repairs( 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. """ + if not self.keep_arrays: + return repairs: dict[int, dict] = {} for index, node in enumerate(trace.nodes): if node.routed_experts is None: @@ -225,7 +234,14 @@ def diff(self) -> list[tuple[str, dict, TraceCursor | None]]: items = getattr(trace, field) sent = cursor.sent[field] if len(items) > sent: - delta[field] = [dump(item) for item in items[sent:]] + exclude = ( + TRAINING_ARRAY_FIELDS + if field == "nodes" and not self.keep_arrays + else None + ) + delta[field] = [ + dump(item, exclude=exclude) for item in items[sent:] + ] cursor.sent[field] = len(items) if field == "nodes": cursor.links.extend( diff --git a/verifiers/v1/serve/server.py b/verifiers/v1/serve/server.py index 8e04bc8069..fee9bc6ce7 100644 --- a/verifiers/v1/serve/server.py +++ b/verifiers/v1/serve/server.py @@ -1,4 +1,5 @@ import asyncio +from collections.abc import Callable import logging import msgpack @@ -8,6 +9,7 @@ from verifiers.v1.clients import ModelContext from verifiers.v1.configs.client import ClientConfig from verifiers.v1.configs.env import EnvConfig +from verifiers.v1.episode import Episode from verifiers.v1.serve.delta import DeltaStreamer, TraceSummary, dump, pack from verifiers.v1.serve.encoding import msgpack_encoder from verifiers.v1.serve.types import ( @@ -25,6 +27,21 @@ logger = logging.getLogger(__name__) +EpisodeHook = Callable[[Episode, SamplingConfig], None] + +_episode_hook: EpisodeHook | None = None + + +def set_episode_hook(hook: EpisodeHook | None) -> None: + """Let the hosting process take each finished episode's per-token training arrays + (routing rows, sampling masks, top-k heads) before the episode leaves the worker. + The hook runs off the event loop with the episode and its sampling config; it may + keep the arrays and leave references in `trace.info`. While a hook is set, the + streamed nodes carry no training arrays.""" + global _episode_hook + _episode_hook = hook + + class EnvServer: def __init__( self, @@ -108,10 +125,16 @@ async def send_delta(delta: dict) -> None: # The gate spans requests: `--max-concurrent` bounds this worker's episodes # in flight the same way the in-process eval's semaphore does. The streamer # ships each trace as it changes; the reply below carries only the rest. - async with DeltaStreamer(lambda: slot.traces, send_delta) as streamer: + hook = _episode_hook + async with DeltaStreamer( + lambda: slot.traces, send_delta, keep_arrays=hook is None + ) as streamer: episode = await self.env.run_slot( slot, ctx, self._gate, on_trace=streamer.watch ) + if hook is not None: + # Before the final flush, so what the hook writes to the traces streams. + await asyncio.to_thread(hook, episode, ctx.sampling) return RunResponse( head=dump(episode, exclude={"traces"}), traces=[ From abbc637fe82e38a1e90ee36c7176b4f3cf6ada48 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Wed, 21 Oct 2026 15:42:01 +0000 Subject: [PATCH 02/10] feat(serve): account for every episode the hook kept The episode hook becomes a `keep` / `discard` protocol with a strict contract: every episode the server passed to `keep` is either delivered in its run's successful reply or handed back to `discard`, including when the run is cancelled while or after `keep` runs, or the reply fails to send. The host can then free what it kept without any sweep. Co-Authored-By: Claude Opus 5.5 --- verifiers/v1/serve/server.py | 40 ++++++++++++++++++++++++++++-------- 1 file changed, 32 insertions(+), 8 deletions(-) diff --git a/verifiers/v1/serve/server.py b/verifiers/v1/serve/server.py index fee9bc6ce7..d551f46aa1 100644 --- a/verifiers/v1/serve/server.py +++ b/verifiers/v1/serve/server.py @@ -1,5 +1,5 @@ import asyncio -from collections.abc import Callable +from typing import Protocol import logging import msgpack @@ -27,17 +27,25 @@ logger = logging.getLogger(__name__) -EpisodeHook = Callable[[Episode, SamplingConfig], None] +class EpisodeHook(Protocol): + """Takes finished episodes' per-token training arrays (routing rows, sampling masks, + top-k heads) before the episodes leave the worker.""" + + def keep(self, episode: Episode, sampling: SamplingConfig) -> None: + """Keep the episode's arrays; may leave references to them in `trace.info`. + Runs off the event loop, before the episode's final flush.""" + + def discard(self, episode: Episode) -> None: + """Free what `keep` kept for an episode that was not delivered.""" + _episode_hook: EpisodeHook | None = None def set_episode_hook(hook: EpisodeHook | None) -> None: - """Let the hosting process take each finished episode's per-token training arrays - (routing rows, sampling masks, top-k heads) before the episode leaves the worker. - The hook runs off the event loop with the episode and its sampling config; it may - keep the arrays and leave references in `trace.info`. While a hook is set, the - streamed nodes carry no training arrays.""" + """Install a hook that keeps finished episodes' training arrays. Every episode it + keeps is either delivered in its run's reply or handed back to `discard`. While a + hook is set, the streamed nodes carry no training arrays.""" global _episode_hook _episode_hook = hook @@ -69,6 +77,8 @@ def __init__( self._gate = asyncio.Semaphore(max_concurrent) if max_concurrent else None # In-flight run tasks by wire request_id, so a `cancel` can abort them self._running: dict[str, asyncio.Task] = {} + # Episodes the hook kept, by run request, until their reply is sent. + self._kept: dict[bytes, tuple[EpisodeHook, Episode]] = {} self.ctx = zmq.asyncio.Context() self.frontend = self.ctx.socket(zmq.ROUTER) self.frontend.setsockopt(zmq.ROUTER_MANDATORY, 1) @@ -134,7 +144,16 @@ async def send_delta(delta: dict) -> None: ) if hook is not None: # Before the final flush, so what the hook writes to the traces streams. - await asyncio.to_thread(hook, episode, ctx.sampling) + keeping = asyncio.ensure_future( + asyncio.to_thread(hook.keep, episode, ctx.sampling) + ) + try: + await asyncio.shield(keeping) + except asyncio.CancelledError: + # The thread finishes regardless; this episode is never delivered. + keeping.add_done_callback(lambda _: hook.discard(episode)) + raise + self._kept[request_id] = (hook, episode) return RunResponse( head=dump(episode, exclude={"traces"}), traces=[ @@ -196,12 +215,17 @@ async def _handle( ).model_dump(), use_bin_type=True, ) + delivered = False try: await self.frontend.send_multipart( [client_id, request_id, b"reply", data], copy=False ) + delivered = response.success except zmq.ZMQError as e: logger.warning("failed to send response: %s", e) + if (kept := self._kept.pop(request_id, None)) is not None and not delivered: + hook, episode = kept + hook.discard(episode) async def run(self) -> None: logger.info( From 71154c1e64e1affa90508f6c065c3b04ff5e6132 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Thu, 22 Oct 2026 14:09:51 +0000 Subject: [PATCH 03/10] style(serve): sort server imports Co-Authored-By: Claude Opus 5.5 --- verifiers/v1/serve/server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/verifiers/v1/serve/server.py b/verifiers/v1/serve/server.py index d551f46aa1..5e72f91198 100644 --- a/verifiers/v1/serve/server.py +++ b/verifiers/v1/serve/server.py @@ -1,6 +1,6 @@ import asyncio -from typing import Protocol import logging +from typing import Protocol import msgpack import zmq From 9aa95cafe99b3ce5f5bfed841bbba539e3cf316d Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Sun, 25 Oct 2026 14:10:15 +0000 Subject: [PATCH 04/10] fix(serve): discard kept arrays when the reply fails to encode A reply that fails to encode goes out as an error, so the client never sees the episode. Count it as undelivered so the hook frees what it kept. Co-Authored-By: Claude Opus 5.5 --- verifiers/v1/serve/server.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/verifiers/v1/serve/server.py b/verifiers/v1/serve/server.py index 1bf16420e3..4d56396119 100644 --- a/verifiers/v1/serve/server.py +++ b/verifiers/v1/serve/server.py @@ -209,12 +209,8 @@ async def _handle( except Exception as e: # Encoding failures must also reply so clients don't wait indefinitely. logger.warning("response encoding failed: %s", e, exc_info=True) - data = msgpack.packb( - BaseResponse( - success=False, error=f"{type(e).__name__}: {e}" - ).model_dump(), - use_bin_type=True, - ) + response = BaseResponse(success=False, error=f"{type(e).__name__}: {e}") + data = msgpack.packb(response.model_dump(), use_bin_type=True) delivered = False try: await self.frontend.send_multipart( From 5c5115fe18adac1d3d01785c0bfef3f4a0a0d68b Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Sun, 25 Oct 2026 16:58:37 +0000 Subject: [PATCH 05/10] feat(serve): never stream training arrays Nodes' routing rows, sampling masks and top-k heads now stay in the worker; an episode hook is the only way they leave it. Drop the keep_arrays switch and the routing-row repairs, which only patched arrays that were already streamed. Co-Authored-By: Claude Opus 5.5 --- tests/unit/serve/test_delta.py | 42 ++------------------ verifiers/v1/serve/delta.py | 72 ++++++---------------------------- verifiers/v1/serve/server.py | 10 ++--- 3 files changed, 18 insertions(+), 106 deletions(-) diff --git a/tests/unit/serve/test_delta.py b/tests/unit/serve/test_delta.py index 9d55fb7eec..fc5443a053 100644 --- a/tests/unit/serve/test_delta.py +++ b/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 ) @@ -272,30 +263,3 @@ def test_finish_refuses_a_gap(): assembly.apply({"trace": "t", "open": {"id": "t"}, "nodes": [{}]}) with pytest.raises(RuntimeError, match="assembled 1 nodes"): assembly.finish({}, [TraceSummary(id="t", nodes=2, calls=0)]) - - -@pytest.mark.asyncio -async def test_streamer_without_arrays_streams_nodes_without_them(): - """A host that keeps the training arrays itself (`set_episode_hook`) streams nodes - without them, and sends no routing repairs.""" - traces: list[vf.Trace] = [] - frames: list[bytes] = [] - - async def send(data: dict) -> None: - frames.append(pack(data)) - - async with DeltaStreamer(lambda: traces, send, keep_arrays=False) as streamer: - trace = make_trace() - traces.append(trace) - streamer.watch(trace) - add_turn(trace, "a1") - trace.nodes[1].routed_experts = np.arange(6, dtype=np.uint8).reshape(3, 2, 1) - await settle() - add_turn(trace, "a2") - trace.nodes[1].routed_experts = np.zeros((3, 2, 1), dtype=np.uint8) - await settle() - deltas = [unpack(frame) for frame in frames] - nodes = [node for delta in deltas for node in delta.get("nodes", [])] - assert len(nodes) == 3 - assert all("routed_experts" not in node for node in nodes) - assert not any("routing_repairs" in delta for delta in deltas) diff --git a/verifiers/v1/serve/delta.py b/verifiers/v1/serve/delta.py index a1ad9c74ac..b5d764ea00 100644 --- a/verifiers/v1/serve/delta.py +++ b/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: @@ -75,7 +74,7 @@ TRAINING_ARRAY_FIELDS = frozenset({"routed_experts", "sampling_mask", "top_logprobs"}) -"""Node fields holding per-token training arrays; streamed only while `keep_arrays`.""" +"""Node fields holding per-token training arrays; never streamed.""" def pack(payload: Any) -> bytes: @@ -100,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([]) @@ -118,12 +115,9 @@ def __init__( self, traces: Callable[[], Iterable[Trace]], send: Callable[[dict], Awaitable[None]], - keep_arrays: bool = True, ) -> None: self.traces = traces self.send = send - self.keep_arrays = keep_arrays - """Stream nodes' training arrays; off when the host keeps them (`set_episode_hook`).""" self.cursors: dict[str, TraceCursor] = {} self._scheduled: asyncio.Handle | None = None self._flushes: set[asyncio.Task] = set() @@ -176,33 +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. - """ - if not self.keep_arrays: - return - 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` @@ -229,20 +196,18 @@ 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: - exclude = ( - TRAINING_ARRAY_FIELDS - if field == "nodes" and not self.keep_arrays - else None - ) # 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, exclude=exclude) + 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) @@ -277,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"): @@ -305,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: diff --git a/verifiers/v1/serve/server.py b/verifiers/v1/serve/server.py index 4d56396119..2b0ffe9dca 100644 --- a/verifiers/v1/serve/server.py +++ b/verifiers/v1/serve/server.py @@ -43,9 +43,9 @@ def discard(self, episode: Episode) -> None: def set_episode_hook(hook: EpisodeHook | None) -> None: - """Install a hook that keeps finished episodes' training arrays. Every episode it - keeps is either delivered in its run's reply or handed back to `discard`. While a - hook is set, the streamed nodes carry no training arrays.""" + """Install a hook that keeps finished episodes' training arrays. Streamed nodes never + carry them, so the hook is the only way they leave the worker. Every episode it + keeps is either delivered in its run's reply or handed back to `discard`.""" global _episode_hook _episode_hook = hook @@ -136,9 +136,7 @@ async def send_delta(delta: dict) -> None: # in flight the same way the in-process eval's semaphore does. The streamer # ships each trace as it changes; the reply below carries only the rest. hook = _episode_hook - async with DeltaStreamer( - lambda: slot.traces, send_delta, keep_arrays=hook is None - ) as streamer: + async with DeltaStreamer(lambda: slot.traces, send_delta) as streamer: episode = await self.env.run_slot( slot, ctx, self._gate, on_trace=streamer.watch ) From e3a53f8037981e18e6f0413d667fd8513d881185 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Sun, 25 Oct 2026 17:38:14 +0000 Subject: [PATCH 06/10] fix(serve): discard on every path keep cannot deliver A keep that raises now hands the episode to discard, and a shutdown that cancels the reply send still discards. A run abandoned while keep runs discards only once keep has returned, from whichever of the two finishes last, so discard never races the arrays keep is storing. Co-Authored-By: Claude Opus 5.5 --- tests/unit/serve/test_episode_hook.py | 58 +++++++++++++++++++++++++++ verifiers/v1/serve/server.py | 52 +++++++++++++++++++----- 2 files changed, 101 insertions(+), 9 deletions(-) create mode 100644 tests/unit/serve/test_episode_hook.py diff --git a/tests/unit/serve/test_episode_hook.py b/tests/unit/serve/test_episode_hook.py new file mode 100644 index 0000000000..d304730b28 --- /dev/null +++ b/tests/unit/serve/test_episode_hook.py @@ -0,0 +1,58 @@ +"""An episode hook's `keep` hands an undelivered episode to `discard` exactly once, and +never while `keep` is still running.""" + +import threading + +import pytest + +from verifiers.v1.serve.server import _Keeping + + +class Hook: + def __init__(self, fail: bool = False) -> None: + self.fail = fail + self.release = threading.Event() + self.release.set() + self.keeping = False + self.calls: list[str] = [] + + def keep(self, episode, sampling) -> None: + self.keeping = True + self.release.wait() + self.calls.append("keep") + self.keeping = False + if self.fail: + raise ValueError("misaligned arrays") + + def discard(self, episode) -> None: + assert not self.keeping, "discard ran while keep was storing" + self.calls.append("discard") + + +def test_a_failed_keep_discards_once(): + hook = Hook(fail=True) + keeping = _Keeping(hook, episode=None) + with pytest.raises(ValueError): + keeping.run(sampling=None) + keeping.abandon() + assert hook.calls == ["keep", "discard"] + + +def test_abandoning_a_running_keep_discards_after_it(): + hook = Hook() + hook.release.clear() + keeping = _Keeping(hook, episode=None) + thread = threading.Thread(target=keeping.run, args=(None,)) + thread.start() + keeping.abandon() + hook.release.set() + thread.join() + assert hook.calls == ["keep", "discard"] + + +def test_abandoning_a_finished_keep_discards_once(): + hook = Hook() + keeping = _Keeping(hook, episode=None) + keeping.run(sampling=None) + keeping.abandon() + assert hook.calls == ["keep", "discard"] diff --git a/verifiers/v1/serve/server.py b/verifiers/v1/serve/server.py index 2b0ffe9dca..7a391696da 100644 --- a/verifiers/v1/serve/server.py +++ b/verifiers/v1/serve/server.py @@ -1,5 +1,6 @@ import asyncio import logging +import threading from typing import Protocol import msgpack @@ -50,6 +51,38 @@ def set_episode_hook(hook: EpisodeHook | None) -> None: _episode_hook = hook +class _Keeping: + """One episode's `keep`, run off the event loop. A failed or abandoned `keep` hands + the episode to `discard` exactly once, and only after `keep` has returned: in the + thread when `keep` finishes last, else in `abandon`.""" + + def __init__(self, hook: EpisodeHook, episode: Episode) -> None: + self.hook = hook + self.episode = episode + self._lock = threading.Lock() + self._kept = False + self._abandoned = False + + def run(self, sampling: SamplingConfig) -> None: + try: + self.hook.keep(self.episode, sampling) + except BaseException: + self.hook.discard(self.episode) + raise + with self._lock: + self._kept = True + discard = self._abandoned + if discard: + self.hook.discard(self.episode) + + def abandon(self) -> None: + with self._lock: + self._abandoned = True + discard = self._kept + if discard: + self.hook.discard(self.episode) + + class EnvServer: def __init__( self, @@ -142,14 +175,12 @@ async def send_delta(delta: dict) -> None: ) if hook is not None: # Before the final flush, so what the hook writes to the traces streams. - keeping = asyncio.ensure_future( - asyncio.to_thread(hook.keep, episode, ctx.sampling) - ) + keeping = _Keeping(hook, episode) try: - await asyncio.shield(keeping) + # The thread runs to the end even when this await is cancelled. + await asyncio.shield(asyncio.to_thread(keeping.run, ctx.sampling)) except asyncio.CancelledError: - # The thread finishes regardless; this episode is never delivered. - keeping.add_done_callback(lambda _: hook.discard(episode)) + keeping.abandon() raise self._kept[request_id] = (hook, episode) return RunResponse( @@ -217,9 +248,12 @@ async def _handle( delivered = response.success except zmq.ZMQError as e: logger.warning("failed to send response: %s", e) - if (kept := self._kept.pop(request_id, None)) is not None and not delivered: - hook, episode = kept - hook.discard(episode) + finally: + # Also on a cancel (shutdown) while the reply is sending. + kept = self._kept.pop(request_id, None) + if kept is not None and not delivered: + hook, episode = kept + hook.discard(episode) async def run(self) -> None: logger.info( From 435979700f477add2777fc3512005b4e9da622c9 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Sun, 25 Oct 2026 17:43:16 +0000 Subject: [PATCH 07/10] fix(serve): discard a delivered episode on a late cancel A client whose caller cancels a run after the reply left drops the reply, so the host's kept arrays had no consumer. The client already sends a cancel; the server now keeps delivered episodes for LATE_CANCEL_SECONDS and discards one when a cancel for it arrives. Co-Authored-By: Claude Opus 5.5 --- tests/unit/serve/test_episode_hook.py | 24 ++++++++++++++++++++- verifiers/v1/serve/server.py | 31 ++++++++++++++++++++++++--- 2 files changed, 51 insertions(+), 4 deletions(-) diff --git a/tests/unit/serve/test_episode_hook.py b/tests/unit/serve/test_episode_hook.py index d304730b28..cbf1ab167d 100644 --- a/tests/unit/serve/test_episode_hook.py +++ b/tests/unit/serve/test_episode_hook.py @@ -1,11 +1,14 @@ """An episode hook's `keep` hands an undelivered episode to `discard` exactly once, and never while `keep` is still running.""" +import asyncio import threading +import msgpack import pytest -from verifiers.v1.serve.server import _Keeping +from verifiers.v1.serve.server import EnvServer, _Keeping +from verifiers.v1.serve.types import CancelRequest class Hook: @@ -56,3 +59,22 @@ def test_abandoning_a_finished_keep_discards_once(): keeping.run(sampling=None) keeping.abandon() assert hook.calls == ["keep", "discard"] + + +def test_a_cancel_after_delivery_discards(): + """The client dropped a reply whose run its caller cancelled; the cancel that follows + frees the delivered episode's arrays.""" + + class Frontend: + async def send_multipart(self, frames, copy=True) -> None: + pass + + server = EnvServer.__new__(EnvServer) + server.frontend = Frontend() + server._running, server._kept, server._delivered = {}, {}, {} + hook = Hook() + server._remember_delivered("run-1", hook, episode=None) + cancel = msgpack.packb(CancelRequest(request_id="run-1").model_dump(mode="json")) + asyncio.run(server._handle(b"client", b"req-2", b"cancel", cancel)) + asyncio.run(server._handle(b"client", b"req-3", b"cancel", cancel)) + assert hook.calls == ["discard"] diff --git a/verifiers/v1/serve/server.py b/verifiers/v1/serve/server.py index 7a391696da..2998a6a30f 100644 --- a/verifiers/v1/serve/server.py +++ b/verifiers/v1/serve/server.py @@ -1,6 +1,7 @@ import asyncio import logging import threading +import time from typing import Protocol import msgpack @@ -42,11 +43,16 @@ def discard(self, episode: Episode) -> None: _episode_hook: EpisodeHook | None = None +LATE_CANCEL_SECONDS = 10.0 +"""How long a delivered episode stays discardable: a client whose caller cancels the run +after the reply left drops the reply and sends a cancel right away.""" + def set_episode_hook(hook: EpisodeHook | None) -> None: """Install a hook that keeps finished episodes' training arrays. Streamed nodes never carry them, so the hook is the only way they leave the worker. Every episode it - keeps is either delivered in its run's reply or handed back to `discard`.""" + keeps is either delivered in its run's reply or handed back to `discard`, and so is + a delivered one whose client cancels the run within `LATE_CANCEL_SECONDS`.""" global _episode_hook _episode_hook = hook @@ -112,6 +118,8 @@ def __init__( self._running: dict[str, asyncio.Task] = {} # Episodes the hook kept, by run request, until their reply is sent. self._kept: dict[bytes, tuple[EpisodeHook, Episode]] = {} + self._delivered: dict[str, tuple[float, EpisodeHook, Episode]] = {} + """Recently delivered kept episodes, oldest first, for a cancel that arrives late.""" self.ctx = zmq.asyncio.Context() self.frontend = self.ctx.socket(zmq.ROUTER) self.frontend.setsockopt(zmq.ROUTER_MANDATORY, 1) @@ -215,6 +223,10 @@ async def _handle( if running is not None: running.cancel() logger.info("cancelled run %s on client request", target) + if (late := self._delivered.pop(target, None)) is not None: + # The client dropped the reply, so nothing will ever take the arrays. + _, hook, episode = late + hook.discard(episode) response = CancelResponse(cancelled=running is not None) else: response = BaseResponse( @@ -251,9 +263,22 @@ async def _handle( finally: # Also on a cancel (shutdown) while the reply is sending. kept = self._kept.pop(request_id, None) - if kept is not None and not delivered: + if kept is not None: hook, episode = kept - hook.discard(episode) + if delivered: + self._remember_delivered(request_id.decode(), hook, episode) + else: + hook.discard(episode) + + def _remember_delivered( + self, request_id: str, hook: EpisodeHook, episode: Episode + ) -> None: + now = time.monotonic() + self._delivered[request_id] = (now, hook, episode) + for old, (delivered_at, _, _) in list(self._delivered.items()): + if now - delivered_at <= LATE_CANCEL_SECONDS: + break + del self._delivered[old] async def run(self) -> None: logger.info( From 098cbdd97cd8a4c9d2c2b75501f348a6d3a457b8 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 26 Oct 2026 17:06:01 +0000 Subject: [PATCH 08/10] fix(serve): close the remaining late-cancel gaps - A run cancelled while the client decodes its reply now sends a cancel too: the server already counted that episode delivered. - A cancel that lands after the run finished but before its reply was sent discards the kept episode. - Delivered episodes expire on a timer instead of on the next delivery, and a cancel never discards an expired one. Co-Authored-By: Claude Opus 5.5 --- tests/unit/serve/test_episode_hook.py | 101 ++++++++++++++++++++++---- verifiers/v1/serve/client.py | 34 ++++++--- verifiers/v1/serve/server.py | 23 ++++-- 3 files changed, 126 insertions(+), 32 deletions(-) diff --git a/tests/unit/serve/test_episode_hook.py b/tests/unit/serve/test_episode_hook.py index cbf1ab167d..c8a81b115e 100644 --- a/tests/unit/serve/test_episode_hook.py +++ b/tests/unit/serve/test_episode_hook.py @@ -7,8 +7,11 @@ import msgpack import pytest +import verifiers.v1 as vf +from verifiers.v1.serve import server as server_module +from verifiers.v1.serve.client import EnvClient from verifiers.v1.serve.server import EnvServer, _Keeping -from verifiers.v1.serve.types import CancelRequest +from verifiers.v1.serve.types import CancelRequest, RunResponse class Hook: @@ -61,20 +64,92 @@ def test_abandoning_a_finished_keep_discards_once(): assert hook.calls == ["keep", "discard"] -def test_a_cancel_after_delivery_discards(): - """The client dropped a reply whose run its caller cancelled; the cancel that follows - frees the delivered episode's arrays.""" +class Frontend: + async def send_multipart(self, frames, copy=True) -> None: + pass - class Frontend: - async def send_multipart(self, frames, copy=True) -> None: - pass +def _server() -> EnvServer: server = EnvServer.__new__(EnvServer) server.frontend = Frontend() server._running, server._kept, server._delivered = {}, {}, {} - hook = Hook() - server._remember_delivered("run-1", hook, episode=None) - cancel = msgpack.packb(CancelRequest(request_id="run-1").model_dump(mode="json")) - asyncio.run(server._handle(b"client", b"req-2", b"cancel", cancel)) - asyncio.run(server._handle(b"client", b"req-3", b"cancel", cancel)) - assert hook.calls == ["discard"] + return server + + +def _cancel(request_id: str) -> bytes: + return msgpack.packb(CancelRequest(request_id=request_id).model_dump(mode="json")) + + +def test_a_cancel_after_delivery_discards_once(): + """The client dropped a reply whose run its caller cancelled; the cancel that follows + frees the delivered episode's arrays.""" + + async def run() -> list[str]: + server, hook = _server(), Hook() + server._remember_delivered("run-1", hook, episode=None) + await server._handle(b"client", b"req-2", b"cancel", _cancel("run-1")) + await server._handle(b"client", b"req-3", b"cancel", _cancel("run-1")) + return hook.calls + + assert asyncio.run(run()) == ["discard"] + + +def test_a_cancel_before_the_reply_leaves_discards_once(): + """A cancel landing after the run finished but before its reply was sent.""" + + async def run() -> list[str]: + server, hook = _server(), Hook() + server._kept[b"run-1"] = (hook, None) + await server._handle(b"client", b"req-2", b"cancel", _cancel("run-1")) + assert not server._kept and not server._delivered + return hook.calls + + assert asyncio.run(run()) == ["discard"] + + +def test_a_stale_delivery_expires_without_discard(monkeypatch): + monkeypatch.setattr(server_module, "LATE_CANCEL_SECONDS", 0.01) + + async def run() -> list[str]: + server, hook = _server(), Hook() + server._remember_delivered("run-1", hook, episode=None) + await asyncio.sleep(0.05) + assert not server._delivered, "expired without waiting for another delivery" + await server._handle(b"client", b"req-2", b"cancel", _cancel("run-1")) + return hook.calls + + assert asyncio.run(run()) == [] + + +def test_a_run_cancelled_while_decoding_its_reply_cancels_on_the_server(): + """The reply arrived, so the server counts the episode delivered; a cancel during + the client's final decode must still reach the server.""" + + async def run() -> bool: + client = EnvClient("tcp://127.0.0.1:1") + cancelled: list[str] = [] + sent: dict[str, str] = {} + + async def request(*args, request_id: str, **kwargs) -> RunResponse: + sent["id"] = request_id + return RunResponse(head={}, traces=[]) + + async def decode(*args): + raise asyncio.CancelledError + + client._request = request + client._decode = decode + client._cancel_run = cancelled.append + try: + with pytest.raises(asyncio.CancelledError): + await client.run( + client=vf.TrainClientConfig(), + model="m", + sampling=vf.SamplingConfig(), + task_data={}, + ) + finally: + await client.close() + return cancelled == [sent["id"]] + + assert asyncio.run(run()) diff --git a/verifiers/v1/serve/client.py b/verifiers/v1/serve/client.py index aaa3d8ffcc..37f37090fc 100644 --- a/verifiers/v1/serve/client.py +++ b/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,18 @@ 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 asyncio.CancelledError: + # The reply already arrived, so the server counts the episode delivered. + self._cancel_run(request_id) + raise async def close(self) -> None: if self._receiver is not None: diff --git a/verifiers/v1/serve/server.py b/verifiers/v1/serve/server.py index 2998a6a30f..bb13e52611 100644 --- a/verifiers/v1/serve/server.py +++ b/verifiers/v1/serve/server.py @@ -223,9 +223,13 @@ async def _handle( if running is not None: running.cancel() logger.info("cancelled run %s on client request", target) - if (late := self._delivered.pop(target, None)) is not None: - # The client dropped the reply, so nothing will ever take the arrays. - _, hook, episode = late + # The client dropped the reply, or will: nothing will take the arrays. + self._expire_delivered() + kept = self._kept.pop(target.encode(), None) + if kept is None and (late := self._delivered.pop(target, None)): + kept = late[1:] + if kept is not None: + hook, episode = kept hook.discard(episode) response = CancelResponse(cancelled=running is not None) else: @@ -273,12 +277,17 @@ async def _handle( def _remember_delivered( self, request_id: str, hook: EpisodeHook, episode: Episode ) -> None: + self._delivered[request_id] = (time.monotonic(), hook, episode) + asyncio.get_running_loop().call_later( + LATE_CANCEL_SECONDS, self._expire_delivered + ) + + def _expire_delivered(self) -> None: now = time.monotonic() - self._delivered[request_id] = (now, hook, episode) - for old, (delivered_at, _, _) in list(self._delivered.items()): - if now - delivered_at <= LATE_CANCEL_SECONDS: + for request_id, (delivered_at, _, _) in list(self._delivered.items()): + if now - delivered_at < LATE_CANCEL_SECONDS: break - del self._delivered[old] + del self._delivered[request_id] async def run(self) -> None: logger.info( From 3dc8572994f7a335d7adad9ea42784034fd1a780 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Mon, 26 Oct 2026 17:24:29 +0000 Subject: [PATCH 09/10] fix(serve): route late cancels through the worker pool The pool broker forgot a run once its reply passed through and answered any later cancel itself, so a worker never discarded the arrays of a run its client abandoned after the reply. The broker now remembers which worker replied for LATE_CANCEL_SECONDS and forwards a late cancel there. The client also cancels when assembling or validating a delivered reply fails, not only when it is cancelled. Co-Authored-By: Claude Opus 5.5 --- tests/unit/serve/test_episode_hook.py | 11 +++-- tests/unit/serve/test_pool_cancel.py | 71 +++++++++++++++++++++++++++ verifiers/v1/serve/client.py | 5 +- verifiers/v1/serve/pool.py | 31 ++++++++++-- 4 files changed, 106 insertions(+), 12 deletions(-) create mode 100644 tests/unit/serve/test_pool_cancel.py diff --git a/tests/unit/serve/test_episode_hook.py b/tests/unit/serve/test_episode_hook.py index c8a81b115e..51004820a3 100644 --- a/tests/unit/serve/test_episode_hook.py +++ b/tests/unit/serve/test_episode_hook.py @@ -121,9 +121,10 @@ async def run() -> list[str]: assert asyncio.run(run()) == [] -def test_a_run_cancelled_while_decoding_its_reply_cancels_on_the_server(): - """The reply arrived, so the server counts the episode delivered; a cancel during - the client's final decode must still reach the server.""" +@pytest.mark.parametrize("error", [asyncio.CancelledError, ValueError]) +def test_a_run_lost_after_its_reply_cancels_on_the_server(error): + """The reply arrived, so the server counts the episode delivered; a run cancelled or + failing while the client assembles and decodes it must still cancel on the server.""" async def run() -> bool: client = EnvClient("tcp://127.0.0.1:1") @@ -135,13 +136,13 @@ async def request(*args, request_id: str, **kwargs) -> RunResponse: return RunResponse(head={}, traces=[]) async def decode(*args): - raise asyncio.CancelledError + raise error client._request = request client._decode = decode client._cancel_run = cancelled.append try: - with pytest.raises(asyncio.CancelledError): + with pytest.raises(error): await client.run( client=vf.TrainClientConfig(), model="m", diff --git a/tests/unit/serve/test_pool_cancel.py b/tests/unit/serve/test_pool_cancel.py new file mode 100644 index 0000000000..340340c16f --- /dev/null +++ b/tests/unit/serve/test_pool_cancel.py @@ -0,0 +1,71 @@ +"""The pool broker routes a cancel to the worker that ran the target run, even after +that run replied: the worker may still hold the run's kept training arrays.""" + +import asyncio +from types import SimpleNamespace + +import msgpack +import zmq + +from verifiers.v1.serve.pool import _CANCEL_MISS, EnvServerPool + + +class FakeWorkerPool(EnvServerPool): + """A pool whose one worker is an in-process ROUTER socket the test drives.""" + + def _spawn_worker(self) -> None: + dealer = self.ctx.socket(zmq.DEALER) + dealer.setsockopt(zmq.LINGER, 0) + dealer.connect("inproc://worker") + self.workers.append( + { + "process": SimpleNamespace(), + "dealer": dealer, + "pipe": SimpleNamespace(), + "active": 0, + "index": 0, + } + ) + if self._poller is not None: + self._poller.register(dealer, zmq.POLLIN) + + +def test_a_late_cancel_reaches_the_worker_that_replied(): + async def run() -> None: + pool = FakeWorkerPool( + {}, max_workers=1, address="tcp://127.0.0.1:0", elastic=False + ) + worker = pool.ctx.socket(zmq.ROUTER) + worker.setsockopt(zmq.LINGER, 0) + worker.bind("inproc://worker") + client = pool.ctx.socket(zmq.DEALER) + client.setsockopt(zmq.LINGER, 0) + client.connect(pool.address) + broker = asyncio.create_task(pool.run()) + + async def recv(socket) -> list[bytes]: + return await asyncio.wait_for(socket.recv_multipart(), 5) + + try: + await client.send_multipart([b"run-1", b"run", b""]) + identity, request_id, method, _ = await recv(worker) + assert (request_id, method) == (b"run-1", b"run") + await worker.send_multipart([identity, b"run-1", b"reply", b"done"]) + assert await recv(client) == [b"run-1", b"reply", b"done"] + + cancel = msgpack.packb({"request_id": "run-1"}) + await client.send_multipart([b"cancel-1", b"cancel", cancel]) + identity, request_id, method, payload = await recv(worker) + assert (request_id, method, payload) == (b"cancel-1", b"cancel", cancel) + + await client.send_multipart( + [b"cancel-2", b"cancel", msgpack.packb({"request_id": "never-ran"})] + ) + assert await recv(client) == [b"cancel-2", b"reply", _CANCEL_MISS] + finally: + worker.close() + client.close() + broker.cancel() + await asyncio.gather(broker, return_exceptions=True) + + asyncio.run(run()) diff --git a/verifiers/v1/serve/client.py b/verifiers/v1/serve/client.py index 37f37090fc..ce41717b1e 100644 --- a/verifiers/v1/serve/client.py +++ b/verifiers/v1/serve/client.py @@ -296,8 +296,9 @@ def apply(delta: dict) -> None: WireEpisode.model_validate, assembly.finish(response.head, response.traces), ) - except asyncio.CancelledError: - # The reply already arrived, so the server counts the episode delivered. + 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 diff --git a/verifiers/v1/serve/pool.py b/verifiers/v1/serve/pool.py index 625e6a827f..8d2efb4aaa 100644 --- a/verifiers/v1/serve/pool.py +++ b/verifiers/v1/serve/pool.py @@ -28,6 +28,7 @@ import signal import tempfile import threading +import time from collections.abc import Callable import msgpack @@ -35,7 +36,7 @@ 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: From 1cee4be7146032b29d97805d4a38dc3cd4298349 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 27 Oct 2026 17:27:00 +0000 Subject: [PATCH 10/10] test: drop the pool-cancel and episode-hook test files Co-Authored-By: Claude Opus 5.5 --- tests/unit/serve/test_episode_hook.py | 156 -------------------------- tests/unit/serve/test_pool_cancel.py | 71 ------------ 2 files changed, 227 deletions(-) delete mode 100644 tests/unit/serve/test_episode_hook.py delete mode 100644 tests/unit/serve/test_pool_cancel.py diff --git a/tests/unit/serve/test_episode_hook.py b/tests/unit/serve/test_episode_hook.py deleted file mode 100644 index 51004820a3..0000000000 --- a/tests/unit/serve/test_episode_hook.py +++ /dev/null @@ -1,156 +0,0 @@ -"""An episode hook's `keep` hands an undelivered episode to `discard` exactly once, and -never while `keep` is still running.""" - -import asyncio -import threading - -import msgpack -import pytest - -import verifiers.v1 as vf -from verifiers.v1.serve import server as server_module -from verifiers.v1.serve.client import EnvClient -from verifiers.v1.serve.server import EnvServer, _Keeping -from verifiers.v1.serve.types import CancelRequest, RunResponse - - -class Hook: - def __init__(self, fail: bool = False) -> None: - self.fail = fail - self.release = threading.Event() - self.release.set() - self.keeping = False - self.calls: list[str] = [] - - def keep(self, episode, sampling) -> None: - self.keeping = True - self.release.wait() - self.calls.append("keep") - self.keeping = False - if self.fail: - raise ValueError("misaligned arrays") - - def discard(self, episode) -> None: - assert not self.keeping, "discard ran while keep was storing" - self.calls.append("discard") - - -def test_a_failed_keep_discards_once(): - hook = Hook(fail=True) - keeping = _Keeping(hook, episode=None) - with pytest.raises(ValueError): - keeping.run(sampling=None) - keeping.abandon() - assert hook.calls == ["keep", "discard"] - - -def test_abandoning_a_running_keep_discards_after_it(): - hook = Hook() - hook.release.clear() - keeping = _Keeping(hook, episode=None) - thread = threading.Thread(target=keeping.run, args=(None,)) - thread.start() - keeping.abandon() - hook.release.set() - thread.join() - assert hook.calls == ["keep", "discard"] - - -def test_abandoning_a_finished_keep_discards_once(): - hook = Hook() - keeping = _Keeping(hook, episode=None) - keeping.run(sampling=None) - keeping.abandon() - assert hook.calls == ["keep", "discard"] - - -class Frontend: - async def send_multipart(self, frames, copy=True) -> None: - pass - - -def _server() -> EnvServer: - server = EnvServer.__new__(EnvServer) - server.frontend = Frontend() - server._running, server._kept, server._delivered = {}, {}, {} - return server - - -def _cancel(request_id: str) -> bytes: - return msgpack.packb(CancelRequest(request_id=request_id).model_dump(mode="json")) - - -def test_a_cancel_after_delivery_discards_once(): - """The client dropped a reply whose run its caller cancelled; the cancel that follows - frees the delivered episode's arrays.""" - - async def run() -> list[str]: - server, hook = _server(), Hook() - server._remember_delivered("run-1", hook, episode=None) - await server._handle(b"client", b"req-2", b"cancel", _cancel("run-1")) - await server._handle(b"client", b"req-3", b"cancel", _cancel("run-1")) - return hook.calls - - assert asyncio.run(run()) == ["discard"] - - -def test_a_cancel_before_the_reply_leaves_discards_once(): - """A cancel landing after the run finished but before its reply was sent.""" - - async def run() -> list[str]: - server, hook = _server(), Hook() - server._kept[b"run-1"] = (hook, None) - await server._handle(b"client", b"req-2", b"cancel", _cancel("run-1")) - assert not server._kept and not server._delivered - return hook.calls - - assert asyncio.run(run()) == ["discard"] - - -def test_a_stale_delivery_expires_without_discard(monkeypatch): - monkeypatch.setattr(server_module, "LATE_CANCEL_SECONDS", 0.01) - - async def run() -> list[str]: - server, hook = _server(), Hook() - server._remember_delivered("run-1", hook, episode=None) - await asyncio.sleep(0.05) - assert not server._delivered, "expired without waiting for another delivery" - await server._handle(b"client", b"req-2", b"cancel", _cancel("run-1")) - return hook.calls - - assert asyncio.run(run()) == [] - - -@pytest.mark.parametrize("error", [asyncio.CancelledError, ValueError]) -def test_a_run_lost_after_its_reply_cancels_on_the_server(error): - """The reply arrived, so the server counts the episode delivered; a run cancelled or - failing while the client assembles and decodes it must still cancel on the server.""" - - async def run() -> bool: - client = EnvClient("tcp://127.0.0.1:1") - cancelled: list[str] = [] - sent: dict[str, str] = {} - - async def request(*args, request_id: str, **kwargs) -> RunResponse: - sent["id"] = request_id - return RunResponse(head={}, traces=[]) - - async def decode(*args): - raise error - - client._request = request - client._decode = decode - client._cancel_run = cancelled.append - try: - with pytest.raises(error): - await client.run( - client=vf.TrainClientConfig(), - model="m", - sampling=vf.SamplingConfig(), - task_data={}, - ) - finally: - await client.close() - return cancelled == [sent["id"]] - - assert asyncio.run(run()) diff --git a/tests/unit/serve/test_pool_cancel.py b/tests/unit/serve/test_pool_cancel.py deleted file mode 100644 index 340340c16f..0000000000 --- a/tests/unit/serve/test_pool_cancel.py +++ /dev/null @@ -1,71 +0,0 @@ -"""The pool broker routes a cancel to the worker that ran the target run, even after -that run replied: the worker may still hold the run's kept training arrays.""" - -import asyncio -from types import SimpleNamespace - -import msgpack -import zmq - -from verifiers.v1.serve.pool import _CANCEL_MISS, EnvServerPool - - -class FakeWorkerPool(EnvServerPool): - """A pool whose one worker is an in-process ROUTER socket the test drives.""" - - def _spawn_worker(self) -> None: - dealer = self.ctx.socket(zmq.DEALER) - dealer.setsockopt(zmq.LINGER, 0) - dealer.connect("inproc://worker") - self.workers.append( - { - "process": SimpleNamespace(), - "dealer": dealer, - "pipe": SimpleNamespace(), - "active": 0, - "index": 0, - } - ) - if self._poller is not None: - self._poller.register(dealer, zmq.POLLIN) - - -def test_a_late_cancel_reaches_the_worker_that_replied(): - async def run() -> None: - pool = FakeWorkerPool( - {}, max_workers=1, address="tcp://127.0.0.1:0", elastic=False - ) - worker = pool.ctx.socket(zmq.ROUTER) - worker.setsockopt(zmq.LINGER, 0) - worker.bind("inproc://worker") - client = pool.ctx.socket(zmq.DEALER) - client.setsockopt(zmq.LINGER, 0) - client.connect(pool.address) - broker = asyncio.create_task(pool.run()) - - async def recv(socket) -> list[bytes]: - return await asyncio.wait_for(socket.recv_multipart(), 5) - - try: - await client.send_multipart([b"run-1", b"run", b""]) - identity, request_id, method, _ = await recv(worker) - assert (request_id, method) == (b"run-1", b"run") - await worker.send_multipart([identity, b"run-1", b"reply", b"done"]) - assert await recv(client) == [b"run-1", b"reply", b"done"] - - cancel = msgpack.packb({"request_id": "run-1"}) - await client.send_multipart([b"cancel-1", b"cancel", cancel]) - identity, request_id, method, payload = await recv(worker) - assert (request_id, method, payload) == (b"cancel-1", b"cancel", cancel) - - await client.send_multipart( - [b"cancel-2", b"cancel", msgpack.packb({"request_id": "never-ran"})] - ) - assert await recv(client) == [b"cancel-2", b"reply", _CANCEL_MISS] - finally: - worker.close() - client.close() - broker.cancel() - await asyncio.gather(broker, return_exceptions=True) - - asyncio.run(run())