From 9c561deb841b9b944099ae399bbd98862a541bb7 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Tue, 1 Sep 2026 10:09:31 -0700 Subject: [PATCH 1/5] feat(monitors): log to the platform through prime-runs Replace the hand-rolled TrainRun (RFT lifecycle, per-step metrics POSTs, every-10th-step Parquet presign->PUT->confirm, atexit failure marking) with prime_runs' pr.init(kind="train"): the SDK owns the uploads on a background thread with retries and backpressure, reports crashed on a process that exits without finalizing, and adds a Prime Traces sink for allowlisted accounts. episodes_to_parquet_bytes and SAMPLE_SCHEMA moved into prime_runs.projection; the kind/subset cohort filter stays here, the step cadence moves into the SDK's training samples sink. TEMPORARY: prime-runs is not on PyPI yet; pinned to the prime repo's feature/prime-runs-train-backend branch (#856 + #873) via uv source. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01AGeqWNdRB8ckrgch1tqVzJ --- pyproject.toml | 5 + src/prime_rl/monitors/prime.py | 374 +++++++-------------------------- uv.lock | 29 ++- 3 files changed, 113 insertions(+), 295 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9d77038a88..2490c6100c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "wandb>=0.26.1", "wandb-workspaces>=0.4.3", "prime>=0.6.19", + "prime-runs[train]>=0.1.0", "pyzmq>=27.1.0", "aiolimiter>=1.2.1", "tenacity>=8.2.0", @@ -221,6 +222,8 @@ flash_attn_3 = false # appear in `uv tree` need to be listed; the resolver ignores entries for # packages it never sees. prime = false +prime-runs = false +prime-traces = false prime-sandboxes = false prime-tunnel = false prime-evals = false @@ -237,6 +240,8 @@ torchao = false prime-kernels = false [tool.uv.sources] +# TEMPORARY: prime-runs is not on PyPI yet (PrimeIntellect-ai/prime #856 + #873) +prime-runs = { git = "https://github.com/PrimeIntellect-ai/prime.git", subdirectory = "packages/prime-runs", branch = "feature/prime-runs-train-backend" } prime-rl-configs = { path = "packages/prime-rl-configs", editable = true } # prime-rl consumes `verifiers.v1` (incl. the built-in harnesses/tasksets); env tasksets # live under deps/{verifiers,prime-envs}/environments. diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index 8bc375d906..ae801c8198 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -1,331 +1,117 @@ from __future__ import annotations import asyncio -import atexit -import io -import json import os -import time -from datetime import datetime, timezone -from typing import Any, Coroutine +from typing import TYPE_CHECKING, Any -import httpx -import pyarrow as pa -import pyarrow.parquet as pq -import verifiers.v1 as vf -from prime_cli.core.config import Config as PrimeConfig -from verifiers.v1.utils.platform import build_samples +import prime_runs as pr from prime_rl.configs.monitors import PrimeMonitorConfig from prime_rl.monitors.base import Kind, Monitor, Subset from prime_rl.utils.config import BaseConfig -from prime_rl.utils.logger import get_logger from prime_rl.utils.utils import sanitize -BASE_URL = "https://api.primeintellect.ai/api/v1/rft" +if TYPE_CHECKING: + import verifiers.v1 as vf + BASE_URL_VAR = "PRIME_API_BASE" -API_KEY_VAR = "PRIME_API_KEY" -SAMPLE_SCHEMA = pa.schema( - [ - ("run_id", pa.string()), - ("step", pa.int64()), - ("tag", pa.string()), - ("problem_id", pa.int64()), - ("sample_id", pa.int64()), - ("prompt", pa.string()), - ("completion", pa.string()), - ("trajectory", pa.string()), - ("answer", pa.string()), - ("env_name", pa.string()), - ("task", pa.string()), - ("info", pa.string()), - ("reward", pa.float64()), - ("advantage", pa.float64()), - ("metrics", pa.string()), - ("timing", pa.string()), - ("num_input_tokens", pa.int64()), - ("num_output_tokens", pa.int64()), - ("created_at", pa.timestamp("us", tz="UTC")), - ] -) + +def _base_url() -> str | None: + """$PRIME_API_BASE historically points at the RFT API root + (``.../api/v1/rft``); the SDK takes the platform base URL and appends + ``/api/v1`` itself, so strip what it re-adds. Unset means the SDK + resolves it (``~/.prime/config.json``, then the production default).""" + base = os.getenv(BASE_URL_VAR) + if not base: + return None + return base.rstrip("/").removesuffix("/rft").removesuffix("/api/v1") class PrimeMonitor(Monitor): - """Logs metrics and episodes to the Prime platform. + """Logs metrics and episodes to the Prime platform through ``prime_runs``. - Uploads are fire-and-forget tasks on the caller's event loop — the prime - monitor only runs in the orchestrator, whose call sites are all async. The - platform calls and the episode serialization run in worker threads via - ``asyncio.to_thread``, so they never stall the loop. + The run handle owns what ``TrainRun`` used to do by hand: the RFT + lifecycle (register or attach, finalize), the per-step metrics POSTs, the + every-10th-step Parquet sample uploads (presign -> PUT -> confirm), and + the terminal status — a process that exits without finalizing is reported + crashed by the SDK's atexit hook, replacing the old ``_mark_failed`` one. + + ``init``/``finish`` do network I/O and run in worker threads; the log + calls are queue puts onto the SDK's uploader thread, which owns retries + and backpressure, so they never stall the loop. """ config: PrimeMonitorConfig + run: pr.Run async def init(self, config: BaseConfig | None = None) -> None: - api_key = os.getenv(API_KEY_VAR) or PrimeConfig().api_key - if not api_key: - raise RuntimeError(f"API key not found - set {API_KEY_VAR} or run `prime login`") - self.run = TrainRun(api_key) + init_kwargs: dict[str, Any] if run_id := os.getenv("RUN_ID"): # A managed launch pre-created the platform run and injected its id - # attach instead of registering a duplicate. The backend owns the run's - # failure marking then; finalize still marks it completed on clean exit. - self.run.id = run_id - self.logger.info(f"Logging metrics and episodes to platform run {run_id} (attached via $RUN_ID)") - return - run_fields: dict[str, Any] = {} - if config is not None: - run_fields = dict( - base_model=config.model.name, - max_steps=config.max_steps or 0, - batch_size=config.batch_size, - rollouts_per_example=config.group_size, - seq_len=config.seq_len, + # failure marking then; a clean finish() still marks it completed. + init_kwargs = {"id": run_id} + elif config is not None: + init_kwargs = dict( + name=self.config.name, + model=config.model.name, environments=[env.env_id for env in config.train.source], - run_config=config.model_dump(exclude_none=True, mode="json"), - wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, + training=pr.TrainingSpec( + max_steps=config.max_steps or 0, + batch_size=config.batch_size, + rollouts_per_example=config.group_size, + seq_len=config.seq_len, + wandb_project=config.monitors.wandb.project if config.monitors.wandb else None, + ), + config=config.model_dump(exclude_none=True, mode="json"), ) - await self.run.create(name=self.config.name, **run_fields) + else: + # The RFT API requires a base model; "unknown" is what the + # pre-SDK register sent for a config-less init. + init_kwargs = {"name": self.config.name, "model": "unknown"} + + # A configured monitor must work (see monitors.setup), so mode="online": + # a missing key ("set PRIME_API_KEY or run `prime login`") or a team + # outside the external-runs allowlist raises here instead of training + # silently untracked. $PRIME_RUNS_MODE=disabled stays the explicit + # opt-out, honoured because mode="online" would override it. + disabled = os.getenv(pr.MODE_ENV, "").strip().lower() == "disabled" + self.run = await asyncio.to_thread( + pr.init, + kind="train", + mode="disabled" if disabled else "online", + base_url=_base_url(), + **init_kwargs, + ) + if self.run.url: + attached = " (attached via $RUN_ID)" if self.run.attached else "" + self.logger.info(f"Logging metrics and episodes to platform run {self.run.id} ({self.run.url}){attached}") + else: + self.logger.info(f"Platform run disabled ({pr.MODE_ENV}=disabled)") async def log_metrics(self, metrics: dict[str, Any], step: int | None) -> None: + # The SDK also drops non-finite values, but silently; sanitize first so + # the dropped paths are named in the log. metrics, dropped = sanitize(metrics) if dropped: self.logger.warning(f"Dropping {len(dropped)} non-finite metric value(s): {', '.join(dropped[:5])}") - # every monitor stamps its own wall time; without it, step=None rows - # (e.g. inference metrics) reach the platform with no time anchor - metrics["_timestamp"] = time.time() - self.run.submit("metrics upload", self.run.log_metrics(metrics)) + # A queue put. The SDK stamps `_timestamp` on every row, so step=None + # rows (e.g. inference metrics) 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: - """Upload one platform sample per episode via the presigned-URL Parquet flow. - Only the trained cohort ships to the platform.""" - # Upload every 10th step, unsampled - the pre-refactor cadence - to not - # overwhelm the platform's ingestion. TODO: Lift once we integrate the - # prime traces SDK. - if kind != "train" or subset != "effective" or not episodes or step % 10 != 0: + """Only the trained cohort ships to the platform. The upload cadence + (every 10th step) and the Parquet encoding live in the SDK's training + samples sink, which reads each episode's dispatch step off ``run.work`` + - the ``TrainRunInfo`` the dispatcher stamps at emit time.""" + if kind != "train" or subset != "effective" or not episodes: return - - async def upload() -> None: - # Serialization dumps every episode's full model - heavy pure-Python work - # that would stall the event loop (and with it dispatch) if run inline. - parquet_bytes = await asyncio.to_thread(episodes_to_parquet_bytes, episodes, self.run.id, step) - if parquet_bytes is not None: - await self.run.upload_samples(parquet_bytes, step) - - self.run.submit(f"episodes upload at step {step}", upload()) + # A queue put that can block briefly under backpressure - off the loop. + await asyncio.to_thread(self.run.log_episodes, episodes) async def finalize(self) -> None: - await self.run.finalize() - - -class TrainRun: - """A training run on the Prime platform's RFT API. - - Owns the HTTP client and the run lifecycle (create, log, finalize) — the - natural seam to be subsumed by the train SDK, and it mirrors ``wandb.Run``'s - exit behavior: a created run that is never finalized is marked failed at - process exit via an atexit hook that ``finalize`` disarms. Fully async, - except the atexit hook itself — at interpreter shutdown there is no event - loop, so it sends one synchronous request. - """ - - def __init__(self, api_key: str): - self.id: str | None = None - self.logger = get_logger() - self._tasks: set[asyncio.Task] = set() - self.base_url = (os.getenv(BASE_URL_VAR) or BASE_URL).rstrip("/") - self.headers = { - "Authorization": f"Bearer {api_key}", - "x-api-key": api_key, - "Content-Type": "application/json", - } - self.client = httpx.AsyncClient( - base_url=self.base_url, - headers=self.headers, - timeout=30, - transport=httpx.AsyncHTTPTransport(retries=3), - ) - - def submit(self, what: str, request: Coroutine[Any, Any, None]) -> None: - """Run a request as a fire-and-forget task; a failure only warns. The task set - keeps strong references - the loop alone won't.""" - - async def guarded() -> None: - try: - await request - except Exception as e: - self.logger.warning(f"Failed {what}: {type(e).__name__}: {e}") - - task = asyncio.get_running_loop().create_task(guarded()) - self._tasks.add(task) - task.add_done_callback(self._tasks.discard) - - async def create( - self, - name: str | None = None, - team_id: str | None = None, - base_model: str = "unknown", - max_steps: int = 0, - batch_size: int | None = None, - rollouts_per_example: int | None = None, - seq_len: int | None = None, - environments: list[str] | None = None, - run_config: dict[str, Any] | None = None, - wandb_project: str | None = None, - ) -> str: - """Register the run with the platform and return its id.""" - prime_config = PrimeConfig() - team_id = team_id or prime_config.team_id - - payload: dict[str, Any] = {"base_model": base_model, "max_steps": max_steps} - if batch_size is not None: - payload["batch_size"] = batch_size - if rollouts_per_example is not None: - payload["rollouts_per_example"] = rollouts_per_example - if seq_len is not None: - payload["seq_len"] = seq_len - if environments is not None: - payload["environments"] = [{"id": env_id} for env_id in environments] - if run_config is not None: - payload["run_config"] = run_config - if wandb_project is not None: - payload["wandb_project"] = wandb_project - if name: - payload["name"] = name - if team_id: - payload["team_id"] = team_id - - response = await self.client.post("/external-runs", json=payload) - if response.status_code != 201: - raise RuntimeError(f"Failed to create platform run (HTTP {response.status_code}): {response.text}") - - self.id = response.json()["run"]["id"] - self._owner_pid = os.getpid() - atexit.register(self._mark_failed) - if prime_config.frontend_url: - self.logger.info( - f"Logging metrics and episodes to platform run {self.id} ({prime_config.frontend_url.rstrip('/')}/dashboard/training/{self.id})" - ) - else: - self.logger.info(f"Logging metrics and episodes to platform run {self.id}") - return self.id - - async def log_metrics(self, metrics: dict[str, Any]) -> None: - (await self.client.post("/metrics", json={"run_id": self.id, "metrics": metrics})).raise_for_status() - - async def upload_samples(self, parquet_bytes: bytes, step: int) -> None: - """Presigned-URL flow: presign -> R2 PUT -> confirm.""" - presign = await self.client.post("/samples/presign", json={"run_id": self.id, "step": step}) - presign.raise_for_status() - data = presign.json()["data"] - # Bare client - the presigned URL rejects the run client's auth headers. - async with httpx.AsyncClient(timeout=30) as client: - put = await client.put( - data["presignedUrl"], content=parquet_bytes, headers={"Content-Type": "application/parquet"} - ) - put.raise_for_status() - confirm = await self.client.post( - "/samples/confirm", json={"run_id": self.id, "step": step, "s3_key": data["s3Key"]} - ) - confirm.raise_for_status() - - async def finalize(self) -> None: - """Finalize the run as completed.""" - self.logger.info(f"Finalizing platform run {self.id}") - # Drain in-flight uploads so the final step's metrics and episodes land - # before the run is marked completed. - await asyncio.gather(*self._tasks, return_exceptions=True) - try: - (await self.client.post("/finalize", json={"run_id": self.id, "summary": {}})).raise_for_status() - except httpx.HTTPError as e: - self.logger.warning(f"Failed to finalize platform run {self.id}: {e}") - await self.set_status(success=True) - atexit.unregister(self._mark_failed) - - def _mark_failed(self) -> None: - # Forked children inherit the atexit table; only the creating process may - # flip the run's status. At interpreter shutdown there is no event loop and - # no executor for async DNS, so this path must stay synchronous. - if os.getpid() != self._owner_pid: - return - self.logger.info(f"Marking platform run {self.id} as failed") - try: - httpx.put( - f"{self.base_url}/external-runs/{self.id}/status", - headers=self.headers, - json={"status": "failed"}, - timeout=30, - ).raise_for_status() - except httpx.HTTPError as e: - self.logger.warning(f"Failed to mark platform run {self.id} as failed: {e}") - - async def set_status(self, success: bool) -> None: - """Mark the run as completed or failed.""" - status = "completed" if success else "failed" - self.logger.info(f"Marking platform run {self.id} as {status}") - try: - put = await self.client.put(f"/external-runs/{self.id}/status", json={"status": status}) - put.raise_for_status() - except httpx.HTTPError as e: - self.logger.warning(f"Failed to mark platform run {self.id} as {status}: {e}") - - -def episodes_to_parquet_bytes(episodes: list[vf.Episode], run_id: str | None, step: int) -> bytes | None: - """One row per episode. Sample construction is shared with verifiers' eval - ``--push`` (``build_samples``: complete native episode in ``info.native_wrapper``, - flat summary from one trainable trace), so a training episode and an eval sample - land on the platform identically; the RFT-only columns (run/step/advantage/ - problem_id/env_name) are layered on here.""" - advantages: dict[str, float | None] = {} - env_names: dict[str, str] = {} - for episode in episodes: - summary_trace = next((trace for trace in episode.traces if trace.agent.trainable), episode.traces[0]) - advantages[episode.id] = summary_trace.info.get("advantage") - env_names[episode.id] = episode.env.id - - now = datetime.now(timezone.utc) - rows = [] - for sample_id, sample in enumerate(build_samples(episodes)): - trajectory = sample["trajectory"] - if not trajectory: # no branches (e.g. an episode that errored before any message) - continue - advantage = advantages.get(sample["episode_id"]) - trajectory = [{**branch, "advantage": advantage} for branch in trajectory] - - try: - problem_id = int(sample["example_id"]) if sample["example_id"] is not None else sample_id - except (TypeError, ValueError): - problem_id = sample_id - - rows.append( - { - "run_id": run_id, - "step": step, - "tag": "", - "problem_id": problem_id, - "sample_id": sample_id, - "prompt": "", - "completion": json.dumps(sample["completion"]), - "trajectory": json.dumps(trajectory), - "answer": "", - "env_name": env_names.get(sample["episode_id"], ""), - "task": json.dumps(sample["task"]), - "info": json.dumps(sample["info"]), - "reward": sample["reward"], - "advantage": advantage, - "metrics": json.dumps(sample["metrics"]), - "timing": json.dumps(sample["timing"]), - "num_input_tokens": trajectory[-1]["num_input_tokens"], - "num_output_tokens": trajectory[-1]["num_output_tokens"], - "created_at": now, - } - ) - - if not rows: - return None - - table = pa.Table.from_pylist(rows, schema=SAMPLE_SCHEMA) - buf = io.BytesIO() - pq.write_table(table, buf, compression="snappy", use_dictionary=True, write_statistics=True) - return buf.getvalue() + # Drains queued uploads so the final step's metrics and episodes land, + # then finalizes (idempotent on the platform side); an attached run's + # failure marking stays with the launcher. + await asyncio.to_thread(self.run.finish) diff --git a/uv.lock b/uv.lock index eece4907bb..bfc855fc24 100644 --- a/uv.lock +++ b/uv.lock @@ -20,12 +20,14 @@ exclude-newer-span = "P7D" vllm = false verifiers = false harbor = "2026-08-11T00:00:00Z" +prime-traces = false prime-pydantic-config = false vllm-router = false dion = false -fastokens = false +prime-runs = false flash-attn-3 = false prime-tunnel = false +fastokens = false deep-gemm = false prime-evals = false torchao = false @@ -4951,6 +4953,7 @@ dependencies = [ { name = "prime" }, { name = "prime-pydantic-config" }, { name = "prime-rl-configs" }, + { name = "prime-runs", extra = ["train"] }, { name = "psutil" }, { name = "pyarrow" }, { name = "pybase64" }, @@ -5113,6 +5116,7 @@ requires-dist = [ { name = "prime-rl", extras = ["gpu"], marker = "extra == 'all'" }, { name = "prime-rl", extras = ["quack"], marker = "extra == 'all'" }, { name = "prime-rl-configs", editable = "packages/prime-rl-configs" }, + { name = "prime-runs", extras = ["train"], git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-runs&branch=feature%2Fprime-runs-train-backend" }, { name = "psutil", specifier = ">=7.0.0" }, { name = "pyarrow", specifier = ">=21.0.0" }, { name = "pybase64", specifier = ">=1.4.2" }, @@ -5182,6 +5186,20 @@ requires-dist = [ { name = "verifiers", specifier = ">=0.3.0" }, ] +[[package]] +name = "prime-runs" +version = "0.1.0" +source = { git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-runs&branch=feature%2Fprime-runs-train-backend#b69d3da61077ca4cd10d348821d73d58daeb0f98" } +dependencies = [ + { name = "httpx" }, + { name = "prime-traces" }, +] + +[package.optional-dependencies] +train = [ + { name = "pyarrow" }, +] + [[package]] name = "prime-sandboxes" version = "0.2.39" @@ -5201,6 +5219,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/db/a5/8a2a1aefecfc58ca926d6b21e8302a127884bf2a45952ed390b5f702eeda/prime_sandboxes-0.2.39-py3-none-any.whl", hash = "sha256:3ea355158473c1697f9ef6ec2a07f3189e9a2555c80d1c88e1f7e22979772e0a", size = 58161, upload-time = "2026-08-19T23:09:26.69Z" }, ] +[[package]] +name = "prime-traces" +version = "0.0.3" +source = { git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-traces&branch=feature%2Fprime-runs-train-backend#b69d3da61077ca4cd10d348821d73d58daeb0f98" } +dependencies = [ + { name = "httpx" }, + { name = "pydantic" }, +] + [[package]] name = "prime-tunnel" version = "0.1.10" From 8d7837e36a1acea00f369a60cdfb53c743fa5d40 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Tue, 1 Sep 2026 10:27:09 -0700 Subject: [PATCH 2/5] chore: update prime-runs --- pyproject.toml | 2 +- uv.lock | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2490c6100c..e84eb2e32a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -241,7 +241,7 @@ prime-kernels = false [tool.uv.sources] # TEMPORARY: prime-runs is not on PyPI yet (PrimeIntellect-ai/prime #856 + #873) -prime-runs = { git = "https://github.com/PrimeIntellect-ai/prime.git", subdirectory = "packages/prime-runs", branch = "feature/prime-runs-train-backend" } +prime-runs = { git = "https://github.com/PrimeIntellect-ai/prime.git", subdirectory = "packages/prime-runs", rev = "9cb688b" } prime-rl-configs = { path = "packages/prime-rl-configs", editable = true } # prime-rl consumes `verifiers.v1` (incl. the built-in harnesses/tasksets); env tasksets # live under deps/{verifiers,prime-envs}/environments. diff --git a/uv.lock b/uv.lock index bfc855fc24..9f126adbb1 100644 --- a/uv.lock +++ b/uv.lock @@ -5116,7 +5116,7 @@ requires-dist = [ { name = "prime-rl", extras = ["gpu"], marker = "extra == 'all'" }, { name = "prime-rl", extras = ["quack"], marker = "extra == 'all'" }, { name = "prime-rl-configs", editable = "packages/prime-rl-configs" }, - { name = "prime-runs", extras = ["train"], git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-runs&branch=feature%2Fprime-runs-train-backend" }, + { name = "prime-runs", extras = ["train"], git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-runs&rev=9cb688b" }, { name = "psutil", specifier = ">=7.0.0" }, { name = "pyarrow", specifier = ">=21.0.0" }, { name = "pybase64", specifier = ">=1.4.2" }, @@ -5189,7 +5189,7 @@ requires-dist = [ [[package]] name = "prime-runs" version = "0.1.0" -source = { git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-runs&branch=feature%2Fprime-runs-train-backend#b69d3da61077ca4cd10d348821d73d58daeb0f98" } +source = { git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-runs&rev=9cb688b#9cb688b86479d49566be7c303be3be79816eaa95" } dependencies = [ { name = "httpx" }, { name = "prime-traces" }, @@ -5222,7 +5222,7 @@ wheels = [ [[package]] name = "prime-traces" version = "0.0.3" -source = { git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-traces&branch=feature%2Fprime-runs-train-backend#b69d3da61077ca4cd10d348821d73d58daeb0f98" } +source = { git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-traces&rev=9cb688b#9cb688b86479d49566be7c303be3be79816eaa95" } dependencies = [ { name = "httpx" }, { name = "pydantic" }, From d1eeeea22e2420f6869d3e1fa1c2719fe40bd322 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Tue, 1 Sep 2026 17:45:08 -0700 Subject: [PATCH 3/5] monitors: pin the current prime-runs, adopt the platform run id, add tests - pyproject/uv.lock: the temporary git source points at the SDK branch head (aba26790) instead of a commit that only a merged PR ref still reaches; it picks up the per-upload sample id ranges, the training sink cooldown, replayable metrics and null handling for non-finite values. Dropped once 0.1.0 is on PyPI. - orchestrator: with an online platform run, episodes carry the platform's run id, so the run's traces can be queried by the id the dashboard shows. W&B keeps the launcher's PRL_RUN_ID either way. - monitor: log_metrics hops off the loop like log_episodes (a queue put can block under backpressure); the PRIME_RUNS_MODE value passes straight through; _base_url only strips /rft (the SDK strips /api/v1); a 60 s finish_timeout bounds the drain on finish and on the atexit crash path. - docs/training.md: the platform-monitoring section describes the SDK, the team requirement, PRIME_RUNS_MODE=disabled and RUN_ID attach. - tests/unit/monitors: first coverage of the monitor (init kwargs, attach, disabled switch, base URL). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Fk5fBAAaTiPfJ5UpNyqxXR --- docs/training.md | 6 +- pyproject.toml | 4 +- src/prime_rl/monitors/prime.py | 33 ++++--- src/prime_rl/orchestrator/orchestrator.py | 11 ++- tests/unit/monitors/__init__.py | 0 tests/unit/monitors/test_prime.py | 102 ++++++++++++++++++++++ uv.lock | 6 +- 7 files changed, 135 insertions(+), 27 deletions(-) create mode 100644 tests/unit/monitors/__init__.py create mode 100644 tests/unit/monitors/test_prime.py diff --git a/docs/training.md b/docs/training.md index d34a446ce1..eb0ddbecee 100644 --- a/docs/training.md +++ b/docs/training.md @@ -371,7 +371,7 @@ prime-rl deliberately logs a **large number of metrics** for maximum observabili ### Platform Monitoring -Register a run on the Prime Intellect platform (Prime Lab) and stream training metrics and episodes to the platform dashboard. Bare flag uses defaults: +Register a run on the Prime Intellect platform and stream training metrics and episodes to its dashboard. Bare flag uses defaults: ```bash uv run rl @ rl.toml --monitors.prime @@ -384,9 +384,9 @@ Or set it in TOML: name = "my-experiment" ``` -Every 10th step the orchestrator uploads the step's episodes (full conversations with rewards and advantages) to the run's sample viewer. +The monitor is a thin layer over the [`prime-runs`](https://github.com/PrimeIntellect-ai/prime/tree/main/packages/prime-runs) SDK (installed as `prime-runs[train]`): it registers the run, streams per-step metrics, uploads every 10th step's episodes (full conversations with rewards and advantages) to the run's sample viewer, and closes the run out. A process that exits without finishing is reported as crashed. Episodes are stamped with the platform run's id, so the run's traces can be queried by the id the dashboard shows; W&B keeps the launcher's `PRL_RUN_ID`. -Requires `PRIME_API_KEY` (set via `prime login` or env var) and an allowlisted team. Currently internal-only. +Requires `PRIME_API_KEY` (`prime login` or the env var) and a team (`PRIME_TEAM_ID`, or the team selected with `prime login`) enabled for external runs. A configured monitor must work: a missing key or a team outside the allowlist fails the launch. `PRIME_RUNS_MODE=disabled` keeps the monitor configured but opens no platform run; `RUN_ID=` attaches to an external run a launcher already created instead of registering a new one. Currently internal-only. ## Rules of Thumb diff --git a/pyproject.toml b/pyproject.toml index fdd3f12758..bbfdaec1dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -240,8 +240,8 @@ torchao = false prime-kernels = false [tool.uv.sources] -# TEMPORARY: prime-runs is not on PyPI yet (PrimeIntellect-ai/prime #856 + #873) -prime-runs = { git = "https://github.com/PrimeIntellect-ai/prime.git", subdirectory = "packages/prime-runs", rev = "9cb688b" } +# TEMPORARY: prime-runs is not on PyPI yet (PrimeIntellect-ai/prime #856); drop once 0.1.0 is published +prime-runs = { git = "https://github.com/PrimeIntellect-ai/prime.git", subdirectory = "packages/prime-runs", rev = "aba26790" } prime-rl-configs = { path = "packages/prime-rl-configs", editable = true } # prime-rl consumes `verifiers.v1` (incl. the built-in harnesses/tasksets); env tasksets # live under deps/{verifiers,prime-envs}/environments. diff --git a/src/prime_rl/monitors/prime.py b/src/prime_rl/monitors/prime.py index ae801c8198..27dfb48147 100644 --- a/src/prime_rl/monitors/prime.py +++ b/src/prime_rl/monitors/prime.py @@ -15,17 +15,17 @@ import verifiers.v1 as vf BASE_URL_VAR = "PRIME_API_BASE" +# How long finish() and the SDK's atexit crash hook let queued uploads drain. The SDK +# default (300 s) is sized for eval sample batches; a crashed training process should +# not linger that long, and a clean finish rarely has more than the last step queued. +FINISH_TIMEOUT = 60.0 def _base_url() -> str | None: - """$PRIME_API_BASE historically points at the RFT API root - (``.../api/v1/rft``); the SDK takes the platform base URL and appends - ``/api/v1`` itself, so strip what it re-adds. Unset means the SDK - resolves it (``~/.prime/config.json``, then the production default).""" + """$PRIME_API_BASE historically points at the RFT API root (``.../api/v1/rft``); + the SDK takes the platform base URL. Unset means the SDK resolves it.""" base = os.getenv(BASE_URL_VAR) - if not base: - return None - return base.rstrip("/").removesuffix("/rft").removesuffix("/api/v1") + return base.rstrip("/").removesuffix("/rft") if base else None class PrimeMonitor(Monitor): @@ -71,17 +71,16 @@ async def init(self, config: BaseConfig | None = None) -> None: # pre-SDK register sent for a config-less init. init_kwargs = {"name": self.config.name, "model": "unknown"} - # A configured monitor must work (see monitors.setup), so mode="online": - # a missing key ("set PRIME_API_KEY or run `prime login`") or a team - # outside the external-runs allowlist raises here instead of training - # silently untracked. $PRIME_RUNS_MODE=disabled stays the explicit - # opt-out, honoured because mode="online" would override it. - disabled = os.getenv(pr.MODE_ENV, "").strip().lower() == "disabled" + # A configured monitor must work (see monitors.setup), so the default is + # mode="online": a missing key or a team outside the external-runs allowlist + # raises here instead of training silently untracked. $PRIME_RUNS_MODE=disabled + # stays the explicit opt-out. self.run = await asyncio.to_thread( pr.init, kind="train", - mode="disabled" if disabled else "online", + mode=os.getenv(pr.MODE_ENV) or "online", base_url=_base_url(), + finish_timeout=FINISH_TIMEOUT, **init_kwargs, ) if self.run.url: @@ -96,9 +95,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. The SDK stamps `_timestamp` on every row, so step=None - # rows (e.g. inference metrics) keep a time anchor. - self.run.log_metrics(metrics, step=step) + # 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) async def log_episodes(self, episodes: list[vf.Episode], step: int, kind: Kind, subset: Subset) -> None: """Only the trained cohort ships to the platform. The upload cadence diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 7f22b5698c..33256987a6 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -220,8 +220,15 @@ async def setup(self) -> None: train_env_names=[env.resolved_name for env in config.train.source], eval_env_names=[source.resolved_name for source in config.eval.source] if config.eval is not None else [], ) - # The launcher-set $PRL_RUN_ID is the run identity; standalone runs mint a local one. - self.run_id = os.environ.get("PRL_RUN_ID") or uuid.uuid4().hex + # The run identity every episode carries. With a platform run it is the platform's id, + # so Prime Traces can be queried by the run the dashboard shows; otherwise the + # launcher-set $PRL_RUN_ID (W&B keeps that either way), or a local one when standalone. + prime = monitors.get(monitors.PrimeMonitor) + self.run_id = ( + prime.run.id + if isinstance(prime, monitors.PrimeMonitor) and prime.run.mode == "online" + else os.environ.get("PRL_RUN_ID") or uuid.uuid4().hex + ) # Base labels for sandboxes created in this process; env-server processes read # the same launcher-set env var themselves. self.run_name = os.environ.get("PRL_RUN_NAME") diff --git a/tests/unit/monitors/__init__.py b/tests/unit/monitors/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/unit/monitors/test_prime.py b/tests/unit/monitors/test_prime.py new file mode 100644 index 0000000000..120fdf1a34 --- /dev/null +++ b/tests/unit/monitors/test_prime.py @@ -0,0 +1,102 @@ +import asyncio +from types import SimpleNamespace + +import prime_runs as pr +import pytest + +from prime_rl.configs.monitors import PrimeMonitorConfig +from prime_rl.monitors.prime import FINISH_TIMEOUT, PrimeMonitor, _base_url + + +@pytest.fixture +def init_calls(monkeypatch): + """Record what the monitor asks the SDK for; hand back a disabled run.""" + real_init, calls, runs = pr.init, [], [] + + def fake_init(**kwargs): + calls.append(kwargs) + runs.append(real_init(kind="train", mode="disabled", id=kwargs.get("id"), model="m")) + return runs[-1] + + monkeypatch.setattr(pr, "init", fake_init) + monkeypatch.delenv("RUN_ID", raising=False) + monkeypatch.delenv(pr.MODE_ENV, raising=False) + yield calls + for run in runs: # else the SDK's atexit hook reports them crashed + run.finish() + + +class FakeConfig(SimpleNamespace): + def model_dump(self, **kwargs): + return {"max_steps": self.max_steps} + + +def orchestrator_config(wandb=None): + return FakeConfig( + model=SimpleNamespace(name="Qwen/Qwen3-8B"), + train=SimpleNamespace(source=[SimpleNamespace(env_id="primeintellect/gsm8k")]), + max_steps=100, + batch_size=64, + group_size=8, + seq_len=4096, + monitors=SimpleNamespace(wandb=wandb), + ) + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + (None, None), + ("https://api.primeintellect.ai/api/v1/rft", "https://api.primeintellect.ai/api/v1"), + ("https://api.primeintellect.ai/api/v1/rft/", "https://api.primeintellect.ai/api/v1"), + ("https://api.primeintellect.ai", "https://api.primeintellect.ai"), + ], +) +def test_base_url_strips_the_rft_root(monkeypatch, value, expected): + if value is None: + monkeypatch.delenv("PRIME_API_BASE", raising=False) + else: + monkeypatch.setenv("PRIME_API_BASE", value) + assert _base_url() == expected + + +def test_init_registers_the_run_from_the_orchestrator_config(init_calls): + monitor = PrimeMonitor(PrimeMonitorConfig(name="exp")) + + asyncio.run(monitor.init(orchestrator_config(wandb=SimpleNamespace(project="proj")))) + + assert init_calls == [ + dict( + kind="train", + mode="online", + base_url=None, + finish_timeout=FINISH_TIMEOUT, + name="exp", + model="Qwen/Qwen3-8B", + environments=["primeintellect/gsm8k"], + training=pr.TrainingSpec( + max_steps=100, batch_size=64, rollouts_per_example=8, seq_len=4096, wandb_project="proj" + ), + config={"max_steps": 100}, + ) + ] + assert monitor.run.kind == "train" + + +def test_init_attaches_to_the_launcher_s_run(init_calls, monkeypatch): + monkeypatch.setenv("RUN_ID", "run-managed") + monitor = PrimeMonitor(PrimeMonitorConfig()) + + asyncio.run(monitor.init(orchestrator_config())) + + (call,) = init_calls + assert call["id"] == "run-managed" and "model" not in call and "training" not in call + assert monitor.run.id == "run-managed" + + +def test_the_disabled_switch_reaches_the_sdk(init_calls, monkeypatch): + monkeypatch.setenv(pr.MODE_ENV, "disabled") + + asyncio.run(PrimeMonitor(PrimeMonitorConfig()).init(orchestrator_config())) + + assert init_calls[0]["mode"] == "disabled" diff --git a/uv.lock b/uv.lock index a54dcbc466..efc2b74f77 100644 --- a/uv.lock +++ b/uv.lock @@ -5115,7 +5115,7 @@ requires-dist = [ { name = "prime-rl", extras = ["gpu"], marker = "extra == 'all'" }, { name = "prime-rl", extras = ["quack"], marker = "extra == 'all'" }, { name = "prime-rl-configs", editable = "packages/prime-rl-configs" }, - { name = "prime-runs", extras = ["train"], git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-runs&rev=9cb688b" }, + { name = "prime-runs", extras = ["train"], git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-runs&rev=aba26790" }, { name = "psutil", specifier = ">=7.0.0" }, { name = "pyarrow", specifier = ">=21.0.0" }, { name = "pybase64", specifier = ">=1.4.2" }, @@ -5188,7 +5188,7 @@ requires-dist = [ [[package]] name = "prime-runs" version = "0.1.0" -source = { git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-runs&rev=9cb688b#9cb688b86479d49566be7c303be3be79816eaa95" } +source = { git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-runs&rev=aba26790#aba26790514749fe484e995f99591fd0d35e2350" } dependencies = [ { name = "httpx" }, { name = "prime-traces" }, @@ -5221,7 +5221,7 @@ wheels = [ [[package]] name = "prime-traces" version = "0.0.3" -source = { git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-traces&rev=9cb688b#9cb688b86479d49566be7c303be3be79816eaa95" } +source = { git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-traces&rev=aba26790#aba26790514749fe484e995f99591fd0d35e2350" } dependencies = [ { name = "httpx" }, { name = "pydantic" }, From 33f842c209f811792825b9c725375aec6928b68a Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Tue, 1 Sep 2026 17:50:59 -0700 Subject: [PATCH 4/5] docs: keep the platform-monitoring intro line as it was Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Fk5fBAAaTiPfJ5UpNyqxXR --- docs/training.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/training.md b/docs/training.md index eb0ddbecee..caa852a600 100644 --- a/docs/training.md +++ b/docs/training.md @@ -371,7 +371,7 @@ prime-rl deliberately logs a **large number of metrics** for maximum observabili ### Platform Monitoring -Register a run on the Prime Intellect platform and stream training metrics and episodes to its dashboard. Bare flag uses defaults: +Register a run on the Prime Intellect platform (Prime Lab) and stream training metrics and episodes to the platform dashboard. Bare flag uses defaults: ```bash uv run rl @ rl.toml --monitors.prime From b00065be3348b44ba7cdef28a7edc64e0676306f Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 2 Sep 2026 15:02:09 -0700 Subject: [PATCH 5/5] orchestrator: keep PRL_RUN_ID as the run identity; the SDK keys uploads Adopting the platform run id for episode provenance coupled the orchestrator's identity to which monitor came up online. prime-runs now sets `run.id`/`run.type` on the uploaded copy of every episode (keeping `work`, the dispatch step), so Prime Traces documents carry the dashboard's id without the orchestrator learning it, and `self.run_id` is back to the launcher's `PRL_RUN_ID` (or a local uuid) as on main. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01FHMjWstGWd1oGKjoj5RgeT --- docs/training.md | 2 +- src/prime_rl/orchestrator/orchestrator.py | 11 ++--------- 2 files changed, 3 insertions(+), 10 deletions(-) diff --git a/docs/training.md b/docs/training.md index caa852a600..88fc9a41cf 100644 --- a/docs/training.md +++ b/docs/training.md @@ -384,7 +384,7 @@ Or set it in TOML: name = "my-experiment" ``` -The monitor is a thin layer over the [`prime-runs`](https://github.com/PrimeIntellect-ai/prime/tree/main/packages/prime-runs) SDK (installed as `prime-runs[train]`): it registers the run, streams per-step metrics, uploads every 10th step's episodes (full conversations with rewards and advantages) to the run's sample viewer, and closes the run out. A process that exits without finishing is reported as crashed. Episodes are stamped with the platform run's id, so the run's traces can be queried by the id the dashboard shows; W&B keeps the launcher's `PRL_RUN_ID`. +The monitor is a thin layer over the [`prime-runs`](https://github.com/PrimeIntellect-ai/prime/tree/main/packages/prime-runs) SDK (installed as `prime-runs[train]`): it registers the run, streams per-step metrics, uploads every 10th step's episodes (full conversations with rewards and advantages) to the run's sample viewer, and closes the run out. A process that exits without finishing is reported as crashed. Uploaded episodes are keyed to the platform run by the SDK; the orchestrator's own run id (the launcher's `PRL_RUN_ID`) stays on W&B and in the local records. Requires `PRIME_API_KEY` (`prime login` or the env var) and a team (`PRIME_TEAM_ID`, or the team selected with `prime login`) enabled for external runs. A configured monitor must work: a missing key or a team outside the allowlist fails the launch. `PRIME_RUNS_MODE=disabled` keeps the monitor configured but opens no platform run; `RUN_ID=` attaches to an external run a launcher already created instead of registering a new one. Currently internal-only. diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 33256987a6..7f22b5698c 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -220,15 +220,8 @@ async def setup(self) -> None: train_env_names=[env.resolved_name for env in config.train.source], eval_env_names=[source.resolved_name for source in config.eval.source] if config.eval is not None else [], ) - # The run identity every episode carries. With a platform run it is the platform's id, - # so Prime Traces can be queried by the run the dashboard shows; otherwise the - # launcher-set $PRL_RUN_ID (W&B keeps that either way), or a local one when standalone. - prime = monitors.get(monitors.PrimeMonitor) - self.run_id = ( - prime.run.id - if isinstance(prime, monitors.PrimeMonitor) and prime.run.mode == "online" - else os.environ.get("PRL_RUN_ID") or uuid.uuid4().hex - ) + # The launcher-set $PRL_RUN_ID is the run identity; standalone runs mint a local one. + self.run_id = os.environ.get("PRL_RUN_ID") or uuid.uuid4().hex # Base labels for sandboxes created in this process; env-server processes read # the same launcher-set env var themselves. self.run_name = os.environ.get("PRL_RUN_NAME")