feat: non-blocking, concurrent prime uploads

pyproject.toml

@@ -16,7 +16,7 @@ dependencies = [
     "transformers==5.19.0",
     "wandb>=0.26.1",
     "wandb-workspaces>=0.4.3",
-    "prime-runs[train]>=0.1.5",
+    "prime-runs[train]>=0.1.6",
     "pyzmq>=27.1.0",
     "aiolimiter>=1.2.1",
     "tenacity>=8.2.0",

src/prime_rl/monitors/prime.py

@@ -199,9 +199,9 @@ async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None:
         metrics, dropped = sanitize(metrics)
         if dropped:
             self.logger.warning(f"Dropping {len(dropped)} non-finite metric value(s): {', '.join(dropped[:5])}")
-        # A queue put that can block briefly under backpressure - off the loop. The SDK
-        # stamps `_timestamp` on every row, so step=None rows keep a time anchor.
-        await asyncio.to_thread(self.run.log_metrics, metrics, step=step)
+        # A queue put that never blocks. The SDK stamps `_timestamp` on every row, so
+        # step=None rows keep a time anchor.
+        self.run.log_metrics(metrics, step=step)
 
     async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None:
         """Every episode ships whole as it arrives (the ``all`` subset), errored and
@@ -211,8 +211,8 @@ async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind,
         if subset != "all" or not episodes:
             return
         if self.run is not None:
-            # A queue put that can block briefly under backpressure - off the loop.
-            await asyncio.to_thread(self.run.log_episodes, episodes)
+            # A queue put that never blocks.
+            self.run.log_episodes(episodes)
             return
         by_env: dict[str, list[vf.Episode]] = {}
         for episode in episodes:
@@ -221,7 +221,7 @@ async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind,
         for env_name, batch in by_env.items():
             run = await self.run_for(env_name, step)
             if run is not None:
-                await asyncio.to_thread(run.log_episodes, batch)  # a queue put, off the loop
+                run.log_episodes(batch)
 
     async def finalize(self, paused: bool = False) -> None:
         # An epoch the run did not finish leaves its evaluation open: close it as cancelled.

tests/unit/monitors/test_monitor_load.py

@@ -1,5 +1,5 @@
-"""Each monitor under an orchestrator-sized load: up to a thousand agentic episodes
-(100 turns, 100k tokens each) and thousands of metric rows, logged through
+"""Each monitor under an orchestrator-sized load: up to a thousand agentic episodes,
+sized like a production RL run's, and thousands of metric rows, logged through
 ``monitors.log`` as the orchestrator does.
 Everything logged must arrive, and the event loop must stay responsive while the
 monitor writes or uploads."""
@@ -8,6 +8,7 @@
 import gc
 import gzip
 import statistics
+import threading
 import time
 
 import numpy as np
@@ -18,6 +19,7 @@
 from prime_runs.sinks import RftMetricsSink, TracesSink
 from prime_runs.worker import UploadWorker
 from prime_traces import TracesClient
+from prime_traces.core.client import GZIP_LEVEL
 from verifiers.v1.graph import MessageNode
 from verifiers.v1.trace import ModelCall
 from verifiers.v1.types import AssistantMessage, Usage, UserMessage
@@ -26,30 +28,42 @@
 from prime_rl.configs.monitors import FileMonitorConfig, PrimeMonitorConfig, WandbMonitorConfig
 from prime_rl.monitors.file.monitor import FileMonitor
 from prime_rl.monitors.file.traces import get_trace_stream
-from prime_rl.monitors.file.traces.chunks import open_chunk
+from prime_rl.monitors.file.traces.chunks import chunk_numbers, open_chunk
 from prime_rl.monitors.prime import PrimeMonitor
 from prime_rl.monitors.wandb import WandbMonitor
 from prime_rl.orchestrator.gc import GarbageCollector
 from prime_rl.utils.pathing import get_file_monitor_dir
 
-TURNS = 100
 TOKENS_PER_MESSAGE = 500
+# Episode sizes (uncompressed JSON) in a production RL run: median 6.4 MB, p90 17 MB,
+# p99 35 MB. A lognormal through the median and p90 also matches its mean (8.5 MB)
+# and p99.
+SIZE_MEDIAN_MB = 6.4
+SIZE_SIGMA = 0.76
+# Prime Traces rejects an episode over 64 MiB.
+MAX_SIZE_MB = 60
+# Episodes share traces at these sizes, which keeps the test's memory small; each
+# episode takes the one nearest its drawn size.
+TRACE_SIZES_MB = np.geomspace(0.5, MAX_SIZE_MB, 24)
 METRIC_ROWS = 2_000
 METRIC_KEYS = 200
 # The orchestrator's own work runs between monitor calls; a stall this long delays
 # dispatch and batch assembly visibly.
 MAX_LAG_S = 0.25
-P99_LAG_S = 0.05
+P99_LAG_S = 0.1
+# A log() call may await the monitor's own work, such as the file monitor's write on its
+# writer thread; the loop keeps serving other tasks meanwhile, which the lag limits cover.
+MAX_LOG_S = 1.0
 TICK_S = 0.005
 
 pytestmark = pytest.mark.slow
 
 
-def _trace(rng: np.random.Generator) -> vf.Trace:
+def _trace(rng: np.random.Generator, turns: int) -> vf.Trace:
     """A multi-turn trace with per-token streams as the env-server wire delivers them."""
     n = TOKENS_PER_MESSAGE
     nodes = []
-    for _ in range(TURNS):
+    for _ in range(turns):
         nodes.append(
             MessageNode(
                 parent=len(nodes) - 1 if nodes else None,
@@ -79,30 +93,36 @@ def _trace(rng: np.random.Generator) -> vf.Trace:
         rewards={"correct": vf.Reward(score=1.0)},
         calls=[
             ModelCall(node=2 * i + 1, usage=Usage(prompt_tokens=n, completion_tokens=n), finish_reason="stop")
-            for i in range(TURNS)
+            for i in range(turns)
         ],
     )
 
 
+def _episode(trace: vf.Trace, i: int) -> vf.Episode:
+    return vf.Episode(
+        env=vf.EnvInfo(id="load", name="load"),
+        task=trace.task,
+        group=vf.GroupInfo(id=f"g{i // 8}"),
+        traces=[trace],
+        ok=True,
+        run=vf.TrainRunInfo(id="run", work=vf.TrainWorkInfo(step=1, policy=vf.PolicySpan(start=0, end=0))),
+    )
+
+
 @pytest.fixture(scope="module")
-def trace() -> vf.Trace:
-    return _trace(np.random.default_rng(0))
-
-
-def _episodes(trace: vf.Trace, count: int) -> list[vf.Episode]:
-    """``count`` episodes that share one trace: every monitor serializes each episode
-    on its own, so the load is the same as distinct traces at a fraction of the memory."""
-    return [
-        vf.Episode(
-            env=vf.EnvInfo(id="load", name="load"),
-            task=trace.task,
-            group=vf.GroupInfo(id=f"g{i // 8}"),
-            traces=[trace],
-            ok=True,
-            run=vf.TrainRunInfo(id="run", work=vf.TrainWorkInfo(step=1, policy=vf.PolicySpan(start=0, end=0))),
-        )
-        for i in range(count)
-    ]
+def traces() -> list[vf.Trace]:
+    """One trace per size in ``TRACE_SIZES_MB``."""
+    rng = np.random.default_rng(0)
+    turn_mb = len(_episode(_trace(rng, 10), 0).to_json()) / 10 / 1e6
+    return [_trace(rng, max(1, round(size / turn_mb))) for size in TRACE_SIZES_MB]
+
+
+def _episodes(traces: list[vf.Trace], count: int) -> list[vf.Episode]:
+    """``count`` episodes with production sizes. Every monitor serializes each episode on
+    its own, so sharing traces keeps the load of distinct ones at a fraction of the memory."""
+    sizes = np.random.default_rng(1).lognormal(np.log(SIZE_MEDIAN_MB), SIZE_SIGMA, count)
+    nearest = np.abs(np.log(sizes)[:, None] - np.log(TRACE_SIZES_MB)).argmin(1)
+    return [_episode(traces[index], i) for i, index in enumerate(nearest)]
 
 
 def _metrics(step: int) -> dict[str, float]:
@@ -152,7 +172,7 @@ def _assert_responsive(name: str, lags: list[float], calls: list[float]) -> None
     )
     assert p99 < P99_LAG_S, f"{name} stalls the event loop: p99 lag {p99 * 1e3:.0f} ms"
     assert max(lags) < MAX_LAG_S, f"{name} stalls the event loop: max lag {max(lags) * 1e3:.0f} ms"
-    assert max(calls) < MAX_LAG_S, f"a {name} log() call took {max(calls) * 1e3:.0f} ms"
+    assert max(calls) < MAX_LOG_S, f"a {name} log() call took {max(calls) * 1e3:.0f} ms"
 
 
 @pytest.fixture
@@ -175,15 +195,18 @@ def registered():
 
 class FakeTracesAPI:
     """The Prime Traces endpoint: gzips each request like the real client and keeps
-    the lines it received."""
+    the lines it received. The client sends several requests at once."""
 
     def __init__(self) -> None:
         self.lines = 0
+        self._lock = threading.Lock()
 
     def upload_batch(self, body, idempotency_key, *, line_format, schema_version=1, context=None, compress=True):
         if compress:
-            gzip.compress(body, mtime=0)
-        self.lines += body.count(b"\n") + (not body.endswith(b"\n"))
+            gzip.compress(body, GZIP_LEVEL, mtime=0)
+        lines = body.count(b"\n") + (not body.endswith(b"\n"))
+        with self._lock:
+            self.lines += lines
         return {"upload_id": idempotency_key.split(":", 1)[1], "status": "committed"}
 
     def close(self) -> None:
@@ -207,8 +230,10 @@ async def _file(tmp_path, monkeypatch):
     await monitor.init(tmp_path)
 
     def delivered() -> tuple[int | None, list[int]]:
-        with open_chunk(get_trace_stream(tmp_path), 0) as chunk:
-            episodes = sum(1 for _ in chunk.read().splitlines())
+        stream, episodes = get_trace_stream(tmp_path), 0
+        for number in chunk_numbers(stream):
+            with open_chunk(stream, number) as chunk:
+                episodes += sum(1 for _ in chunk)
         rows = (get_file_monitor_dir(tmp_path) / "metrics.jsonl").read_bytes().splitlines()
         return episodes, [orjson.loads(row)["step"] for row in rows]
 
@@ -257,27 +282,16 @@ def counting_log(data, *args, **kwargs):
     return monitor, lambda: (None, steps)
 
 
-# TODO: fix and drop this mark. The SDK uploads one episode at a time, and gzips each
-# at level 9 (about 70 ms here). A thousand episodes logged at once fill its queue, and
-# log() then waits for room (up to 5 s) on the orchestrator's arrival path, then drops.
-SDK_BACKPRESSURE = pytest.mark.xfail(reason="prime_runs blocks log() while its upload queue is full", strict=True)
-
-
 @pytest.mark.parametrize(
     "start,count",
     [
-        pytest.param(
-            start,
-            count,
-            id=f"{start.__name__[1:]}-{count}",
-            marks=SDK_BACKPRESSURE if (start, count) == (_prime, 1000) else (),
-        )
+        pytest.param(start, count, id=f"{start.__name__[1:]}-{count}")
         for start in (_file, _prime, _wandb)
         for count in (10, 100, 1000)
     ],
 )
-def test_monitor_under_load(start, count, trace, tmp_path, registered, orchestrator_gc, monkeypatch):
-    episodes = _episodes(trace, count)
+def test_monitor_under_load(start, count, traces, tmp_path, registered, orchestrator_gc, monkeypatch):
+    episodes = _episodes(traces, count)
 
     async def run():
         monitor, delivered = await start(tmp_path, monkeypatch)