diff --git a/pyproject.toml b/pyproject.toml index ca8e6db0a6..8752d26311 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index 4f4b8f5dda..a708eb3b70 100644 --- a/src/prime_rl/monitors/prime.py +++ b/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. diff --git a/tests/unit/monitors/test_monitor_load.py b/tests/unit/monitors/test_monitor_load.py index 0f64fd061a..9c91b1fc3f 100644 --- a/tests/unit/monitors/test_monitor_load.py +++ b/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) diff --git a/uv.lock b/uv.lock index 9484af262d..d86fb8f01d 100644 --- a/uv.lock +++ b/uv.lock @@ -5357,7 +5357,7 @@ requires-dist = [ { name = "prime-rl", extras = ["quack"], marker = "extra == 'all'" }, { name = "prime-rl", extras = ["sonic-moe"], marker = "extra == 'all'" }, { name = "prime-rl-configs", editable = "packages/prime-rl-configs" }, - { name = "prime-runs", extras = ["train"], specifier = ">=0.1.5" }, + { name = "prime-runs", extras = ["train"], specifier = ">=0.1.6" }, { name = "prometheus-client", specifier = ">=0.20.0" }, { name = "psutil", specifier = ">=7.0.0" }, { name = "pybase64", specifier = ">=1.4.2" }, @@ -5431,16 +5431,16 @@ requires-dist = [ [[package]] name = "prime-runs" -version = "0.1.5" +version = "0.1.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "httpx" }, { name = "orjson" }, { name = "prime-traces" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3b/85/78a77c31ab994605b433f9594b2dedfdd92c6e0e0c015db3a3bdcd229cad/prime_runs-0.1.5.tar.gz", hash = "sha256:feddd70d3deaf83f242747e8b845d00ef30b15a6cdc3dcc78eaa299a51d9a57b", size = 67663, upload-time = "2026-10-29T17:56:49.718Z" } +sdist = { url = "https://files.pythonhosted.org/packages/44/76/a6810b874ed23b16f5210feb50c034595bd199a21a6db2619eef968c487b/prime_runs-0.1.6.tar.gz", hash = "sha256:60c2c656c9c076bd8995a3943b060fdaa893a248f9f627d5f2b26040821e829e", size = 68575, upload-time = "2026-11-02T23:44:48.27Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9f/a3/43c9773f431f0d7a9f3af6b4cf2c69a3ba9af7abf189f29cd2a493419ab5/prime_runs-0.1.5-py3-none-any.whl", hash = "sha256:f2bf5099a942eb254f98ffaf5c53f3fbc91b3b5a4a655914e99a4f06e5ed3ead", size = 46374, upload-time = "2026-10-29T17:56:48.14Z" }, + { url = "https://files.pythonhosted.org/packages/14/22/8745496681a8ec9884f049acaf848642adea0fb529a817d55a1d05a4eebf/prime_runs-0.1.6-py3-none-any.whl", hash = "sha256:78934a05fa80c0abf3ccaee4b37bfe89c822a72b10231f9e667ce37586a1cbf6", size = 46844, upload-time = "2026-11-02T23:44:46.867Z" }, ] [package.optional-dependencies] @@ -5469,16 +5469,16 @@ wheels = [ [[package]] name = "prime-traces" -version = "0.0.11" +version = "0.0.12" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "httpx" }, { name = "orjson" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e7/6c/0c1be7958d7e98b7168bcc48ce51c2743caae2daff92cab717c6b83cb331/prime_traces-0.0.11.tar.gz", hash = "sha256:a5933843e5ed2bf18279225f03539d6a85adeb86f9f906df4c8eff819f79b874", size = 60137, upload-time = "2026-10-29T17:56:49.936Z" } +sdist = { url = "https://files.pythonhosted.org/packages/69/50/de56a11d243fddb0b7c9feb47e512681d5a1a499ef7c42ba33bbb34e3f86/prime_traces-0.0.12.tar.gz", hash = "sha256:bc34ab96fd7bb9bed3ab39f28546fa936bdebf2e900523c8aee2a4be6ae3ea37", size = 61727, upload-time = "2026-11-02T23:44:51.943Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/10/a4/82298907c004527930892f096fdb0f539434666091e898492caae65b17a4/prime_traces-0.0.11-py3-none-any.whl", hash = "sha256:eefaf69409bbf639ca4972b7248eee0dcd31b17b204aefcc7fc20204fa6b234f", size = 43734, upload-time = "2026-10-29T17:56:48.691Z" }, + { url = "https://files.pythonhosted.org/packages/d4/35/19558b051a460662db16e630492d6aa70ce714afa2421c3d11f4883b6f01/prime_traces-0.0.12-py3-none-any.whl", hash = "sha256:9f6bba4fb0f6fcf09662475bdf86e81f42df94cce27409ca7acba3b7d1e90025", size = 44478, upload-time = "2026-11-02T23:44:50.652Z" }, ] [[package]]