feat: stream served episodes turn by turn

docs/v1/evaluation.md

@@ -42,8 +42,6 @@ The output from evaluations are written into `outputs/<env>--<model>--<harness>/
 - `num_rollouts` — rollouts per task
 - `verbose` — log at debug instead of info
 - `shuffle` — samples the task order (fixed seed); an error on an infinite taskset
-- `serve` — on by default: rollouts run through an elastic env-server worker pool
-  (`[serve]` sizes it, e.g. `serve.pool.type`); `--no-serve` runs them in-process
 - `rich` — the live dashboard (default); `--no-rich` streams logs to the console and
   prints each trace as JSON at the end
 - `rich.show_logs` — replace the dashboard's per-rollout rows with a live tail of the

tests/v1/conftest.py

@@ -5,7 +5,8 @@
 not unit tests of individual components. They need a model API key (`PRIME_API_KEY`);
 without one the `e2e`-marked tests skip (config parsing still runs).
 
-`run_v1` mirrors the eval CLI's in-process path (`run_eval` with `--no-serve`). Placement coverage (harness x harness runtime x tool
+`run_v1` mirrors the eval CLI (`run_eval`, in-process); `run_v1_server` drives the same env
+through an env-server worker pool, the path prime-rl trains through. Placement coverage (harness x harness runtime x tool
 server runtime) is PAIRWISE, not a full cross product: each test carries a curated list of
 combinations (in test_e2e.py) that hits every axis value and the cross-boundary pairs with
 distinct networking. The full cross bought flake exposure and CI minutes, not coverage — add
@@ -32,6 +33,7 @@
 CI runs deterministic tests across the Python matrix and the remaining live E2Es once.
 """
 
+import asyncio
 import os
 from pathlib import Path
 
@@ -130,9 +132,7 @@ def _eval_config(
     taskset_overrides: dict | None = None,
     runtime: dict | None = None,
     env: dict | None = None,
-    pool: dict | None = None,
     reasoning_effort: str | None = None,
-    server: bool = False,
 ) -> EvalConfig:
     """Build the smallest `EvalConfig` that still exercises the path, shared by the in-process
     (`run_v1`) and env-server (`run_v1_server`) fixtures. `taskset_overrides` merges onto the
@@ -184,7 +184,6 @@ def _eval_config(
             "reasoning_effort": reasoning_effort,
         },
         rich=None,
-        serve=({"pool": pool} if pool else {}) if server else None,
         output_dir=output_dir.parent,
         run={"dir": output_dir.name},
         model=CI_MODEL,
@@ -194,8 +193,7 @@ def _eval_config(
 
 @pytest.fixture
 def run_v1():
-    """Run a v1 taskset end-to-end in-process (`run_eval` with `--no-serve`) and return
-    its traces."""
+    """Run a v1 taskset end-to-end in-process (`run_eval`) and return its traces."""
 
     async def _run(taskset: str, **kwargs) -> list[Trace]:
         config = _eval_config(taskset, **kwargs)
@@ -208,17 +206,56 @@ async def _run(taskset: str, **kwargs) -> list[Trace]:
 
 @pytest.fixture
 def run_v1_server():
-    """Run a v1 taskset through the env-server worker pool (`run_eval`'s default path) —
-    the path a CLI run and prime-rl training both take. Spawns the broker + a worker, so
-    it's the only fixture that exercises serving resources (shared tool servers,
-    interception pool) being stood up by the *server* rather than the in-process runner.
-    Pinned to a single static worker for determinism."""
+    """Run a v1 taskset through an env-server worker pool — the path prime-rl trains
+    through. Spawns the broker + a worker, so it's the only fixture that exercises
+    serving resources (shared tool servers, interception pool) being stood up by the
+    *server* rather than the in-process runner. Pinned to a single static worker for
+    determinism."""
+    import multiprocessing as mp
+
+    from verifiers.v1.serve import EnvClient, env_config_data, serve_env
+    from verifiers.v1.utils.loaders import load_taskset
 
     async def _run(taskset: str, **kwargs) -> list[Trace]:
-        kwargs.setdefault("pool", {"type": "static", "num_workers": 1})
-        config = _eval_config(taskset, server=True, **kwargs)
-        records = await run_eval(config)
-        return [t for r in records for t in r.traces]
+        config = _eval_config(taskset, **kwargs)
+        tasks = list(load_taskset(config.env.taskset).head(config.num_tasks))
+        mpctx = mp.get_context("spawn")
+        address_queue: mp.Queue = mpctx.Queue()
+        proc = mpctx.Process(
+            target=serve_env,
+            kwargs={
+                "max_workers": 1,
+                "elastic": False,
+                "address": "tcp://127.0.0.1:0",
+                "address_queue": address_queue,
+                "config_data": env_config_data(config.env),
+                "max_concurrent": config.max_concurrent,
+            },
+        )
+        proc.start()
+        try:
+            address = await asyncio.to_thread(address_queue.get, timeout=600)
+            client = EnvClient(address=address)
+            try:
+                await client.wait_for_server_startup(timeout=600)
+                episodes = await asyncio.gather(
+                    *(
+                        client.run(
+                            client=config.client,
+                            model=config.model,
+                            sampling=config.sampling,
+                            task_data=task.data.model_dump(mode="json"),
+                        )
+                        for task in tasks
+                        for _ in range(config.num_rollouts)
+                    )
+                )
+            finally:
+                await client.close()
+        finally:
+            proc.terminate()
+            await asyncio.to_thread(proc.join, 10)
+        return [t for e in episodes for t in e.traces]
 
     return _run
 

tests/v1/test_serve_delta.py

@@ -0,0 +1,253 @@
+"""The env-serve delta stream: a trace's changes leave once each and reassemble into
+the episode the worker finished with."""
+
+import asyncio
+import copy
+
+import pytest
+
+import verifiers.v1 as vf
+from verifiers.v1.episode import EnvInfo, Episode, WireEpisode
+from verifiers.v1.graph import MessageNode
+from verifiers.v1.semantic import ParentLink
+from verifiers.v1.serve.delta import (
+    DeltaStreamer,
+    EpisodeAssembly,
+    TraceSummary,
+    dump,
+    pack,
+    unpack,
+)
+from verifiers.v1.trace import ModelCall, TimeSpan, TraceTask
+from verifiers.v1.types import AssistantMessage, Usage, UserMessage
+
+
+class MyTask(vf.TaskData):
+    prompt: str
+    answer: str
+
+
+class Slot:
+    def __init__(self) -> None:
+        self.traces: list[vf.Trace] = []
+
+
+async def settle() -> None:
+    """Let a notified change reach the wire: the flush is a callback that starts a task."""
+    for _ in range(3):
+        await asyncio.sleep(0)
+
+
+def make_trace() -> vf.Trace:
+    return vf.Trace(
+        agent=vf.AgentInfo(config=vf.AgentConfig()),
+        task=TraceTask(type="MyTask", data=MyTask(idx=0, prompt="q", answer="a")),
+    )
+
+
+def add_turn(trace: vf.Trace, reply: str) -> None:
+    parent = len(trace.nodes) - 1 if trace.nodes else None
+    if parent is None:
+        trace.nodes.append(MessageNode(parent=None, message=UserMessage(content="q")))
+        parent = 0
+    trace.nodes.append(
+        MessageNode(
+            parent=parent,
+            message=AssistantMessage(content=reply),
+            sampled=True,
+            token_ids=[1, 2, 3],
+            mask=[False, True, True],
+            logprobs=[-0.1, -0.2],
+        )
+    )
+    trace.calls.append(
+        ModelCall(
+            node=len(trace.nodes) - 1,
+            usage=Usage(prompt_tokens=3, completion_tokens=2),
+            time=TimeSpan(start=1.0, end=2.0),
+        )
+    )
+    trace.notify()
+
+
+def test_delta_fields_cover_every_serialized_trace_field():
+    """A new Trace field must be routed through the stream, or it silently vanishes."""
+    from verifiers.v1.serve.delta import HEADER_FIELDS, LIST_FIELDS, SCALAR_FIELDS
+
+    streamed = set(HEADER_FIELDS) | set(LIST_FIELDS) | set(SCALAR_FIELDS)
+    serialized = {
+        name for name, info in vf.Trace.model_fields.items() if not info.exclude
+    }
+    assert streamed == serialized
+
+
+@pytest.mark.asyncio
+async def test_failed_send_is_diffed_again():
+    """A delta the wire refused is not lost: its cursor stays and the next flush
+    carries the same content, so the client still assembles the whole trace."""
+    slot = Slot()
+    frames: list[bytes] = []
+    fail = {"on": False}
+
+    async def send(data: bytes) -> None:
+        if fail["on"]:
+            raise OSError("host unreachable")
+        frames.append(data)
+
+    async with DeltaStreamer(slot, send) as streamer:
+        trace = make_trace()
+        slot.traces.append(trace)
+        streamer.watch(trace)
+        await settle()
+        fail["on"] = True
+        add_turn(trace, "a1")
+        await settle()
+        fail["on"] = False
+        add_turn(trace, "a2")
+        await settle()
+        slot.traces = [trace]
+    assembly = EpisodeAssembly()
+    for frame in frames:
+        assembly.apply(unpack(frame))
+    assert len(assembly.traces[trace.id]["nodes"]) == 3
+    assert len(assembly.traces[trace.id]["calls"]) == 2
+
+
+@pytest.mark.asyncio
+async def test_pending_preview_streams_and_clears_on_commit():
+    slot = Slot()
+    frames: list[bytes] = []
+
+    async def send(data: bytes) -> None:
+        frames.append(data)
+
+    async with DeltaStreamer(slot, send) as streamer:
+        trace = make_trace()
+        slot.traces.append(trace)
+        streamer.watch(trace)
+        add_turn(trace, "a1")
+        await settle()
+        # the harness sends its next request: the tool result is previewed at once
+        trace.preview([UserMessage(content="tool says 42")])
+        await settle()
+        preview = unpack(frames[-1])
+        assert preview["pending"][0]["content"] == "tool says 42"
+        assert "nodes" not in preview
+        # the model answers: the turn commits and the preview goes with it
+        trace.nodes.append(
+            MessageNode(parent=1, message=UserMessage(content="tool says 42"))
+        )
+        trace.clear_preview()
+        add_turn(trace, "a2")
+        await settle()
+        committed = unpack(frames[-1])
+        assert len(committed["nodes"]) == 2 and "pending" not in committed
+        trace.stop("agent_completed")
+        trace.ok = True
+        episode = Episode(
+            env=EnvInfo(id="my-env"), task=trace.task, ok=True, traces=[trace]
+        )
+        slot.traces = list(episode.traces)
+    assembly = EpisodeAssembly()
+    for frame in frames:
+        assembly.apply(unpack(frame))
+        if unpack(frame).get("pending"):
+            assert assembly.traces[trace.id]["pending"][0]["content"] == "tool says 42"
+    assert assembly.traces[trace.id]["pending"] == []
+    summaries = [
+        TraceSummary(id=trace.id, nodes=len(trace.nodes), calls=len(trace.calls))
+    ]
+    record = assembly.finish(dump(episode, exclude={"traces"}), summaries)
+    assert "pending" not in record["traces"][0]
+    assert WireEpisode.model_validate(unpack(pack(record))).traces[0].id == trace.id
+
+
+@pytest.mark.asyncio
+async def test_deltas_stream_once_and_reassemble_the_episode():
+    slot = Slot()
+    frames: list[bytes] = []
+
+    async def send(data: bytes) -> None:
+        frames.append(data)
+
+    async with DeltaStreamer(slot, send) as streamer:
+        trace = make_trace()
+        slot.traces.append(trace)
+        streamer.watch(trace)
+        trace.timing.boot.start = 1.0
+        trace.notify()
+        await settle()
+        assert len(frames) == 1, "the mint and the boot span coalesce into one delta"
+        add_turn(trace, "a1")
+        await settle()
+        add_turn(trace, "a2")
+        trace.nodes[0].semantic_parents.append(ParentLink(node=1, type="reply"))
+        await settle()
+        trace.record_reward("match", 1.0)
+        trace.stop("agent_completed")
+        trace.ok = True
+        episode = Episode(
+            env=EnvInfo(id="my-env"), task=trace.task, ok=True, traces=[trace]
+        )
+        slot.traces = list(episode.traces)
+    # The exit flushed the tail (reward, stop, ok) without re-sending any node.
+    deltas = [unpack(frame) for frame in frames]
+    sent_nodes = sum(len(delta.get("nodes", [])) for delta in deltas)
+    sent_calls = sum(len(delta.get("calls", [])) for delta in deltas)
+    assert (sent_nodes, sent_calls) == (3, 2)
+    assert sum("open" in delta for delta in deltas) == 1
+    assert any("links" in delta for delta in deltas)
+    assert deltas[-1]["set"]["stop_condition"] == "agent_completed"
+
+    handed = copy.deepcopy(deltas)
+    assembly = EpisodeAssembly()
+    for delta in deltas:
+        assembly.apply(delta)
+    assert deltas == handed, "assembling never mutates a delta the caller holds"
+    head = dump(episode, exclude={"traces"})
+    summaries = [
+        TraceSummary(id=trace.id, nodes=len(trace.nodes), calls=len(trace.calls))
+    ]
+    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
+    assert (
+        rebuilt.traces[0].nodes[0].semantic_parents == trace.nodes[0].semantic_parents
+    )
+    assert rebuilt.traces[0].rewards["match"].value == 1.0
+    assert rebuilt.traces[0].stop_condition == "agent_completed"
+    assert len(rebuilt.traces[0].calls) == 2
+
+
+@pytest.mark.asyncio
+async def test_discarded_attempt_drops_its_trace():
+    slot = Slot()
+    frames: list[bytes] = []
+
+    async def send(data: bytes) -> None:
+        frames.append(data)
+
+    async with DeltaStreamer(slot, send) as streamer:
+        first = make_trace()
+        slot.traces.append(first)
+        streamer.watch(first)
+        add_turn(first, "a1")
+        await settle()
+        slot.traces = []  # the attempt is retried
+        second = make_trace()
+        slot.traces.append(second)
+        streamer.watch(second)
+        await settle()
+    deltas = [unpack(frame) for frame in frames]
+    assert {"trace": first.id, "discard": True} in deltas
+    assembly = EpisodeAssembly()
+    for delta in deltas:
+        assembly.apply(delta)
+    assert list(assembly.traces) == [second.id]
+
+
+def test_finish_refuses_a_gap():
+    assembly = EpisodeAssembly()
+    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)])