diff --git a/tests/unit/serve/test_delta.py b/tests/unit/serve/test_delta.py index cb2b9ef5a2..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 ) 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/client.py b/verifiers/v1/serve/client.py index aaa3d8ffcc..ce41717b1e 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,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: diff --git a/verifiers/v1/serve/delta.py b/verifiers/v1/serve/delta.py index 24b67c543c..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: @@ -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: 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: diff --git a/verifiers/v1/serve/server.py b/verifiers/v1/serve/server.py index e13e2ff221..bb13e52611 100644 --- a/verifiers/v1/serve/server.py +++ b/verifiers/v1/serve/server.py @@ -1,5 +1,8 @@ import asyncio import logging +import threading +import time +from typing import Protocol import msgpack import zmq @@ -9,6 +12,7 @@ from verifiers.v1.configs.client import ClientConfig from verifiers.v1.configs.env import EnvConfig from verifiers.v1.configs.sampling import SamplingConfig +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 +29,66 @@ logger = logging.getLogger(__name__) +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 + +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`, and so is + a delivered one whose client cancels the run within `LATE_CANCEL_SECONDS`.""" + global _episode_hook + _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, @@ -52,6 +116,10 @@ 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._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) @@ -108,10 +176,21 @@ 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. + hook = _episode_hook async with DeltaStreamer(lambda: slot.traces, send_delta) 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. + keeping = _Keeping(hook, episode) + try: + # 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: + keeping.abandon() + raise + self._kept[request_id] = (hook, episode) return RunResponse( head=dump(episode, exclude={"traces"}), traces=[ @@ -144,6 +223,14 @@ async def _handle( if running is not None: running.cancel() logger.info("cancelled run %s on client request", target) + # 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: response = BaseResponse( @@ -167,18 +254,40 @@ 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( [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) + finally: + # Also on a cancel (shutdown) while the reply is sending. + kept = self._kept.pop(request_id, None) + if kept is not None: + hook, episode = kept + 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: + 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() + for request_id, (delivered_at, _, _) in list(self._delivered.items()): + if now - delivered_at < LATE_CANCEL_SECONDS: + break + del self._delivered[request_id] async def run(self) -> None: logger.info(