Repository navigation
prime-runs: RftBackend for training runs - #873
Merged
kcoopermiller merged 5 commits intoSep 1, 2026
Merged
Conversation
What prime-rl's monitors/prime.py::TrainRun did by hand, in the SDK, so the
monitor can become a thin adapter over pr.init(kind="train"). Additive: the
eval path is unchanged apart from run.type now following the run's kind.
init(kind="train", model=<base model>, environments=[hub ids],
training=TrainingSpec(max_steps, batch_size, rollouts_per_example, seq_len,
wandb_*), config=<run_config>, team_id=...) registers an external run. A team
is required: the platform creates external runs for a team and enables them
per team, so a team outside the allowlist gets the platform's 403 reason as
a ForbiddenError from init() and the producer decides what that means.
init(kind="train", id=$RUN_ID) attaches to a run a managed launch created:
nothing is registered, the platform keeps the run's failure marking (a
non-completed finish is logged and left to the launcher), and a clean
finish() still completes it — prime-rl's contract.
Terminal status: completed -> POST /rft/finalize {exit_code: 0}, which is a
compare-and-set on the server with "already finalized" as success, so it is
replayed on an ambiguous failure. failed/cancelled/crashed -> PUT
/external-runs/{id}/status {status: failed, error_message} with the finer
reason in the message (the RFT vocabulary is completed | failed); a 409 means
someone closed it first and is not an error. update() is a no-op: run_config
travels with create and there is no summary document.
Run.log_metrics(values, step=) queues one row per call on a second uploader
(a metrics dict is not a record, and a slow sample upload must not hold a
step's metrics back); rows are stamped with _timestamp, step is set unless
present, non-finite values dropped. RftMetricsSink posts one row per request
(the endpoint takes a single step; 60/min per token, Retry-After honoured).
RftSamplesSink is the training viewer's Parquet table: episodes whose
run.work is training work (verifiers TrainRunInfo, which prime-rl stamps on
every dispatched episode) are grouped by step, encoded on every 10th step
(prime-rl's cadence, configurable) and sent through presign -> PUT -> confirm.
Presign is replayable, the storage PUT is retried, confirm is not replayed.
The encoder is prime-rl's episodes_to_parquet_bytes moved into
prime_runs.projection, duck-typed like the rest; pyarrow is an optional
extra (prime-runs[train]) and without it the sink turns itself off with a
warning while traces still flow. Eval-work episodes, bare traces and JSON
episodes have no row here and reach Prime Traces only.
stamp_run takes the run type from the sink context, so bare dicts logged to
a training run carry run.type="train" — the traces service's vocabulary.
64 new tests (RFT backend, both sinks, init wiring, attach semantics, the
real Parquet encoder under importorskip). Version unchanged: this ships in
the release after 0.1.0.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GBkJRDCFihXMjU1YpJTt14
kcoopermiller
requested review from
DamianB-BitFlipper,
JannikSt,
JohannesHa,
burnpiro,
kennethnym and
willccbb
as code owners
August 28, 2026 23:09
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6c57662995
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
RftSamplesSink lazily opens its own httpx.Client for the presigned PUTs and did not register with _fork, so a forked child inherited the parent's socket and its atexit finish() would have closed it (sending close_notify on the parent's connection). Register the sink; reset_after_fork drops an owned client (rebuilt on the next upload) and marks an injected one unusable, the way TracesSink and PlatformClient already do. Also spell out the one-object-per-step contract: a step's episodes must be logged in one call (prime-rl does), since the platform keeps one object key per (run, step). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GBkJRDCFihXMjU1YpJTt14
Member
Author
|
Note on CI: |
kcoopermiller
changed the title
kennethnym
previously approved these changes
Sep 1, 2026
…recover - RftSamplesSink hands the encoder a per-upload sample_id offset (stride 2**32) so a step uploaded as several objects keeps (step, sample_id) unique; the confirm POST is replayable; the module docstring and README describe the additive per-step contract (dispatched-at-N vs trained-at-N) - RftMetricsSink posts are replayable: the platform merges rows by step - UploadWorker(retire_cooldown=): a sink that struck out on transient failures is paused and tried again after the cooldown; init() sets a 300 s cooldown for kind="train", evals keep the permanent retirement - attach() targets external runs (the v1 monitoring endpoints answer 400 for a hosted run): documented on init(), RftBackend.attach and the README Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Fk5fBAAaTiPfJ5UpNyqxXR
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit ba770fd. Configure here.
…nt's A forked child inherits the training samples sink's per-step upload counters, so it would hand out the same sample_id range as its parent for the next object of a step. After reset_after_fork the child draws its ranges at random from the upper half of the range space, which the parent's counter never reaches; every id stays below 2**53 since the viewer parses sample_id as a JSON number. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Fk5fBAAaTiPfJ5UpNyqxXR
Merged
kcoopermiller
added a commit
that referenced
this pull request
Sep 4, 2026
## Why
Eval and training jobs need one call that opens a run on the platform, streams records while they execute, and marks the run finished, failed, cancelled, or crashed when they stop. verifiers and prime-rl each hand-roll that lifecycle against the raw API today — two copies of the same thing — and nobody uploads the config a run was actually launched with, which is why the Config tab shows a v0 blob instead of anyone's `eval.toml`.
## What
A new leaf package, `prime-runs` (`httpx` + `prime-traces`; `pyarrow` behind a `train` extra), where the run is an object:
```python
import prime_runs as pr
run = pr.init(name="gsm8k-qwen3-8b", environments=["gsm8k"],
model=config.model, framework="verifiers",
config="eval.toml") # stored byte-for-byte
print(run.url)
for episode in rollouts:
run.log_episodes([episode]) # a queue put; streams on a background thread
run.finish(summary=pr.metrics.from_episodes(episodes))
```
Training runs use the same handle over the RFT external-runs API (folded in from #873):
```python
run = pr.init(kind="train", model=base_model, environments=[env_id],
training=pr.TrainingSpec(max_steps=1000, batch_size=64),
config=cfg.model_dump(), team_id=team) # or id=os.environ["RUN_ID"] to attach
run.log_episodes(step_episodes) # every 10th step -> Parquet presign / PUT / confirm
run.log_metrics(metrics, step=step) # POST /rft/metrics, on its own uploader
run.finish() # idempotent finalize; failed/cancelled/crashed -> status PUT
```
`Run` exposes `id`, `url`, `kind`, `log_episodes()` / `log_traces()` / `log_metrics()`, `update_summary()`, `flush()`, `finish()`, `fail()`, a context manager (an exception → `failed`, Ctrl-C → `cancelled`, an exit without `finish()` → `crashed` via atexit), and loss counters (`dropped_records`, `failed_records`, `errors`).
Records are keyed to the run at upload: the traces sink sets `run.id` / `run.type` on the uploaded copy of every trace and episode (member traces included) and keeps the producer's `name` / `work`, so verifiers and prime-rl keep their own local run ids and never learn the platform's.
Every record is written to two places during the transition: Prime Traces (the system of record; an account outside the beta has that sink retire quietly) and the table today's viewer reads — the v0 sample table for evals, per-step Parquet objects for training. Once the viewers read traces natively, the legacy sink is one line to drop and no producer changes.
How it degrades: queued batches coalesce so the request rate tracks upload latency; a full queue drops (counted) rather than stalling the run; a transient failure costs its batch and three in a row retire the sink — a training run pauses it for five minutes instead, since it outlives any platform deploy; non-finite numbers become `null` in the sample tables rather than losing the row. `POST /evaluations/`, `/rft/external-runs` and `/samples` are never replayed after an ambiguous failure (a retry could duplicate); presign, confirm, `/rft/metrics` (one row per step server-side) and `/rft/finalize` are.
Integrations: [verifiers#2415](PrimeIntellect-ai/verifiers#2415) (eval), [prime-rl#3459](PrimeIntellect-ai/prime-rl#3459) (training).
## Known gaps
- **No producer-facing way to mark an evaluation FAILED** in the API as deployed today: `finalize` only moves PROCESSING → COMPLETED, so a crashed or interrupted eval stays RUNNING on the dashboard. The SDK records status/error under `metadata.prime_runs` and warns. platform#4740 adds `PUT /evaluations/{id} {status: FAILED|CANCELLED, error_message}` with a 409 guard; once deployed, `EvalsBackend.finalize` should use it and keep the metadata fallback for older backends (`RunStatus.CANCELLED` already exists for that).
- **`config_source` is stored, not yet shown.** `init(config=<path>)` keeps the launched file byte-for-byte under `metadata.config_source` (`run_config.config_source` for training). The eval Config tab renders it once platform#4740 lands; verifiers#2415 sends only the three v0 keys for now and adds the file in a follow-up, and the training Config tab needs the same branch.
- **Attach targets external runs.** `init(kind="train", id=...)` works for runs created through `POST /rft/external-runs`; the v1 monitoring endpoints answer 400 for a hosted run's id (the same endpoints prime-rl's old monitor used, so nothing regresses).
- **Training samples are keyed by dispatch step.** Objects are additive per step on the platform, so an off-policy episode landing in a later batch adds an object for its dispatch step; each object numbers its `sample_id`s in its own range (a forked child in a random band). "Step N" in the sample viewer means dispatched at N, while the metrics at N describe the batch trained at N.
## Before merge
1. ~~prime-traces 0.0.3 on PyPI (#872), then `prime-traces>=0.0.3` here.~~ Done.
2. A PyPI **pending publisher** for `prime-runs` (owner `PrimeIntellect-ai`, repo `prime`, workflow `release-runs.yml`, environment `pypi-prod`). The name is unclaimed; the first run of the release workflow creates the project. The workflow publishes before it tags, so a failed first publish leaves no tag and re-runs from `main` via `workflow_dispatch`.
3. ~~Coalesce queued sample batches, add `RunStatus.CANCELLED`, fold #873 in so 0.1.0 carries the training backend.~~ Done.
4. Squash-merge (`main` requires linear history); `release-runs.yml` publishes 0.1.0 with eval + train. verifiers#2415 and prime-rl#3459 then drop their `[tool.uv.sources]` git pins and relock.
<!-- CURSOR_SUMMARY -->
---
> [!NOTE]
> **Medium Risk**
> New client code talks to production eval/RFT/traces APIs with nuanced retry and upload semantics; mistakes could duplicate runs, drop samples, or mis-report terminal status, though behavior is heavily tested and defaults to warn-not-raise.
>
> **Overview**
> Introduces **`prime-runs`**, a new leaf Python package (`httpx` + `prime-traces`, optional `pyarrow` for training) that centralizes platform run lifecycle: **`init()`** registers or attaches to a run, **`log_episodes` / `log_traces` / `log_metrics`** stream records on background uploaders, and **`finish()`** drains queues and closes the run (with context-manager, atexit crash reporting, and fork-safe HTTP/traces clients).
>
> **Eval** runs use **`EvalsBackend`** (`/evaluations`, hub environment resolve) and dual sinks to Prime Traces plus the legacy v0 sample table; **training** runs use **`RftBackend`** (`/rft/external-runs`, attach via `id=`), Parquet sample uploads on a step cadence, and a separate metrics uploader. The SDK stamps platform **`run.id` / `run.type`** on records, stores launch **`config_source`** verbatim, and implements deliberate retry/idempotency (no replay on ambiguous `POST` creates or append-only samples).
>
> **Repo wiring:** CI gains **`test-runs`** (Python 3.11–3.13) and version-bump enforcement for `packages/prime-runs`; **`release-runs.yml`** publishes to PyPI and tags `prime-runs-v*`. Hermetic tests cover HTTP, backends, sinks, projection, and end-to-end `init()` flows.
>
> <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 1b65b9a. 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 Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EYdQmKAVnGteT1tkimJb3g
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
The training half of the SDK:
pr.init(kind="train", ...)opens an external training run over/api/v1/rft/external-runs— what prime-rl'smonitors/prime.py::TrainRundoes by hand today (its docstring calls itself "the natural seam to be subsumed by the train SDK"). Additive: the eval path is unchanged apart fromrun.typefollowing the run's kind.RftBackendbackend.pyupdate()is a no-op (no config/summary document on RFT runs)Run.log_metrics(values, step=)run.pyUploadWorkerso a slow sample upload never holds metrics back;_timestampstamped, non-finite droppedRftMetricsSinksinks/metrics.pyPOST /rft/metricsper row (endpoint takes one step; 60/min, Retry-After honoured)RftSamplesSinksinks/train_samples.pyrun.work.step, every 10th step (prime-rl's cadence, configurable), presign → PUT → confirmepisodes_to_parquet_bytesprojection.pyprime-runs[train], sink turns itself off with a warning without itTrainingSpec,RunKind,RunSpec.kind/trainingmodels.pyinit(kind=, id=, training=)run.pyid=attaches to a managed launch's runDecisions worth a look
completed→POST /rft/finalize {exit_code: 0}: a server-side compare-and-set with "already finalized" as success, so it is replayed on an ambiguous failure (unlike the eval finalize).failed/cancelled/crashed→PUT /external-runs/{id}/status {status: "failed", error_message}with the finer reason in the message, since the RFT vocabulary iscompleted | failed. A 409 there means someone closed it first — logged, not raised.POST /external-runsrefuses without ateam_idand returns 403 for teams outsiderft_external_runs_allowed_team_ids. The SDK raisesConfigurationErrorbefore the request when there is no team, and passes the platform's 403 reason through asForbiddenErrorfrominit(). The producer decides whether that means "run locally" (the way verifiers treats a missing key) — same contract as evals, no silent downgrade insideinit().id=$RUN_ID). Mirrors prime-rl: nothing registered, a non-completedfinish()is logged and left to the launcher (which marks its own run failed when the process goes away), a cleanfinish()still finalizes.train/effectivestep; the SDK applies the step cadence (step_interval=10) and reads the work kind off the episode, so the producer just logs everything and Prime Traces gets it all. Theeffective-subset filter stays in prime-rl (the SDK cannot know it).finish(summary=)is not sent for training runs — there is nowhere to put it (prime-rl sendssummary: {}and the backend ignores it). Documented; final numbers go throughlog_metrics.Not in this PR
TrainRunwith this (afterprime-runsis on PyPI; prime-rl also needs itsdeps/verifierssubmodule past verifiers#2415, which makesimport prime_runspart of its import path anyway).dashboard/training/[id]/ConfigTab.tsxrenderingrun_config.config_sourceverbatim (platform follow-up, same shape as platform#4740's eval Config tab).release/*only).Note
Medium Risk
Touches run lifecycle, multi-day training uploads, and external RFT APIs (team allowlist, metrics rate limits, object storage); eval path is largely unchanged but shared upload/metrics plumbing grew more complex.
Overview
Adds training runs to
prime-runsalongside eval:pr.init(kind="train", ...)registers or attaches to external runs on/api/v1/rft/external-runs, withTrainingSpec, requiredteam_id, andRun.log_metricson a separate upload worker so metrics are not blocked by sample uploads.New platform plumbing:
RftBackendhandles create/attach/finalize (RFTcompleted|failedmapping, idempotent finalize);RftMetricsSinkposts per-step metrics;RftSamplesSinkgroups training-work episodes byrun.work.step, uploads Parquet on a configurable cadence (default every 10th step) via presign/PUT/confirm, with fork-safesample_idranges. Parquet encoding lives inprojection.episodes_to_parquet_bytesbehind optionalprime-runs[train](pyarrow).Resilience for long runs: training sinks use a 5-minute pause-and-retry after transient failures instead of permanent retirement; attached runs (
id=) skip registration and non-success terminal status updates so launchers own failure marking. Eval behavior stays the same aside from stampingrun.typefrom kind.Reviewed by Cursor Bugbot for commit c000e19. Bugbot is set up for automated code reviews on this repo. Configure here.