Skip to content

feat: encode episodes for the file and prime monitors one node at a time - #4366

Merged
mikasenghaas merged 21 commits into
mainfrom
feat/episode-json
Nov 3, 2026
Merged

mikasenghaas merged 21 commits into
mainfrom
feat/episode-json

Conversation

@mikasenghaas

@mikasenghaas mikasenghaas commented Nov 2, 2026 •

Copy link
Copy Markdown
Member

Summary

Companion to PrimeIntellect-ai/verifiers#3042 and PrimeIntellect-ai/prime#1140 (released in prime-runs 0.1.6). The file monitor and the Prime monitor encode each episode in short steps, so a large episode no longer stalls the orchestrator's event loop for as long as encoding it takes.

  • Pins.
    • The verifiers submodule moves to verifiers main, which includes verifiers#3042: to_json encodes one node at a time and releases the GIL every millisecond, and record_run records the run on every trace.
    • prime-runs 0.1.6, already on main, contains prime#1140: the SDK uploads episode.to_json() as it is when the episode and its traces name the upload's run. The dispatcher's record_run makes that true for RL episodes.
    • Verifiers main also includes feat: remove the eval's platform push and the prime-runs dependency verifiers#3046, which removes vf-eval's platform push and verifiers' prime-runs dependency. In uv.lock (uv 0.12.0), only the verifiers metadata changes.
  • The file monitor writes episode.to_json().
    • The stream line is byte for byte what orjson.dumps(episode.to_record()) gave before.
    • The index row comes from index_record(episode). It dumps everything but the nodes' payloads, one trace at a time, and gives the same row.
    • ChunkedJsonl.append takes a line in parts, so the newline is written without copying a long line.
  • The load test's lag limits drop from p99 100 ms and max 250 ms to 25 ms and 100 ms, so a regression shows. Across runs at production sizes, the highest values were a p99 of 12.5 ms and a max of 22 ms.
  • A hosted orchestrator stamps the platform's run id. A hosted launch sets only $RUN_ID, not $PRL_RUN_ID, so the orchestrator stamped episodes with a random id, and the SDK re-stamped every upload. It now falls back to $RUN_ID, the Prime run id.

Verification

Headline: event-loop lag at 1,000 production-sized episodes (~8 GB of JSON), monitor load test, dev box:

Main (before) This PR
prime: lag p99 / max 14 / 17 ms 3.1 / 3.5 ms
file: lag p99 / max 47 / 93 ms 4.8 / 15 ms
file: log() max 177 ms 195 ms
  • A second run: inside the full unit suite, with more load on the box, it measured prime 11 / 18 ms, file 5.2 / 22 ms, and file log() max 224 ms.
  • Why the GIL release matters: encoding node by node without it left the file monitor at 15–18 / 31–66 ms, varying from run to run. A busy encoder thread keeps the GIL until CPython's 5 ms switch interval forces a handoff, and the event loop waits that out several times per wakeup.

The stall no longer grows with the episode. One thread encodes episodes of 10 to 10,000 turns (1,000 tokens per turn) while a 5 ms ticker on the event loop records the worst stall. Medians of 5 runs, under the orchestrator's GC policy:

Turns Episode Encode CPU, before / after Worst stall, before Worst stall, after
10 0.2 MB 0.4 / 0.5 ms 5.5 ms 5.2 ms
100 2.2 MB 4.0 / 4.5 ms 8.0 ms 1.6 ms
1,000 22 MB 50 / 51 ms 40 ms 3.1 ms
3,000 65 MB 152 / 167 ms 100 ms 5.2 ms
10,000 217 MB 508 / 571 ms 326 ms 15 ms
  • Before: the stall grows linearly with the episode.
  • After: it stays around 2–5 ms up to 3,000 turns, about the GIL switch interval. The 10-turn case is that interval alone, since each episode encodes in under a millisecond.
  • The rest at 10,000 turns: the trace header's pydantic dump, which holds every model call (one per turn). Alone it stalls 4.4 / 8.4 / 15 ms at 1,000 / 3,000 / 10,000 turns. Agentic episodes in production have a few hundred turns, so this part stays small. Encoding the calls one at a time, like the nodes, would remove it.
  • CPU: encoding costs up to ~12% more, from the extra calls and GIL releases.

Encoding: test_received_episode_encodes_as_its_record_did sends a two-trace episode over the real delta wire into a WireEpisode. Then:

  • the file monitor's stream line equals orjson.dumps(to_record()), byte for byte;
  • its index row equals the row from the full record;
  • after record_run, the SDK takes the episode's own to_json(), which carries the same record as the old stamped line.

Tests: pytest tests/unit -m "not gpu": 1,077 passed, 7 failed. The same 7 fail on main on this dev box: an attention kernel without sm_120 support, and no ninja for the CPU Adam extension.

CI: on the CI runner, the file monitor at 1,000 production-sized episodes measured event-loop lag p99 6.0 ms and max 12 ms, against 63 / 152 ms on main. Its slowest log() was 382 ms, under the 1 s limit.

🤖 Generated with Claude Code


Note

Medium Risk
Changes affect how every episode is persisted and indexed and how run identity is chosen on hosted launches; behavior is contract-tested but touches high-volume I/O on the training path.

Overview
Large episodes no longer block the orchestrator event loop during trace logging. The file monitor writes stream lines with episode.to_json() (incremental, GIL-friendly encoding from verifiers) instead of a single orjson.dumps(to_record()), and builds index rows via new index_record() so summarization skips heavy node payloads without materializing the full record first. ChunkedJsonl.append accepts byte parts so the newline is appended without copying megabyte-scale lines.

The orchestrator run_id now falls back to $RUN_ID when $PRL_RUN_ID is unset, so hosted runs align with the platform id and Prime uploads are not re-stamped with a random id.

Monitor load tests use stricter event-loop lag thresholds to catch regressions. A new unit test asserts on-disk bytes and SDK upload JSON still match the legacy to_record() encoding.

Reviewed by Cursor Bugbot for commit f83e849. Bugbot is set up for automated code reviews on this repo. Configure here.

mikasenghaas and others added 6 commits November 2, 2026 18:39
PrimeIntellect-ai/prime#1136 gzips at level 1, keeps up to eight
upload requests in flight, and bounds the upload queue by records
with a submit that never blocks. prime-runs moves into
override-dependencies next to prime-traces, because verifiers pins
it to an older commit.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The SDK's submit never blocks now, so the prime monitor calls
run.log_metrics and run.log_episodes on the event loop. Each
asyncio.to_thread hop had to win the GIL from the uploader thread
several times, so under load a log() call took up to 87 ms.

The load test drops its expected failure for 1,000 episodes. Its
fake Traces endpoint gzips at the SDK's level and counts lines under
a lock, since the client now sends several requests at once.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The load test used one 2.2 MB episode. A production RL run's episodes
are larger and vary a lot: median 6.4 MB, p90 17 MB, p99 35 MB,
uncompressed. Each episode now draws its size from a lognormal through
that median and p90, which also matches the run's mean and p99. Sizes
are capped at 60 MB, under Prime Traces' 64 MiB line limit. The
episodes share 24 traces at log-spaced sizes, so the test's memory
stays small. A thousand episodes is now about 8 GB of JSON, close to
one production step of 1,024 samples.

The file monitor's check counted only the first trace chunk, read
whole. It now streams every chunk, since 8 GB spans two.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
PrimeIntellect-ai/verifiers#3042 encodes episodes and traces to JSON
one node at a time, with an optional run stamped on every trace.
PrimeIntellect-ai/prime#1140 uploads that encoding as it is. The
verifiers submodule moves to verifiers main plus #3042; main's own
changes since the old pin are harness and dependency bumps.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The file monitor encoded each episode with to_record() and one orjson
call, both of which hold the GIL for the whole episode, so a large
episode stalled the orchestrator's event loop from the writer thread.
It now writes Episode.to_json(), which encodes one node at a time. The
stream line is byte for byte the same.

The index row came from the same record. It now comes from
index_record, which dumps everything but the nodes' payloads, a trace
at a time. ChunkedJsonl.append takes a line in parts, so the newline
is written without copying a long line.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
PrimeIntellect-ai/verifiers#3042 now releases the GIL after each
millisecond of encoding, so a busy event loop does not wait out
CPython's switch interval several times per wakeup.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
mikasenghaas and others added 2 commits November 2, 2026 20:25
PrimeIntellect-ai/verifiers#3042 records the run on an episode and
each of its traces (Episode.record_run), so to_json carries it.
PrimeIntellect-ai/prime#1140 uploads to_json as it is when the episode
and its traces name the upload's run, and stamps any other record.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A hosted launch sets only the platform's $RUN_ID, not $PRL_RUN_ID, so
the orchestrator stamped episodes with a random run id. The SDK then
re-stamped every upload with the Prime run. The orchestrator now falls
back to $RUN_ID, so hosted episodes name their Prime run and upload
their own to_json encoding.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
mikasenghaas and others added 11 commits November 2, 2026 20:42
A log() call may await the monitor's own work: the file monitor awaits
its write on a worker thread, which encodes the episode. At production
sizes that takes up to ~400 ms on CI for a 60 MB episode, while the
event loop keeps serving other tasks. The lag limits cover the loop,
so log() calls get a looser limit of 1 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A log() call may await the monitor's own work: the file monitor awaits
its write on a worker thread, which encodes the episode. At production
sizes that takes up to ~400 ms on CI for a 60 MB episode, while the
event loop keeps serving other tasks. The lag limits cover the loop,
so log() calls get a looser limit of 1 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Encoding episodes one node at a time with GIL releases keeps the
monitors' event-loop lag at a p99 of up to ~13 ms and a max of up to
~22 ms at production sizes. The limits drop from 50 ms p99 and
250 ms max to 25 ms and 100 ms, so a regression in either shows.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
# Conflicts:
#	tests/unit/monitors/test_monitor_load.py
At production sizes the file monitor encodes each episode in one call,
which stalls the event loop with a p99 of ~60 ms on CI. The limit
covers that until the encoder works in short steps, which lowers it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
# Conflicts:
#	tests/unit/monitors/test_monitor_load.py
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…episode-json

# Conflicts:
#	pyproject.toml
#	uv.lock
Verifiers main requires prime-runs from PyPI instead of pinning a git
commit, so with the submodule on main the SDK overrides go away. This
also moves prime-sandboxes to 0.5.4, which verifiers main requires.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
# Conflicts:
#	deps/verifiers
#	pyproject.toml
#	uv.lock
…episode-json

# Conflicts:
#	deps/verifiers
#	uv.lock
mikasenghaas added a commit that referenced this pull request Nov 3, 2026
## Summary

Prime uploads keep up with the orchestrator, and logging to Prime never
blocks it. Companion to PrimeIntellect-ai/prime#1136.

- **SDK pin.** `prime-runs[train]>=0.1.6`, which pulls `prime-traces`
0.0.12 from PyPI. These releases contain PrimeIntellect-ai/prime#1136:
gzip level 1, up to 8 upload requests in flight, an upload queue bounded
by records whose `submit` never blocks, and lazily built upload records.
- `uv.lock` is written with uv 0.12.0 (lockfile revision 3); only the
two SDK packages change.
- **No thread hop for Prime.** `PrimeMonitor` calls `run.log_metrics`
and `run.log_episodes` directly on the event loop, not through
`asyncio.to_thread`. A queue put that never blocks needs no thread, and
each hop had to win the GIL back from the uploader thread.
- **A production-sized load test.**
`tests/unit/monitors/test_monitor_load.py`:
- **Episode sizes:** episodes draw their size from a lognormal fitted to
a production RL run (median 6.4 MB, p90 17 MB; it also matches the mean
of 8.5 MB and the p99 of 35 MB), capped at 60 MB. 1,000 episodes is ~8
GB of JSON, close to one production step of 1,024 samples.
- **Memory:** episodes share 24 traces at log-spaced sizes, so the
test's own memory stays small.
- **Fake endpoint:** the fake Traces endpoint gzips at the SDK's level
and counts lines under a lock, because the client now sends several
requests at once.
- **File monitor check:** it streams every trace chunk, since 8 GB spans
two 5 GiB chunks.
- **No expected failure:** the prime monitor at 1,000 episodes is no
longer marked as one.
- **A separate `log()` limit of 1 s.** A `log()` call may await the
monitor's own work: the file monitor awaits its write, which encodes the
episode on a worker thread, while the event loop keeps serving other
tasks. The lag limits still cover the loop.
- **The p99 lag limit rises from 50 to 100 ms.** At production sizes the
file monitor encodes each episode in one call, which stalls the event
loop with a p99 of ~60 ms on CI. Main's code shows the same stall with
this test. #4366 makes the encoder work in short steps and tightens the
limits to 25 ms p99 and 100 ms max.

## Verification

**Headline: Prime monitor, 1,000 production-sized episodes (~8 GB of
JSON), main against this branch.** Main ran this branch's test, with the
fake endpoint gzipping at main's level 9.

| | main | This branch |
|---|---|---|
| Episodes delivered | 725 of 1,000 | **1,000 of 1,000** |
| Lost to a full upload queue | 27 | 0 |
| Lost when the 60 s drain ran out | 248 | 0 |
| Slowest `log()` | up to 5 s (blocked on the full queue) | **0.6 ms** |
| Event-loop lag p99 / max | not reached (failed on delivery) | 14 / 17
ms |
| Test time | 217 s | 28 s |

**Scale: 10,000 episodes of 2.2 MB (~22 GB)**, from a local copy of the
test with the count raised:

| Setup | Result |
|---|---|
| No network time | 0 dropped; `log()` max 1.4 ms; event-loop lag p99 14
ms, max 21 ms; 570 MB peak RSS |
| 100 ms + 100 MB/s per request, 1 request in flight | 2,595 delivered
within the 60 s drain |
| Same, 4 in flight | 9,822 delivered |
| Same, 8 in flight (the default) | **10,000 delivered**; event-loop lag
p99 12 ms; 860 MB peak RSS |

The upload rate is ~180 episodes/s (~400 MB/s of JSON). Encoding the
records is the limit.

**Full load-test grid at production sizes, this branch (dev box):** 9
passed.

| 1,000 episodes | Event-loop lag p99 / max | `log()` max |
|---|---|---|
| prime | 14 / 17 ms | 0.6 ms |
| file | 47 / 93 ms | 177 ms |
| wandb | 0.9 / 0.9 ms | 0.6 ms |

**CI:** on the CI runner, the file monitor at 1,000 episodes measured
event-loop lag p99 60 ms and max 157 ms, within this PR's limits (100 /
250 ms), and `log()` up to 299 ms (limit 1 s). With #4366 the same case
measures p99 6 ms and max 12 ms.

Other checks:
- `pytest tests/unit/train/sft/test_prime_monitor.py
tests/unit/utils/test_links.py
tests/unit/orchestrator/test_episode_records.py tests/unit/train/rl -m
"not gpu"`: 29 passed.
- `uvx uv@0.12.0 lock --check` passes.

## Merge order

1. ~~PrimeIntellect-ai/prime#1136~~ and
~~PrimeIntellect-ai/prime#1140~~: merged and released as `prime-runs`
0.1.6 and `prime-traces` 0.0.12.
2. **#4363 (this PR)**: ready to merge. It does not need
PrimeIntellect-ai/verifiers#3042.
3. PrimeIntellect-ai/verifiers#3042
4. #4366

🤖 Generated with [Claude Code](https://claude.com/claude-code)


<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> **Medium Risk**
> Changes the hot path for RL training telemetry (episode/metrics
enqueue) and bumps external SDK versions; misbehavior could drop or
delay uploads under load, though load tests now enforce delivery and
loop responsiveness.
> 
> **Overview**
> **Prime logging no longer blocks the orchestrator event loop** by
pinning `prime-runs[train]>=0.1.6` (and `prime-traces` 0.0.12 in the
lockfile), which uses a non-blocking upload queue, concurrent uploads,
and lighter gzip.
> 
> **`PrimeMonitor`** now calls `run.log_metrics` and `run.log_episodes`
directly instead of `asyncio.to_thread`, matching SDK behavior where
enqueue never blocks.
> 
> The **monitor load test** is reworked to use **production-like episode
sizes** (lognormal ~6.4 MB median), shared trace templates for memory, a
thread-safe fake Traces API, and multi-chunk file delivery checks. The
**Prime monitor at 1,000 episodes is no longer expected to fail**;
responsiveness asserts split event-loop lag (p99 100 ms) from
per-`log()` duration (max 1 s).
> 
> <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit
11ae629. Bugbot is set up for automated
code reviews on this repo. Configure
[here](https://www.cursor.com/dashboard/bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Base automatically changed from feat/sdk-upload-perf to main November 3, 2026 00:34
# Conflicts:
#	pyproject.toml
#	tests/unit/monitors/test_monitor_load.py
mikasenghaas added a commit to PrimeIntellect-ai/verifiers that referenced this pull request Nov 3, 2026
…races (#3042)

## Summary

- **`Episode.to_json` and `Trace.to_json` encode one node at a time.**
- Pydantic's dump and orjson each hold the GIL for a whole call. Before,
each method made one call over the whole record, so encoding a large
episode stalled every other thread of the process, such as an asyncio
event loop, for as long as it took.
- Each call now covers one node, and the encoder releases the GIL after
each millisecond of encoding (`time.sleep(0)`). Without that release, a
busy thread keeps the GIL until CPython's 5 ms switch interval forces a
handoff, and an event loop waits that long several times per wakeup.
Releasing after every node instead made encoding ~3x slower.
- The output is byte for byte `orjson.dumps(to_record(),
option=OPT_SERIALIZE_NUMPY)`. Each trace and node is dumped as its
declared field type, as the parent's dump does. The parts are joined
once at the end, so the record is copied once.
- **`Trace.run`.** `Trace` gets a `run` field, and
`Episode.record_run(run)` sets it on the episode and on each of its
traces.
- Prime Traces reads the run from each trace (`trace.run.id`) and
derives an episode's run from its members. So `to_record()` and
`to_json()` now carry what the Prime SDK used to stamp, and the SDK can
upload `to_json()` as it is (PrimeIntellect-ai/prime#1140).
- The run types (`PolicySpan`, `TrainWorkInfo`, `EvalWorkInfo`,
`WorkInfo`, `TrainRunInfo`, `EvalRunInfo`, `RunInfo`) move from
`episode.py` to `trace.py`, which `Trace` needs. `verifiers.v1` exports
them as before.
- The delta stream sends `run` with the trace header. Env servers send
`None`; the consumer stamps the run when the episode arrives.

## Verification

- `pytest tests/unit`: 245 passed, 10 skipped (no harbor, no
`PRIME_API_KEY`).
- `to_json` equals `orjson.dumps(to_record())` byte for byte: for a
trace and an episode, with and without float rounding, with nested
`exclude` specs (`__all__` plus a per-trace index), and with a
`TaskData` subclass.
- `record_run` puts the run on the episode and every trace, `to_json`
carries it, and the record validates back to the same run.
- Worst event-loop stall while another thread encodes, and CPU per
episode. The episodes are synthetic agentic episodes at the sizes of a
production RL run:

  | Episode | Before: CPU / worst stall | After: CPU / worst stall |
  |---|---|---|
  | 8.4 MB (production mean) | 20 ms / 15 ms | 21 ms / 2.3 ms |
  | 35 MB (production p99) | 81 ms / 55 ms | 91 ms / 5.2 ms |
  | 133 MB (production max) | 319 ms / 191 ms | 359 ms / 12 ms |

- End-to-end numbers are in PrimeIntellect-ai/prime-rl#4366, which
writes and uploads episodes with this encoder under prime-rl's monitor
load test.

## Merge order

1. ~~PrimeIntellect-ai/prime#1136~~ and
~~PrimeIntellect-ai/prime#1140~~: merged and released as `prime-runs`
0.1.6 and `prime-traces` 0.0.12.
2. **#3042 (this PR)** and PrimeIntellect-ai/prime-rl#4363, in either
order.
3. PrimeIntellect-ai/prime-rl#4366

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> **Medium Risk**
> Serialization logic is substantially rewritten (though
contract-tested); wire JSON gains optional `trace.run`, and code that
imported run types from `verifiers.v1.episode` must import from `trace`
or `vf` instead.
> 
> **Overview**
> **`Trace.to_json` and `Episode.to_json`** no longer dump the full
record in one pydantic/orjson call. They build JSON in pieces
(`splice_json`, `trace_json_parts`), encoding each message node
separately and yielding the GIL about every millisecond so large
rollouts don’t block asyncio and other threads. Output stays
**byte-identical** to `orjson.dumps(to_record(), OPT_SERIALIZE_NUMPY)`,
including nested `exclude` specs and custom task types.
> 
> **Run metadata moves onto each trace.** `Trace` gains optional `run`,
and **`Episode.record_run`** now stamps the same `RunInfo` on the
episode and every trace in `traces`. The run-related models
(`PolicySpan`, `TrainRunInfo`, `EvalRunInfo`, etc.) live in
**`trace.py`** instead of `episode.py`; **`import verifiers.v1 as vf`**
still exposes them. The live trace delta stream sends **`run` in the
trace header** (often `None` until the consumer stamps it).
> 
> Tests assert `to_json` ↔ `to_record` parity and that stamped runs
round-trip through JSON validation.
> 
> <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit
dabe65d. Bugbot is set up for automated
code reviews on this repo. Configure
[here](https://www.cursor.com/dashboard/bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
@mikasenghaas
mikasenghaas marked this pull request as ready for review November 3, 2026 00:45
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@mikasenghaas
mikasenghaas merged commit d07cd84 into main Nov 3, 2026
10 of 11 checks passed
@mikasenghaas
mikasenghaas deleted the feat/episode-json branch November 3, 2026 00:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants