verifiers 0.3.2.dev99__py3-none-any.whl → 0.3.2.dev101__py3-none-any.whl

verifiers/v1/interception/server.py

@@ -483,6 +483,7 @@ class InterceptionServer(Interception):
                 acp=acp,
             )
         )
+        session.trace.notify()
 
     async def handle_request(
         self, request: web.Request, dialect: Dialect
@@ -660,6 +661,9 @@ class InterceptionServer(Interception):
             return web.json_response(dialect.error_body(str(error)), status=400)
         except RolloutError as error:
             return self._fail(session, dialect, error)
+        # The tail is what the harness added since the last turn (tool results, user
+        # turns): live watchers see it now rather than with the model's reply.
+        session.trace.preview(turn, turn.tail)
 
         inspect_response = bool(session.response_interceptors or session.response_stops)
         if relay_streaming:
@@ -782,6 +786,8 @@ class InterceptionServer(Interception):
                     error = e
                     raise
             finally:
+                if node is None:
+                    turn.abandon()
                 # The turn's one per-exchange record: settings, timing, outcome, and
                 # the error that ended it (if any).
                 self.record_call(
@@ -1047,6 +1053,8 @@ class InterceptionServer(Interception):
                 error = e
             raise
         finally:
+            if node is None:
+                turn.abandon()
             # The turn's one per-exchange record: settings, timing, outcome, and the
             # error that ended it (if any).
             self.record_call(

verifiers/v1/rollout.py

@@ -180,6 +180,7 @@ class Rollout:
         proceed; a setup failure is captured onto the trace."""
         self._opened = True
         self.trace.timing.boot.start = time.time()
+        self.trace.notify()
         if self._borrowed_runtime is None:
             self.runtime = make_runtime(self.runtime_config, name=self.trace.id)
         elif self._borrowed_runtime is not None and self._borrowed_runtime.stopped:
@@ -219,6 +220,7 @@ class Rollout:
             now = time.time()
             self.trace.timing.boot.end = now
             self.trace.timing.setup.start = now
+            self.trace.notify()
             # Task setup and harness provisioning share one setup-stage deadline.
             setup_deadline = (
                 None
@@ -342,6 +344,7 @@ class Rollout:
         now = time.time()
         self.trace.timing.setup.end = now
         self.trace.timing.agent.start = now
+        self.trace.notify()
         return not self._session.stopped
 
     async def step(self, messages: Messages | None = None) -> bool:
@@ -464,6 +467,7 @@ class Rollout:
             finally:
                 if trace.timing.agent.start and not trace.timing.agent.end:
                     trace.timing.agent.end = time.time()
+                trace.notify()
             if not self._failed and self._opened:
                 assert runtime is not None
                 trace.timing.finalize.start = time.time()
@@ -479,6 +483,7 @@ class Rollout:
                 now = time.time()
                 trace.timing.finalize.end = now
                 trace.timing.scoring.start = now
+                trace.notify()
                 async with boundary(TaskError, "scoring"):
                     # Cross-trace judgement runs later, after the runtime is gone.
                     await asyncio.wait_for(
@@ -489,6 +494,7 @@ class Rollout:
                         self._timeouts.scoring,
                     )
                 trace.timing.scoring.end = time.time()
+                trace.notify()
         except Exception as e:  # noqa: BLE001 - finalize boundary records every rollout failure
             self.fail(e)
         finally:

verifiers/v1/serve/__init__.py

@@ -1,4 +1,5 @@
 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.types import (
@@ -12,6 +13,7 @@ __all__ = [
     "EnvClient",
     "EnvServer",
     "EnvServerPool",
+    "EpisodeAssembly",
     "HealthRequest",
     "HealthResponse",
     "RunRequest",

verifiers/v1/serve/client.py

@@ -1,11 +1,13 @@
 """ZMQ client for the env server.
 
-A DEALER socket + msgpack, with a single receive loop matching responses to
-per-request futures by `request_id`. Speaks the typed pydantic request/response
-models (`serve/types.py`) end-to-end: a request is `model_dump`ed onto the wire
-and the reply is `model_validate`d back — `Trace`s come back typed as
-`Trace[WireTaskData]` (non-strict task, so env fields survive without importing the
-env). Health is just another request (no dedicated probe thread).
+A DEALER socket + msgpack, with a single receive loop matching replies to
+per-request futures by `request_id` and routing a run's `delta` frames to its
+assembly (`serve/delta.py`). Speaks the typed pydantic request/response models
+(`serve/types.py`) end-to-end: a request is `model_dump`ed onto the wire; a run's
+traces stream back turn by turn and the assembled record is `model_validate`d into a
+`WireEpisode` once the reply lands — `Trace`s typed as `Trace[WireTaskData]`
+(non-strict task, so env fields survive without importing the env). Health is just
+another request (no dedicated probe thread).
 """
 
 import asyncio
@@ -13,6 +15,7 @@ import contextlib
 import logging
 import time
 import uuid
+from collections.abc import Callable
 from typing import TypeVar
 
 import msgpack
@@ -21,6 +24,7 @@ import zmq.asyncio
 
 from verifiers.v1.configs.client import ClientConfig
 from verifiers.v1.episode import WireEpisode
+from verifiers.v1.serve.delta import EpisodeAssembly, unpack
 from verifiers.v1.serve.types import (
     BaseRequest,
     BaseResponse,
@@ -47,6 +51,7 @@ class EnvClient:
         self.socket.setsockopt(zmq.RCVHWM, 0)
         self.socket.connect(address)
         self._pending: dict[str, asyncio.Future[bytes]] = {}
+        self._deltas: dict[str, Callable[[bytes], None]] = {}
         # Strong refs to in-flight fire-and-forget cancels: the loop only
         # holds weak references to tasks, so an unreferenced one can be
         # garbage-collected before it ever sends
@@ -59,27 +64,48 @@ class EnvClient:
             self._receiver = asyncio.create_task(self._receive_loop())
 
     async def _receive_loop(self) -> None:
+        # The one receiver serves every in-flight request, so a bad frame or a delta
+        # handler that raises is logged and skipped rather than allowed to end the loop.
         while True:
             try:
-                request_id_bytes, data = await self.socket.recv_multipart()
+                frames = await self.socket.recv_multipart()
             except asyncio.CancelledError:
                 break
-            future = self._pending.pop(request_id_bytes.decode(), None)
-            if future is not None and not future.done():
-                future.set_result(data)
+            try:
+                if len(frames) != 3:
+                    raise ValueError(
+                        f"expected [request_id, kind, data], got {len(frames)} frames - "
+                        "is the env server speaking the same serve protocol?"
+                    )
+                request_id_bytes, kind, data = frames
+                request_id = request_id_bytes.decode()
+                if kind == b"delta":
+                    on_delta = self._deltas.get(request_id)
+                    if on_delta is not None:
+                        on_delta(data)
+                    continue
+                future = self._pending.pop(request_id, None)
+                if future is not None and not future.done():
+                    future.set_result(data)
+            except Exception:  # keep receiving for the other requests
+                logger.warning("dropping a malformed env-server frame", exc_info=True)
 
     async def _request(
         self,
         request: BaseRequest,
         response_type: type[ResponseT],
         timeout: float | None = None,
+        on_delta: Callable[[bytes], None] | 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."""
+        `timeout` is only used for health polling — rollouts run untimed. `on_delta`
+        receives each `delta` frame the request streams before its reply."""
         self._ensure_receiver()
         request_id = uuid.uuid4().hex
         future: asyncio.Future[bytes] = asyncio.get_running_loop().create_future()
         self._pending[request_id] = future
+        if on_delta is not None:
+            self._deltas[request_id] = on_delta
         payload = msgpack.packb(request.model_dump(mode="json"), use_bin_type=True)
         await self.socket.send_multipart(
             [request_id.encode(), request.method.encode(), payload]
@@ -101,31 +127,30 @@ class EnvClient:
                 self._cancel_tasks.add(task)
                 task.add_done_callback(self._cancel_tasks.discard)
             raise
-        if response_type is HealthResponse:
-            response = response_type.model_validate(msgpack.unpackb(data, raw=False))
-        else:
-            # Keep large trace replies compact on the loop and expand only one at a time.
-            await self._decode_slots.acquire()
-            decoding = asyncio.create_task(
-                asyncio.to_thread(
-                    lambda: response_type.model_validate(
-                        msgpack.unpackb(data, raw=False)
-                    )
-                )
-            )
-            # Hold the slot until the worker finishes so cancellation cannot overlap decodes.
-            decoding.add_done_callback(lambda _: self._decode_slots.release())
-            try:
-                response = await asyncio.shield(decoding)
-            except asyncio.CancelledError:
-                decoding.add_done_callback(
-                    lambda task: None if task.cancelled() else task.exception()
-                )
-                raise
+        finally:
+            self._deltas.pop(request_id, None)
+        response = response_type.model_validate(unpack(data))
         if not response.success:
             raise RuntimeError(response.error or "env server request failed")
         return response
 
+    async def _validate_episode(self, record: dict) -> WireEpisode:
+        """Type the assembled record off the loop, one episode at a time: a long
+        trace is a lot of token spans to validate."""
+        await self._decode_slots.acquire()
+        decoding = asyncio.create_task(
+            asyncio.to_thread(WireEpisode.model_validate, record)
+        )
+        # Hold the slot until the worker finishes so cancellation cannot overlap decodes.
+        decoding.add_done_callback(lambda _: self._decode_slots.release())
+        try:
+            return await asyncio.shield(decoding)
+        except asyncio.CancelledError:
+            decoding.add_done_callback(
+                lambda task: None if task.cancelled() else task.exception()
+            )
+            raise
+
     async def _send_cancel(self, run_request_id: str) -> None:
         """Best-effort server-side abort of an abandoned run."""
         with contextlib.suppress(Exception):
@@ -164,10 +189,22 @@ class EnvClient:
         model: str,
         sampling: SamplingConfig,
         task_data: dict,
+        on_delta: Callable[[dict], None] | None = None,
     ) -> WireEpisode:
         """Run one rollout; return its episode record — flat traces (typed
         `Trace[WireTaskData]`) plus the shared stamp. The server takes the task
-        itself (`task_data`, its dumped `TaskData`)."""
+        itself (`task_data`, its dumped `TaskData`). The traces stream in as the
+        rollout runs, one delta per turn or phase change (`serve.delta`); `on_delta`
+        sees each as it lands, so a caller can relay or persist the stream. The
+        episode returned is assembled from the same deltas."""
+        assembly = EpisodeAssembly()
+
+        def apply(data: bytes) -> None:
+            delta = unpack(data)
+            assembly.apply(delta)
+            if on_delta is not None:
+                on_delta(delta)
+
         response = await self._request(
             RunRequest(
                 task_data=task_data,
@@ -176,8 +213,12 @@ class EnvClient:
                 sampling=sampling,
             ),
             RunResponse,
+            on_delta=apply,
+        )
+        assert response.head is not None
+        return await self._validate_episode(
+            assembly.finish(response.head, response.traces)
         )
-        return response.episode
 
     async def close(self) -> None:
         if self._receiver is not None:

verifiers/v1/serve/delta.py

@@ -0,0 +1,281 @@
+"""Turn-by-turn episode deltas over the env-serve wire.
+
+The worker streams a served episode as it grows. Each trace announces its own changes
+(`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 landing 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 (a turn's nodes are committed complete, with their
+tokens), so apart from the preview, which a committed turn repeats, 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
+send that fails is simply diffed again at the next flush.
+
+The client applies deltas into raw dicts (`EpisodeAssembly`) — cheap enough to hand a
+caller after every turn — and validates the assembled record into a `WireEpisode` once,
+when the run's reply lands.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import copy
+import logging
+from collections.abc import Awaitable, Callable
+from typing import TYPE_CHECKING, Any, Self
+
+import msgpack
+from pydantic import BaseModel
+
+from verifiers.v1.serve.encoding import msgpack_encoder
+
+if TYPE_CHECKING:
+    from verifiers.v1.env import RunSlot
+    from verifiers.v1.trace import Trace
+
+logger = logging.getLogger(__name__)
+
+HEADER_FIELDS = ("version", "id", "verifiers", "task")
+"""Sent once, when a trace first appears."""
+
+LIST_FIELDS = (
+    "nodes",
+    "calls",
+    "errors",
+    "extra_usage",
+    "request_rewrites",
+    "response_rewrites",
+)
+"""Append-only on the worker: each delta carries the items past the sent count."""
+
+SCALAR_FIELDS = (
+    "agent",
+    "tools",
+    "mm_token_type_id_map",
+    "rewards",
+    "metrics",
+    "info",
+    "root_reply",
+    "is_completed",
+    "ok",
+    "stop_condition",
+    "timing",
+)
+"""Re-sent whole whenever their dumped value changes."""
+
+
+def pack(payload: Any) -> bytes:
+    return msgpack.packb(payload, default=msgpack_encoder, use_bin_type=True)
+
+
+def unpack(data: bytes) -> Any:
+    # `links` is keyed by node index (int).
+    return msgpack.unpackb(data, raw=False, strict_map_key=False)
+
+
+def dump(model: BaseModel, **kwargs: Any) -> dict:
+    # Task data and agent configs are subclasses of their declared types; dump the
+    # runtime type so env-specific fields survive the wire.
+    return model.model_dump(mode="python", serialize_as_any=True, **kwargs)
+
+
+class TraceCursor:
+    """What of one trace has been sent."""
+
+    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.scalars: dict[str, bytes] = {}
+        self.pending: bytes = pack([])
+
+
+class DeltaStreamer:
+    """Sends a `RunSlot`'s trace changes as deltas, one flush per burst of changes.
+
+    `watch` subscribes a minted trace (pass it as `run_slot`'s `on_trace`); each
+    `Trace.notify` schedules a flush, and changes landing in the same loop iteration
+    ride one flush. Use as an async context manager around the rollout: a clean exit
+    flushes the final state (the slot then holds the finished episode's traces), a
+    cancelled rollout flushes nothing — its client has already gone."""
+
+    def __init__(self, slot: RunSlot, send: Callable[[bytes], Awaitable[None]]) -> None:
+        self.slot = slot
+        self.send = send
+        self.cursors: dict[str, TraceCursor] = {}
+        self._scheduled: asyncio.Handle | None = None
+        self._flushes: set[asyncio.Task] = set()
+        self._lock = asyncio.Lock()
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        if self._scheduled is not None:
+            self._scheduled.cancel()
+            self._scheduled = None
+        if exc_type is None:
+            # let in-flight flushes finish, then send whatever the run left behind
+            await asyncio.gather(*self._flushes, return_exceptions=True)
+            await self.flush()
+        else:
+            for task in self._flushes:
+                task.cancel()
+
+    def watch(self, trace: Trace) -> None:
+        trace.watch(self._changed)
+        self._changed(trace)
+
+    def _changed(self, trace: Trace) -> None:
+        if self._scheduled is not None:
+            return
+        self._scheduled = asyncio.get_running_loop().call_soon(self._start_flush)
+
+    def _start_flush(self) -> None:
+        self._scheduled = None
+        task = asyncio.create_task(self.flush())
+        self._flushes.add(task)
+        task.add_done_callback(self._flushes.discard)
+
+    async def flush(self) -> None:
+        # Serialized: deltas must leave in diff order, and the reply after the last.
+        async with self._lock:
+            for trace_id, delta, cursor in self.diff():
+                try:
+                    await self.send(pack(delta))
+                except Exception:  # the cursor stays put: the next flush diffs it again
+                    logger.warning(
+                        "failed to send delta for %s", trace_id, exc_info=True
+                    )
+                    continue
+                if cursor is None:
+                    self.cursors.pop(trace_id, None)
+                else:
+                    self.cursors[trace_id] = cursor
+
+    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`
+        stores a cursor only after its delta left."""
+        traces = list(self.slot.traces)
+        live = {trace.id for trace in traces}
+        deltas: list[tuple[str, dict, TraceCursor | None]] = []
+        for trace_id in [trace_id for trace_id in self.cursors if trace_id not in live]:
+            # A retried attempt abandons its traces; the client drops them too.
+            deltas.append((trace_id, {"trace": trace_id, "discard": True}, None))
+        for trace in traces:
+            delta: dict[str, Any] = {"trace": trace.id}
+            sent = self.cursors.get(trace.id)
+            cursor = copy.deepcopy(sent) if sent is not None else TraceCursor()
+            if sent is None:
+                delta["open"] = dump(trace, include=set(HEADER_FIELDS))
+            links: dict[int, list[dict]] = {}
+            for index in range(cursor.sent["nodes"]):
+                node_links = trace.nodes[index].semantic_parents
+                if len(node_links) > cursor.links[index]:
+                    links[index] = [
+                        dump(link) for link in node_links[cursor.links[index] :]
+                    ]
+                    cursor.links[index] = len(node_links)
+            if links:
+                delta["links"] = links
+            for field in LIST_FIELDS:
+                items = getattr(trace, field)
+                sent = cursor.sent[field]
+                if len(items) > sent:
+                    delta[field] = [dump(item) for item in items[sent:]]
+                    cursor.sent[field] = len(items)
+                    if field == "nodes":
+                        cursor.links.extend(
+                            len(node.semantic_parents) for node in items[sent:]
+                        )
+            changed: dict[str, Any] = {}
+            for field, value in dump(trace, include=set(SCALAR_FIELDS)).items():
+                packed = pack(value)
+                if cursor.scalars.get(field) != packed:
+                    cursor.scalars[field] = packed
+                    changed[field] = value
+            if changed:
+                delta["set"] = changed
+            pending = [dump(message) for message in trace.pending]
+            packed = pack(pending)
+            if packed != cursor.pending:
+                cursor.pending = packed
+                # nodes landing clear the preview on the client by themselves
+                if pending or "nodes" not in delta:
+                    delta["pending"] = pending
+            if len(delta) > 1:
+                deltas.append((trace.id, delta, cursor))
+        return deltas
+
+
+class EpisodeAssembly:
+    """The client's growing picture of one served episode: raw trace dicts in arrival
+    order, each shaped like a dumped `Trace` plus its `pending` preview."""
+
+    def __init__(self) -> None:
+        self.traces: dict[str, dict] = {}
+
+    def apply(self, delta: dict) -> None:
+        trace_id = delta["trace"]
+        if delta.get("discard"):
+            self.traces.pop(trace_id, None)
+            return
+        trace = self.traces.get(trace_id)
+        if trace is None:
+            trace = self.traces[trace_id] = {
+                **delta["open"],
+                **{field: [] for field in LIST_FIELDS},
+                "pending": [],
+            }
+        for index, links in (delta.get("links") or {}).items():
+            trace["nodes"][int(index)]["semantic_parents"].extend(links)
+        # 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:
+            trace["nodes"].extend(
+                {**node, "semantic_parents": list(node.get("semantic_parents") or [])}
+                for node in delta["nodes"]
+            )
+        for field in LIST_FIELDS:
+            if field in delta and field != "nodes":
+                trace[field].extend(delta[field])
+        trace.update(delta.get("set") or {})
+        # a committed turn absorbs the preview; an explicit preview replaces it
+        if "nodes" in delta:
+            trace["pending"] = []
+        if "pending" in delta:
+            trace["pending"] = delta["pending"]
+
+    def finish(self, head: dict, summaries: list[TraceSummary]) -> dict:
+        """The full episode record: `head` (the episode without its traces) joined with
+        the assembled traces in the server's order. A count mismatch means a delta was
+        lost, which the wire never does in a live connection — fail loudly."""
+        traces = []
+        for summary in summaries:
+            trace = self.traces.get(summary.id)
+            if trace is None:
+                raise RuntimeError(f"served episode is missing trace {summary.id}")
+            trace = {key: value for key, value in trace.items() if key != "pending"}
+            if (
+                len(trace["nodes"]) != summary.nodes
+                or len(trace["calls"]) != summary.calls
+            ):
+                raise RuntimeError(
+                    f"served trace {summary.id} assembled {len(trace['nodes'])} nodes / "
+                    f"{len(trace['calls'])} calls, server reports {summary.nodes} / {summary.calls}"
+                )
+            traces.append(trace)
+        return {**head, "traces": traces}
+
+
+class TraceSummary(BaseModel):
+    """A finished trace's identity and size, for the client's assembly check."""
+
+    id: str
+    nodes: int
+    calls: int

verifiers/v1/serve/pool.py

@@ -193,7 +193,7 @@ class EnvServerPool:
                     ) = await self.frontend.recv_multipart()
                     if method == b"health":
                         await self.frontend.send_multipart(
-                            [client_id, request_id, _HEALTH]
+                            [client_id, request_id, b"reply", _HEALTH]
                         )
                     elif method == b"cancel":
                         # Route the cancel to the worker holding the target run;
@@ -208,7 +208,7 @@ class EnvServerPool:
                         target_entry = pending.get(target)
                         if target_entry is None:
                             await self.frontend.send_multipart(
-                                [client_id, request_id, _CANCEL_MISS]
+                                [client_id, request_id, b"reply", _CANCEL_MISS]
                             )
                         else:
                             worker = target_entry["worker"]
@@ -238,16 +238,25 @@ class EnvServerPool:
                             self._maybe_scale_up(in_flight)
                 for w in self.workers:
                     if w["dealer"] in events:
-                        request_id, data = await w["dealer"].recv_multipart(copy=False)
-                        # Copy only the routing key; relay the response Frames unchanged.
-                        entry = pending.pop(request_id.bytes, None)
-                        if entry is None:
-                            continue
-                        entry["worker"]["active"] -= 1
-                        in_flight -= 1
+                        request_id, kind, data = await w["dealer"].recv_multipart(
+                            copy=False
+                        )
+                        # Copy only the routing key; relay the Frames unchanged. A
+                        # `delta` is one of many for its request — only the `reply`
+                        # closes the request out.
+                        if kind.bytes == b"reply":
+                            entry = pending.pop(request_id.bytes, None)
+                            if entry is None:
+                                continue
+                            entry["worker"]["active"] -= 1
+                            in_flight -= 1
+                        else:
+                            entry = pending.get(request_id.bytes)
+                            if entry is None:
+                                continue
                         with contextlib.suppress(zmq.ZMQError):
                             await self.frontend.send_multipart(
-                                [entry["client_id"], request_id, data], copy=False
+                                [entry["client_id"], request_id, kind, data], copy=False
                             )
         except (asyncio.CancelledError, KeyboardInterrupt):
             pass