From 0dcb62a8167f98c71180bb0a7925d6780e1eb829 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 13:52:53 -0700 Subject: [PATCH 01/39] feat(runs): add the prime-runs SDK for eval and training runs prime-evals models a run as three stateless calls over Dict[str, Any], so neither of its intended consumers uses it on their main path: verifiers reimplemented create -> batch -> finalize inline in v1/utils/platform.py, and prime-rl hits a different API family entirely from utils/monitor/prime.py. Producers have objects, a long-running loop, steps, ranks, forks and crashes; each of them solved that privately, twice. prime-runs makes the run an object instead: run = pr.init(name=..., environments=["gsm8k"], model=..., framework=...) run.log_traces([episode]); run.log({"reward": r}, step=step) run.finish(summary=pr.metrics.from_episodes(episodes)) Identity: init() is called before rollouts start and the ID it returns is the run ID everywhere, local archive included. Producers already stamp the run onto their traces, so nothing is re-stamped and no producer record is rewritten. The join key is run.id inside the trace document -- an indexed ClickHouse column with a delete-by-run path -- not an upload-scoped context key. Backends and sinks are independent axes. Backends own lifecycle (EvalsBackend, OfflineBackend); sinks own transport (TracesSink, plus EvalSamplesSink for the viewer's flat table). Both sinks run during the transition, because Prime Traces is gated to an account allowlist and a traces-only client would leave everyone else with an empty dashboard. When the Viewer API reads traces natively, the default sink list drops one entry and no producer changes. The SDK owns the operational work: streaming instead of buffering, a bounded upload queue that drops and counts rather than stalling a training run, fork safety via register_at_fork, contained errors (on_error="warn" by default), terminal status through the context manager / atexit / signals, rank awareness, and an offline mode that is a real run -- which is what lets producers delete their --no-push branching. trace_to_sample / build_samples move here from verifiers: it is knowledge about a platform wire format, and prime-rl currently reaches across a repo boundary to import it from a module path that has already drifted. Known platform gap: there is no producer-facing way to mark an evaluation failed (finalize only goes PROCESSING -> COMPLETED, UpdateEvaluationRequest has no status). EvalsBackend calls the status endpoint it needs, latches on 404 so it probes once, and falls back to recording the terminal state in metadata while warning that the run will keep showing as running. The fallback stops firing on its own once the endpoint ships. Leaf package by construction -- httpx, pydantic, tenacity, prime-traces and nothing else -- because the prime CLI depends on verifiers, so verifiers can never depend on prime. verifiers already takes prime-tunnel and prime-sandboxes on the same terms. 107 tests, hermetic (httpx.MockTransport + tmp dirs). Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/ci.yml | 28 + .github/workflows/release-runs.yml | 78 ++ packages/prime-runs/README.md | 153 ++++ packages/prime-runs/pyproject.toml | 70 ++ .../prime-runs/src/prime_runs/__init__.py | 96 +++ packages/prime-runs/src/prime_runs/_http.py | 231 ++++++ .../src/prime_runs/backends/__init__.py | 13 + .../src/prime_runs/backends/base.py | 68 ++ .../src/prime_runs/backends/evals.py | 259 +++++++ .../src/prime_runs/backends/offline.py | 160 ++++ packages/prime-runs/src/prime_runs/config.py | 75 ++ .../prime-runs/src/prime_runs/exceptions.py | 84 +++ packages/prime-runs/src/prime_runs/metrics.py | 49 ++ packages/prime-runs/src/prime_runs/models.py | 110 +++ .../prime-runs/src/prime_runs/projection.py | 194 +++++ packages/prime-runs/src/prime_runs/py.typed | 0 packages/prime-runs/src/prime_runs/run.py | 704 ++++++++++++++++++ .../src/prime_runs/sinks/__init__.py | 14 + .../prime-runs/src/prime_runs/sinks/base.py | 65 ++ .../src/prime_runs/sinks/offline.py | 90 +++ .../src/prime_runs/sinks/samples.py | 99 +++ .../prime-runs/src/prime_runs/sinks/traces.py | 194 +++++ packages/prime-runs/src/prime_runs/worker.py | 250 +++++++ packages/prime-runs/tests/_fakes.py | 115 +++ packages/prime-runs/tests/conftest.py | 148 ++++ .../prime-runs/tests/test_evals_backend.py | 169 +++++ packages/prime-runs/tests/test_http.py | 110 +++ packages/prime-runs/tests/test_init.py | 231 ++++++ packages/prime-runs/tests/test_metrics.py | 77 ++ packages/prime-runs/tests/test_projection.py | 141 ++++ packages/prime-runs/tests/test_run.py | 295 ++++++++ .../prime-runs/tests/test_samples_sink.py | 61 ++ packages/prime-runs/tests/test_traces_sink.py | 152 ++++ packages/prime-runs/tests/test_worker.py | 167 +++++ pyproject.toml | 2 + uv.lock | 31 + 36 files changed, 4783 insertions(+) create mode 100644 .github/workflows/release-runs.yml create mode 100644 packages/prime-runs/README.md create mode 100644 packages/prime-runs/pyproject.toml create mode 100644 packages/prime-runs/src/prime_runs/__init__.py create mode 100644 packages/prime-runs/src/prime_runs/_http.py create mode 100644 packages/prime-runs/src/prime_runs/backends/__init__.py create mode 100644 packages/prime-runs/src/prime_runs/backends/base.py create mode 100644 packages/prime-runs/src/prime_runs/backends/evals.py create mode 100644 packages/prime-runs/src/prime_runs/backends/offline.py create mode 100644 packages/prime-runs/src/prime_runs/config.py create mode 100644 packages/prime-runs/src/prime_runs/exceptions.py create mode 100644 packages/prime-runs/src/prime_runs/metrics.py create mode 100644 packages/prime-runs/src/prime_runs/models.py create mode 100644 packages/prime-runs/src/prime_runs/projection.py create mode 100644 packages/prime-runs/src/prime_runs/py.typed create mode 100644 packages/prime-runs/src/prime_runs/run.py create mode 100644 packages/prime-runs/src/prime_runs/sinks/__init__.py create mode 100644 packages/prime-runs/src/prime_runs/sinks/base.py create mode 100644 packages/prime-runs/src/prime_runs/sinks/offline.py create mode 100644 packages/prime-runs/src/prime_runs/sinks/samples.py create mode 100644 packages/prime-runs/src/prime_runs/sinks/traces.py create mode 100644 packages/prime-runs/src/prime_runs/worker.py create mode 100644 packages/prime-runs/tests/_fakes.py create mode 100644 packages/prime-runs/tests/conftest.py create mode 100644 packages/prime-runs/tests/test_evals_backend.py create mode 100644 packages/prime-runs/tests/test_http.py create mode 100644 packages/prime-runs/tests/test_init.py create mode 100644 packages/prime-runs/tests/test_metrics.py create mode 100644 packages/prime-runs/tests/test_projection.py create mode 100644 packages/prime-runs/tests/test_run.py create mode 100644 packages/prime-runs/tests/test_samples_sink.py create mode 100644 packages/prime-runs/tests/test_traces_sink.py create mode 100644 packages/prime-runs/tests/test_worker.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2c090d84c..9072262e1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -41,6 +41,7 @@ jobs: "packages/prime-sandboxes/src/prime_sandboxes/__init__.py" "packages/prime-evals/src/prime_evals/__init__.py" "packages/prime-traces/src/prime_traces/__init__.py" + "packages/prime-runs/src/prime_runs/__init__.py" ) BUMPED_FILES="" @@ -160,6 +161,33 @@ jobs: working-directory: packages/prime-traces run: uv run pytest tests/ -v + test-runs: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.11", "3.12", "3.13"] + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + + - name: Install uv + uses: astral-sh/setup-uv@v4 + with: + version: "latest" + + - name: Install dependencies + working-directory: packages/prime-runs + run: uv sync --all-extras + + # No secrets: the runs tests are hermetic (httpx.MockTransport + tmp dirs). + - name: Run tests + working-directory: packages/prime-runs + run: uv run pytest tests/ -v + test-prime: runs-on: ubuntu-latest strategy: diff --git a/.github/workflows/release-runs.yml b/.github/workflows/release-runs.yml new file mode 100644 index 000000000..432f4b57c --- /dev/null +++ b/.github/workflows/release-runs.yml @@ -0,0 +1,78 @@ +name: Release prime-runs + +on: + push: + branches: [ "main" ] + paths: + - 'packages/prime-runs/**' + - '.github/workflows/release-runs.yml' + workflow_dispatch: + inputs: + force_release: + description: 'Force release even if tag exists' + required: false + default: 'false' + +jobs: + release: + runs-on: ubuntu-latest + environment: pypi-prod + permissions: + contents: write + packages: write + id-token: write + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Get version + id: version + working-directory: packages/prime-runs + run: | + VERSION=$(grep -E "^__version__\s*=\s*" src/prime_runs/__init__.py | sed -E 's/^__version__[[:space:]]*=[[:space:]]*"([^"]+)".*$/\1/') + echo "version=$VERSION" >> $GITHUB_OUTPUT + echo "Version: $VERSION" + + - name: Check for existing tag + id: check_tag + run: | + TAG="prime-runs-v${{ steps.version.outputs.version }}" + if git tag -l "$TAG" | grep -q .; then + echo "exists=true" >> $GITHUB_OUTPUT + echo "Tag $TAG already exists" + else + echo "exists=false" >> $GITHUB_OUTPUT + echo "Tag $TAG does not exist" + fi + + - name: Set up Python + if: steps.check_tag.outputs.exists != 'true' || inputs.force_release == 'true' + uses: actions/setup-python@v6 + with: + python-version: '3.11' + + - name: Install uv + if: steps.check_tag.outputs.exists != 'true' || inputs.force_release == 'true' + uses: astral-sh/setup-uv@v4 + + - name: Build package + if: steps.check_tag.outputs.exists != 'true' || inputs.force_release == 'true' + working-directory: packages/prime-runs + run: | + uv build --out-dir dist + + - name: Create tag + if: steps.check_tag.outputs.exists != 'true' + run: | + TAG="prime-runs-v${{ steps.version.outputs.version }}" + git config user.name 'github-actions[bot]' + git config user.email 'github-actions[bot]@users.noreply.github.com' + git tag -a "$TAG" -m "Release $TAG" + git push origin "$TAG" + + - name: Publish to PyPI + if: steps.check_tag.outputs.exists != 'true' || inputs.force_release == 'true' + uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # v1.14.2 + with: + packages-dir: packages/prime-runs/dist diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md new file mode 100644 index 000000000..fb1f3d09b --- /dev/null +++ b/packages/prime-runs/README.md @@ -0,0 +1,153 @@ +# Prime Runs SDK + +Track eval and training runs on the Prime Intellect platform. + +```bash +pip install prime-runs +``` + +## Quick start + +```python +import prime_runs as pr + +run = pr.init( + name="gsm8k-qwen3-8b", + environments=["gsm8k"], + model="Qwen/Qwen3-8B", + framework="verifiers", + config={"num_rollouts": 4, "max_tokens": 2048}, +) +print(run.url) # https://app.primeintellect.ai/dashboard/evaluations/eval-... + +for episode in rollouts: # episodes carry run.id — see "Identity" below + run.log_traces([episode]) + run.log({"reward": episode.reward}) + +run.finish(summary=pr.metrics.from_episodes(episodes)) +``` + +`init()` opens the run and returns a handle carrying its ID and dashboard URL. +Records stream out on a background thread while the run proceeds — the dashboard +fills in as rollouts land, rather than all at once at the end. `finish()` closes +the run out. + +Prefer a `with` block, which also reports a terminal status on the paths where +you never reach `finish()`: + +```python +with pr.init(name="gsm8k-qwen3-8b", environments=["gsm8k"]) as run: + ... +``` + +## Identity + +`init()` is called **before** the first rollout, and the ID it returns is *the* +run ID everywhere — including inside every trace document you write, and +including the local archive. Stamp it once at rollout time: + +```python +run = pr.init(...) +trace.record_run(EvalRunInfo(id=run.id)) # verifiers +``` + +Nothing is re-stamped afterwards and no record of yours is rewritten. The +ingestion service extracts `run.id` from the trace body into an indexed column +with a delete-by-run path, so "every trace for this run" is a fast query rather +than a scan over upload metadata. + +`init()` also exports `PRIME_RUN_ID`, so forked workers and subprocess launchers +join the run their parent opened instead of each opening their own. + +## Modes + +| mode | what happens | +| --- | --- | +| `online` | the run lives on the platform (default when an API key is present) | +| `offline` | the run lives in a local directory, ready to sync later | +| `disabled` | every call is a no-op, with the same object shape | + +An offline run is a real run: a real ID, a status, a config, a summary, a metrics +stream, and records written in the JSONL wire format the traces service accepts. +That is why producers do not need a `--no-push` branch — the call sites are +identical either way, and a missing API key degrades to offline rather than +skipping the run. + +Set the mode explicitly, or through `$PRIME_RUNS_MODE`. + +## What the run handle does for you + +- **Streams instead of buffering.** Records go out as they are produced, so a + run with a hundred thousand episodes does not hold them all in memory. +- **Contains its own errors.** With the default `on_error="warn"`, nothing the + platform raises escapes into your loop. Use `on_error="raise"` in tests and CI, + where a silent upload failure is the bug. +- **Applies backpressure.** The upload queue is bounded; if a producer durably + outruns the uploader, records are dropped and counted (`run.dropped_records`) + rather than stalling the run. +- **Survives forks.** A forked child gets a fresh uploader instead of inheriting + the parent's queue and locks. +- **Reports a terminal status.** Context manager, `atexit` and signal handlers + all route to the same idempotent `finish()`, so a killed process is recorded as + crashed rather than left running forever. +- **Knows about ranks.** Rank 0 owns creation and finalization; other ranks join + through `PRIME_RUN_ID` and upload their own records. + +## Configuration + +Resolved from environment variables first, then `~/.prime/config.json`: + +| setting | env var | default | +| --- | --- | --- | +| API key | `PRIME_API_KEY` | — | +| team | `PRIME_TEAM_ID` | — | +| platform API | `PRIME_API_BASE_URL` | `https://api.primeintellect.ai` | +| dashboard | `PRIME_FRONTEND_URL` | `https://app.primeintellect.ai` | +| traces service | `PRIME_TRACES_URL` | resolved by `prime-traces` | +| offline runs | `PRIME_RUNS_DIR` | `./prime-runs` | + +Or pass them to `init()` directly (`api_key=`, `base_url=`, `team_id=`, `dir=`). + +## Backends and sinks + +Two independent axes: + +- **Backends** own run *lifecycle* — `EvalsBackend` (`/api/v1/evaluations/*`), + `OfflineBackend` (a local directory). Selected by `kind`. +- **Sinks** own sample *transport* — `TracesSink` (primary; streaming, + episode-aware, content-addressed and therefore idempotent on retry) and + `EvalSamplesSink` (the flat v0 sample table today's viewer reads). + +Both sinks run during the transition, because Prime Traces is in closed beta and +a traces-only client would leave non-allowlisted accounts with an empty +dashboard. When the Viewer API reads traces natively, the default sink list drops +one entry — and no producer changes. + +Turn either off with `pr.init(traces=False)` / `pr.init(samples=False)`, or pass +`sinks=[...]` to supply your own. + +## Also here + +`prime_runs.projection` holds `trace_to_sample` / `build_samples`, the projection +from native traces onto the platform's v0 eval-sample format. It lives here +because it is knowledge about a platform wire format, not about any one eval +framework. `prime_runs.metrics.from_episodes` is the run-level aggregation the +eval dashboard reads — opt-in, because what a run's headline number means is a +judgement that belongs next to the producer. + +Both are duck-typed: verifiers `Trace`/`Episode` and prime-rl `Rollout` satisfy +them structurally, and no producer package is imported. This is a leaf package by +design — the `prime` CLI depends on `verifiers`, so verifiers can never depend on +`prime`. + +## Status + +Eval runs are supported. Training runs (`kind="train"`, over +`/api/v1/rft/external-runs`) are next; `pr.init(kind="train")` raises a clear +error until then. + +One platform gap is worth knowing about: there is currently no producer-facing +way to mark an evaluation **failed**. The SDK calls the status endpoint it needs, +treats its absence as expected, and records the terminal state in the run's +metadata as a fallback — a failed run will keep showing as running on the +dashboard, and the SDK says so in a warning. diff --git a/packages/prime-runs/pyproject.toml b/packages/prime-runs/pyproject.toml new file mode 100644 index 000000000..058d90bdc --- /dev/null +++ b/packages/prime-runs/pyproject.toml @@ -0,0 +1,70 @@ +[project] +name = "prime-runs" +# Version is single-sourced from src/prime_runs/__init__.py via Hatch +dynamic = ["version"] +description = "Prime Intellect Runs SDK - Track eval and training runs on the Prime platform" +readme = "README.md" +requires-python = ">=3.10" +license = {text = "MIT"} +authors = [ + { name = "Prime Intellect", email = "contact@primeintellect.ai" } +] +# Deliberately a leaf package. The `prime` CLI depends on `verifiers`, so +# verifiers can never depend on `prime` — and verifiers is a first-class +# consumer of this SDK. Nothing here may pull in prime, verifiers, typer, +# rich or textual, directly or transitively. +dependencies = [ + "httpx>=0.25.0", + "pydantic>=2.0.0", + "tenacity>=9.1.2", + "prime-traces>=0.0.2", +] +keywords = ["evals", "evaluations", "training", "runs", "traces", "observability"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Software Development :: Libraries :: Python Modules", + "Topic :: Scientific/Engineering :: Artificial Intelligence" +] + +[project.urls] +Homepage = "https://github.com/PrimeIntellect-ai/prime" +Documentation = "https://github.com/PrimeIntellect-ai/prime/tree/main/packages/prime-runs" +Repository = "https://github.com/PrimeIntellect-ai/prime.git" + +[project.optional-dependencies] +dev = [ + "pytest>=7.0.0", + "pytest-asyncio>=0.23.0", + "ruff>=0.13.1", +] + +[tool.uv.sources] +prime-traces = { workspace = true } + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.version] +path = "src/prime_runs/__init__.py" + +[tool.hatch.build.targets.wheel] +packages = ["src/prime_runs"] + +[tool.pytest.ini_options] +addopts = "-v" +testpaths = ["tests"] + +[tool.ruff] +line-length = 100 +target-version = "py310" + +[tool.ruff.lint] +extend-select = ["E", "F", "I"] diff --git a/packages/prime-runs/src/prime_runs/__init__.py b/packages/prime-runs/src/prime_runs/__init__.py new file mode 100644 index 000000000..d295c9194 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/__init__.py @@ -0,0 +1,96 @@ +"""Prime Intellect Runs SDK. + +Track eval and training runs on the Prime platform:: + + import prime_runs as pr + + run = pr.init(name="gsm8k-qwen3-8b", environments=["gsm8k"], model=model) + print(run.url) + + for episode in rollouts: # stamp run.id onto the traces + run.log_traces([episode]) + run.log({"reward": episode.reward}, step=step) + + run.finish(summary=pr.metrics.from_episodes(episodes)) + +``init()`` opens the run and returns a handle carrying its ID and dashboard +URL; records stream out on a background thread as the run proceeds; ``finish()`` +closes it out. A ``with`` block does the last part for you, including on the +paths where the producer never gets to it. + +This is a leaf package on purpose. The ``prime`` CLI depends on ``verifiers``, +so verifiers can never depend on ``prime`` — and verifiers is one of the two +producers this SDK exists to serve. Nothing here imports a producer package; +records are duck-typed through ``to_record()``. +""" + +from . import metrics, projection +from .backends import Backend, EvalsBackend, OfflineBackend +from .config import Config +from .exceptions import ( + ConfigurationError, + EnvironmentResolutionError, + NotFoundError, + PaymentRequiredError, + PrimeRunsError, + RetryableAPIError, + RunAPIError, + RunFinishedError, + TransportError, + UnauthorizedError, +) +from .models import ( + EnvironmentRef, + Mode, + OnError, + RunHandle, + RunKind, + RunSpec, + RunStatus, +) +from .projection import build_samples, trace_to_sample +from .run import MODE_ENV, RUN_ID_ENV, Run, init +from .sinks import EvalSamplesSink, OfflineSink, Sink, TracesSink + +__version__ = "0.1.0" + +__all__ = [ + # The surface almost every caller needs + "init", + "Run", + "metrics", + "projection", + # Types + "Config", + "EnvironmentRef", + "Mode", + "OnError", + "RunHandle", + "RunKind", + "RunSpec", + "RunStatus", + "MODE_ENV", + "RUN_ID_ENV", + # Backends & sinks, for callers assembling their own + "Backend", + "EvalsBackend", + "OfflineBackend", + "Sink", + "EvalSamplesSink", + "OfflineSink", + "TracesSink", + # The v0 sample projection, moved here from verifiers + "build_samples", + "trace_to_sample", + # Exceptions + "PrimeRunsError", + "ConfigurationError", + "EnvironmentResolutionError", + "RunAPIError", + "RunFinishedError", + "NotFoundError", + "PaymentRequiredError", + "RetryableAPIError", + "TransportError", + "UnauthorizedError", +] diff --git a/packages/prime-runs/src/prime_runs/_http.py b/packages/prime-runs/src/prime_runs/_http.py new file mode 100644 index 000000000..a56970bcb --- /dev/null +++ b/packages/prime-runs/src/prime_runs/_http.py @@ -0,0 +1,231 @@ +"""Shared HTTP client for the platform run APIs. + +One client for all backends and the legacy samples sink, because they share a +host, a credential and a retry policy. Two things it does that a bare +``httpx.Client`` does not: + +- maps status codes onto :mod:`prime_runs.exceptions` so callers branch on a + type rather than on a message; +- retries 429/502/503/504 and transport failures with exponential backoff, + honouring ``Retry-After`` when the server sends one. + +Retries are safe here because every call it makes is either idempotent (PUT, +GET) or create-shaped and guarded upstream: run creation happens exactly once +per ``init()``, and sample POSTs that get retried after a lost response are the +known duplicate-append case the traces sink exists to replace. +""" + +import json +import sys +import time +from typing import Any, Dict, Mapping, Optional, Union + +import httpx + +from .exceptions import ( + NotFoundError, + PaymentRequiredError, + RetryableAPIError, + RunAPIError, + TransportError, + UnauthorizedError, +) + +DEFAULT_TIMEOUT = httpx.Timeout(60.0, connect=10.0) +# Sample batches are megabytes and the platform fans them out to storage before +# answering, so uploads get their own, much longer budget. +UPLOAD_TIMEOUT = httpx.Timeout(300.0, connect=10.0) +RETRY_STATUS = frozenset({429, 502, 503, 504}) +DEFAULT_MAX_ATTEMPTS = 5 +MAX_BACKOFF_SECONDS = 16.0 + + +def _user_agent() -> str: + from . import __version__ + + py = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" + return f"prime-runs/{__version__} python/{py}" + + +def retry_delay(attempt: int, retry_after: Optional[float]) -> float: + """Seconds to wait before ``attempt`` (1-based). Server wins if it spoke.""" + if retry_after is not None and retry_after >= 0: + return min(retry_after, MAX_BACKOFF_SECONDS) + return min(2.0 ** (attempt - 1), MAX_BACKOFF_SECONDS) + + +def _parse_retry_after(response: httpx.Response) -> Optional[float]: + raw = response.headers.get("retry-after") + if not raw: + return None + try: + return float(raw) + except ValueError: + # HTTP-date form. Not worth parsing for a backoff hint — fall back to + # the exponential schedule rather than guessing a clock skew. + return None + + +def encode_json(value: Any) -> bytes: + """Compact UTF-8 JSON, matching the encoding used to size batches. + + ``allow_nan=False`` matters: a NaN reward serialized as JavaScript's bare + ``NaN`` is rejected by strict JSON parsers server-side, and the failure + surfaces as an opaque 400 on a payload the producer cannot inspect. + """ + return json.dumps(value, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode( + "utf-8" + ) + + +class PlatformClient: + """Minimal authenticated client for ``{base_url}/api/v1``.""" + + def __init__( + self, + *, + api_key: str, + base_url: str, + timeout: httpx.Timeout = DEFAULT_TIMEOUT, + max_attempts: int = DEFAULT_MAX_ATTEMPTS, + client: Optional[httpx.Client] = None, + ) -> None: + self.base_url = base_url.rstrip("/") + self.api_prefix = f"{self.base_url}/api/v1" + self.max_attempts = max(1, max_attempts) + self._owns_client = client is None + self._client = client or httpx.Client( + headers={ + "Authorization": f"Bearer {api_key}", + "User-Agent": _user_agent(), + }, + follow_redirects=True, + timeout=timeout, + ) + + def request( + self, + method: str, + path: str, + *, + json_body: Optional[Mapping[str, Any]] = None, + content: Optional[bytes] = None, + params: Optional[Mapping[str, Any]] = None, + timeout: Union[httpx.Timeout, float, None] = None, + max_attempts: Optional[int] = None, + ) -> Dict[str, Any]: + """Send one request, retrying transient failures. Returns the JSON body.""" + url = f"{self.api_prefix}{path}" + body = content if content is not None else (encode_json(json_body) if json_body else None) + headers = {"Content-Type": "application/json"} if body is not None else None + attempts = max_attempts or self.max_attempts + + last_error: Optional[Exception] = None + for attempt in range(1, attempts + 1): + try: + response = self._client.request( + method, + url, + content=body, + headers=headers, + params=dict(params) if params else None, + timeout=timeout, + ) + except httpx.TimeoutException as exc: + last_error = TransportError(f"{method} {path} timed out: {exc}") + except httpx.RequestError as exc: + last_error = TransportError(f"{method} {path} failed: {type(exc).__name__}: {exc}") + else: + if response.status_code in RETRY_STATUS: + last_error = RetryableAPIError( + _error_message(response), + status_code=response.status_code, + code=_error_code(response), + retry_after=_parse_retry_after(response), + ) + elif response.is_error: + raise _map_error(response) + else: + return _decode(response) + + if attempt == attempts: + break + after = getattr(last_error, "retry_after", None) + time.sleep(retry_delay(attempt, after)) + + assert last_error is not None + raise last_error + + def get(self, path: str, **kwargs: Any) -> Dict[str, Any]: + return self.request("GET", path, **kwargs) + + def post(self, path: str, **kwargs: Any) -> Dict[str, Any]: + return self.request("POST", path, **kwargs) + + def put(self, path: str, **kwargs: Any) -> Dict[str, Any]: + return self.request("PUT", path, **kwargs) + + def close(self) -> None: + if self._owns_client: + self._client.close() + + def __enter__(self) -> "PlatformClient": + return self + + def __exit__(self, *exc_info: Any) -> None: + self.close() + + +def _decode(response: httpx.Response) -> Dict[str, Any]: + if not response.content: + return {} + try: + payload = response.json() + except ValueError as exc: + raise RunAPIError( + f"{response.request.method} {response.request.url.path} returned non-JSON " + f"({response.status_code}): {response.text[:200]!r}" + ) from exc + return payload if isinstance(payload, dict) else {"data": payload} + + +def _error_body(response: httpx.Response) -> Dict[str, Any]: + try: + payload = response.json() + except ValueError: + return {} + return payload if isinstance(payload, dict) else {} + + +def _error_code(response: httpx.Response) -> Optional[str]: + body = _error_body(response) + code = body.get("code") or body.get("error_code") + return str(code) if code else None + + +def _error_message(response: httpx.Response) -> str: + body = _error_body(response) + detail = body.get("detail") or body.get("message") or body.get("error") + if detail is None: + detail = response.text[:200] or response.reason_phrase + return ( + f"HTTP {response.status_code} from " + f"{response.request.method} {response.request.url.path}: {detail}" + ) + + +def _map_error(response: httpx.Response) -> RunAPIError: + message = _error_message(response) + code = _error_code(response) + status = response.status_code + if status == 401: + return UnauthorizedError( + f"{message} — check PRIME_API_KEY or run `prime login`.", + status_code=status, + code=code, + ) + if status == 402: + return PaymentRequiredError(message, status_code=status, code=code) + if status == 404: + return NotFoundError(message, status_code=status, code=code) + return RunAPIError(message, status_code=status, code=code) diff --git a/packages/prime-runs/src/prime_runs/backends/__init__.py b/packages/prime-runs/src/prime_runs/backends/__init__.py new file mode 100644 index 000000000..1b8693ade --- /dev/null +++ b/packages/prime-runs/src/prime_runs/backends/__init__.py @@ -0,0 +1,13 @@ +"""Run lifecycle backends, one per platform run system.""" + +from .base import Backend +from .evals import EvalsBackend +from .offline import OfflineBackend, default_dir, new_run_id + +__all__ = [ + "Backend", + "EvalsBackend", + "OfflineBackend", + "default_dir", + "new_run_id", +] diff --git a/packages/prime-runs/src/prime_runs/backends/base.py b/packages/prime-runs/src/prime_runs/backends/base.py new file mode 100644 index 000000000..ad4d14ea0 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/backends/base.py @@ -0,0 +1,68 @@ +"""The contract a run backend implements. + +A backend owns one thing: the *lifecycle* of a run — bringing it into +existence, updating what is known about it, and closing it out with a terminal +status. It does not move samples; that is a sink's job (see +:mod:`prime_runs.sinks`). Keeping the two axes independent is what lets the +eval and training APIs — which agree on almost nothing at the wire level — +share a single ``Run`` handle, and what lets the sample transport change +underneath without touching either. +""" + +from typing import Any, Dict, Optional, Protocol, runtime_checkable + +from ..models import RunHandle, RunSpec, RunStatus + + +@runtime_checkable +class Backend(Protocol): + """Lifecycle operations for one family of runs.""" + + kind: str + """The ``RunKind`` this backend serves.""" + + supports_step_metrics: bool + """Whether ``log_metrics`` records a point per step. + + ``False`` means the API has no time series and the run keeps a last-value + summary instead. The ``Run`` handle reads this to decide whether + ``log(..., step=)`` is a real write or a summary merge, so producers get + the same call either way. + """ + + def create(self, spec: RunSpec) -> RunHandle: + """Open a new run and return its platform identity.""" + ... + + def attach(self, run_id: str) -> RunHandle: + """Re-acquire an existing run, for resume and for non-primary ranks.""" + ... + + def update( + self, + run_id: str, + *, + config: Optional[Dict[str, Any]] = None, + summary: Optional[Dict[str, Any]] = None, + ) -> None: + """Persist config (inputs) and/or summary (outputs) mid-run.""" + ... + + def log_metrics(self, run_id: str, metrics: Dict[str, Any], step: Optional[int] = None) -> None: + """Append one point to the run's time series. No-op when unsupported.""" + ... + + def finalize( + self, + run_id: str, + *, + status: RunStatus, + summary: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + ) -> None: + """Close the run out. Called exactly once per run.""" + ... + + def close(self) -> None: + """Release transport resources.""" + ... diff --git a/packages/prime-runs/src/prime_runs/backends/evals.py b/packages/prime-runs/src/prime_runs/backends/evals.py new file mode 100644 index 000000000..26c4caba2 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/backends/evals.py @@ -0,0 +1,259 @@ +"""Eval runs over ``/api/v1/evaluations/*``. + +Wraps the endpoints verifiers previously called inline, plus environment +resolution through the hub's get-or-create so a local run uploads without a +prior ``prime env push``. + +One gap is worth stating plainly, because it shapes the code below: **the eval +API has no producer-facing way to mark a run failed.** ``finalize`` moves a run +PROCESSING -> COMPLETED, ``UpdateEvaluationRequest`` carries no ``status``, and +FAILED is written only when an internal Cloud Task trigger fails. So a crashed +run stays RUNNING forever. ``_report_failure`` calls the status endpoint this +SDK needs, treats its absence as expected, and falls back to recording the +terminal state in ``metadata`` so the failure is at least visible and machine +readable. When the endpoint lands, the fallback stops firing on its own. +""" + +import logging +import uuid +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional + +from .._http import PlatformClient +from ..exceptions import ( + ConfigurationError, + EnvironmentResolutionError, + NotFoundError, + RunAPIError, +) +from ..models import EnvironmentRef, RunHandle, RunSpec, RunStatus + +logger = logging.getLogger(__name__) + +# Statuses the platform's EvaluationStatus enum uses, keyed by ours. +_PLATFORM_STATUS = { + RunStatus.COMPLETED: "COMPLETED", + RunStatus.FAILED: "FAILED", + RunStatus.CRASHED: "FAILED", +} + + +class EvalsBackend: + """Lifecycle for evaluation runs.""" + + kind = "eval" + # The evaluations API stores a single metrics blob, not a time series. + supports_step_metrics = False + + def __init__( + self, + client: PlatformClient, + *, + frontend_url: str, + team_id: Optional[str] = None, + ) -> None: + self._client = client + self._frontend_url = frontend_url.rstrip("/") + self._team_id = team_id + self._status_endpoint_missing = False + + # ------------------------------------------------------------------ create + + def create(self, spec: RunSpec) -> RunHandle: + environment_ids = self._resolve_environments(spec.environments) + if not environment_ids: + raise ConfigurationError( + "An eval run needs at least one environment. Pass " + 'environments=["my-env"] to init().' + ) + + run_name: str = spec.name or _default_name(spec) + payload: Dict[str, Any] = { + "name": run_name, + "environments": [{"id": environment_id} for environment_id in environment_ids], + "tags": list(spec.tags), + } + _set_if(payload, "model_name", spec.model) + _set_if(payload, "dataset", spec.dataset or _first_environment_name(spec)) + _set_if(payload, "framework", spec.framework) + _set_if(payload, "description", spec.description) + _set_if(payload, "metadata", spec.config or None) + _set_if(payload, "metrics", spec.summary or None) + _set_if(payload, "team_id", spec.team_id or self._team_id) + + response = self._client.post("/evaluations/", json_body=payload) + run_id = response.get("evaluation_id") + if not run_id: + raise RunAPIError( + f"POST /evaluations/ returned no evaluation_id (keys: {sorted(response)})" + ) + return RunHandle( + id=run_id, + name=str(response.get("name") or run_name), + url=response.get("viewer_url") or self.url_for(run_id), + raw=response, + ) + + def attach(self, run_id: str) -> RunHandle: + try: + response = self._client.get(f"/evaluations/{run_id}") + except NotFoundError: + raise + except RunAPIError as exc: + # Attach is a convenience — a resume or a non-primary rank joining. + # Losing the run's name to a transient read is not worth failing on; + # the ID is what everything downstream actually needs. + logger.debug("Could not read evaluation %s on attach: %s", run_id, exc) + return RunHandle(id=run_id, url=self.url_for(run_id)) + return RunHandle( + id=run_id, + name=response.get("name"), + url=response.get("viewer_url") or self.url_for(run_id), + raw=response, + ) + + def url_for(self, run_id: str) -> str: + return f"{self._frontend_url}/dashboard/evaluations/{run_id}" + + # ------------------------------------------------------------------ update + + def update( + self, + run_id: str, + *, + config: Optional[Dict[str, Any]] = None, + summary: Optional[Dict[str, Any]] = None, + ) -> None: + payload: Dict[str, Any] = {} + _set_if(payload, "metadata", config or None) + _set_if(payload, "metrics", summary or None) + if not payload: + return + self._client.put(f"/evaluations/{run_id}", json_body=payload) + + def log_metrics(self, run_id: str, metrics: Dict[str, Any], step: Optional[int] = None) -> None: + """No-op: see ``supports_step_metrics``. + + The run keeps a last-value summary and flushes it through ``update``, + which is the only shape this API can store. + """ + + # ---------------------------------------------------------------- finalize + + def finalize( + self, + run_id: str, + *, + status: RunStatus, + summary: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + ) -> None: + if status is RunStatus.COMPLETED: + body: Dict[str, Any] = {} + _set_if(body, "metrics", summary or None) + self._client.post(f"/evaluations/{run_id}/finalize", json_body=body or {"metrics": {}}) + return + self._report_failure(run_id, status=status, summary=summary, error=error) + + def _report_failure( + self, + run_id: str, + *, + status: RunStatus, + summary: Optional[Dict[str, Any]], + error: Optional[str], + ) -> None: + """Mark a run failed, or record why we could not.""" + terminal = { + "status": status.value, + "finished_at": datetime.now(timezone.utc).isoformat(), + } + if error: + terminal["error"] = error + + if not self._status_endpoint_missing: + try: + self._client.post( + f"/evaluations/{run_id}/status", + json_body={"status": _PLATFORM_STATUS[status], "error": error}, + max_attempts=1, + ) + return + except NotFoundError: + # Expected until the status endpoint ships. Latch so a run that + # fails repeatedly does not pay for the probe every time. + self._status_endpoint_missing = True + logger.debug( + "Platform has no /evaluations/{id}/status endpoint; " + "recording terminal state in metadata instead" + ) + except RunAPIError as exc: + if exc.status_code not in (405, 422): + raise + self._status_endpoint_missing = True + logger.debug("Status endpoint rejected the request (%s); using metadata", exc) + + # Fallback: the run cannot be moved out of RUNNING, but the failure is + # at least recorded where an operator and the dashboard can both read it. + self.update(run_id, config={"prime_runs": terminal}, summary=summary) + logger.warning( + "Run %s %s, but the platform has no way to mark an evaluation failed; " + "it will keep showing as running. Recorded the failure in metadata.prime_runs.", + run_id, + status.value, + ) + + def close(self) -> None: + self._client.close() + + # ----------------------------------------------------------- environments + + def _resolve_environments(self, refs: List[EnvironmentRef]) -> List[str]: + """Environment IDs for the hub references a producer named. + + Unlike the old client, a reference that cannot be resolved raises + instead of being skipped: dropping one silently produces a run attached + to the wrong environments, which looks like a successful upload and is + found much later. + """ + resolved: List[str] = [] + for ref in refs: + if ref.id: + resolved.append(ref.id) + continue + body: Dict[str, Any] = {"name": ref.name} + _set_if(body, "team_id", self._team_id) + try: + response = self._client.post("/environmentshub/resolve", json_body=body) + except RunAPIError as exc: + raise EnvironmentResolutionError( + f"Could not resolve environment {ref.name!r}: {exc}" + ) from exc + environment_id = (response.get("data") or {}).get("id") + if not environment_id: + raise EnvironmentResolutionError(f"Hub returned no id for environment {ref.name!r}") + resolved.append(environment_id) + return resolved + + +def _set_if(payload: Dict[str, Any], key: str, value: Any) -> None: + if value is not None: + payload[key] = value + + +def _first_environment_name(spec: RunSpec) -> Optional[str]: + for ref in spec.environments: + if ref.name: + return ref.name + return None + + +def _default_name(spec: RunSpec) -> str: + """A name for producers that did not pick one. + + The API requires a name, so the alternative to generating one is a 422 at + the worst possible moment. Leads with the environment so runs sort together + in the dashboard list. + """ + stem = _first_environment_name(spec) or spec.framework or spec.kind + return f"{stem}-{uuid.uuid4().hex[:8]}" diff --git a/packages/prime-runs/src/prime_runs/backends/offline.py b/packages/prime-runs/src/prime_runs/backends/offline.py new file mode 100644 index 000000000..2b571e5ab --- /dev/null +++ b/packages/prime-runs/src/prime_runs/backends/offline.py @@ -0,0 +1,160 @@ +"""Offline runs: a local directory that looks exactly like a platform run. + +This is not a debugging affordance, it is the reason producers can delete their +``--no-push`` branching. A run that never reaches the network still has an ID, +a status, a config, a summary and a metrics stream, so the call sites above it +are identical whether or not anyone is logged in. The locally issued ID is used +as the run ID everywhere — including inside the trace documents — so a later +sync attaches the archive to a platform run without rewriting a single record. + +Layout, one directory per run:: + + //run.json spec + status + timestamps + //metrics.jsonl one JSON object per log() call + //records/ whatever the offline sink wrote +""" + +import json +import logging +import os +import uuid +from dataclasses import asdict +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, Optional, Union + +from ..models import RunHandle, RunSpec, RunStatus + +logger = logging.getLogger(__name__) + +DEFAULT_DIR_ENV = "PRIME_RUNS_DIR" +DEFAULT_DIR = "prime-runs" + + +def default_dir() -> Path: + return Path(os.getenv(DEFAULT_DIR_ENV) or DEFAULT_DIR) + + +def new_run_id() -> str: + """A locally issued run ID, visibly distinct from a platform one.""" + return f"offline-{uuid.uuid4().hex[:16]}" + + +class OfflineBackend: + """Run lifecycle recorded on the local filesystem.""" + + kind = "offline" + supports_step_metrics = True + + def __init__(self, directory: Union[str, Path, None] = None) -> None: + self.directory = Path(directory) if directory is not None else default_dir() + + def run_dir(self, run_id: str) -> Path: + return self.directory / run_id + + def create(self, spec: RunSpec) -> RunHandle: + run_id = new_run_id() + path = self.run_dir(run_id) + path.mkdir(parents=True, exist_ok=True) + run_name: str = spec.name or run_id + state: Dict[str, Any] = { + "id": run_id, + "name": run_name, + "kind": spec.kind, + "status": RunStatus.RUNNING.value, + "created_at": _now(), + "spec": _spec_to_json(spec), + } + self._write_state(run_id, state) + return RunHandle(id=run_id, name=run_name, url=str(path.resolve()), raw=state) + + def attach(self, run_id: str) -> RunHandle: + path = self.run_dir(run_id) + path.mkdir(parents=True, exist_ok=True) + state = self._read_state(run_id) + return RunHandle( + id=run_id, + name=str(state.get("name") or run_id), + url=str(path.resolve()), + raw=state, + ) + + def update( + self, + run_id: str, + *, + config: Optional[Dict[str, Any]] = None, + summary: Optional[Dict[str, Any]] = None, + ) -> None: + state = self._read_state(run_id) + if config: + state.setdefault("config", {}).update(config) + if summary: + state.setdefault("summary", {}).update(summary) + state["updated_at"] = _now() + self._write_state(run_id, state) + + def log_metrics(self, run_id: str, metrics: Dict[str, Any], step: Optional[int] = None) -> None: + line = {"step": step, "timestamp": _now(), **metrics} + path = self.run_dir(run_id) / "metrics.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(line, ensure_ascii=False, default=str) + "\n") + + def finalize( + self, + run_id: str, + *, + status: RunStatus, + summary: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + ) -> None: + state = self._read_state(run_id) + state["status"] = status.value + state["finished_at"] = _now() + if error: + state["error"] = error + if summary: + state.setdefault("summary", {}).update(summary) + self._write_state(run_id, state) + + def close(self) -> None: + """Nothing to release — every write is already flushed to disk.""" + + # ------------------------------------------------------------------ state + + def _state_path(self, run_id: str) -> Path: + return self.run_dir(run_id) / "run.json" + + def _read_state(self, run_id: str) -> Dict[str, Any]: + path = self._state_path(run_id) + if not path.exists(): + return {"id": run_id} + try: + state = json.loads(path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError, UnicodeDecodeError) as exc: + logger.warning("Could not read %s (%s); starting a fresh record", path, exc) + return {"id": run_id} + return state if isinstance(state, dict) else {"id": run_id} + + def _write_state(self, run_id: str, state: Dict[str, Any]) -> None: + path = self._state_path(run_id) + path.parent.mkdir(parents=True, exist_ok=True) + # Write-then-rename: a crash mid-write must not leave the run's own + # record truncated, since it is the only description of what ran. + temp = path.with_suffix(".json.tmp") + temp.write_text(json.dumps(state, indent=2, ensure_ascii=False, default=str), "utf-8") + temp.replace(path) + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _spec_to_json(spec: RunSpec) -> Dict[str, Any]: + data = asdict(spec) + data["environments"] = [ + {key: value for key, value in env.items() if value is not None} + for env in data.get("environments", []) + ] + return data diff --git a/packages/prime-runs/src/prime_runs/config.py b/packages/prime-runs/src/prime_runs/config.py new file mode 100644 index 000000000..e84ec2985 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/config.py @@ -0,0 +1,75 @@ +"""Lightweight configuration for the Prime Runs SDK. + +Same shape as the other prime SDK packages: reads ``~/.prime/config.json`` +plus environment variables, env taking precedence. Adds ``frontend_url`` +(where run URLs point) on top of the shared platform settings. +""" + +import json +import os +from pathlib import Path +from typing import Optional + + +class Config: + """Minimal configuration class for SDK packages. + + Reads from ~/.prime/config.json and environment variables. + """ + + DEFAULT_BASE_URL: str = "https://api.primeintellect.ai" + DEFAULT_FRONTEND_URL: str = "https://app.primeintellect.ai" + + def __init__(self) -> None: + self.config_dir = Path.home() / ".prime" + self.config_file = self.config_dir / "config.json" + self._load_config() + + def _load_config(self) -> None: + """Load configuration from file.""" + config_data: object = {} + if self.config_file.exists(): + try: + config_data = json.loads(self.config_file.read_text()) + except (json.JSONDecodeError, OSError, UnicodeDecodeError): + config_data = {} + # Valid JSON that is not an object (a list, a bare string) must degrade + # the same way invalid JSON does: every accessor below assumes a dict. + self.config = config_data if isinstance(config_data, dict) else {} + + @staticmethod + def _strip_api_v1(url: str) -> str: + return url.rstrip("/").removesuffix("/api/v1") + + @property + def api_key(self) -> str: + """API key with precedence: env > file > empty.""" + return os.getenv("PRIME_API_KEY") or self.config.get("api_key", "") + + @property + def team_id(self) -> Optional[str]: + """Team ID with precedence: env > file > None.""" + team_id = os.getenv("PRIME_TEAM_ID") + if team_id is not None: + return team_id + return self.config.get("team_id") or None + + @property + def base_url(self) -> str: + """Platform API base URL with precedence: env > file > default.""" + env_val = os.getenv("PRIME_API_BASE_URL") or os.getenv("PRIME_BASE_URL") + if env_val: + return self._strip_api_v1(env_val) + return self._strip_api_v1(self.config.get("base_url", self.DEFAULT_BASE_URL)) + + @property + def frontend_url(self) -> str: + """Dashboard base URL, used to build the run URL a producer prints. + + The platform returns a ``viewer_url`` on create; this is the fallback + for responses that omit it and for offline/legacy paths. + """ + env_val = os.getenv("PRIME_FRONTEND_URL") + if env_val: + return env_val.rstrip("/") + return str(self.config.get("frontend_url") or self.DEFAULT_FRONTEND_URL).rstrip("/") diff --git a/packages/prime-runs/src/prime_runs/exceptions.py b/packages/prime-runs/src/prime_runs/exceptions.py new file mode 100644 index 000000000..412a85da4 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/exceptions.py @@ -0,0 +1,84 @@ +"""Exceptions for the Prime Runs SDK. + +Producers run for hours; the default posture is that nothing here escapes into +a training loop (``on_error="warn"``). These types exist so that callers who +opt into ``on_error="raise"`` — tests, CI, hosted workers — can branch on what +actually failed instead of matching log strings. +""" + +from typing import Optional + + +class PrimeRunsError(Exception): + """Base exception for the Prime Runs SDK.""" + + +class ConfigurationError(PrimeRunsError): + """The SDK was asked to do something the local configuration cannot support. + + Missing API key, an unknown ``kind``, ``mode="online"`` with no way to reach + the platform. Raised before any request is made. + """ + + +class RunAPIError(PrimeRunsError): + """An HTTP error response from a run backend.""" + + def __init__( + self, + message: str, + *, + status_code: Optional[int] = None, + code: Optional[str] = None, + ): + self.status_code = status_code + self.code = code + super().__init__(message) + + +class UnauthorizedError(RunAPIError): + """401 — the credential was rejected. Stop rather than retry.""" + + +class PaymentRequiredError(RunAPIError): + """402 — payment required. Check billing status.""" + + +class NotFoundError(RunAPIError): + """404 — the run, environment or evaluation does not exist for this owner.""" + + +class RetryableAPIError(RunAPIError): + """429/5xx — retry the same request after ``retry_after`` seconds.""" + + def __init__( + self, + message: str, + *, + status_code: Optional[int] = None, + code: Optional[str] = None, + retry_after: Optional[float] = None, + ): + super().__init__(message, status_code=status_code, code=code) + self.retry_after = retry_after + + +class TransportError(RunAPIError): + """The request failed below HTTP — connection refused, TLS failure, timeout.""" + + +class EnvironmentResolutionError(PrimeRunsError): + """An environment named in ``init()`` could not be resolved to a hub ID. + + Distinct from a generic API error because it is usually a typo or a + permissions problem on the environment, not an outage, and because an eval + run cannot be created without at least one resolved environment. + """ + + +class RunFinishedError(PrimeRunsError): + """A finished run was written to again. + + Terminal status is reported once. Logging after ``finish()`` is a producer + bug — the data would land on a run the platform has already closed out. + """ diff --git a/packages/prime-runs/src/prime_runs/metrics.py b/packages/prime-runs/src/prime_runs/metrics.py new file mode 100644 index 000000000..25ec6d888 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/metrics.py @@ -0,0 +1,49 @@ +"""Run-level aggregates over native episodes. + +Opt-in, not automatic. The SDK does not decide what a run's headline number is +— that judgement belongs next to the producer, which knows which agents are +being scored and what counts as an error. This module ships the aggregation +verifiers already used, so callers migrating off ``verifiers.v1.utils.platform`` +keep byte-identical dashboard numbers, and anyone else can pass their own dict +to ``run.finish(summary=...)``. + +Duck-typed like :mod:`prime_runs.projection`: no producer package is imported. +""" + +from typing import Any, Dict, Optional, Sequence + + +def from_episodes( + episodes: Sequence[Any], traces: Optional[Sequence[Any]] = None +) -> Dict[str, Any]: + """Run-level aggregates in the shape the eval dashboard reads. + + Rewards and metrics aggregate over the trainable traces only — fixed agents + (a judge, a modeled user) often carry no rewards and would dilute every mean + with structural zeros — falling back to all traces when none are trainable, + the same rule the dashboard applies. ``avg_error`` is the share of EPISODES + that aren't ok: a hook failure counts even when its traces are clean or it + left none behind. + """ + if traces is None: + traces = [trace for episode in episodes for trace in episode.traces] + scored = [trace for trace in traces if trace.agent.trainable] or list(traces) + + sums: Dict[str, float] = {} + counts: Dict[str, int] = {} + for trace in scored: + scores = { + name: reward.score for name, reward in trace.rewards.items() if reward is not None + } + metrics = {name: value for name, value in trace.metrics.items() if value is not None} + for name, value in {**scores, **metrics}.items(): + sums[name] = sums.get(name, 0.0) + value + counts[name] = counts.get(name, 0) + 1 + + n = len(scored) + avg_error = sum(not episode.ok for episode in episodes) / len(episodes) if episodes else 0.0 + return { + "avg_reward": sum(trace.reward for trace in scored) / n if n else 0.0, + "avg_metrics": {name: sums[name] / counts[name] for name in sums}, + "avg_error": avg_error, + } diff --git a/packages/prime-runs/src/prime_runs/models.py b/packages/prime-runs/src/prime_runs/models.py new file mode 100644 index 000000000..e783f140a --- /dev/null +++ b/packages/prime-runs/src/prime_runs/models.py @@ -0,0 +1,110 @@ +"""Types shared across backends, sinks and the ``Run`` handle. + +Only the values that cross a module boundary live here. Response bodies are +deliberately *not* modeled: the platform returns more fields than any producer +reads, and freezing them in pydantic here would make every backend addition a +breaking SDK release. Backends pull the two or three fields they need and hand +back a ``RunHandle``. +""" + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Dict, List, Literal, Optional + +RunKind = Literal["eval", "train"] +"""Which run system owns the lifecycle. Selects the backend.""" + +Mode = Literal["online", "offline", "disabled"] +"""``online`` talks to the platform, ``offline`` writes a local run directory +that can be synced later, ``disabled`` makes every call a no-op while keeping +the same object shape so producer code needs no branching.""" + +OnError = Literal["warn", "raise"] + + +class RunStatus(str, Enum): + """Terminal state a producer can report. + + ``crashed`` is distinct from ``failed``: ``failed`` means the producer + decided the run failed, ``crashed`` means the process exited without ever + saying. Only the second one is inferred by the SDK (atexit / signal), and + the distinction is what tells an operator whether to look at the run's own + error or at the machine it ran on. + """ + + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CRASHED = "crashed" + + def is_terminal(self) -> bool: + return self is not RunStatus.RUNNING + + +@dataclass +class EnvironmentRef: + """An environment as a producer names it, before hub resolution. + + ``id`` short-circuits resolution; ``name`` goes through the hub's + get-or-create so a local run uploads without a prior ``prime env push``. + """ + + name: Optional[str] = None + id: Optional[str] = None + version_id: Optional[str] = None + + @classmethod + def coerce(cls, value: Any) -> "EnvironmentRef": + if isinstance(value, EnvironmentRef): + return value + if isinstance(value, str): + return cls(name=value) + if isinstance(value, dict): + return cls( + name=value.get("name"), + id=value.get("id"), + version_id=value.get("version_id"), + ) + raise TypeError( + "environments entries must be a str, dict or EnvironmentRef, " + f"got {type(value).__name__}" + ) + + def __post_init__(self) -> None: + if not self.name and not self.id: + raise ValueError("EnvironmentRef needs a name or an id") + + +@dataclass +class RunSpec: + """Everything a backend needs to open a run, in producer vocabulary. + + This is the argument surface of ``init()`` after normalization — backends + translate it into whatever their API family calls these things, which is + the whole reason eval and training runs can share one handle. + """ + + name: Optional[str] = None + kind: RunKind = "eval" + environments: List[EnvironmentRef] = field(default_factory=list) + model: Optional[str] = None + framework: Optional[str] = None + dataset: Optional[str] = None + description: Optional[str] = None + tags: List[str] = field(default_factory=list) + team_id: Optional[str] = None + # W&B's split, which maps cleanly onto the platform's existing columns: + # `config` is what you set going in (-> metadata), `summary` is what came + # out (-> metrics). + config: Dict[str, Any] = field(default_factory=dict) + summary: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class RunHandle: + """What a backend returns once the run exists on the other side.""" + + id: str + name: Optional[str] = None + url: Optional[str] = None + raw: Dict[str, Any] = field(default_factory=dict) diff --git a/packages/prime-runs/src/prime_runs/projection.py b/packages/prime-runs/src/prime_runs/projection.py new file mode 100644 index 000000000..de50348f4 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/projection.py @@ -0,0 +1,194 @@ +"""Projection of native traces onto the platform's v0 eval-sample format. + +Moved here from ``verifiers.v1.utils.platform``. This is knowledge about a +*platform wire format*, so it belongs in the client for that wire — not in an +eval framework that prime-rl then has to reach across a repo boundary to import +(``from verifiers.v1.push import trace_to_sample``, which had already drifted +from the module's real path). + +Everything is duck-typed. Verifiers ``Trace``/``Episode`` and prime-rl +``Rollout`` satisfy it structurally, and none of them is imported: the leaf +package that both producers depend on cannot depend back on either. + +The projection exists for the *current* viewer, which reads the flat sample +table. Once the Viewer API reads traces natively this module stops being on +the default path — which is why it is a standalone function rather than +something woven through the run lifecycle. +""" + +import logging +from typing import Any, Dict, Iterable, List, Optional, Sequence + +from ._http import encode_json + +logger = logging.getLogger(__name__) + +# Repeated /samples posts append; this is the platform's per-request ceiling. +MAX_SAMPLES_PAYLOAD_BYTES = 25 * 1024 * 1024 +# The bytes an empty {"samples":[]} envelope costs, counted against every batch. +ENVELOPE_BYTES = len(b'{"samples":[]}') + + +def json_bytes(value: Any) -> int: + """Serialized size of ``value`` under the encoding actually sent.""" + return len(encode_json(value)) + + +def _dump(items: Iterable[Any]) -> List[Dict[str, Any]]: + return [item.model_dump(mode="json", exclude_none=True) for item in items] + + +def is_episode(record: Any) -> bool: + """Whether a record is an episode (a group of traces) rather than a trace.""" + return hasattr(record, "traces") and not hasattr(record, "branches") + + +def summary_trace_index(episode: Any) -> int: + """Index of the trace whose flat projection represents the episode. + + The first trainable trace, else the first trace. Shared by the projection + and the rollout counter so the two can never disagree about which trace + (and therefore which ``task.data.idx``) an episode is numbered under. + """ + return next( + (index for index, trace in enumerate(episode.traces) if trace.agent.trainable), + 0, + ) + + +def trace_to_sample( + trace: Any, rollout_number: int = 1, episode_id: Optional[str] = None +) -> Dict[str, Any]: + """One trace -> the platform's sample dict (the v0 eval-sample format). + + The hub table stays flat — one row per trace; its episode is denormalized + onto the row (``episode_id`` from the envelope, plus the trace's own + ``agent``/``trainable``), so a multi-trace rollout's grouping travels with + each row without a nested schema. No prompt/completion split (meaningless + mid-branch): ``completion`` is the final branch's messages, ``trajectory`` + one message list per branch. + """ + task = trace.task.data.model_dump(mode="json", exclude_none=True) + branches = trace.branches + sample: Dict[str, Any] = { + "sample_id": trace.id, + "example_id": trace.task.data.idx, + "rollout_number": rollout_number, + "episode_id": episode_id, + "agent": trace.agent.name, + "trainable": trace.agent.trainable, + "task": task, + "prompt": [], + "completion": _dump(branches[-1].messages) if branches else [], + "answer": task.get("answer"), + # Keyed `tool_defs` because the v0 sample format already carries it there. + "tool_defs": _dump(trace.tools) if trace.tools else None, + "reward": trace.reward, + "timing": trace.timing.model_dump(mode="json", exclude_none=True), + "is_completed": trace.is_completed, + "is_truncated": trace.is_truncated, + "metrics": trace.metrics, + "error": trace.last_error.model_dump(mode="json", exclude_none=True) + if trace.last_error + else None, + "stop_condition": trace.stop_condition, + "trajectory": [ + { + "messages": _dump(branch.messages), + "num_input_tokens": branch.num_input_tokens, + "num_output_tokens": branch.num_output_tokens, + } + for branch in branches + ], + "token_usage": trace.usage.model_dump(mode="json", exclude_none=True) + if trace.usage + else None, + "info": dict(trace.info) or None, + } + # Flatten sub-rewards to top-level keys the way v0 does (raw scores, as v0's + # per-function outputs were); env metrics stay nested. + for name, reward in trace.rewards.items(): + if reward is not None: + sample.setdefault(name, reward.score) + return sample + + +def episode_to_samples(episode: Any, rollout_number: int) -> List[Dict[str, Any]]: + """One episode -> the sample rows the platform should store for it. + + Normally a single row: the native episode in ``info.native_wrapper`` is + authoritative and carries every trace, while one trainable trace (or the + first) supplies the flat summary older consumers read, identified by + ``native_trace_index``. An episode too large for one request falls back to + one projected row per trace, which loses the native wrapper but keeps the + run visible rather than dropping it. + """ + if not episode.traces: + return [] + summary_index = summary_trace_index(episode) + summary_trace = episode.traces[summary_index] + sample = trace_to_sample(summary_trace, rollout_number, episode.id) + sample["sample_id"] = episode.id + sample["info"] = { + **(sample["info"] or {}), + "native_wrapper": episode.to_record(), + "native_trace_index": summary_index, + } + if ENVELOPE_BYTES + json_bytes(sample) <= MAX_SAMPLES_PAYLOAD_BYTES: + return [sample] + + logger.warning( + "Episode %s exceeds the platform sample limit; uploading projected traces", + episode.id, + ) + return [trace_to_sample(trace, rollout_number, episode.id) for trace in episode.traces] + + +def build_samples( + episodes: Sequence[Any], rollout_numbers: Optional[Dict[Any, int]] = None +) -> List[Dict[str, Any]]: + """Project episodes to platform samples. + + ``rollout_numbers`` carries the per-example counter across calls. Streaming + producers pass the same dict every time so rollout numbering stays + consistent when episodes arrive in batches instead of one final list; a + one-shot caller can ignore it. + """ + counts = rollout_numbers if rollout_numbers is not None else {} + samples: List[Dict[str, Any]] = [] + for episode in episodes: + if not episode.traces: + continue + idx = episode.traces[summary_trace_index(episode)].task.data.idx + counts[idx] = number = counts.get(idx, 0) + 1 + samples.extend(episode_to_samples(episode, number)) + return samples + + +def batch_samples(samples: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any]]]: + """Split samples into request-sized batches. + + Raises ``ValueError`` on a sample too large to send alone: silently dropping + it would report a successful run that is missing rows. + """ + batches: List[List[Dict[str, Any]]] = [] + batch: List[Dict[str, Any]] = [] + payload_bytes = ENVELOPE_BYTES + for index, sample in enumerate(samples): + sample_bytes = json_bytes(sample) + if ENVELOPE_BYTES + sample_bytes > MAX_SAMPLES_PAYLOAD_BYTES: + raise ValueError( + f"sample {index} is too large to upload " + f"({ENVELOPE_BYTES + sample_bytes} > {MAX_SAMPLES_PAYLOAD_BYTES} bytes)" + ) + # The +1 is the comma that joins this sample to the previous one. + next_bytes = payload_bytes + (1 if batch else 0) + sample_bytes + if batch and next_bytes > MAX_SAMPLES_PAYLOAD_BYTES: + batches.append(batch) + batch = [] + next_bytes = ENVELOPE_BYTES + sample_bytes + batch.append(sample) + payload_bytes = next_bytes + if batch: + batches.append(batch) + return batches diff --git a/packages/prime-runs/src/prime_runs/py.typed b/packages/prime-runs/src/prime_runs/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py new file mode 100644 index 000000000..13ab0d482 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/run.py @@ -0,0 +1,704 @@ +"""The run handle, and ``init()`` that produces one. + +A run is a long-lived thing with a status, so it is an object, not three +stateless calls. That single change is what lets the SDK take on the work every +producer was doing privately: streaming instead of buffering, containing its +own errors, reporting a terminal status when the process dies, and behaving the +same on rank 3 of a training job as on a laptop. + +The identity rule matters most, so it is worth stating once. ``init()`` is +called *before* rollouts start, and the ID it returns is *the* run ID +everywhere — including inside every trace document the producer writes, and +including the local archive. Nothing is re-stamped afterwards and no producer +record is rewritten. Verifiers already stamps ``EvalRunInfo(id=config.run.id)`` +at rollout time; the only change is where that ID comes from. Offline runs get +a locally issued ID through the same path, so there is one code path, not two. +""" + +import atexit +import logging +import math +import os +import signal +import threading +import time +from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Union + +from ._http import DEFAULT_TIMEOUT, PlatformClient +from .backends import Backend, EvalsBackend, OfflineBackend +from .config import Config +from .exceptions import ConfigurationError, RunFinishedError +from .models import EnvironmentRef, Mode, OnError, RunHandle, RunKind, RunSpec, RunStatus +from .sinks import EvalSamplesSink, OfflineSink, Sink, TracesSink +from .worker import MetricItem, UploadWorker, WriteItem + +logger = logging.getLogger(__name__) + +RUN_ID_ENV = "PRIME_RUN_ID" +MODE_ENV = "PRIME_RUNS_MODE" +#: Rank variables, in the order prime-rl sets them. Rank 0 owns the lifecycle. +RANK_ENV_VARS = ("RANK", "DP_RANK", "LOCAL_RANK") +DEFAULT_SUMMARY_FLUSH_SECONDS = 10.0 + + +class Run: + """A live run: an ID, a URL, somewhere to put metrics, somewhere to put traces. + + Every method that touches the network is contained. With the default + ``on_error="warn"`` nothing raised by the platform escapes into a producer's + loop — a run that has been going for six hours does not get killed by a 502 + on a telemetry call. ``on_error="raise"`` inverts that for tests and CI, + where a silent upload failure is the bug. + """ + + def __init__( + self, + *, + backend: Backend, + handle: RunHandle, + spec: RunSpec, + sinks: Optional[List[Sink]] = None, + mode: Mode = "online", + on_error: OnError = "warn", + is_primary: bool = True, + owns_lifecycle: bool = True, + summary_flush_seconds: float = DEFAULT_SUMMARY_FLUSH_SECONDS, + queue_size: Optional[int] = None, + ) -> None: + self._backend = backend + self._handle = handle + self._spec = spec + self._mode: Mode = mode + self._on_error: OnError = on_error + self._is_primary = is_primary + # A non-primary rank shares the run but must not create or close it: + # eight ranks racing to finalize produce seven confusing failures and + # one winner. + self._owns_lifecycle = owns_lifecycle and is_primary + self._status = RunStatus.RUNNING + + self.config: Dict[str, Any] = dict(spec.config) + self.summary: Dict[str, Any] = dict(spec.summary) + self.errors: List[str] = [] + + self._summary_flush_seconds = summary_flush_seconds + self._last_summary_flush = time.monotonic() + self._summary_dirty = False + self._config_dirty = False + self._finish_lock = threading.RLock() + self._finished = False + + sinks = sinks or [] + worker_kwargs: Dict[str, Any] = {} + if queue_size is not None: + worker_kwargs["max_queue_size"] = queue_size + self._worker = UploadWorker( + sinks, + on_error=self._record_sink_error, + metric_writer=self._write_metrics if backend.supports_step_metrics else None, + **worker_kwargs, + ) + context = _sink_context(spec, handle) + for sink in sinks: + try: + sink.start(handle.id, context) + except Exception as exc: # noqa: BLE001 - a bad sink is not a bad run + sink.enabled = False + self._report(f"starting sink {getattr(sink, 'name', sink)}", exc) + + self._atexit_hook = self._on_process_exit + atexit.register(self._atexit_hook) + self._previous_signal_handlers: Dict[int, Any] = {} + + # -------------------------------------------------------------- identity + + @property + def id(self) -> str: + """The run ID. Stamp this onto every trace the run produces.""" + return self._handle.id + + @property + def name(self) -> Optional[str]: + return self._handle.name + + @property + def url(self) -> Optional[str]: + """Where to open this run — a dashboard URL, or a local path offline.""" + return self._handle.url + + @property + def kind(self) -> RunKind: + return self._spec.kind + + @property + def mode(self) -> Mode: + return self._mode + + @property + def status(self) -> RunStatus: + return self._status + + @property + def is_primary(self) -> bool: + """Whether this process owns the run's lifecycle (rank 0, or single-process).""" + return self._is_primary + + @property + def finished(self) -> bool: + return self._finished + + @property + def dropped_records(self) -> int: + """Records the uploader could not keep up with. Should be zero.""" + return self._worker.dropped + + def __repr__(self) -> str: + return ( + f"" + ) + + # ------------------------------------------------------------------- log + + def log( + self, + metrics: Mapping[str, Any], + *, + step: Optional[int] = None, + commit: bool = True, + ) -> None: + """Record scalar metrics. + + Values always land in ``summary`` last-value-wins, the way W&B's + implicit summary works. Whether they *also* become a time series + depends on the backend: the training API stores one, the evaluations + API stores a single metrics blob, and rather than making producers care, + a backend without a time series simply keeps the summary — flushed on a + timer so a tight loop does not turn into one PUT per step. + + ``commit=False`` stages values without scheduling a write, for callers + assembling a step from several places. + """ + self._require_live("log") + cleaned = _clean_metrics(metrics) + if not cleaned: + return + self.summary.update(cleaned) + self._summary_dirty = True + if not commit: + return + if self._backend.supports_step_metrics: + self._worker.submit(MetricItem(metrics=cleaned, step=step)) + return + self._maybe_flush_summary() + + def log_traces( + self, + records: Iterable[Any], + *, + line_format: Optional[str] = None, + step: Optional[int] = None, + ) -> None: + """Hand traces or episodes to the sinks. Returns immediately. + + Accepts whatever the producer already has: verifiers ``Trace`` / + ``Episode``, prime-rl ``Rollout``, or plain JSON mappings. Nothing is + buffered until the end of the run — call this as rollouts complete and + the dashboard fills in while the run is still going. + """ + self._require_live("log_traces") + batch = list(records) + if not batch: + return + self._worker.submit(WriteItem(records=batch, line_format=line_format, step=step)) + + # Producers that think in episodes rather than traces; same path. + log_episodes = log_traces + + def log_samples(self, records: Iterable[Any], *, step: Optional[int] = None) -> None: + """Alias for :meth:`log_traces`, matching prime-rl's ``Monitor`` vocabulary.""" + self.log_traces(records, step=step) + + def update_config(self, values: Mapping[str, Any]) -> None: + """Merge into the run's config (its inputs). Flushed with the summary.""" + self._require_live("update_config") + self.config.update(values) + self._config_dirty = True + + # ---------------------------------------------------------------- finish + + def finish( + self, + summary: Optional[Mapping[str, Any]] = None, + *, + status: Union[RunStatus, str] = RunStatus.COMPLETED, + error: Optional[str] = None, + ) -> None: + """Flush everything and close the run out. Idempotent. + + Safe to call from ``__exit__``, an atexit hook and a signal handler at + once — whichever gets there first reports the status, and the rest + return. + """ + with self._finish_lock: + if self._finished: + return + self._finished = True + + resolved = RunStatus(status) if not isinstance(status, RunStatus) else status + if summary: + self.summary.update(_clean_metrics(summary)) + self._status = resolved + + # Order matters: records first, so a dashboard that reacts to the + # terminal status never sees a finished run with samples still landing. + self._worker.flush(timeout=60.0) + self._worker.close() + + if self._owns_lifecycle: + self._report_guarded( + "updating the run", + lambda: self._backend.update( + self.id, + config=self.config if (self._config_dirty or self.config) else None, + summary=self.summary or None, + ), + ) + self._report_guarded( + "finalizing the run", + lambda: self._backend.finalize( + self.id, + status=resolved, + summary=self.summary or None, + error=error or (self.errors[0] if self.errors else None), + ), + ) + self._report_guarded("closing the backend", self._backend.close) + + atexit.unregister(self._atexit_hook) + self._restore_signal_handlers() + + if self._worker.dropped: + logger.warning( + "Run %s finished with %d dropped record(s)", self.id, self._worker.dropped + ) + + def fail(self, error: Union[str, BaseException]) -> None: + """Close the run out as failed.""" + self.finish(status=RunStatus.FAILED, error=_describe(error)) + + def flush(self, timeout: Optional[float] = 30.0) -> bool: + """Block until queued records have been written. Mostly for tests.""" + flushed = self._worker.flush(timeout=timeout) + self._flush_summary() + return flushed + + # -------------------------------------------------------- context manager + + def __enter__(self) -> "Run": + return self + + def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: + if exc_type is None: + self.finish() + elif isinstance(exc, KeyboardInterrupt): + # An interrupt is a decision, not a fault. Recording it as FAILED + # would put every cancelled run in the same bucket as broken ones. + self.finish(status=RunStatus.FAILED, error="interrupted") + else: + self.finish(status=RunStatus.FAILED, error=_describe(exc)) + return False + + # ------------------------------------------------------------- internals + + def install_signal_handlers(self) -> None: + """Report a terminal status when the process is killed. + + Only installed on the main thread, and only over a *default* handler: + replacing a handler the application chose would be worse than missing a + status. The previous handler is always called afterwards, so SIGINT + still raises ``KeyboardInterrupt`` and SIGTERM still terminates. + """ + if threading.current_thread() is not threading.main_thread(): + return + for signum in (signal.SIGINT, signal.SIGTERM): + try: + current = signal.getsignal(signum) + except (ValueError, OSError): # pragma: no cover - platform dependent + continue + if current not in (signal.SIG_DFL, signal.default_int_handler): + continue + try: + signal.signal(signum, self._handle_signal) + except (ValueError, OSError): # pragma: no cover + continue + self._previous_signal_handlers[signum] = current + + def _handle_signal(self, signum: int, frame: Any) -> None: + name = signal.Signals(signum).name + if not self._finished: + self.finish(status=RunStatus.FAILED, error=f"received {name}") + previous = self._previous_signal_handlers.get(signum, signal.SIG_DFL) + signal.signal(signum, previous) + if callable(previous): + previous(signum, frame) + else: + os.kill(os.getpid(), signum) + + def _restore_signal_handlers(self) -> None: + for signum, previous in self._previous_signal_handlers.items(): + try: + if signal.getsignal(signum) is self._handle_signal: + signal.signal(signum, previous) + except (ValueError, OSError): # pragma: no cover + continue + self._previous_signal_handlers.clear() + + def _on_process_exit(self) -> None: + """Last resort: the process is exiting and nobody called ``finish()``. + + Reported as CRASHED rather than FAILED — the producer never said the run + failed, it just stopped existing, and the distinction tells an operator + whether to read the run's error or go look at the machine. + """ + if self._finished: + return + logger.warning("Run %s was never finished; reporting it as crashed", self.id) + self.finish(status=RunStatus.CRASHED, error="process exited without finishing the run") + + def _require_live(self, operation: str) -> None: + if self._finished: + raise RunFinishedError( + f"{operation}() was called on run {self.id}, which is already finished. " + "The platform has closed this run out; start a new one." + ) + + def _write_metrics(self, metrics: Dict[str, Any], step: Optional[int]) -> None: + self._backend.log_metrics(self.id, metrics, step) + + def _maybe_flush_summary(self) -> None: + now = time.monotonic() + if now - self._last_summary_flush < self._summary_flush_seconds: + return + self._flush_summary() + + def _flush_summary(self) -> None: + if not (self._summary_dirty or self._config_dirty) or not self._owns_lifecycle: + return + config = self.config if self._config_dirty else None + summary = self.summary if self._summary_dirty else None + self._last_summary_flush = time.monotonic() + self._summary_dirty = False + self._config_dirty = False + self._report_guarded( + "flushing run metrics", + lambda: self._backend.update(self.id, config=config, summary=summary), + ) + + def _record_sink_error(self, sink_name: str, exc: Exception) -> None: + self._report(f"writing to the {sink_name} sink", exc) + + def _report_guarded(self, what: str, call: Any) -> None: + try: + call() + except Exception as exc: # noqa: BLE001 - routed through the error policy + self._report(what, exc) + + def _report(self, what: str, exc: Exception) -> None: + message = f"{what} failed: {type(exc).__name__}: {exc}" + self.errors.append(message) + if self._on_error == "raise": + raise exc + logger.warning("Run %s: %s", self._handle.id, message) + + +# --------------------------------------------------------------------- init + + +def init( + *, + name: Optional[str] = None, + kind: RunKind = "eval", + environments: Optional[Sequence[Any]] = None, + model: Optional[str] = None, + framework: Optional[str] = None, + dataset: Optional[str] = None, + description: Optional[str] = None, + tags: Optional[Sequence[str]] = None, + config: Optional[Mapping[str, Any]] = None, + summary: Optional[Mapping[str, Any]] = None, + id: Optional[str] = None, + mode: Optional[Mode] = None, + dir: Optional[str] = None, + team_id: Optional[str] = None, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + traces_url: Optional[str] = None, + traces: bool = True, + samples: bool = True, + sinks: Optional[List[Sink]] = None, + on_error: OnError = "warn", + handle_signals: bool = True, + queue_size: Optional[int] = None, +) -> Run: + """Start a run and return a handle to it. + + Call this *before* the first rollout: the ID it returns is what every trace + in the run should carry, and the URL it returns is what a producer prints so + someone can watch the run fill in. + + ``mode`` selects where the run lives. Left unset it is read from + ``$PRIME_RUNS_MODE``, and failing that inferred: online when there is an API + key, offline when there is not. Offline is a real run with a real ID and a + real directory, which is why producers no longer need a ``--no-push`` + branch — the call sites are identical either way. + + ``id`` attaches to an existing run instead of creating one, for resuming + after a crash and for non-primary ranks joining a run rank 0 created. + """ + resolved_config = Config() + api_key = api_key if api_key is not None else resolved_config.api_key + base_url = base_url or resolved_config.base_url + team_id = team_id if team_id is not None else resolved_config.team_id + + spec = RunSpec( + name=name, + kind=kind, + environments=[EnvironmentRef.coerce(entry) for entry in (environments or [])], + model=model, + framework=framework, + dataset=dataset, + description=description, + tags=list(tags or []), + team_id=team_id, + config=dict(config or {}), + summary=dict(summary or {}), + ) + + is_primary = _is_primary_rank() + inherited_id = id or os.getenv(RUN_ID_ENV) or None + resolved_mode = _resolve_mode(mode, api_key=api_key, is_primary=is_primary, run_id=inherited_id) + + if resolved_mode == "disabled": + backend: Backend = _DisabledBackend() + handle = RunHandle(id=inherited_id or _local_id(), name=name) + return _build( + spec, backend, handle, [], resolved_mode, on_error, is_primary, False, queue_size + ) + + if resolved_mode == "offline": + offline = OfflineBackend(dir) + handle = offline.attach(inherited_id) if inherited_id else offline.create(spec) + run_sinks = sinks if sinks is not None else [OfflineSink(offline.directory)] + run = _build( + spec, + offline, + handle, + run_sinks, + resolved_mode, + on_error, + is_primary, + True, + queue_size, + ) + _announce(run, handle_signals) + return run + + if kind != "eval": + raise ConfigurationError( + f"kind={kind!r} is not supported yet — training runs arrive with the RFT backend. " + 'Use kind="eval", or mode="offline" to record the run locally.' + ) + if not api_key: + raise ConfigurationError( + 'mode="online" needs an API key. Set PRIME_API_KEY, run `prime login`, ' + 'or pass mode="offline".' + ) + + client = PlatformClient(api_key=api_key, base_url=base_url, timeout=DEFAULT_TIMEOUT) + backend = EvalsBackend(client, frontend_url=resolved_config.frontend_url, team_id=team_id) + owns_lifecycle = inherited_id is None + handle = backend.attach(inherited_id) if inherited_id else backend.create(spec) + + if sinks is None: + run_sinks = [] + if traces: + run_sinks.append(TracesSink(api_key=api_key, traces_url=traces_url, team_id=team_id)) + if samples: + # Both transports run during the transition: traces is the system of + # record, the sample table is what today's viewer reads, and Prime + # Traces is still gated to an account allowlist. + run_sinks.append(EvalSamplesSink(client)) + else: + run_sinks = list(sinks) + + run = _build( + spec, + backend, + handle, + run_sinks, + resolved_mode, + on_error, + is_primary, + owns_lifecycle, + queue_size, + ) + _announce(run, handle_signals) + return run + + +def _build( + spec: RunSpec, + backend: Backend, + handle: RunHandle, + sinks: List[Sink], + mode: Mode, + on_error: OnError, + is_primary: bool, + owns_lifecycle: bool, + queue_size: Optional[int], +) -> Run: + return Run( + backend=backend, + handle=handle, + spec=spec, + sinks=sinks, + mode=mode, + on_error=on_error, + is_primary=is_primary, + owns_lifecycle=owns_lifecycle, + queue_size=queue_size, + ) + + +def _announce(run: Run, handle_signals: bool) -> None: + """Publish the run ID to child processes and arm crash reporting. + + Exporting ``PRIME_RUN_ID`` is how forked workers and subprocess launchers + join the run their parent created instead of each opening their own — the + same trick prime-rl's monitor used with ``RUN_ID``, generalized so every + producer gets it. + """ + os.environ.setdefault(RUN_ID_ENV, run.id) + if handle_signals: + run.install_signal_handlers() + if run.url: + logger.info("Run %s: %s", run.id, run.url) + + +class _DisabledBackend: + """No-op lifecycle, so ``mode="disabled"`` needs no branching upstream.""" + + kind = "disabled" + supports_step_metrics = False + + def create(self, spec: RunSpec) -> RunHandle: + return RunHandle(id=_local_id()) + + def attach(self, run_id: str) -> RunHandle: + return RunHandle(id=run_id) + + def update(self, run_id: str, **kwargs: Any) -> None: + return None + + def log_metrics(self, run_id: str, metrics: Dict[str, Any], step: Optional[int] = None) -> None: + return None + + def finalize(self, run_id: str, **kwargs: Any) -> None: + return None + + def close(self) -> None: + return None + + +def _local_id() -> str: + from .backends.offline import new_run_id + + return new_run_id() + + +def _is_primary_rank() -> bool: + """Whether this process should own the run's lifecycle. + + Any rank variable set to something other than 0 means a peer process is + rank 0 and owns creation and finalization. Non-primary ranks still upload + their own records — the point is that eight processes contribute to one run + rather than creating eight. + """ + for name in RANK_ENV_VARS: + value = os.getenv(name) + if value and value.strip() not in ("0", ""): + return False + return True + + +def _resolve_mode( + mode: Optional[Mode], *, api_key: str, is_primary: bool, run_id: Optional[str] +) -> Mode: + if mode is None: + env_mode = os.getenv(MODE_ENV) + if env_mode: + mode = env_mode.strip().lower() # type: ignore[assignment] + if mode not in (None, "online", "offline", "disabled"): + raise ConfigurationError(f"mode={mode!r} is not one of 'online', 'offline' or 'disabled'") + if mode is None: + if api_key: + mode = "online" + else: + logger.warning( + "No API key found (set PRIME_API_KEY or run `prime login`); " + "recording this run offline instead." + ) + mode = "offline" + if mode == "online" and not is_primary and not run_id: + # A non-primary rank with no run to join would create a second run for + # the same job. Recording nothing is better than that. + logger.debug("Non-primary rank with no %s; disabling this run handle", RUN_ID_ENV) + return "disabled" + return mode # type: ignore[return-value] + + +def _sink_context(spec: RunSpec, handle: RunHandle) -> Dict[str, str]: + """Upload-scoped provenance. + + Not the join key — that is ``run.id`` inside the trace document, which the + ingestion service extracts into an indexed column. What goes here is what + you would want when looking at an upload and asking where it came from. + """ + context = {"source": "prime-runs", "run_kind": spec.kind} + if spec.framework: + context["framework"] = spec.framework + if spec.model: + context["model"] = spec.model + return context + + +def _describe(error: Union[str, BaseException]) -> str: + if isinstance(error, BaseException): + return f"{type(error).__name__}: {error}" + return str(error) + + +def _clean_metrics(metrics: Mapping[str, Any]) -> Dict[str, Any]: + """Drop values JSON cannot carry. + + NaN and infinity are the ones that matter: a diverged loss serializes as + JavaScript's bare ``NaN``, which strict JSON rejects, and the failure + surfaces as an opaque 400 on a payload nobody can inspect. Dropping the key + loses one point; sending it loses the request. + """ + cleaned: Dict[str, Any] = {} + for key, value in metrics.items(): + if isinstance(value, float) and not math.isfinite(value): + logger.debug("Dropping non-finite metric %s=%r", key, value) + continue + if isinstance(value, Mapping): + nested = _clean_metrics(value) + if nested: + cleaned[key] = nested + continue + cleaned[key] = value + return cleaned + + +__all__ = ["Run", "init", "RUN_ID_ENV", "MODE_ENV"] diff --git a/packages/prime-runs/src/prime_runs/sinks/__init__.py b/packages/prime-runs/src/prime_runs/sinks/__init__.py new file mode 100644 index 000000000..115536d25 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/__init__.py @@ -0,0 +1,14 @@ +"""Sample transports. Independent of backends, and of each other.""" + +from .base import Sink, to_mapping +from .offline import OfflineSink +from .samples import EvalSamplesSink +from .traces import TracesSink + +__all__ = [ + "Sink", + "to_mapping", + "EvalSamplesSink", + "OfflineSink", + "TracesSink", +] diff --git a/packages/prime-runs/src/prime_runs/sinks/base.py b/packages/prime-runs/src/prime_runs/sinks/base.py new file mode 100644 index 000000000..e6a4771e1 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/base.py @@ -0,0 +1,65 @@ +"""The contract a sample sink implements. + +A sink moves *records* — traces, episodes, rollouts — to wherever they are +stored. It knows nothing about run lifecycle; a backend closing a run and a +sink flushing its last batch are separate events on purpose. + +Sinks are independent of backends and of each other. During the transition both +the traces sink and the legacy eval-samples sink run at once, so the dashboard +keeps working for accounts outside the traces beta while traces becomes the +system of record. When the Viewer API reads traces natively, the default sink +list drops one entry — and no producer changes. + +Every sink must be *degradable*: a sink that cannot write sets ``enabled = +False`` and says why, once. A run whose traces are gated is still a valid run. +""" + +from typing import Any, Mapping, Optional, Protocol, Sequence, runtime_checkable + + +@runtime_checkable +class Sink(Protocol): + """A destination for run records.""" + + name: str + enabled: bool + + def start(self, run_id: str, context: Mapping[str, str]) -> None: + """Bind the sink to a run before the first write.""" + ... + + def write( + self, + records: Sequence[Any], + *, + line_format: Optional[str] = None, + step: Optional[int] = None, + ) -> None: + """Send one batch. Called from the uploader thread, never inline.""" + ... + + def flush(self) -> None: + """Block until everything handed over so far has been written.""" + ... + + def close(self) -> None: + """Release transport resources.""" + ... + + +def to_mapping(record: Any) -> Mapping[str, Any]: + """The JSON mapping for a record, whatever shape the producer handed us. + + Mirrors ``prime_traces.SupportsToRecord``: verifiers ``Trace``/``Episode`` + and prime-rl ``Rollout`` all implement ``to_record()``, and plain dicts pass + straight through. + """ + if isinstance(record, Mapping): + return record + to_record = getattr(record, "to_record", None) + if callable(to_record): + value = to_record() + if isinstance(value, Mapping): + return value + raise TypeError(f"{type(record).__name__}.to_record() must return a mapping") + raise TypeError(f"{type(record).__name__} is not a mapping and has no to_record()") diff --git a/packages/prime-runs/src/prime_runs/sinks/offline.py b/packages/prime-runs/src/prime_runs/sinks/offline.py new file mode 100644 index 000000000..73f65392c --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/offline.py @@ -0,0 +1,90 @@ +"""Local JSONL sink, written in the wire format Prime Traces accepts. + +Deliberately not a debug dump: the files this writes are valid trace/episode +JSONL, so ``prime_traces.TracesClient.upload_file`` can send them later +untouched. That is what makes an offline run a *deferred* run rather than a +different one — the run ID stamped into the records was issued at ``init()`` +and does not change on sync. +""" + +import json +import logging +from pathlib import Path +from typing import Any, Mapping, Optional, Sequence, TextIO, Union + +from .base import Sink, to_mapping + +logger = logging.getLogger(__name__) + + +class OfflineSink(Sink): + """Appends records to ``//records/.jsonl``.""" + + name = "offline" + + def __init__(self, directory: Union[str, Path], *, stamp_run: bool = True) -> None: + self.enabled = True + self.directory = Path(directory) + self._stamp_run = stamp_run + self._run_id: Optional[str] = None + self._run_kind: Optional[str] = None + self._handles: dict[str, TextIO] = {} + self.records_written = 0 + + def start(self, run_id: str, context: Mapping[str, str]) -> None: + self._run_id = run_id + self._run_kind = context.get("run_kind") + self._records_dir.mkdir(parents=True, exist_ok=True) + + @property + def _records_dir(self) -> Path: + return self.directory / (self._run_id or "unknown") / "records" + + def write( + self, + records: Sequence[Any], + *, + line_format: Optional[str] = None, + step: Optional[int] = None, + ) -> None: + if not self.enabled or not records: + return + name = str(getattr(line_format, "value", line_format) or _infer_format(records[0])) + handle = self._handle(name) + for record in records: + mapping = dict(to_mapping(record)) + if self._stamp_run and not mapping.get("run") and self._run_id: + run: dict[str, Any] = {"id": self._run_id} + if self._run_kind: + run["type"] = self._run_kind + mapping["run"] = run + handle.write( + json.dumps(mapping, ensure_ascii=False, separators=(",", ":"), default=str) + "\n" + ) + self.records_written += 1 + + def _handle(self, name: str) -> TextIO: + handle = self._handles.get(name) + if handle is None: + self._records_dir.mkdir(parents=True, exist_ok=True) + handle = (self._records_dir / f"{name}.jsonl").open("a", encoding="utf-8") + self._handles[name] = handle + return handle + + def flush(self) -> None: + for handle in self._handles.values(): + handle.flush() + + def close(self) -> None: + for handle in self._handles.values(): + try: + handle.close() + except OSError as exc: # pragma: no cover - teardown must not raise + logger.debug("Error closing an offline record file: %s", exc) + self._handles.clear() + + +def _infer_format(record: Any) -> str: + if isinstance(record, Mapping): + return "episode" if "traces" in record else "trace" + return "episode" if hasattr(record, "traces") else "trace" diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py new file mode 100644 index 000000000..55a1755f2 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -0,0 +1,99 @@ +"""Legacy sink: the flat eval-sample table behind today's viewer. + +This exists so the migration is a refactor rather than a regression. The viewer +reads the v0 sample table; Prime Traces is in closed beta on an account +allowlist. Shipping traces-only would leave every non-allowlisted account +staring at an empty dashboard — so both sinks run, and this one retires when +the Viewer API reads traces natively. Retiring it is a one-line change to the +default sink list, with nothing to do in verifiers or prime-rl. + +Its known weakness is why traces is the primary: ``POST /samples`` *appends*, +so a retried request whose response was lost duplicates rows. Content-addressed +uploads do not have that problem, which is exactly the property the traces sink +was built on. +""" + +import logging +from typing import Any, Dict, List, Mapping, Optional, Sequence + +from .._http import UPLOAD_TIMEOUT, PlatformClient, encode_json +from ..projection import batch_samples, build_samples +from .base import Sink + +logger = logging.getLogger(__name__) + + +class EvalSamplesSink(Sink): + """Projects episodes to v0 samples and pushes them to the evaluations API.""" + + name = "eval_samples" + + def __init__(self, client: PlatformClient) -> None: + self.enabled = True + self._client = client + self._run_id: Optional[str] = None + # Carried across calls so a streaming producer numbers rollouts the same + # way a one-shot upload does: the Nth episode for an example is rollout N, + # whether it arrived alone or in a batch of five hundred. + self._rollout_numbers: Dict[Any, int] = {} + self.samples_written = 0 + + def start(self, run_id: str, context: Mapping[str, str]) -> None: + self._run_id = run_id + + def write( + self, + records: Sequence[Any], + *, + line_format: Optional[str] = None, + step: Optional[int] = None, + ) -> None: + if not self.enabled or not records: + return + if self._run_id is None: + raise RuntimeError("EvalSamplesSink.write called before start()") + + samples = self._to_samples(records) + if not samples: + return + for batch in batch_samples(samples): + self._client.post( + f"/evaluations/{self._run_id}/samples", + content=encode_json({"samples": batch}), + timeout=UPLOAD_TIMEOUT, + ) + self.samples_written += len(batch) + + def _to_samples(self, records: Sequence[Any]) -> List[Dict[str, Any]]: + """Split a batch into episodes to project and samples to pass through. + + A producer that already speaks the v0 sample format (a dict with + ``sample_id``) sends it unchanged; anything with ``traces`` is a native + episode and gets projected. Anything else is skipped loudly rather than + posted as a malformed row the API would reject for the whole batch. + """ + episodes: List[Any] = [] + passthrough: List[Dict[str, Any]] = [] + for record in records: + if isinstance(record, Mapping): + if "sample_id" in record: + passthrough.append(dict(record)) + elif "traces" in record: + logger.debug( + "Skipping a pre-serialized episode: this sink projects native " + "episode objects, not their JSON records" + ) + else: + logger.debug("Skipping a record with no sample_id and no traces") + continue + if hasattr(record, "traces"): + episodes.append(record) + else: + logger.debug("Skipping %s: not an episode", type(record).__name__) + return build_samples(episodes, self._rollout_numbers) + passthrough + + def flush(self) -> None: + """Writes are synchronous; the uploader thread owns the asynchrony.""" + + def close(self) -> None: + """The platform client is shared with the backend, which closes it.""" diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py new file mode 100644 index 000000000..8709aaedb --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -0,0 +1,194 @@ +"""Streaming sink over Prime Traces — the primary sample transport. + +Records go out as they are produced, in content-addressed JSONL batches. Two +properties of that transport are why the run handle can promise what it does: +uploads are idempotent (the same bytes resolve to the same upload ID, so a +retry after a lost response replays rather than duplicates), and they are +episode-aware, so a multi-trace rollout keeps its grouping instead of being +flattened into one summary row. + +**The join key is ``run.id`` inside the trace document, not an upload context +key.** The ingestion service extracts ``run.id`` into an indexed column with a +delete-by-run path; ``context`` is an upload-scoped map that answers a +different question. Producers already stamp the run onto their traces, and +``init()`` returns the ID they stamp — so this sink adds nothing to the join +and uses ``context`` only for provenance. +""" + +import logging +from typing import Any, Dict, Mapping, Optional, Sequence + +from .base import Sink + +logger = logging.getLogger(__name__) + + +class TracesSink(Sink): + """Uploads records through the Prime Traces service.""" + + name = "traces" + + def __init__( + self, + *, + client: Optional[Any] = None, + api_key: Optional[str] = None, + traces_url: Optional[str] = None, + team_id: Optional[str] = None, + stamp_run: bool = True, + compress: bool = True, + ) -> None: + self.enabled = True + self._client = client + # Left unset, prime-traces resolves its own endpoint. That matters: + # the service has its own URL (PRIME_TRACES_URL / config `traces_url`) + # which is not necessarily the platform API's, and passing the + # platform base URL through here would quietly override it. + self._client_kwargs: Dict[str, Any] = {} + if api_key is not None: + self._client_kwargs["api_key"] = api_key + if traces_url is not None: + self._client_kwargs["base_url"] = traces_url + if team_id is not None: + self._client_kwargs["team_id"] = team_id + self._stamp_run = stamp_run + self._compress = compress + self._run_id: Optional[str] = None + self._run_kind: Optional[str] = None + self._context: Dict[str, str] = {} + self.receipts: list = [] + + # ------------------------------------------------------------------ setup + + def start(self, run_id: str, context: Mapping[str, str]) -> None: + self._run_id = run_id + self._run_kind = context.get("run_kind") + self._context = {key: str(value) for key, value in context.items() if value is not None} + if self._client is None: + try: + from prime_traces import TracesClient + except ImportError as exc: # pragma: no cover - dependency is declared + self._disable(f"prime-traces is not installed ({exc})") + return + try: + self._client = TracesClient(**self._client_kwargs) + except Exception as exc: # noqa: BLE001 - construction must not kill a run + self._disable(f"could not construct the traces client ({exc})") + + # ------------------------------------------------------------------ write + + def write( + self, + records: Sequence[Any], + *, + line_format: Optional[str] = None, + step: Optional[int] = None, + ) -> None: + if not self.enabled or not records or self._client is None: + return + + from prime_traces import LineFormat + + resolved = _resolve_line_format(line_format, records, LineFormat) + context = dict(self._context) + if step is not None: + context["step"] = str(step) + + payload = [self._prepare(record) for record in records] + try: + receipts = self._client.upload_records( + payload, + line_format=resolved, + context=context or None, + compress=self._compress, + ) + except Exception as exc: # noqa: BLE001 - classified below + if self._is_gated(exc): + self._disable( + f"Prime Traces is not enabled for this account ({exc}); " + "falling back to the remaining sinks" + ) + return + raise + self.receipts.extend(receipts) + + def _prepare(self, record: Any) -> Any: + """Stamp the run onto plain mappings that do not already carry one. + + Producer objects are passed through untouched — verifiers and prime-rl + both stamp the run themselves at rollout time, and rewriting a caller's + object to add something it already has is how two sources of truth for + the run ID appear. A bare dict has no such convention, so filling in + the indexed field is the difference between a queryable run and an + orphaned upload. + """ + if not self._stamp_run or not isinstance(record, Mapping): + return record + if record.get("run"): + return record + run: Dict[str, Any] = {"id": self._run_id} + if self._run_kind: + run["type"] = self._run_kind + return {**record, "run": run} + + def flush(self) -> None: + """Uploads are synchronous, so nothing is held back here. + + Batching happens inside ``upload_records``; the asynchrony a producer + cares about lives one level up, in the uploader thread. + """ + + def close(self) -> None: + client = self._client + self._client = None + if client is not None and hasattr(client, "close"): + try: + client.close() + except Exception as exc: # noqa: BLE001 - teardown must not raise + logger.debug("Error closing the traces client: %s", exc) + + # ----------------------------------------------------------- degradation + + def _disable(self, reason: str) -> None: + if self.enabled: + logger.warning("Traces sink disabled: %s", reason) + self.enabled = False + + @staticmethod + def _is_gated(exc: Exception) -> bool: + """Whether this failure means "not allowed", not "try again". + + Prime Traces is in closed beta: a non-allowlisted account gets 403 + ``service_not_enabled``, and a write-only hosted-eval token gets 403 + ``forbidden`` on anything it may not do. Neither is fixable at runtime, + so the sink turns itself off instead of retrying for the rest of the run. + """ + try: + from prime_traces.exceptions import ForbiddenError + except ImportError: # pragma: no cover - dependency is declared + return False + return isinstance(exc, ForbiddenError) + + +def _resolve_line_format(line_format: Optional[str], records: Sequence[Any], enum: Any) -> Any: + """Pick the wire format, preferring what the caller said. + + The default is inferred from the records themselves: anything carrying + ``traces`` is an episode. Guessing wrong is not cosmetic — the same bytes + submitted under a different format are rejected as a conflict — so the + inference only ever looks at the first record's shape, which is stable + within a batch a producer handed over as a unit. + """ + if line_format is not None: + return enum(line_format) if not isinstance(line_format, enum) else line_format + first = records[0] + mapping = _try_mapping(first) + if mapping is not None: + return enum.EPISODE if "traces" in mapping else enum.TRACE + return enum.EPISODE if hasattr(first, "traces") else enum.TRACE + + +def _try_mapping(record: Any) -> Optional[Mapping[str, Any]]: + if isinstance(record, Mapping): + return record + return None diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py new file mode 100644 index 000000000..1d6a73ba8 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -0,0 +1,250 @@ +"""Background uploader: the thread that keeps the network off the rollout loop. + +Three things a producer should never have to think about, handled once here: + +**Backpressure.** The queue is bounded. Verifiers' uploader held every episode +of a run in memory and posted them all at the end, which is fine at a hundred +episodes and is an OOM at a hundred thousand. A bounded queue trades that for a +short block, and — past the block — a counted drop, because stalling a training +run to protect telemetry is the wrong trade in the other direction. + +**Fork safety.** Hosted evals fork after the SDK is initialized. A forked child +inherits the queue's *memory* but not the thread that drains it, so anything +already queued would sit there forever and any lock held mid-write stays held. +The child therefore starts over with an empty queue and a fresh thread, and +drops what it inherited: those records belong to the parent, which is still +running and will upload them itself. + +**Containment.** A sink that raises is retried once, then disabled for the rest +of the run with the error reported through the run's error handler. The upload +thread never propagates into the producer, and never dies quietly either. +""" + +import logging +import os +import queue +import threading +from dataclasses import dataclass, field +from typing import Any, Callable, List, Optional, Sequence + +logger = logging.getLogger(__name__) + +DEFAULT_QUEUE_SIZE = 256 +DEFAULT_PUT_TIMEOUT = 5.0 + + +@dataclass +class WriteItem: + """One batch of records destined for every enabled sink.""" + + records: Sequence[Any] + line_format: Optional[str] = None + step: Optional[int] = None + + +@dataclass +class MetricItem: + """One ``log()`` call destined for a backend that stores a time series.""" + + metrics: dict + step: Optional[int] = None + + +@dataclass +class _Flush: + """A barrier the caller waits on.""" + + event: threading.Event = field(default_factory=threading.Event) + + +class UploadWorker: + """Drains a bounded queue into a list of sinks on one daemon thread.""" + + def __init__( + self, + sinks: List[Any], + *, + max_queue_size: int = DEFAULT_QUEUE_SIZE, + put_timeout: float = DEFAULT_PUT_TIMEOUT, + on_error: Optional[Callable[[str, Exception], None]] = None, + metric_writer: Optional[Callable[[dict, Optional[int]], None]] = None, + ) -> None: + self.sinks = sinks + self.max_queue_size = max_queue_size + self.put_timeout = put_timeout + self._on_error = on_error + # Set when the backend stores a real time series. Metrics then ride the + # same queue as records, so a per-step log() in a training loop costs a + # queue put rather than an HTTP round trip. + self._metric_writer = metric_writer + self._queue: "queue.Queue[Any]" = queue.Queue(maxsize=max_queue_size) + self._thread: Optional[threading.Thread] = None + self._stopping = threading.Event() + self._lock = threading.Lock() + self.dropped = 0 + self._pid = os.getpid() + self._register_fork_hook() + + # ----------------------------------------------------------------- thread + + def start(self) -> None: + with self._lock: + if self._thread is not None and self._thread.is_alive(): + return + self._stopping.clear() + self._thread = threading.Thread( + target=self._run, name="prime-runs-uploader", daemon=True + ) + self._thread.start() + + def _run(self) -> None: + while True: + item = self._queue.get() + try: + if item is None: + return + if isinstance(item, _Flush): + self._flush_sinks() + item.event.set() + continue + if isinstance(item, MetricItem): + self._write_metrics(item) + continue + self._dispatch(item) + except Exception as exc: # noqa: BLE001 - the thread must outlive one bad batch + logger.debug("Uploader iteration failed: %s", exc) + finally: + self._queue.task_done() + + def _dispatch(self, item: WriteItem) -> None: + for sink in self.sinks: + if not getattr(sink, "enabled", True): + continue + try: + sink.write(item.records, line_format=item.line_format, step=item.step) + except Exception as exc: # noqa: BLE001 - one sink failing must not stop the others + self._fail_sink(sink, exc) + + def _write_metrics(self, item: MetricItem) -> None: + if self._metric_writer is None: + return + try: + self._metric_writer(item.metrics, item.step) + except Exception as exc: # noqa: BLE001 - metrics must not kill the uploader + logger.warning( + "Dropped metrics for step %s: %s: %s", item.step, type(exc).__name__, exc + ) + if self._on_error is not None: + try: + self._on_error("metrics", exc) + except Exception: # noqa: BLE001 + logger.debug("Error handler raised while reporting metrics", exc_info=True) + + def _flush_sinks(self) -> None: + for sink in self.sinks: + if not getattr(sink, "enabled", True): + continue + try: + sink.flush() + except Exception as exc: # noqa: BLE001 + self._fail_sink(sink, exc) + + def _fail_sink(self, sink: Any, exc: Exception) -> None: + """Disable a sink that raised, and report it exactly once. + + Not retried here: the transports already retry internally (traces on + content-addressed uploads, the platform client on 429/5xx), so an error + that reaches this point has already exhausted its budget. Continuing to + call a sink in that state produces one log line per batch for the rest + of the run and hides whatever failed first. + """ + name = getattr(sink, "name", type(sink).__name__) + sink.enabled = False + logger.warning("Sink %s disabled after an error: %s: %s", name, type(exc).__name__, exc) + if self._on_error is not None: + try: + self._on_error(name, exc) + except Exception: # noqa: BLE001 - the handler is the caller's problem + logger.debug("Error handler raised while reporting a sink failure", exc_info=True) + + # ------------------------------------------------------------------ queue + + def submit(self, item: Any) -> bool: + """Hand a batch or a metric point to the uploader. + + ``False`` means it was dropped: the queue stayed full for the whole + timeout, so the producer is durably outrunning the uploader. Blocking + further would turn a telemetry backlog into a stalled training run. + """ + if self._stopping.is_set(): + return False + thread = self._thread + if thread is None or not thread.is_alive(): + self.start() + try: + self._queue.put(item, timeout=self.put_timeout) + return True + except queue.Full: + count = len(item.records) if isinstance(item, WriteItem) else 1 + self.dropped += count + logger.warning( + "Upload queue full after %.1fs; dropped %d item(s) (%d total). " + "The producer is outrunning the uploader.", + self.put_timeout, + count, + self.dropped, + ) + return False + + def flush(self, timeout: Optional[float] = None) -> bool: + """Block until everything queued so far has been written.""" + if self._thread is None or not self._thread.is_alive(): + self._flush_sinks() + return True + barrier = _Flush() + try: + self._queue.put(barrier, timeout=self.put_timeout) + except queue.Full: + logger.warning("Could not enqueue a flush barrier; the queue is saturated") + return False + return barrier.event.wait(timeout) + + def close(self, timeout: Optional[float] = 30.0) -> None: + """Drain, stop the thread, and close every sink.""" + self._stopping.set() + thread = self._thread + if thread is not None and thread.is_alive(): + try: + self._queue.put(None, timeout=self.put_timeout) + except queue.Full: + logger.warning("Upload queue saturated at close; some records may be lost") + thread.join(timeout) + if thread.is_alive(): + logger.warning("Uploader did not stop within %ss; abandoning it", timeout) + self._thread = None + for sink in self.sinks: + try: + sink.close() + except Exception as exc: # noqa: BLE001 - teardown must not raise + logger.debug("Error closing sink %s: %s", getattr(sink, "name", sink), exc) + + # ------------------------------------------------------------------- fork + + def _register_fork_hook(self) -> None: + if not hasattr(os, "register_at_fork"): # pragma: no cover - Windows + return + os.register_at_fork(after_in_child=self._reinit_after_fork) + + def _reinit_after_fork(self) -> None: + """Give the child a clean uploader. + + Everything queued at fork time belongs to the parent, which still has a + live thread and will send it. Inheriting that queue would upload each + record twice; inheriting the lock could deadlock the child on its first + write. + """ + self._pid = os.getpid() + self._queue = queue.Queue(maxsize=self.max_queue_size) + self._thread = None + self._stopping = threading.Event() + self._lock = threading.Lock() diff --git a/packages/prime-runs/tests/_fakes.py b/packages/prime-runs/tests/_fakes.py new file mode 100644 index 000000000..4cef94549 --- /dev/null +++ b/packages/prime-runs/tests/_fakes.py @@ -0,0 +1,115 @@ +"""Stand-ins for verifiers ``Trace``/``Episode``. + +The SDK never imports a producer package — records are duck-typed — so the +tests must not either, or they would be testing an import rather than the +protocol. These objects implement exactly the surface +:mod:`prime_runs.projection` and :mod:`prime_runs.metrics` touch, which makes +that surface explicit and makes an accidental widening of it fail here. +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + + +class Dumpable: + """Anything the projection calls ``model_dump()`` on.""" + + def __init__(self, **values: Any) -> None: + self._values = values + + def model_dump(self, mode: str = "python", exclude_none: bool = False) -> Dict[str, Any]: + if exclude_none: + return {key: value for key, value in self._values.items() if value is not None} + return dict(self._values) + + +@dataclass +class Reward: + score: float + + +@dataclass +class Agent: + name: str = "solver" + trainable: bool = True + + +@dataclass +class Branch: + messages: List[Dumpable] = field(default_factory=list) + num_input_tokens: int = 10 + num_output_tokens: int = 5 + + +class TaskData(Dumpable): + def __init__(self, idx: int = 0, answer: str = "42", **extra: Any) -> None: + super().__init__(idx=idx, answer=answer, **extra) + self.idx = idx + + +@dataclass +class Task: + data: TaskData = field(default_factory=TaskData) + + +@dataclass +class Trace: + id: str = "trace-1" + task: Task = field(default_factory=Task) + agent: Agent = field(default_factory=Agent) + branches: List[Branch] = field(default_factory=list) + tools: Optional[List[Dumpable]] = None + reward: float = 1.0 + timing: Dumpable = field(default_factory=lambda: Dumpable(total_ms=120)) + is_completed: bool = True + is_truncated: bool = False + metrics: Dict[str, float] = field(default_factory=dict) + last_error: Optional[Dumpable] = None + stop_condition: Optional[str] = "stop" + usage: Optional[Dumpable] = None + info: Dict[str, Any] = field(default_factory=dict) + rewards: Dict[str, Optional[Reward]] = field(default_factory=dict) + + def to_record(self) -> Dict[str, Any]: + return {"id": self.id, "reward": self.reward} + + +@dataclass +class Episode: + id: str = "episode-1" + traces: List[Trace] = field(default_factory=list) + ok: bool = True + + def to_record(self) -> Dict[str, Any]: + return {"id": self.id, "traces": [trace.to_record() for trace in self.traces]} + + +def make_trace( + *, + trace_id: str = "trace-1", + idx: int = 0, + trainable: bool = True, + reward: float = 1.0, + agent: str = "solver", + rewards: Optional[Dict[str, Optional[Reward]]] = None, + metrics: Optional[Dict[str, float]] = None, + branches: int = 1, +) -> Trace: + return Trace( + id=trace_id, + task=Task(data=TaskData(idx=idx)), + agent=Agent(name=agent, trainable=trainable), + branches=[ + Branch(messages=[Dumpable(role="assistant", content=f"branch {n}")]) + for n in range(branches) + ], + reward=reward, + metrics=metrics or {}, + rewards=rewards or {}, + ) + + +def make_episode( + episode_id: str = "episode-1", traces: Optional[List[Trace]] = None, ok: bool = True +) -> Episode: + return Episode(id=episode_id, traces=traces if traces is not None else [make_trace()], ok=ok) diff --git a/packages/prime-runs/tests/conftest.py b/packages/prime-runs/tests/conftest.py new file mode 100644 index 000000000..745ba07d4 --- /dev/null +++ b/packages/prime-runs/tests/conftest.py @@ -0,0 +1,148 @@ +"""Shared fixtures. Every test is hermetic: no network, no real ~/.prime.""" + +from pathlib import Path +from typing import Any, Callable, Dict, List + +import httpx +import pytest + +from prime_runs._http import PlatformClient + +_PRIME_ENV_VARS = ( + "PRIME_API_KEY", + "PRIME_TEAM_ID", + "PRIME_API_BASE_URL", + "PRIME_BASE_URL", + "PRIME_TRACES_URL", + "PRIME_FRONTEND_URL", + "PRIME_RUN_ID", + "PRIME_RUNS_MODE", + "PRIME_RUNS_DIR", + "RANK", + "DP_RANK", + "LOCAL_RANK", +) + + +@pytest.fixture(autouse=True) +def isolated_prime_config(monkeypatch, tmp_path): + """Never read the developer's real ~/.prime, env vars or rank.""" + monkeypatch.setattr(Path, "home", lambda: tmp_path) + for name in _PRIME_ENV_VARS: + monkeypatch.delenv(name, raising=False) + return tmp_path + + +@pytest.fixture(autouse=True) +def no_sleep(monkeypatch): + """Record retry backoff instead of waiting it out.""" + sleeps: List[float] = [] + monkeypatch.setattr("prime_runs._http.time.sleep", sleeps.append) + return sleeps + + +class RecordingHandler: + """A MockTransport handler that records requests and replies from a route map.""" + + def __init__(self, routes: Dict[str, Any]) -> None: + self.routes = routes + self.requests: List[httpx.Request] = [] + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.requests.append(request) + key = f"{request.method} {request.url.path}" + route = self.routes.get(key) + if route is None: + return httpx.Response(404, json={"detail": f"no route for {key}"}) + if callable(route): + return route(request) + return httpx.Response(200, json=route) + + def paths(self) -> List[str]: + return [f"{r.method} {r.url.path}" for r in self.requests] + + def bodies_for(self, path: str) -> List[Any]: + import json + + return [json.loads(r.content) for r in self.requests if r.url.path == path and r.content] + + +@pytest.fixture +def make_platform_client() -> Callable[..., PlatformClient]: + def _make(handler: Callable[[httpx.Request], httpx.Response], **kwargs: Any) -> PlatformClient: + client = httpx.Client( + base_url="http://testserver", + transport=httpx.MockTransport(handler), + headers={"Authorization": "Bearer test-key"}, + ) + return PlatformClient( + api_key="test-key", base_url="http://testserver", client=client, **kwargs + ) + + return _make + + +@pytest.fixture +def eval_routes() -> Dict[str, Any]: + """The happy path for an eval run: resolve env, create, samples, finalize.""" + return { + "POST /api/v1/environmentshub/resolve": {"data": {"id": "env-123"}}, + "POST /api/v1/evaluations/": lambda request: httpx.Response( + 201, + json={ + "evaluation_id": "eval-abc", + "name": "test-run", + "status": "RUNNING", + "eval_type": "environment", + "viewer_url": "https://app.example/dashboard/evaluations/eval-abc", + "created_at": "2026-08-19T00:00:00Z", + }, + ), + "POST /api/v1/evaluations/eval-abc/samples": { + "evaluation_id": "eval-abc", + "samples_pushed": 1, + "status": "RUNNING", + }, + "POST /api/v1/evaluations/eval-abc/finalize": { + "evaluation_id": "eval-abc", + "status": "PROCESSING", + }, + "PUT /api/v1/evaluations/eval-abc": { + "evaluation_id": "eval-abc", + "name": "test-run", + "status": "RUNNING", + "updated_at": "2026-08-19T00:00:00Z", + }, + } + + +class FakeSink: + """A sink that records what it was given, and can be told to fail.""" + + def __init__(self, name: str = "fake", fail_on_write: bool = False) -> None: + self.name = name + self.enabled = True + self.fail_on_write = fail_on_write + self.started: List[Any] = [] + self.batches: List[Any] = [] + self.flushes = 0 + self.closed = False + + def start(self, run_id: str, context: Dict[str, str]) -> None: + self.started.append((run_id, dict(context))) + + def write(self, records, *, line_format=None, step=None) -> None: + if self.fail_on_write: + raise RuntimeError("sink is broken") + self.batches.append((list(records), line_format, step)) + + def flush(self) -> None: + self.flushes += 1 + + def close(self) -> None: + self.closed = True + + +@pytest.fixture +def fake_sink() -> Callable[..., FakeSink]: + return FakeSink diff --git a/packages/prime-runs/tests/test_evals_backend.py b/packages/prime-runs/tests/test_evals_backend.py new file mode 100644 index 000000000..076369795 --- /dev/null +++ b/packages/prime-runs/tests/test_evals_backend.py @@ -0,0 +1,169 @@ +"""Eval run lifecycle against ``/api/v1/evaluations/*``.""" + +import httpx +import pytest +from conftest import RecordingHandler + +from prime_runs.backends import EvalsBackend +from prime_runs.exceptions import ConfigurationError, EnvironmentResolutionError +from prime_runs.models import EnvironmentRef, RunSpec, RunStatus + + +def make_backend(make_platform_client, routes, **kwargs): + handler = RecordingHandler(routes) + client = make_platform_client(handler) + backend = EvalsBackend(client, frontend_url="https://app.example", **kwargs) + return backend, handler + + +def test_create_resolves_environment_names_through_the_hub(make_platform_client, eval_routes): + backend, handler = make_backend(make_platform_client, eval_routes) + spec = RunSpec( + name="test-run", + environments=[EnvironmentRef(name="gsm8k")], + model="Qwen/Qwen3-8B", + framework="verifiers", + ) + + handle = backend.create(spec) + + assert handler.paths()[0] == "POST /api/v1/environmentshub/resolve" + created = handler.bodies_for("/api/v1/evaluations/")[0] + assert created["environments"] == [{"id": "env-123"}] + assert created["model_name"] == "Qwen/Qwen3-8B" + assert created["framework"] == "verifiers" + # The environment name doubles as the dataset, as the old uploader did. + assert created["dataset"] == "gsm8k" + assert handle.id == "eval-abc" + assert handle.url == "https://app.example/dashboard/evaluations/eval-abc" + + +def test_an_explicit_environment_id_skips_the_hub(make_platform_client, eval_routes): + backend, handler = make_backend(make_platform_client, eval_routes) + + backend.create(RunSpec(name="r", environments=[EnvironmentRef(id="env-999")])) + + assert "POST /api/v1/environmentshub/resolve" not in handler.paths() + assert handler.bodies_for("/api/v1/evaluations/")[0]["environments"] == [{"id": "env-999"}] + + +def test_an_unresolvable_environment_fails_the_run_rather_than_being_skipped( + make_platform_client, eval_routes +): + """Silently dropping it produces a run attached to the wrong environments — + an upload that looks successful and is discovered wrong much later.""" + routes = dict(eval_routes) + routes["POST /api/v1/environmentshub/resolve"] = lambda request: httpx.Response( + 404, json={"detail": "no such environment"} + ) + backend, _ = make_backend(make_platform_client, routes) + + with pytest.raises(EnvironmentResolutionError, match="gsm8k"): + backend.create(RunSpec(name="r", environments=[EnvironmentRef(name="gsm8k")])) + + +def test_a_run_with_no_environments_is_rejected_before_any_request( + make_platform_client, eval_routes +): + backend, handler = make_backend(make_platform_client, eval_routes) + + with pytest.raises(ConfigurationError, match="at least one environment"): + backend.create(RunSpec(name="r")) + assert handler.requests == [] + + +def test_a_run_without_a_name_gets_one(make_platform_client, eval_routes): + """The API requires a name; the alternative to generating one is a 422 at + the worst possible moment.""" + backend, handler = make_backend(make_platform_client, eval_routes) + + backend.create(RunSpec(environments=[EnvironmentRef(name="gsm8k")])) + + assert handler.bodies_for("/api/v1/evaluations/")[0]["name"].startswith("gsm8k-") + + +def test_team_id_is_forwarded_to_the_hub_and_the_run(make_platform_client, eval_routes): + backend, handler = make_backend(make_platform_client, eval_routes, team_id="team-7") + + backend.create(RunSpec(name="r", environments=[EnvironmentRef(name="gsm8k")])) + + assert handler.bodies_for("/api/v1/environmentshub/resolve")[0]["team_id"] == "team-7" + assert handler.bodies_for("/api/v1/evaluations/")[0]["team_id"] == "team-7" + + +def test_finalizing_a_completed_run_posts_its_metrics(make_platform_client, eval_routes): + backend, handler = make_backend(make_platform_client, eval_routes) + + backend.finalize("eval-abc", status=RunStatus.COMPLETED, summary={"avg_reward": 0.75}) + + assert handler.bodies_for("/api/v1/evaluations/eval-abc/finalize")[0] == { + "metrics": {"avg_reward": 0.75} + } + + +def test_a_failed_run_falls_back_to_metadata_when_the_status_endpoint_is_missing( + make_platform_client, eval_routes, caplog +): + """The platform has no producer-facing way to fail an evaluation yet. + + Until it does, the run cannot leave RUNNING — but the failure must still be + recorded somewhere an operator and the dashboard can both read it, and the + SDK must say plainly that the run will keep showing as running. + """ + backend, handler = make_backend(make_platform_client, eval_routes) + + with caplog.at_level("WARNING"): + backend.finalize("eval-abc", status=RunStatus.FAILED, error="boom") + + assert "POST /api/v1/evaluations/eval-abc/status" in handler.paths() + terminal = handler.bodies_for("/api/v1/evaluations/eval-abc")[0]["metadata"]["prime_runs"] + assert terminal["status"] == "failed" + assert terminal["error"] == "boom" + assert "keep showing as running" in caplog.text + + +def test_the_missing_status_endpoint_is_probed_once_per_backend(make_platform_client, eval_routes): + backend, handler = make_backend(make_platform_client, eval_routes) + + backend.finalize("eval-abc", status=RunStatus.FAILED, error="one") + backend.finalize("eval-abc", status=RunStatus.CRASHED, error="two") + + assert handler.paths().count("POST /api/v1/evaluations/eval-abc/status") == 1 + + +def test_a_status_endpoint_that_exists_is_used_instead_of_the_fallback( + make_platform_client, eval_routes +): + routes = dict(eval_routes) + routes["POST /api/v1/evaluations/eval-abc/status"] = {"evaluation_id": "eval-abc"} + backend, handler = make_backend(make_platform_client, routes) + + backend.finalize("eval-abc", status=RunStatus.FAILED, error="boom") + + assert handler.bodies_for("/api/v1/evaluations/eval-abc/status")[0] == { + "status": "FAILED", + "error": "boom", + } + assert "PUT /api/v1/evaluations/eval-abc" not in handler.paths() + + +def test_update_sends_nothing_when_there_is_nothing_to_send(make_platform_client, eval_routes): + backend, handler = make_backend(make_platform_client, eval_routes) + + backend.update("eval-abc") + + assert handler.requests == [] + + +def test_attach_survives_a_read_failure(make_platform_client, eval_routes): + """Losing a run's name to a transient read is not worth failing a resume on.""" + routes = dict(eval_routes) + routes["GET /api/v1/evaluations/eval-abc"] = lambda request: httpx.Response( + 500, json={"detail": "nope"} + ) + backend, _ = make_backend(make_platform_client, routes) + + handle = backend.attach("eval-abc") + + assert handle.id == "eval-abc" + assert handle.url == "https://app.example/dashboard/evaluations/eval-abc" diff --git a/packages/prime-runs/tests/test_http.py b/packages/prime-runs/tests/test_http.py new file mode 100644 index 000000000..79e4b0118 --- /dev/null +++ b/packages/prime-runs/tests/test_http.py @@ -0,0 +1,110 @@ +"""Transport behaviour: error mapping and retry.""" + +import httpx +import pytest + +from prime_runs._http import PlatformClient, encode_json, retry_delay +from prime_runs.exceptions import ( + NotFoundError, + PaymentRequiredError, + RetryableAPIError, + RunAPIError, + TransportError, + UnauthorizedError, +) + + +def client_for(handler, **kwargs) -> PlatformClient: + return PlatformClient( + api_key="test-key", + base_url="http://testserver", + client=httpx.Client(transport=httpx.MockTransport(handler)), + **kwargs, + ) + + +@pytest.mark.parametrize( + "status,expected", + [ + (401, UnauthorizedError), + (402, PaymentRequiredError), + (404, NotFoundError), + (400, RunAPIError), + (422, RunAPIError), + ], +) +def test_status_codes_map_to_types_callers_can_branch_on(status, expected): + client = client_for(lambda request: httpx.Response(status, json={"detail": "nope"})) + + with pytest.raises(expected) as caught: + client.get("/evaluations/x") + + assert caught.value.status_code == status + assert "nope" in str(caught.value) + + +def test_an_unauthorized_error_says_what_to_do_about_it(): + client = client_for(lambda request: httpx.Response(401, json={"detail": "bad token"})) + + with pytest.raises(UnauthorizedError, match="PRIME_API_KEY"): + client.get("/evaluations/x") + + +def test_retryable_statuses_are_retried_then_surface(no_sleep): + attempts = [] + + def handler(request): + attempts.append(request) + return httpx.Response(503, json={"code": "ingest_unavailable"}) + + with pytest.raises(RetryableAPIError) as caught: + client_for(handler, max_attempts=3).get("/evaluations/x") + + assert len(attempts) == 3 + assert caught.value.code == "ingest_unavailable" + assert no_sleep == [1.0, 2.0] + + +def test_a_retry_succeeds_without_bothering_the_caller(no_sleep): + responses = [httpx.Response(429), httpx.Response(200, json={"ok": True})] + + client = client_for(lambda request: responses.pop(0)) + + assert client.get("/evaluations/x") == {"ok": True} + + +def test_retry_after_beats_the_exponential_schedule(): + assert retry_delay(1, 7.5) == 7.5 + assert retry_delay(1, None) == 1.0 + assert retry_delay(4, None) == 8.0 + # Never wait longer than the ceiling, whatever the server asked for. + assert retry_delay(1, 900.0) == 16.0 + + +def test_transport_failures_are_retried_and_typed(no_sleep): + def handler(request): + raise httpx.ConnectError("refused", request=request) + + with pytest.raises(TransportError): + client_for(handler, max_attempts=2).get("/evaluations/x") + + assert len(no_sleep) == 1 + + +def test_an_empty_body_is_a_valid_response(): + client = client_for(lambda request: httpx.Response(204)) + + assert client.post("/evaluations/x/finalize", json_body={"metrics": {}}) == {} + + +def test_a_non_json_body_names_the_request_that_produced_it(): + client = client_for(lambda request: httpx.Response(200, text="gateway")) + + with pytest.raises(RunAPIError, match="non-JSON"): + client.get("/evaluations/x") + + +def test_encoding_refuses_values_json_cannot_carry(): + """Bare ``NaN`` is JavaScript, not JSON; it comes back as an opaque 400.""" + with pytest.raises(ValueError): + encode_json({"reward": float("nan")}) diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py new file mode 100644 index 000000000..1bf8f6a83 --- /dev/null +++ b/packages/prime-runs/tests/test_init.py @@ -0,0 +1,231 @@ +"""``init()``: mode resolution, offline runs, online runs, rank handling.""" + +import json +import os + +import pytest +from _fakes import make_episode, make_trace +from conftest import RecordingHandler + +import prime_runs as pr +from prime_runs.exceptions import ConfigurationError +from prime_runs.models import RunStatus +from prime_runs.run import RUN_ID_ENV + +# ------------------------------------------------------------------- offline + + +def test_an_offline_run_is_a_real_run(tmp_path): + """The reason producers can delete their ``--no-push`` branch: same ID, same + status, same calls — just a different destination.""" + run = pr.init(name="local", environments=["gsm8k"], mode="offline", dir=str(tmp_path)) + + assert run.id.startswith("offline-") + assert run.url == str((tmp_path / run.id).resolve()) + assert run.mode == "offline" + + run.log({"reward": 0.5}, step=1) + run.log_traces([{"id": "t1"}]) + run.finish(summary={"avg_reward": 0.5}) + + state = json.loads((tmp_path / run.id / "run.json").read_text()) + assert state["status"] == RunStatus.COMPLETED.value + assert state["summary"]["avg_reward"] == 0.5 + assert state["spec"]["environments"] == [{"name": "gsm8k"}] + + +def test_offline_records_are_written_in_the_wire_format(tmp_path): + """The archive is a deferred upload, not a debug dump: these bytes are what + ``TracesClient.upload_file`` sends, with the run already stamped.""" + run = pr.init(mode="offline", dir=str(tmp_path)) + run.log_traces([{"id": "t1"}, {"id": "t2"}]) + run.finish() + + lines = (tmp_path / run.id / "records" / "trace.jsonl").read_text().splitlines() + records = [json.loads(line) for line in lines] + + assert [record["id"] for record in records] == ["t1", "t2"] + assert all(record["run"]["id"] == run.id for record in records) + + +def test_episodes_are_written_to_their_own_file(tmp_path): + run = pr.init(mode="offline", dir=str(tmp_path)) + run.log_traces([make_episode("ep-1", [make_trace()])]) + run.finish() + + assert (tmp_path / run.id / "records" / "episode.jsonl").exists() + + +def test_offline_metrics_are_a_time_series(tmp_path): + run = pr.init(mode="offline", dir=str(tmp_path)) + run.log({"loss": 2.0}, step=1) + run.log({"loss": 1.0}, step=2) + run.flush() + run.finish() + + lines = (tmp_path / run.id / "metrics.jsonl").read_text().splitlines() + assert [json.loads(line)["loss"] for line in lines] == [2.0, 1.0] + + +def test_a_record_that_already_names_a_run_is_left_alone(tmp_path): + """Producers stamp the run themselves; two sources of truth for the run ID + is how traces end up on the wrong run.""" + run = pr.init(mode="offline", dir=str(tmp_path)) + run.log_traces([{"id": "t1", "run": {"id": "someone-elses-run"}}]) + run.finish() + + record = json.loads((tmp_path / run.id / "records" / "trace.jsonl").read_text()) + assert record["run"]["id"] == "someone-elses-run" + + +# -------------------------------------------------------------------- modes + + +def test_no_api_key_degrades_to_offline_rather_than_skipping_the_run(tmp_path, caplog): + with caplog.at_level("WARNING"): + run = pr.init(name="local", dir=str(tmp_path)) + + assert run.mode == "offline" + assert "offline" in caplog.text + run.finish() + + +def test_the_mode_can_be_set_from_the_environment(monkeypatch, tmp_path): + monkeypatch.setenv("PRIME_RUNS_MODE", "disabled") + + run = pr.init(name="local", api_key="test-key", dir=str(tmp_path)) + + assert run.mode == "disabled" + run.finish() + + +def test_an_unknown_mode_is_rejected(tmp_path): + with pytest.raises(ConfigurationError, match="not one of"): + pr.init(mode="sideways", dir=str(tmp_path)) + + +def test_a_disabled_run_still_answers_every_call(tmp_path): + """Same object shape, so producer code needs no branching.""" + run = pr.init(mode="disabled", dir=str(tmp_path)) + + run.log({"reward": 1.0}, step=1) + run.log_traces([{"id": "t1"}]) + run.finish(summary={"avg_reward": 1.0}) + + assert run.id + assert run.status is RunStatus.COMPLETED + assert not list(tmp_path.iterdir()) + + +def test_training_runs_are_not_supported_yet(tmp_path): + with pytest.raises(ConfigurationError, match="training runs"): + pr.init(kind="train", api_key="test-key", environments=["gsm8k"]) + + +# --------------------------------------------------------------------- rank + + +def test_a_non_primary_rank_with_no_run_to_join_records_nothing(monkeypatch, tmp_path): + """Otherwise rank 3 creates a second run for the same job.""" + monkeypatch.setenv("RANK", "3") + + run = pr.init(name="local", api_key="test-key", environments=["gsm8k"]) + + assert run.mode == "disabled" + assert run.is_primary is False + run.finish() + + +def test_a_run_id_in_the_environment_is_joined_not_recreated(monkeypatch, tmp_path): + monkeypatch.setenv("DP_RANK", "2") + monkeypatch.setenv(RUN_ID_ENV, "offline-shared") + + run = pr.init(mode="offline", dir=str(tmp_path)) + + assert run.id == "offline-shared" + assert run.is_primary is False + run.finish() + + +def test_init_publishes_the_run_id_for_child_processes(tmp_path): + """Forked workers and subprocess launchers join the run their parent opened + instead of each opening their own.""" + run = pr.init(mode="offline", dir=str(tmp_path)) + + assert os.environ[RUN_ID_ENV] == run.id + run.finish() + + +# ------------------------------------------------------------------- online + + +@pytest.fixture +def online(monkeypatch, make_platform_client, eval_routes): + """``init(mode="online")`` wired to a MockTransport.""" + + def _init(routes=None, **kwargs): + handler = RecordingHandler(routes or eval_routes) + client = make_platform_client(handler) + monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: client) + run = pr.init( + name="test-run", + environments=["gsm8k"], + model="Qwen3-8B", + framework="verifiers", + api_key="test-key", + traces=False, + **kwargs, + ) + return run, handler + + return _init + + +def test_an_online_run_returns_the_platforms_id_and_viewer_url(online): + run, _ = online() + + assert run.id == "eval-abc" + assert run.url == "https://app.example/dashboard/evaluations/eval-abc" + assert run.mode == "online" + run.finish() + + +def test_episodes_stream_to_the_sample_table_while_the_run_is_going(online): + run, handler = online() + + run.log_traces([make_episode("ep-1", [make_trace()])]) + run.flush() + + posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples") + assert len(posted) == 1 + assert posted[0]["samples"][0]["sample_id"] == "ep-1" + # Still running: the point is that the dashboard fills in as rollouts land. + assert "POST /api/v1/evaluations/eval-abc/finalize" not in handler.paths() + run.finish() + + +def test_finishing_an_online_run_finalizes_it_with_its_metrics(online): + run, handler = online() + + run.finish(summary={"avg_reward": 0.9}) + + body = handler.bodies_for("/api/v1/evaluations/eval-abc/finalize")[0] + assert body["metrics"]["avg_reward"] == 0.9 + + +def test_the_end_to_end_shape_a_producer_writes(online): + """The whole surface, in the order verifiers will call it.""" + from prime_runs import metrics + + episodes = [make_episode(f"ep-{n}", [make_trace(idx=n, reward=float(n))]) for n in range(3)] + + run, handler = online() + for episode in episodes: + run.log_traces([episode]) + run.finish(summary=metrics.from_episodes(episodes)) + + assert handler.paths().count("POST /api/v1/evaluations/eval-abc/samples") == 3 + finalize = handler.bodies_for("/api/v1/evaluations/eval-abc/finalize")[0] + assert finalize["metrics"]["avg_reward"] == 1.0 + assert run.status is RunStatus.COMPLETED + assert run.errors == [] diff --git a/packages/prime-runs/tests/test_metrics.py b/packages/prime-runs/tests/test_metrics.py new file mode 100644 index 000000000..3673cc5d9 --- /dev/null +++ b/packages/prime-runs/tests/test_metrics.py @@ -0,0 +1,77 @@ +"""Run-level aggregates, moved here from verifiers. + +The aggregation rules encode judgements about what a run's headline number +means. They are asserted here because a silent change to any of them changes +every number on the dashboard without changing a line of producer code. +""" + +from _fakes import Reward, make_episode, make_trace + +from prime_runs import metrics + + +def test_only_trainable_traces_are_scored(): + """A judge scoring 0 must not halve a solver's average. + + Fixed agents often carry no rewards at all, so including them dilutes every + mean with structural zeros. + """ + episode = make_episode( + "ep-1", + [ + make_trace(trace_id="solver", reward=1.0), + make_trace(trace_id="judge", agent="judge", trainable=False, reward=0.0), + ], + ) + + assert metrics.from_episodes([episode])["avg_reward"] == 1.0 + + +def test_all_traces_count_when_none_are_trainable(): + """The dashboard's fallback: an eval of a frozen model still has a score.""" + episode = make_episode( + "ep-1", + [ + make_trace(trace_id="a", trainable=False, reward=1.0), + make_trace(trace_id="b", trainable=False, reward=0.0), + ], + ) + + assert metrics.from_episodes([episode])["avg_reward"] == 0.5 + + +def test_sub_rewards_and_env_metrics_average_together(): + episodes = [ + make_episode("ep-1", [make_trace(rewards={"format": Reward(1.0)}, metrics={"turns": 4.0})]), + make_episode("ep-2", [make_trace(rewards={"format": Reward(0.0)}, metrics={"turns": 6.0})]), + ] + + avg = metrics.from_episodes(episodes)["avg_metrics"] + + assert avg["format"] == 0.5 + assert avg["turns"] == 5.0 + + +def test_a_metric_present_on_only_some_traces_averages_over_those_traces(): + """Counts are per-key, not per-run: a metric only some environments emit + must not be divided by traces that never reported it.""" + episodes = [ + make_episode("ep-1", [make_trace(metrics={"partial": 1.0})]), + make_episode("ep-2", [make_trace(metrics={})]), + ] + + assert metrics.from_episodes(episodes)["avg_metrics"]["partial"] == 1.0 + + +def test_avg_error_counts_episodes_not_traces(): + """A hook failure counts even when its traces are clean or it left none.""" + episodes = [ + make_episode("ok", [make_trace(), make_trace(trace_id="t2")], ok=True), + make_episode("broken", [], ok=False), + ] + + assert metrics.from_episodes(episodes)["avg_error"] == 0.5 + + +def test_empty_run_produces_zeros_rather_than_dividing_by_zero(): + assert metrics.from_episodes([]) == {"avg_reward": 0.0, "avg_metrics": {}, "avg_error": 0.0} diff --git a/packages/prime-runs/tests/test_projection.py b/packages/prime-runs/tests/test_projection.py new file mode 100644 index 000000000..3929a426a --- /dev/null +++ b/packages/prime-runs/tests/test_projection.py @@ -0,0 +1,141 @@ +"""The v0 eval-sample projection, moved here from verifiers. + +These assertions are the contract the current viewer reads. They exist so the +move is provably a relocation and not a rewrite: a run pushed through this SDK +must produce the same rows verifiers produced. +""" + +import pytest +from _fakes import Reward, make_episode, make_trace + +from prime_runs.projection import ( + MAX_SAMPLES_PAYLOAD_BYTES, + batch_samples, + build_samples, + is_episode, + summary_trace_index, + trace_to_sample, +) + + +def test_trace_to_sample_carries_the_flat_row(): + trace = make_trace(trace_id="t1", idx=7, reward=0.5, metrics={"tokens": 12.0}) + + sample = trace_to_sample(trace, rollout_number=3, episode_id="ep-9") + + assert sample["sample_id"] == "t1" + assert sample["example_id"] == 7 + assert sample["rollout_number"] == 3 + assert sample["episode_id"] == "ep-9" + assert sample["agent"] == "solver" + assert sample["trainable"] is True + assert sample["reward"] == 0.5 + assert sample["metrics"] == {"tokens": 12.0} + # No prompt/completion split mid-branch: completion is the final branch. + assert sample["prompt"] == [] + assert sample["completion"] == [{"role": "assistant", "content": "branch 0"}] + + +def test_trace_to_sample_flattens_sub_rewards_to_top_level(): + trace = make_trace(rewards={"format": Reward(1.0), "correct": Reward(0.0), "skip": None}) + + sample = trace_to_sample(trace) + + assert sample["format"] == 1.0 + assert sample["correct"] == 0.0 + assert "skip" not in sample + + +def test_trace_to_sample_does_not_let_a_sub_reward_clobber_a_real_field(): + """``setdefault`` semantics: a sub-reward named ``reward`` must not win. + + Sub-rewards are flattened into the same namespace as the row's own columns, + so an environment that names a reward function after one of them would + otherwise silently overwrite the value the dashboard reads. + """ + trace = make_trace(reward=0.25, rewards={"reward": Reward(9.0)}) + + assert trace_to_sample(trace)["reward"] == 0.25 + + +def test_trajectory_keeps_one_entry_per_branch(): + trace = make_trace(branches=3) + + trajectory = trace_to_sample(trace)["trajectory"] + + assert len(trajectory) == 3 + assert trajectory[1]["messages"] == [{"role": "assistant", "content": "branch 1"}] + assert trajectory[0]["num_input_tokens"] == 10 + + +def test_build_samples_emits_one_row_per_episode_with_the_native_wrapper(): + episode = make_episode("ep-1", [make_trace(trace_id="t1")]) + + samples = build_samples([episode]) + + assert len(samples) == 1 + assert samples[0]["sample_id"] == "ep-1" + assert samples[0]["info"]["native_wrapper"] == episode.to_record() + assert samples[0]["info"]["native_trace_index"] == 0 + + +def test_summary_trace_is_the_first_trainable_one(): + """A judge or modeled user must not become the row the dashboard shows.""" + episode = make_episode( + "ep-1", + [ + make_trace(trace_id="judge", agent="judge", trainable=False, reward=0.0), + make_trace(trace_id="solver", reward=1.0), + ], + ) + + assert summary_trace_index(episode) == 1 + sample = build_samples([episode])[0] + assert sample["info"]["native_trace_index"] == 1 + assert sample["reward"] == 1.0 + + +def test_rollout_numbers_continue_across_streaming_calls(): + """Streaming uploads must number rollouts the way one final upload did. + + The old code built every sample in a single call, so its per-example counter + lived in a local. Streaming means several calls, and without a carried + counter every batch would restart at rollout 1. + """ + counters: dict = {} + first = build_samples([make_episode("ep-1", [make_trace(idx=4)])], counters) + second = build_samples([make_episode("ep-2", [make_trace(idx=4)])], counters) + other_example = build_samples([make_episode("ep-3", [make_trace(idx=5)])], counters) + + assert first[0]["rollout_number"] == 1 + assert second[0]["rollout_number"] == 2 + assert other_example[0]["rollout_number"] == 1 + + +def test_episodes_without_traces_are_skipped(): + assert build_samples([make_episode("empty", [])]) == [] + + +def test_batch_samples_splits_on_the_payload_ceiling(): + big = {"sample_id": "x", "blob": "a" * (MAX_SAMPLES_PAYLOAD_BYTES // 3)} + batches = batch_samples([dict(big, sample_id=str(n)) for n in range(4)]) + + assert len(batches) > 1 + assert sum(len(batch) for batch in batches) == 4 + + +def test_batch_samples_refuses_a_sample_that_cannot_be_sent(): + """Dropping it would report a complete run that is missing rows.""" + oversized = {"sample_id": "x", "blob": "a" * (MAX_SAMPLES_PAYLOAD_BYTES + 1)} + + with pytest.raises(ValueError, match="too large"): + batch_samples([oversized]) + + +def test_batch_samples_returns_nothing_for_no_samples(): + assert batch_samples([]) == [] + + +def test_is_episode_distinguishes_episodes_from_traces(): + assert is_episode(make_episode()) + assert not is_episode(make_trace()) diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py new file mode 100644 index 000000000..3c22f733a --- /dev/null +++ b/packages/prime-runs/tests/test_run.py @@ -0,0 +1,295 @@ +"""The run handle: lifecycle, containment, ranks, terminal status.""" + +from typing import Any, Dict, List, Optional + +import pytest +from conftest import FakeSink + +from prime_runs.exceptions import RunFinishedError +from prime_runs.models import RunHandle, RunSpec, RunStatus +from prime_runs.run import Run + + +class FakeBackend: + def __init__(self, supports_step_metrics: bool = False, fail_on: Optional[str] = None) -> None: + self.kind = "eval" + self.supports_step_metrics = supports_step_metrics + self.fail_on = fail_on + self.updates: List[Dict[str, Any]] = [] + self.points: List[Any] = [] + self.finalized: List[Dict[str, Any]] = [] + self.closed = False + + def create(self, spec: RunSpec) -> RunHandle: + return RunHandle(id="run-1", name=spec.name, url="https://app.example/run-1") + + def attach(self, run_id: str) -> RunHandle: + return RunHandle(id=run_id) + + def update(self, run_id, *, config=None, summary=None) -> None: + if self.fail_on == "update": + raise RuntimeError("update exploded") + self.updates.append({"config": config, "summary": summary}) + + def log_metrics(self, run_id, metrics, step=None) -> None: + self.points.append((metrics, step)) + + def finalize(self, run_id, *, status, summary=None, error=None) -> None: + if self.fail_on == "finalize": + raise RuntimeError("finalize exploded") + self.finalized.append({"status": status, "summary": summary, "error": error}) + + def close(self) -> None: + self.closed = True + + +def make_run(backend=None, sinks=None, **kwargs) -> Run: + backend = backend or FakeBackend() + spec = RunSpec(name="test-run", kind="eval", framework="verifiers", model="Qwen3-8B") + return Run( + backend=backend, + handle=backend.create(spec), + spec=spec, + sinks=sinks if sinks is not None else [], + **kwargs, + ) + + +def test_the_handle_exposes_what_a_producer_prints(): + run = make_run() + + assert run.id == "run-1" + assert run.url == "https://app.example/run-1" + assert run.status is RunStatus.RUNNING + assert run.is_primary is True + run.finish() + + +def test_sinks_are_started_with_the_run_id_and_provenance(): + sink = FakeSink() + run = make_run(sinks=[sink]) + + run_id, context = sink.started[0] + + assert run_id == "run-1" + # Provenance only — the join key is run.id inside the trace document. + assert context["source"] == "prime-runs" + assert context["framework"] == "verifiers" + assert "evaluation_id" not in context + run.finish() + + +def test_traces_reach_the_sinks_while_the_run_is_still_going(): + sink = FakeSink() + run = make_run(sinks=[sink]) + + run.log_traces([{"id": "t1"}], step=2) + run.flush() + + assert sink.batches[0][0] == [{"id": "t1"}] + assert sink.batches[0][2] == 2 + assert not run.finished + run.finish() + + +def test_an_empty_batch_is_not_sent(): + sink = FakeSink() + run = make_run(sinks=[sink]) + + run.log_traces([]) + run.flush() + + assert sink.batches == [] + run.finish() + + +def test_metrics_land_in_the_summary_when_the_backend_has_no_time_series(): + backend = FakeBackend(supports_step_metrics=False) + run = make_run(backend, summary_flush_seconds=0.0) + + run.log({"reward": 0.5}, step=1) + run.log({"reward": 0.75}, step=2) + + assert run.summary["reward"] == 0.75 + assert backend.points == [] + assert backend.updates, "the summary was flushed" + run.finish() + + +def test_metrics_become_a_time_series_when_the_backend_has_one(): + backend = FakeBackend(supports_step_metrics=True) + run = make_run(backend) + + run.log({"loss": 2.0}, step=1) + run.flush() + + assert backend.points == [({"loss": 2.0}, 1)] + run.finish() + + +def test_commit_false_stages_without_writing(): + backend = FakeBackend(supports_step_metrics=True) + run = make_run(backend) + + run.log({"loss": 2.0}, step=1, commit=False) + run.flush() + + assert backend.points == [] + assert run.summary["loss"] == 2.0 + run.finish() + + +def test_non_finite_metrics_are_dropped_rather_than_failing_the_request(): + """A diverged loss serializes as bare ``NaN``, which strict JSON rejects — + the whole request fails on a payload nobody can inspect.""" + run = make_run(summary_flush_seconds=0.0) + + run.log({"loss": float("nan"), "grad": float("inf"), "reward": 0.5}) + + assert run.summary == {"reward": 0.5} + run.finish() + + +def test_finish_flushes_records_before_reporting_the_terminal_status(): + """A dashboard reacting to the status must never see a finished run with + samples still landing.""" + order: List[str] = [] + + class OrderedSink(FakeSink): + def write(self, records, *, line_format=None, step=None) -> None: + order.append("write") + super().write(records, line_format=line_format, step=step) + + class OrderedBackend(FakeBackend): + def finalize(self, run_id, *, status, summary=None, error=None) -> None: + order.append("finalize") + super().finalize(run_id, status=status, summary=summary, error=error) + + run = make_run(OrderedBackend(), sinks=[OrderedSink()]) + run.log_traces([{"id": "t1"}]) + run.finish() + + assert order == ["write", "finalize"] + + +def test_finish_is_idempotent(): + backend = FakeBackend() + run = make_run(backend) + + run.finish(summary={"avg_reward": 1.0}) + run.finish(status=RunStatus.FAILED) + + assert len(backend.finalized) == 1 + assert backend.finalized[0]["status"] is RunStatus.COMPLETED + assert backend.finalized[0]["summary"] == {"avg_reward": 1.0} + + +def test_logging_after_finish_is_a_producer_bug(): + run = make_run() + run.finish() + + with pytest.raises(RunFinishedError): + run.log({"reward": 1.0}) + with pytest.raises(RunFinishedError): + run.log_traces([{"id": "t1"}]) + + +def test_the_context_manager_completes_a_clean_run(): + backend = FakeBackend() + with make_run(backend) as run: + run.log({"reward": 1.0}) + + assert backend.finalized[0]["status"] is RunStatus.COMPLETED + + +def test_an_exception_inside_the_block_fails_the_run_and_still_propagates(): + backend = FakeBackend() + + with pytest.raises(ValueError): + with make_run(backend): + raise ValueError("rollout blew up") + + assert backend.finalized[0]["status"] is RunStatus.FAILED + assert "rollout blew up" in backend.finalized[0]["error"] + + +def test_an_interrupt_is_recorded_as_a_decision_not_a_fault(): + backend = FakeBackend() + + with pytest.raises(KeyboardInterrupt): + with make_run(backend): + raise KeyboardInterrupt + + assert backend.finalized[0]["error"] == "interrupted" + + +def test_a_process_that_exits_without_finishing_reports_crashed(): + """The producer never said the run failed — it stopped existing. The + distinction tells an operator where to look.""" + backend = FakeBackend() + run = make_run(backend) + + run._on_process_exit() + + assert backend.finalized[0]["status"] is RunStatus.CRASHED + + +def test_a_backend_failure_does_not_escape_into_the_producer_by_default(): + """Six hours of rollouts must not be lost to a 502 on a telemetry call.""" + backend = FakeBackend(fail_on="finalize") + run = make_run(backend) + + run.finish() + + assert run.errors and "finalize exploded" in run.errors[0] + + +def test_on_error_raise_surfaces_the_failure_for_tests_and_ci(): + backend = FakeBackend(fail_on="finalize") + run = make_run(backend, on_error="raise") + + with pytest.raises(RuntimeError, match="finalize exploded"): + run.finish() + + +def test_a_sink_error_is_recorded_on_the_run(): + run = make_run(sinks=[FakeSink("broken", fail_on_write=True)]) + + run.log_traces([{"id": "t1"}]) + run.flush() + + assert any("broken" in error for error in run.errors) + run.finish() + + +def test_a_non_primary_rank_does_not_close_the_shared_run(): + """Eight ranks racing to finalize produce seven confusing failures.""" + backend = FakeBackend() + run = make_run(backend, is_primary=False) + + run.log({"reward": 1.0}) + run.finish() + + assert backend.finalized == [] + assert backend.updates == [] + assert run.is_primary is False + + +def test_a_non_primary_rank_still_uploads_its_own_records(): + """The point of eight ranks is that they contribute to one run.""" + sink = FakeSink() + run = make_run(sinks=[sink], is_primary=False) + + run.log_traces([{"id": "t1"}]) + run.flush() + + assert sink.batches + run.finish() + + +def test_dropped_records_are_reported_on_the_handle(): + run = make_run() + run._worker.dropped = 3 + + assert run.dropped_records == 3 + run.finish() diff --git a/packages/prime-runs/tests/test_samples_sink.py b/packages/prime-runs/tests/test_samples_sink.py new file mode 100644 index 000000000..2e95737e4 --- /dev/null +++ b/packages/prime-runs/tests/test_samples_sink.py @@ -0,0 +1,61 @@ +"""The legacy sample sink that keeps today's viewer working.""" + +from _fakes import make_episode, make_trace +from conftest import RecordingHandler + +from prime_runs.sinks import EvalSamplesSink + + +def make_sink(make_platform_client, eval_routes): + handler = RecordingHandler(eval_routes) + sink = EvalSamplesSink(make_platform_client(handler)) + sink.start("eval-abc", {}) + return sink, handler + + +def test_episodes_are_projected_and_posted(make_platform_client, eval_routes): + sink, handler = make_sink(make_platform_client, eval_routes) + + sink.write([make_episode("ep-1", [make_trace()])]) + + body = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")[0] + assert body["samples"][0]["sample_id"] == "ep-1" + assert sink.samples_written == 1 + + +def test_rollout_numbering_is_continuous_across_streamed_batches(make_platform_client, eval_routes): + """Streaming must produce the numbering the old one-shot upload produced.""" + sink, handler = make_sink(make_platform_client, eval_routes) + + sink.write([make_episode("ep-1", [make_trace(idx=0)])]) + sink.write([make_episode("ep-2", [make_trace(idx=0)])]) + + posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples") + assert [body["samples"][0]["rollout_number"] for body in posted] == [1, 2] + + +def test_a_producer_that_already_speaks_v0_is_passed_through(make_platform_client, eval_routes): + sink, handler = make_sink(make_platform_client, eval_routes) + + sink.write([{"sample_id": "s1", "reward": 1.0}]) + + body = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")[0] + assert body["samples"] == [{"sample_id": "s1", "reward": 1.0}] + + +def test_records_this_sink_cannot_project_are_skipped_not_posted(make_platform_client, eval_routes): + """A malformed row would be rejected for the whole batch, taking the valid + rows with it.""" + sink, handler = make_sink(make_platform_client, eval_routes) + + sink.write([{"unrelated": True}]) + + assert handler.requests == [] + + +def test_an_empty_batch_makes_no_request(make_platform_client, eval_routes): + sink, handler = make_sink(make_platform_client, eval_routes) + + sink.write([]) + + assert handler.requests == [] diff --git a/packages/prime-runs/tests/test_traces_sink.py b/packages/prime-runs/tests/test_traces_sink.py new file mode 100644 index 000000000..612e99988 --- /dev/null +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -0,0 +1,152 @@ +"""The primary sample transport, and how it degrades.""" + +import pytest +from _fakes import make_episode, make_trace +from prime_traces import LineFormat +from prime_traces.exceptions import ForbiddenError, RetryableAPIError + +from prime_runs.sinks import TracesSink + + +class FakeTracesClient: + def __init__(self, raises: Exception = None) -> None: + self.raises = raises + self.calls = [] + self.closed = False + + def upload_records(self, records, **kwargs): + self.calls.append((list(records), kwargs)) + if self.raises is not None: + raise self.raises + return ["receipt"] + + def close(self) -> None: + self.closed = True + + +def make_sink(client=None, **kwargs) -> TracesSink: + sink = TracesSink(client=client or FakeTracesClient(), **kwargs) + sink.start("run-1", {"source": "prime-runs", "run_kind": "eval", "framework": "verifiers"}) + return sink + + +def test_records_go_out_with_provenance_but_not_the_join_key(): + """``run.id`` inside the document is the indexed column; ``context`` is an + upload-scoped map that answers a different question.""" + client = FakeTracesClient() + sink = make_sink(client) + + sink.write([{"id": "t1", "run": {"id": "run-1"}}], step=4) + + _, kwargs = client.calls[0] + assert kwargs["context"] == { + "source": "prime-runs", + "run_kind": "eval", + "framework": "verifiers", + "step": "4", + } + assert "run_id" not in kwargs["context"] + + +def test_the_line_format_is_inferred_from_the_records(): + client = FakeTracesClient() + sink = make_sink(client) + + sink.write([make_trace()]) + sink.write([make_episode()]) + + assert client.calls[0][1]["line_format"] is LineFormat.TRACE + assert client.calls[1][1]["line_format"] is LineFormat.EPISODE + + +def test_an_explicit_line_format_wins(): + client = FakeTracesClient() + sink = make_sink(client) + + sink.write([{"id": "t1"}], line_format="episode") + + assert client.calls[0][1]["line_format"] is LineFormat.EPISODE + + +def test_a_bare_mapping_gets_the_run_stamped_onto_a_copy(): + """A dict has no stamping convention, and an upload with no ``run.id`` is + orphaned — unqueryable and undeletable by run.""" + client = FakeTracesClient() + sink = make_sink(client) + original = {"id": "t1"} + + sink.write([original]) + + assert client.calls[0][0][0]["run"] == {"id": "run-1", "type": "eval"} + assert original == {"id": "t1"}, "the caller's dict was not mutated" + + +def test_producer_objects_are_passed_through_untouched(): + """Verifiers and prime-rl stamp the run at rollout time; rewriting their + objects here is how a second source of truth appears.""" + client = FakeTracesClient() + sink = make_sink(client) + trace = make_trace() + + sink.write([trace]) + + assert client.calls[0][0][0] is trace + + +def test_a_gated_account_disables_the_sink_instead_of_failing_the_run(caplog): + """Prime Traces is in closed beta; no runtime action fixes a 403, so + retrying it for the rest of the run only produces noise.""" + client = FakeTracesClient( + raises=ForbiddenError("not in beta", status_code=403, code="service_not_enabled") + ) + sink = make_sink(client) + + with caplog.at_level("WARNING"): + sink.write([{"id": "t1"}]) + + assert sink.enabled is False + assert "not enabled" in caplog.text + + +def test_a_transient_failure_is_raised_so_the_worker_can_report_it(): + """Unlike a 403, this one is about the moment, not the account.""" + client = FakeTracesClient(raises=RetryableAPIError("busy", status_code=503)) + sink = make_sink(client) + + with pytest.raises(RetryableAPIError): + sink.write([{"id": "t1"}]) + assert sink.enabled is True + + +def test_a_disabled_sink_stops_calling_the_service(): + client = FakeTracesClient() + sink = make_sink(client) + sink.enabled = False + + sink.write([{"id": "t1"}]) + + assert client.calls == [] + + +def test_closing_the_sink_closes_the_client(): + client = FakeTracesClient() + sink = make_sink(client) + + sink.close() + + assert client.closed is True + + +def test_a_missing_traces_client_disables_the_sink_rather_than_raising(monkeypatch, caplog): + """Construction failures must not take down a run that has not started.""" + + def explode(**kwargs): + raise RuntimeError("no credentials") + + monkeypatch.setattr("prime_traces.TracesClient", explode) + sink = TracesSink() + + with caplog.at_level("WARNING"): + sink.start("run-1", {"run_kind": "eval"}) + + assert sink.enabled is False diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py new file mode 100644 index 000000000..d3673a8bf --- /dev/null +++ b/packages/prime-runs/tests/test_worker.py @@ -0,0 +1,167 @@ +"""The background uploader: backpressure, containment, fork safety.""" + +import queue +import threading + +from conftest import FakeSink + +from prime_runs.worker import MetricItem, UploadWorker, WriteItem + + +def drain(worker: UploadWorker) -> None: + assert worker.flush(timeout=5.0) + + +def test_records_reach_every_enabled_sink(): + sinks = [FakeSink("a"), FakeSink("b")] + worker = UploadWorker(sinks) + + worker.submit(WriteItem(records=[{"id": 1}], line_format="trace", step=3)) + drain(worker) + + for sink in sinks: + assert sink.batches == [([{"id": 1}], "trace", 3)] + worker.close() + + +def test_a_disabled_sink_is_skipped(): + live, dead = FakeSink("live"), FakeSink("dead") + dead.enabled = False + worker = UploadWorker([live, dead]) + + worker.submit(WriteItem(records=[{"id": 1}])) + drain(worker) + + assert live.batches and not dead.batches + worker.close() + + +def test_one_sink_failing_does_not_stop_the_others(): + broken, healthy = FakeSink("broken", fail_on_write=True), FakeSink("healthy") + reported = [] + worker = UploadWorker([broken, healthy], on_error=lambda name, exc: reported.append(name)) + + worker.submit(WriteItem(records=[{"id": 1}])) + drain(worker) + + assert healthy.batches + assert broken.enabled is False + assert reported == ["broken"] + worker.close() + + +def test_a_failed_sink_is_not_called_again(): + """One log line per batch for the rest of a run hides whatever failed first.""" + broken = FakeSink("broken", fail_on_write=True) + reported = [] + worker = UploadWorker([broken], on_error=lambda name, exc: reported.append(name)) + + for _ in range(3): + worker.submit(WriteItem(records=[{"id": 1}])) + drain(worker) + + assert reported == ["broken"] + worker.close() + + +def test_a_full_queue_drops_rather_than_blocking_the_producer(): + """Stalling a training run to protect telemetry is the wrong trade.""" + + class BlockingSink(FakeSink): + def __init__(self) -> None: + super().__init__("blocking") + self.entered = threading.Event() + self.released = threading.Event() + + def write(self, records, *, line_format=None, step=None) -> None: + self.entered.set() + self.released.wait(5.0) + super().write(records, line_format=line_format, step=step) + + sink = BlockingSink() + worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.05) + + # First item is picked up and wedges the uploader inside sink.write(). + assert worker.submit(WriteItem(records=[{"id": 0}])) + assert sink.entered.wait(5.0), "the uploader never reached the sink" + # Second fills the one-slot queue; third has nowhere to go. + worker.submit(WriteItem(records=[{"id": 1}])) + accepted = worker.submit(WriteItem(records=[{"id": 2}, {"id": 3}])) + + assert accepted is False + assert worker.dropped == 2 + + sink.released.set() + worker.close() + + +def test_metrics_ride_the_same_queue_when_the_backend_stores_a_time_series(): + points = [] + worker = UploadWorker([], metric_writer=lambda metrics, step: points.append((metrics, step))) + + worker.submit(MetricItem(metrics={"loss": 0.5}, step=7)) + drain(worker) + + assert points == [({"loss": 0.5}, 7)] + worker.close() + + +def test_a_metric_write_that_raises_does_not_kill_the_uploader(): + sink = FakeSink() + + def explode(metrics, step): + raise RuntimeError("nope") + + worker = UploadWorker([sink], metric_writer=explode) + worker.submit(MetricItem(metrics={"loss": 0.5}, step=1)) + worker.submit(WriteItem(records=[{"id": 1}])) + drain(worker) + + assert sink.batches, "the uploader survived the metric failure" + worker.close() + + +def test_close_drains_then_closes_every_sink(): + sink = FakeSink() + worker = UploadWorker([sink]) + + worker.submit(WriteItem(records=[{"id": 1}])) + worker.close() + + assert sink.batches + assert sink.closed is True + + +def test_submitting_after_close_is_refused(): + worker = UploadWorker([FakeSink()]) + worker.close() + + assert worker.submit(WriteItem(records=[{"id": 1}])) is False + + +def test_flush_without_a_running_thread_still_flushes_the_sinks(): + sink = FakeSink() + worker = UploadWorker([sink]) + + assert worker.flush(timeout=1.0) is True + assert sink.flushes == 1 + + +def test_a_forked_child_starts_over_instead_of_re_uploading_the_parents_queue(): + """The queued records belong to the parent, which still has a live thread. + + Inheriting them would upload each record twice; inheriting the lock could + deadlock the child on its first write. + """ + sink = FakeSink() + worker = UploadWorker([sink], max_queue_size=4) + worker._queue.put(WriteItem(records=[{"id": "parents"}])) + old_queue = worker._queue + + worker._reinit_after_fork() + + assert worker._queue is not old_queue + assert worker._queue.empty() + assert worker._thread is None + assert isinstance(worker._queue, queue.Queue) + assert isinstance(worker._lock, type(threading.Lock())) or worker._lock is not None diff --git a/pyproject.toml b/pyproject.toml index ceff429ca..fb068c938 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ prime-sandboxes = false prime-tunnel = false prime-evals = false prime-traces = false +prime-runs = false verifiers = false tasksets = false harnesses = false @@ -50,6 +51,7 @@ pythonpath = [ "packages/prime-sandboxes/src", "packages/prime-evals/src", "packages/prime-traces/src", + "packages/prime-runs/src", ] testpaths = [ "packages/*/tests", diff --git a/uv.lock b/uv.lock index 680749c86..6092621c5 100644 --- a/uv.lock +++ b/uv.lock @@ -15,6 +15,7 @@ exclude-newer-span = "P7D" verifiers = false tasksets = false prime-traces = false +prime-runs = false prime-tunnel = false prime-sandboxes = false harnesses = false @@ -25,6 +26,7 @@ prime = false members = [ "prime", "prime-evals", + "prime-runs", "prime-sandboxes", "prime-traces", "prime-tunnel", @@ -1445,6 +1447,35 @@ toml = [ { name = "tomli" }, ] +[[package]] +name = "prime-runs" +source = { editable = "packages/prime-runs" } +dependencies = [ + { name = "httpx" }, + { name = "prime-traces" }, + { name = "pydantic" }, + { name = "tenacity" }, +] + +[package.optional-dependencies] +dev = [ + { name = "pytest" }, + { name = "pytest-asyncio" }, + { name = "ruff" }, +] + +[package.metadata] +requires-dist = [ + { name = "httpx", specifier = ">=0.25.0" }, + { name = "prime-traces", editable = "packages/prime-traces" }, + { name = "pydantic", specifier = ">=2.0.0" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, + { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.23.0" }, + { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.13.1" }, + { name = "tenacity", specifier = ">=9.1.2" }, +] +provides-extras = ["dev"] + [[package]] name = "prime-sandboxes" source = { editable = "packages/prime-sandboxes" } From 14bbac596790316c24160b91c234dac5423e1583 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 14:08:25 -0700 Subject: [PATCH 02/39] fix(runs): address Bugbot review findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All six were real. Verified against the backend where the claim depended on service behaviour. Run identity (high). init() exported PRIME_RUN_ID and then read it back, so a second init() in the same process silently attached to the first run and never created or finalized one of its own. Exports now record the exporting PID, and only an ID from a *different* PID counts as inherited — a forked child sees the parent's PID and joins, a re-init sees its own and opens a fresh run. Lifecycle ownership follows intent rather than "was an ID present": an explicit id= is a deliberate resume and finalizes, an ID picked up from the environment belongs to whoever exported it. finish() also stops advertising the run it owned. Fork safety (high). The child inherited the parent's httpx pools and buffered file handles. Two processes writing one socket interleave into a single HTTP stream, and a duplicated write buffer gets flushed twice. Connection and file holders now reset in the child through a single process-wide hook in _fork.py: dropped, never closed (closing sends close_notify down a socket the parent is still reading) and never flushed (the buffer holds records the parent will write itself). The hook is registered once rather than per instance, because os.register_at_fork cannot be undone — the old per-worker registration pinned every run the process ever opened. Metadata replacement (medium). The service writes metadata with {"$set": {"metadata": ...}}, a document-level replace. The failure fallback PUT carried only {"prime_runs": ...}, erasing the config finish() had just written. finalize() now receives the run's full config and merges into it. Abandoned uploads (medium). finish() waited 60s for a flush, ignored the result, then close() joined 30s — while a single sample POST is allowed 300s. The budget is now derived from the upload timeout, a flush that does not drain warns instead of passing silently, and close() leaves sinks open when the thread is still alive rather than pulling a client out from under a live request. Signal status (medium). The handler reported FAILED while atexit, RunStatus and the README all said CRASHED. Signals and Ctrl-C now report CRASHED: the producer never said the run failed, it was stopped from outside its control flow. FAILED stays for what the producer itself reports. Environment version pinning (medium). EnvironmentRef accepted version_id and dropped it, even though the API's EnvironmentReference carries it — the run attached to whatever version the hub resolved that day. Found while testing the signal fix: _handle_signal read the displaced handler *after* finish(), which restores and clears that table, so chaining always fell back to SIG_DFL — re-raising the signal at default disposition and killing the process instead of running the handler the application installed. Captured before finish() now. 122 tests (up from 107), including a real os.fork() end-to-end check that the child joins the parent's run and no record is written twice. Co-Authored-By: Claude Opus 5 (1M context) --- packages/prime-runs/README.md | 16 +++- packages/prime-runs/src/prime_runs/_fork.py | 60 ++++++++++++ packages/prime-runs/src/prime_runs/_http.py | 32 +++++-- .../src/prime_runs/backends/base.py | 15 ++- .../src/prime_runs/backends/evals.py | 69 ++++++++----- .../src/prime_runs/backends/offline.py | 3 + packages/prime-runs/src/prime_runs/run.py | 92 +++++++++++++++--- .../src/prime_runs/sinks/offline.py | 11 +++ .../prime-runs/src/prime_runs/sinks/traces.py | 45 ++++++--- packages/prime-runs/src/prime_runs/worker.py | 26 +++-- packages/prime-runs/tests/conftest.py | 6 ++ .../prime-runs/tests/test_evals_backend.py | 41 ++++++++ packages/prime-runs/tests/test_init.py | 96 +++++++++++++++++++ packages/prime-runs/tests/test_run.py | 55 ++++++++++- packages/prime-runs/tests/test_worker.py | 70 +++++++++++++- 15 files changed, 564 insertions(+), 73 deletions(-) create mode 100644 packages/prime-runs/src/prime_runs/_fork.py diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index fb1f3d09b..66629359c 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -85,11 +85,19 @@ Set the mode explicitly, or through `$PRIME_RUNS_MODE`. - **Applies backpressure.** The upload queue is bounded; if a producer durably outruns the uploader, records are dropped and counted (`run.dropped_records`) rather than stalling the run. -- **Survives forks.** A forked child gets a fresh uploader instead of inheriting - the parent's queue and locks. +- **Waits for its own uploads.** `finish()` gives queued records the same budget + a single upload gets (300s, `finish_timeout=`) and says so in a warning if + they do not drain, rather than finalizing over records still in flight. +- **Survives forks.** A forked child gets a fresh uploader, a fresh connection + pool and fresh file handles instead of writing the parent's — which would + interleave two processes into one HTTP stream and flush the parent's buffered + records a second time. It also joins the parent's run rather than opening + its own. - **Reports a terminal status.** Context manager, `atexit` and signal handlers - all route to the same idempotent `finish()`, so a killed process is recorded as - crashed rather than left running forever. + all route to the same idempotent `finish()`. A run the producer decided had + failed is `failed`; one stopped from outside its control flow — Ctrl-C, + SIGTERM, an exit that never reached `finish()` — is `crashed`. Neither is + left running forever. - **Knows about ranks.** Rank 0 owns creation and finalization; other ranks join through `PRIME_RUN_ID` and upload their own records. diff --git a/packages/prime-runs/src/prime_runs/_fork.py b/packages/prime-runs/src/prime_runs/_fork.py new file mode 100644 index 000000000..77b44cf36 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/_fork.py @@ -0,0 +1,60 @@ +"""One process-wide ``os.register_at_fork`` hook, shared by everything stateful. + +Hosted evals fork after the SDK is initialized, and a forked child inherits far +more than the queue: it gets copies of every open socket and every buffered file +handle. Using those is not merely untidy — two processes writing the same TCP +connection interleave bytes into one HTTP stream, and a duplicated write buffer +gets flushed twice, once from each side. + +So anything holding a connection or a file registers here and gets told to start +over in the child. Two rules for a ``reset_after_fork`` implementation: + +- **Drop, do not close.** Closing an inherited transport can send bytes — a TLS + ``close_notify``, an HTTP ``Connection: close`` — down a socket the parent is + still using. Release the reference and let the child's copies of the + descriptors go when it exits. +- **Do not flush.** A buffer inherited from the parent holds records the parent + has not written yet and will write itself. Flushing it in the child writes + them a second time. + +Registration is weak and the hook is installed once. A per-object +``register_at_fork`` call cannot be undone, so registering per instance would +pin every run the process ever opened in memory and re-run hooks for runs that +finished hours ago. +""" + +import logging +import os +import threading +import weakref +from typing import Any + +logger = logging.getLogger(__name__) + +_registry: "weakref.WeakSet[Any]" = weakref.WeakSet() +_lock = threading.Lock() +_installed = False + + +def register(obj: Any) -> None: + """Have ``obj.reset_after_fork()`` called in any child forked from here.""" + global _installed + with _lock: + _registry.add(obj) + if _installed or not hasattr(os, "register_at_fork"): # pragma: no cover - Windows + return + os.register_at_fork(after_in_child=_reset_all) + _installed = True + + +def _reset_all() -> None: + # A fresh lock: the child inherits the parent's, which may have been held by + # a thread that does not exist here. Nothing else runs at this point, so the + # swap is safe. + global _lock + _lock = threading.Lock() + for obj in list(_registry): + try: + obj.reset_after_fork() + except Exception as exc: # noqa: BLE001 - a fork hook must never raise + logger.debug("reset_after_fork failed for %r: %s", obj, exc) diff --git a/packages/prime-runs/src/prime_runs/_http.py b/packages/prime-runs/src/prime_runs/_http.py index a56970bcb..07cb8edc1 100644 --- a/packages/prime-runs/src/prime_runs/_http.py +++ b/packages/prime-runs/src/prime_runs/_http.py @@ -22,6 +22,7 @@ import httpx +from . import _fork from .exceptions import ( NotFoundError, PaymentRequiredError, @@ -94,15 +95,34 @@ def __init__( self.api_prefix = f"{self.base_url}/api/v1" self.max_attempts = max(1, max_attempts) self._owns_client = client is None - self._client = client or httpx.Client( - headers={ - "Authorization": f"Bearer {api_key}", - "User-Agent": _user_agent(), - }, + self._headers = { + "Authorization": f"Bearer {api_key}", + "User-Agent": _user_agent(), + } + self._timeout = timeout + self._client = client or self._new_client() + if self._owns_client: + # An injected client belongs to the caller (tests, the CLI); only a + # pool we opened ourselves is ours to rebuild after a fork. + _fork.register(self) + + def _new_client(self) -> httpx.Client: + return httpx.Client( + headers=dict(self._headers), follow_redirects=True, - timeout=timeout, + timeout=self._timeout, ) + def reset_after_fork(self) -> None: + """Rebuild the connection pool in a forked child. + + The inherited pool's sockets are the parent's: writing them would + interleave two processes' requests into one HTTP stream. The old client + is dropped rather than closed, because closing can send ``close_notify`` + on a connection the parent is still reading. + """ + self._client = self._new_client() + def request( self, method: str, diff --git a/packages/prime-runs/src/prime_runs/backends/base.py b/packages/prime-runs/src/prime_runs/backends/base.py index ad4d14ea0..f8fc7809d 100644 --- a/packages/prime-runs/src/prime_runs/backends/base.py +++ b/packages/prime-runs/src/prime_runs/backends/base.py @@ -45,7 +45,12 @@ def update( config: Optional[Dict[str, Any]] = None, summary: Optional[Dict[str, Any]] = None, ) -> None: - """Persist config (inputs) and/or summary (outputs) mid-run.""" + """Persist config (inputs) and/or summary (outputs) mid-run. + + ``config`` is the run's *whole* config, not a patch. The evaluations API + stores metadata with a document-level ``$set``, so a partial write + replaces whatever was there — every caller must send the full picture. + """ ... def log_metrics(self, run_id: str, metrics: Dict[str, Any], step: Optional[int] = None) -> None: @@ -59,8 +64,14 @@ def finalize( status: RunStatus, summary: Optional[Dict[str, Any]] = None, error: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, ) -> None: - """Close the run out. Called exactly once per run.""" + """Close the run out. Called exactly once per run. + + ``config`` is passed so a backend that has to record the terminal state + *inside* metadata can merge it into the full config rather than + replacing the document with one key. + """ ... def close(self) -> None: diff --git a/packages/prime-runs/src/prime_runs/backends/evals.py b/packages/prime-runs/src/prime_runs/backends/evals.py index 26c4caba2..0680506ad 100644 --- a/packages/prime-runs/src/prime_runs/backends/evals.py +++ b/packages/prime-runs/src/prime_runs/backends/evals.py @@ -60,8 +60,8 @@ def __init__( # ------------------------------------------------------------------ create def create(self, spec: RunSpec) -> RunHandle: - environment_ids = self._resolve_environments(spec.environments) - if not environment_ids: + environments = self._resolve_environments(spec.environments) + if not environments: raise ConfigurationError( "An eval run needs at least one environment. Pass " 'environments=["my-env"] to init().' @@ -70,7 +70,7 @@ def create(self, spec: RunSpec) -> RunHandle: run_name: str = spec.name or _default_name(spec) payload: Dict[str, Any] = { "name": run_name, - "environments": [{"id": environment_id} for environment_id in environment_ids], + "environments": environments, "tags": list(spec.tags), } _set_if(payload, "model_name", spec.model) @@ -124,6 +124,12 @@ def update( config: Optional[Dict[str, Any]] = None, summary: Optional[Dict[str, Any]] = None, ) -> None: + """Persist config and/or summary. + + ``config`` must be the run's whole config: the service writes metadata + with ``{"$set": {"metadata": ...}}``, which replaces the stored document + rather than merging into it. + """ payload: Dict[str, Any] = {} _set_if(payload, "metadata", config or None) _set_if(payload, "metrics", summary or None) @@ -147,13 +153,14 @@ def finalize( status: RunStatus, summary: Optional[Dict[str, Any]] = None, error: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, ) -> None: if status is RunStatus.COMPLETED: body: Dict[str, Any] = {} _set_if(body, "metrics", summary or None) self._client.post(f"/evaluations/{run_id}/finalize", json_body=body or {"metrics": {}}) return - self._report_failure(run_id, status=status, summary=summary, error=error) + self._report_failure(run_id, status=status, summary=summary, error=error, config=config) def _report_failure( self, @@ -162,6 +169,7 @@ def _report_failure( status: RunStatus, summary: Optional[Dict[str, Any]], error: Optional[str], + config: Optional[Dict[str, Any]] = None, ) -> None: """Mark a run failed, or record why we could not.""" terminal = { @@ -195,7 +203,14 @@ def _report_failure( # Fallback: the run cannot be moved out of RUNNING, but the failure is # at least recorded where an operator and the dashboard can both read it. - self.update(run_id, config={"prime_runs": terminal}, summary=summary) + # The terminal block is merged into the full config because this PUT + # replaces the stored metadata document — sending it alone would erase + # everything finish() just wrote. + self.update( + run_id, + config={**(config or {}), "prime_runs": terminal}, + summary=summary, + ) logger.warning( "Run %s %s, but the platform has no way to mark an evaluation failed; " "it will keep showing as running. Recorded the failure in metadata.prime_runs.", @@ -208,33 +223,41 @@ def close(self) -> None: # ----------------------------------------------------------- environments - def _resolve_environments(self, refs: List[EnvironmentRef]) -> List[str]: - """Environment IDs for the hub references a producer named. + def _resolve_environments(self, refs: List[EnvironmentRef]) -> List[Dict[str, Any]]: + """Hub references as the API's ``EnvironmentReference`` objects. + + ``version_id`` is carried through when the producer pinned one — the + API accepts it, and dropping it would silently attach the run to + whatever version the hub resolves today, which is the difference + between a reproducible eval and one that quietly moved. Unlike the old client, a reference that cannot be resolved raises instead of being skipped: dropping one silently produces a run attached to the wrong environments, which looks like a successful upload and is found much later. """ - resolved: List[str] = [] + resolved: List[Dict[str, Any]] = [] for ref in refs: - if ref.id: - resolved.append(ref.id) - continue - body: Dict[str, Any] = {"name": ref.name} - _set_if(body, "team_id", self._team_id) - try: - response = self._client.post("/environmentshub/resolve", json_body=body) - except RunAPIError as exc: - raise EnvironmentResolutionError( - f"Could not resolve environment {ref.name!r}: {exc}" - ) from exc - environment_id = (response.get("data") or {}).get("id") - if not environment_id: - raise EnvironmentResolutionError(f"Hub returned no id for environment {ref.name!r}") - resolved.append(environment_id) + entry: Dict[str, Any] = {"id": ref.id or self._lookup_environment(ref)} + _set_if(entry, "version_id", ref.version_id) + resolved.append(entry) return resolved + def _lookup_environment(self, ref: EnvironmentRef) -> str: + """Resolve one environment name to a hub ID (get-or-create).""" + body: Dict[str, Any] = {"name": ref.name} + _set_if(body, "team_id", self._team_id) + try: + response = self._client.post("/environmentshub/resolve", json_body=body) + except RunAPIError as exc: + raise EnvironmentResolutionError( + f"Could not resolve environment {ref.name!r}: {exc}" + ) from exc + environment_id = (response.get("data") or {}).get("id") + if not environment_id: + raise EnvironmentResolutionError(f"Hub returned no id for environment {ref.name!r}") + return str(environment_id) + def _set_if(payload: Dict[str, Any], key: str, value: Any) -> None: if value is not None: diff --git a/packages/prime-runs/src/prime_runs/backends/offline.py b/packages/prime-runs/src/prime_runs/backends/offline.py index 2b571e5ab..e0c508dde 100644 --- a/packages/prime-runs/src/prime_runs/backends/offline.py +++ b/packages/prime-runs/src/prime_runs/backends/offline.py @@ -108,8 +108,11 @@ def finalize( status: RunStatus, summary: Optional[Dict[str, Any]] = None, error: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, ) -> None: state = self._read_state(run_id) + if config: + state.setdefault("config", {}).update(config) state["status"] = status.value state["finished_at"] = _now() if error: diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 13ab0d482..87ef09562 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -24,7 +24,7 @@ import time from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Union -from ._http import DEFAULT_TIMEOUT, PlatformClient +from ._http import DEFAULT_TIMEOUT, UPLOAD_TIMEOUT, PlatformClient from .backends import Backend, EvalsBackend, OfflineBackend from .config import Config from .exceptions import ConfigurationError, RunFinishedError @@ -39,6 +39,18 @@ #: Rank variables, in the order prime-rl sets them. Rank 0 owns the lifecycle. RANK_ENV_VARS = ("RANK", "DP_RANK", "LOCAL_RANK") DEFAULT_SUMMARY_FLUSH_SECONDS = 10.0 +#: How long ``finish()`` waits for queued records. Derived from the upload +#: timeout rather than picked: a single in-flight sample POST is allowed 300s, +#: so a shorter budget here would routinely abandon an upload that was about to +#: succeed and then finalize the run without it. +DEFAULT_FINISH_TIMEOUT = float(UPLOAD_TIMEOUT.read or 300.0) + +#: Run IDs this process exported into ``PRIME_RUN_ID``, mapped to the PID that +#: exported them. The PID is the whole point: it is what distinguishes "my +#: parent opened this run and I should join it" from "I opened this run a moment +#: ago and the variable is still lying around". A forked child sees a different +#: PID and correctly treats the entry as inherited. +_exported_run_ids: Dict[str, int] = {} class Run: @@ -63,6 +75,7 @@ def __init__( is_primary: bool = True, owns_lifecycle: bool = True, summary_flush_seconds: float = DEFAULT_SUMMARY_FLUSH_SECONDS, + finish_timeout: float = DEFAULT_FINISH_TIMEOUT, queue_size: Optional[int] = None, ) -> None: self._backend = backend @@ -82,6 +95,7 @@ def __init__( self.errors: List[str] = [] self._summary_flush_seconds = summary_flush_seconds + self._finish_timeout = finish_timeout self._last_summary_flush = time.monotonic() self._summary_dirty = False self._config_dirty = False @@ -252,8 +266,14 @@ def finish( # Order matters: records first, so a dashboard that reacts to the # terminal status never sees a finished run with samples still landing. - self._worker.flush(timeout=60.0) - self._worker.close() + if not self._worker.flush(timeout=self._finish_timeout): + logger.warning( + "Run %s: uploads did not drain within %ss; finalizing anyway. " + "Some records may be missing from this run.", + self.id, + self._finish_timeout, + ) + self._worker.close(timeout=self._finish_timeout) if self._owns_lifecycle: self._report_guarded( @@ -271,12 +291,14 @@ def finish( status=resolved, summary=self.summary or None, error=error or (self.errors[0] if self.errors else None), + config=self.config or None, ), ) self._report_guarded("closing the backend", self._backend.close) atexit.unregister(self._atexit_hook) self._restore_signal_handlers() + self._retract_run_id() if self._worker.dropped: logger.warning( @@ -302,9 +324,10 @@ def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: if exc_type is None: self.finish() elif isinstance(exc, KeyboardInterrupt): - # An interrupt is a decision, not a fault. Recording it as FAILED - # would put every cancelled run in the same bucket as broken ones. - self.finish(status=RunStatus.FAILED, error="interrupted") + # An interrupt is a decision, not a fault, so it must not land in + # the same bucket as broken ones. Matches the SIGINT handler, which + # normally gets there first when signal handling is on. + self.finish(status=RunStatus.CRASHED, error="interrupted") else: self.finish(status=RunStatus.FAILED, error=_describe(exc)) return False @@ -336,9 +359,16 @@ def install_signal_handlers(self) -> None: def _handle_signal(self, signum: int, frame: Any) -> None: name = signal.Signals(signum).name - if not self._finished: - self.finish(status=RunStatus.FAILED, error=f"received {name}") + # Read the displaced handler *before* finishing: finish() restores and + # then clears this table, so looking it up afterwards always yields + # SIG_DFL — which re-raises the signal at its default disposition and + # kills the process instead of running the handler the app installed. previous = self._previous_signal_handlers.get(signum, signal.SIG_DFL) + if not self._finished: + # CRASHED, not FAILED: the producer never said the run failed, it was + # stopped from outside its own control flow. Same bucket as the + # atexit path, and deliberately not the bucket a broken eval lands in. + self.finish(status=RunStatus.CRASHED, error=f"received {name}") signal.signal(signum, previous) if callable(previous): previous(signum, frame) @@ -366,6 +396,19 @@ def _on_process_exit(self) -> None: logger.warning("Run %s was never finished; reporting it as crashed", self.id) self.finish(status=RunStatus.CRASHED, error="process exited without finishing the run") + def _retract_run_id(self) -> None: + """Stop advertising a finished run to processes started from here. + + Only retracts what this run published: if the value now points somewhere + else, another run owns it and clearing it would orphan that one's + children. + """ + if not self._owns_lifecycle: + return + if os.environ.get(RUN_ID_ENV) == self.id: + os.environ.pop(RUN_ID_ENV, None) + _exported_run_ids.pop(self.id, None) + def _require_live(self, operation: str) -> None: if self._finished: raise RunFinishedError( @@ -476,7 +519,13 @@ def init( ) is_primary = _is_primary_rank() - inherited_id = id or os.getenv(RUN_ID_ENV) or None + joined_id = _inherited_run_id() + inherited_id = id or joined_id + # Owning the lifecycle means "this call is responsible for creating and + # finalizing the run". An explicit `id=` is a deliberate resume, so it owns. + # An ID picked up from the environment belongs to whoever exported it, so it + # does not. (A non-primary rank never owns either way — see Run.__init__.) + owns_lifecycle = id is not None or joined_id is None resolved_mode = _resolve_mode(mode, api_key=api_key, is_primary=is_primary, run_id=inherited_id) if resolved_mode == "disabled": @@ -498,7 +547,7 @@ def init( resolved_mode, on_error, is_primary, - True, + owns_lifecycle, queue_size, ) _announce(run, handle_signals) @@ -517,7 +566,6 @@ def init( client = PlatformClient(api_key=api_key, base_url=base_url, timeout=DEFAULT_TIMEOUT) backend = EvalsBackend(client, frontend_url=resolved_config.frontend_url, team_id=team_id) - owns_lifecycle = inherited_id is None handle = backend.attach(inherited_id) if inherited_id else backend.create(spec) if sinks is None: @@ -577,15 +625,33 @@ def _announce(run: Run, handle_signals: bool) -> None: Exporting ``PRIME_RUN_ID`` is how forked workers and subprocess launchers join the run their parent created instead of each opening their own — the same trick prime-rl's monitor used with ``RUN_ID``, generalized so every - producer gets it. + producer gets it. The PID is recorded alongside so that *this* process does + not later mistake its own export for a parent's. """ - os.environ.setdefault(RUN_ID_ENV, run.id) + os.environ[RUN_ID_ENV] = run.id + _exported_run_ids[run.id] = os.getpid() if handle_signals: run.install_signal_handlers() if run.url: logger.info("Run %s: %s", run.id, run.url) +def _inherited_run_id() -> Optional[str]: + """A run ID this process should join, or ``None`` to open a fresh run. + + ``PRIME_RUN_ID`` set by an ancestor means "join that run". The same variable + set by an earlier ``init()`` *in this process* means nothing of the sort — + without this check, a second eval in one process would silently attach to + the first, and would never create or finalize a run of its own. + """ + value = os.getenv(RUN_ID_ENV) + if not value: + return None + if _exported_run_ids.get(value) == os.getpid(): + return None + return value + + class _DisabledBackend: """No-op lifecycle, so ``mode="disabled"`` needs no branching upstream.""" diff --git a/packages/prime-runs/src/prime_runs/sinks/offline.py b/packages/prime-runs/src/prime_runs/sinks/offline.py index 73f65392c..e7d35f9d3 100644 --- a/packages/prime-runs/src/prime_runs/sinks/offline.py +++ b/packages/prime-runs/src/prime_runs/sinks/offline.py @@ -12,6 +12,7 @@ from pathlib import Path from typing import Any, Mapping, Optional, Sequence, TextIO, Union +from .. import _fork from .base import Sink, to_mapping logger = logging.getLogger(__name__) @@ -30,6 +31,16 @@ def __init__(self, directory: Union[str, Path], *, stamp_run: bool = True) -> No self._run_kind: Optional[str] = None self._handles: dict[str, TextIO] = {} self.records_written = 0 + _fork.register(self) + + def reset_after_fork(self) -> None: + """Abandon inherited file handles; ``_handle`` reopens on next write. + + Dropped without flushing or closing: the inherited buffer holds records + the parent has not written yet and will write itself, so flushing it + here would put every one of them in the file twice. + """ + self._handles = {} def start(self, run_id: str, context: Mapping[str, str]) -> None: self._run_id = run_id diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py index 8709aaedb..8e02d24df 100644 --- a/packages/prime-runs/src/prime_runs/sinks/traces.py +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -18,6 +18,7 @@ import logging from typing import Any, Dict, Mapping, Optional, Sequence +from .. import _fork from .base import Sink logger = logging.getLogger(__name__) @@ -40,6 +41,7 @@ def __init__( ) -> None: self.enabled = True self._client = client + self._injected_client = client is not None # Left unset, prime-traces resolves its own endpoint. That matters: # the service has its own URL (PRIME_TRACES_URL / config `traces_url`) # which is not necessarily the platform API's, and passing the @@ -57,6 +59,7 @@ def __init__( self._run_kind: Optional[str] = None self._context: Dict[str, str] = {} self.receipts: list = [] + _fork.register(self) # ------------------------------------------------------------------ setup @@ -64,16 +67,36 @@ def start(self, run_id: str, context: Mapping[str, str]) -> None: self._run_id = run_id self._run_kind = context.get("run_kind") self._context = {key: str(value) for key, value in context.items() if value is not None} - if self._client is None: - try: - from prime_traces import TracesClient - except ImportError as exc: # pragma: no cover - dependency is declared - self._disable(f"prime-traces is not installed ({exc})") - return - try: - self._client = TracesClient(**self._client_kwargs) - except Exception as exc: # noqa: BLE001 - construction must not kill a run - self._disable(f"could not construct the traces client ({exc})") + self._ensure_client() + + def _ensure_client(self) -> bool: + """Build the traces client if we do not have one. Lazy so that a fork + reset — which drops the inherited client — is repaired on next write.""" + if self._client is not None: + return True + if self._injected_client: + # The caller handed us a client and a fork took it away. Rebuilding + # would silently swap their transport for a default one. + return False + try: + from prime_traces import TracesClient + except ImportError as exc: # pragma: no cover - dependency is declared + self._disable(f"prime-traces is not installed ({exc})") + return False + try: + self._client = TracesClient(**self._client_kwargs) + except Exception as exc: # noqa: BLE001 - construction must not kill a run + self._disable(f"could not construct the traces client ({exc})") + return False + return True + + def reset_after_fork(self) -> None: + """Drop the inherited traces client; the next write builds a fresh one. + + Not closed: the child's copy of the socket is the parent's connection, + and shutting it down here would cut the parent off mid-upload. + """ + self._client = None # ------------------------------------------------------------------ write @@ -84,7 +107,7 @@ def write( line_format: Optional[str] = None, step: Optional[int] = None, ) -> None: - if not self.enabled or not records or self._client is None: + if not self.enabled or not records or not self._ensure_client(): return from prime_traces import LineFormat diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index 1d6a73ba8..022ac3464 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -27,6 +27,8 @@ from dataclasses import dataclass, field from typing import Any, Callable, List, Optional, Sequence +from . import _fork + logger = logging.getLogger(__name__) DEFAULT_QUEUE_SIZE = 256 @@ -83,7 +85,7 @@ def __init__( self._lock = threading.Lock() self.dropped = 0 self._pid = os.getpid() - self._register_fork_hook() + _fork.register(self) # ----------------------------------------------------------------- thread @@ -220,7 +222,16 @@ def close(self, timeout: Optional[float] = 30.0) -> None: logger.warning("Upload queue saturated at close; some records may be lost") thread.join(timeout) if thread.is_alive(): - logger.warning("Uploader did not stop within %ss; abandoning it", timeout) + # Closing the sinks now would pull an httpx client, or a file + # handle, out from under a request that is still running on that + # thread — turning a slow upload into a crash inside a daemon + # thread nobody is watching. Leave them to the interpreter. + logger.warning( + "Uploader still running after %ss; leaving it and its sinks open. " + "Records still in flight may not finish before the process exits.", + timeout, + ) + return self._thread = None for sink in self.sinks: try: @@ -230,18 +241,15 @@ def close(self, timeout: Optional[float] = 30.0) -> None: # ------------------------------------------------------------------- fork - def _register_fork_hook(self) -> None: - if not hasattr(os, "register_at_fork"): # pragma: no cover - Windows - return - os.register_at_fork(after_in_child=self._reinit_after_fork) - - def _reinit_after_fork(self) -> None: + def reset_after_fork(self) -> None: """Give the child a clean uploader. Everything queued at fork time belongs to the parent, which still has a live thread and will send it. Inheriting that queue would upload each record twice; inheriting the lock could deadlock the child on its first - write. + write. The sinks reset themselves through the same hook — their sockets + and file buffers are the parent's too, and using those from two + processes interleaves one HTTP stream or writes one buffer twice. """ self._pid = os.getpid() self._queue = queue.Queue(maxsize=self.max_queue_size) diff --git a/packages/prime-runs/tests/conftest.py b/packages/prime-runs/tests/conftest.py index 745ba07d4..96e4b995e 100644 --- a/packages/prime-runs/tests/conftest.py +++ b/packages/prime-runs/tests/conftest.py @@ -107,6 +107,12 @@ def eval_routes() -> Dict[str, Any]: "evaluation_id": "eval-abc", "status": "PROCESSING", }, + "GET /api/v1/evaluations/eval-abc": { + "evaluation_id": "eval-abc", + "name": "test-run", + "status": "RUNNING", + "viewer_url": "https://app.example/dashboard/evaluations/eval-abc", + }, "PUT /api/v1/evaluations/eval-abc": { "evaluation_id": "eval-abc", "name": "test-run", diff --git a/packages/prime-runs/tests/test_evals_backend.py b/packages/prime-runs/tests/test_evals_backend.py index 076369795..fa3c68e87 100644 --- a/packages/prime-runs/tests/test_evals_backend.py +++ b/packages/prime-runs/tests/test_evals_backend.py @@ -167,3 +167,44 @@ def test_attach_survives_a_read_failure(make_platform_client, eval_routes): assert handle.id == "eval-abc" assert handle.url == "https://app.example/dashboard/evaluations/eval-abc" + + +def test_a_pinned_environment_version_reaches_the_api(make_platform_client, eval_routes): + """The API's EnvironmentReference carries version_id. Dropping it attaches + the run to whatever version the hub resolves today — the difference between + a reproducible eval and one that quietly moved.""" + backend, handler = make_backend(make_platform_client, eval_routes) + + backend.create(RunSpec(name="r", environments=[EnvironmentRef(id="env-1", version_id="v-7")])) + + assert handler.bodies_for("/api/v1/evaluations/")[0]["environments"] == [ + {"id": "env-1", "version_id": "v-7"} + ] + + +def test_a_version_pin_survives_hub_resolution(make_platform_client, eval_routes): + backend, handler = make_backend(make_platform_client, eval_routes) + + backend.create(RunSpec(name="r", environments=[EnvironmentRef(name="gsm8k", version_id="v-7")])) + + assert handler.bodies_for("/api/v1/evaluations/")[0]["environments"] == [ + {"id": "env-123", "version_id": "v-7"} + ] + + +def test_the_failure_fallback_preserves_the_run_config(make_platform_client, eval_routes): + """The service writes metadata with a document-level $set, so a PUT carrying + only the terminal block would erase everything finish() just wrote.""" + backend, handler = make_backend(make_platform_client, eval_routes) + + backend.finalize( + "eval-abc", + status=RunStatus.FAILED, + error="boom", + config={"num_rollouts": 4, "model": "Qwen3-8B"}, + ) + + metadata = handler.bodies_for("/api/v1/evaluations/eval-abc")[0]["metadata"] + assert metadata["num_rollouts"] == 4 + assert metadata["model"] == "Qwen3-8B" + assert metadata["prime_runs"]["status"] == "failed" diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index 1bf8f6a83..0af0d074c 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -229,3 +229,99 @@ def test_the_end_to_end_shape_a_producer_writes(online): assert finalize["metrics"]["avg_reward"] == 1.0 assert run.status is RunStatus.COMPLETED assert run.errors == [] + + +# ------------------------------------------------------- run id inheritance + + +def test_a_second_init_in_one_process_opens_its_own_run(tmp_path): + """init() exports PRIME_RUN_ID for child processes. Reading our own export + back would silently attach the second eval to the first, and it would never + create or finalize a run of its own.""" + first = pr.init(mode="offline", dir=str(tmp_path)) + first.finish() + + second = pr.init(mode="offline", dir=str(tmp_path)) + second.finish() + + assert second.id != first.id + assert (tmp_path / second.id / "run.json").exists() + + +def test_a_finished_run_stops_advertising_itself(tmp_path): + run = pr.init(mode="offline", dir=str(tmp_path)) + assert os.environ[RUN_ID_ENV] == run.id + + run.finish() + + assert RUN_ID_ENV not in os.environ + + +def test_an_id_inherited_from_a_parent_process_is_joined(monkeypatch, tmp_path): + """The env var without a matching PID belongs to an ancestor.""" + monkeypatch.setenv(RUN_ID_ENV, "offline-from-parent") + + run = pr.init(mode="offline", dir=str(tmp_path)) + + assert run.id == "offline-from-parent" + run.finish() + + +def test_an_explicit_id_is_a_resume_and_still_finalizes(online): + """Resuming after a crash has to be able to close the run out; only an ID + picked up from the environment belongs to someone else.""" + run, handler = online(id="eval-abc") + + run.finish(summary={"avg_reward": 1.0}) + + assert "POST /api/v1/evaluations/eval-abc/finalize" in handler.paths() + + +def test_an_id_inherited_from_the_environment_does_not_finalize(monkeypatch, online): + monkeypatch.setenv(RUN_ID_ENV, "eval-abc") + + run, handler = online() + run.finish() + + assert "POST /api/v1/evaluations/eval-abc/finalize" not in handler.paths() + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork is POSIX-only") +# Forking a threaded process is exactly the situation under test — hosted evals +# do it, and the uploader thread is why the SDK needs a fork hook at all. +@pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning") +def test_a_forked_child_joins_the_run_without_duplicating_the_parents_records(tmp_path): + """The end-to-end shape hosted evals actually hit. + + At fork time the parent has records in the upload queue and bytes in the + sink's write buffer. The child inherits copies of both; writing them would + put every one of those records in the file twice, and opening its own run + would split one job across two. + """ + import json + + run = pr.init(mode="offline", dir=str(tmp_path), handle_signals=False) + run.log_traces([{"id": f"parent-{n}"} for n in range(5)]) + + pid = os.fork() + if pid == 0: # pragma: no cover - asserted through the child's exit code + code = 0 + try: + child = pr.init(mode="offline", dir=str(tmp_path), handle_signals=False) + if child.id != run.id: + code = 1 + child.log_traces([{"id": "child-1"}]) + child.finish() + except BaseException: + code = 2 + finally: + os._exit(code) + + _, status = os.waitpid(pid, 0) + assert os.waitstatus_to_exitcode(status) == 0, "the child did not join the parent's run" + run.finish() + + lines = (tmp_path / run.id / "records" / "trace.jsonl").read_text().splitlines() + ids = [json.loads(line)["id"] for line in lines] + + assert sorted(ids) == sorted([f"parent-{n}" for n in range(5)] + ["child-1"]) diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index 3c22f733a..a9e8cb80c 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -1,5 +1,6 @@ """The run handle: lifecycle, containment, ranks, terminal status.""" +import signal from typing import Any, Dict, List, Optional import pytest @@ -34,10 +35,12 @@ def update(self, run_id, *, config=None, summary=None) -> None: def log_metrics(self, run_id, metrics, step=None) -> None: self.points.append((metrics, step)) - def finalize(self, run_id, *, status, summary=None, error=None) -> None: + def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> None: if self.fail_on == "finalize": raise RuntimeError("finalize exploded") - self.finalized.append({"status": status, "summary": summary, "error": error}) + self.finalized.append( + {"status": status, "summary": summary, "error": error, "config": config} + ) def close(self) -> None: self.closed = True @@ -161,9 +164,9 @@ def write(self, records, *, line_format=None, step=None) -> None: super().write(records, line_format=line_format, step=step) class OrderedBackend(FakeBackend): - def finalize(self, run_id, *, status, summary=None, error=None) -> None: + def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> None: order.append("finalize") - super().finalize(run_id, status=status, summary=summary, error=error) + super().finalize(run_id, status=status, summary=summary, error=error, config=config) run = make_run(OrderedBackend(), sinks=[OrderedSink()]) run.log_traces([{"id": "t1"}]) @@ -214,15 +217,59 @@ def test_an_exception_inside_the_block_fails_the_run_and_still_propagates(): def test_an_interrupt_is_recorded_as_a_decision_not_a_fault(): + """Ctrl-C must not land in the same bucket as a broken eval — and it must + agree with the SIGINT handler, which normally gets there first.""" backend = FakeBackend() with pytest.raises(KeyboardInterrupt): with make_run(backend): raise KeyboardInterrupt + assert backend.finalized[0]["status"] is RunStatus.CRASHED assert backend.finalized[0]["error"] == "interrupted" +def test_a_termination_signal_reports_crashed_like_atexit_does(): + """The producer never said the run failed; it was stopped from outside.""" + backend = FakeBackend() + run = make_run(backend) + chained = [] + # Stand in for the handler the SDK displaced. Anything but SIG_DFL, which + # would re-raise the signal and take the test runner down with it. + run._previous_signal_handlers[signal.SIGTERM] = lambda *a: chained.append(a) + + run._handle_signal(signal.SIGTERM, None) + + assert backend.finalized[0]["status"] is RunStatus.CRASHED + assert "SIGTERM" in backend.finalized[0]["error"] + assert chained, "the displaced handler still runs" + signal.signal(signal.SIGTERM, signal.SIG_DFL) + + +def test_finish_hands_the_full_config_to_finalize(): + """The evaluations API replaces metadata wholesale, so a backend recording + terminal state inside it needs the whole picture to merge into.""" + backend = FakeBackend() + run = make_run(backend) + run.update_config({"num_rollouts": 4}) + + run.finish(status=RunStatus.FAILED, error="boom") + + assert backend.finalized[0]["config"]["num_rollouts"] == 4 + + +def test_finish_warns_when_uploads_do_not_drain(caplog): + """Finalizing over an unfinished upload silently drops records.""" + run = make_run(sinks=[FakeSink()]) + run._finish_timeout = 0.01 + run._worker.flush = lambda timeout=None: False + + with caplog.at_level("WARNING"): + run.finish() + + assert "did not drain" in caplog.text + + def test_a_process_that_exits_without_finishing_reports_crashed(): """The producer never said the run failed — it stopped existing. The distinction tells an operator where to look.""" diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index d3673a8bf..83733ffc9 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -158,10 +158,78 @@ def test_a_forked_child_starts_over_instead_of_re_uploading_the_parents_queue(): worker._queue.put(WriteItem(records=[{"id": "parents"}])) old_queue = worker._queue - worker._reinit_after_fork() + worker.reset_after_fork() assert worker._queue is not old_queue assert worker._queue.empty() assert worker._thread is None assert isinstance(worker._queue, queue.Queue) assert isinstance(worker._lock, type(threading.Lock())) or worker._lock is not None + + +def test_close_leaves_sinks_open_when_the_uploader_will_not_stop(caplog): + """Closing them would pull an httpx client or a file handle out from under + a request still running on that thread.""" + + class WedgedSink(FakeSink): + def __init__(self) -> None: + super().__init__("wedged") + self.released = threading.Event() + + def write(self, records, *, line_format=None, step=None) -> None: + self.released.wait(10.0) + + sink = WedgedSink() + worker = UploadWorker([sink]) + worker.submit(WriteItem(records=[{"id": 1}])) + + with caplog.at_level("WARNING"): + worker.close(timeout=0.2) + + assert sink.closed is False + assert "leaving it and its sinks open" in caplog.text + sink.released.set() + + +def test_a_forked_child_resets_every_registered_holder_of_a_connection(): + """One process-wide hook, not one per object: a per-instance + register_at_fork can never be undone, so it would pin every run the process + ever opened and re-run hooks for runs that finished hours ago.""" + from prime_runs import _fork + + class Holder: + def __init__(self) -> None: + self.reset = 0 + + def reset_after_fork(self) -> None: + self.reset += 1 + + holder = Holder() + _fork.register(holder) + + _fork._reset_all() + + assert holder.reset == 1 + + +def test_one_registered_object_raising_does_not_block_the_others(): + from prime_runs import _fork + + class Boom: + def reset_after_fork(self) -> None: + raise RuntimeError("nope") + + class Fine: + def __init__(self) -> None: + self.reset = 0 + + def reset_after_fork(self) -> None: + self.reset += 1 + + fine = Fine() + _fork.register(Boom()) + _fork.register(fine) + + _fork._reset_all() + + assert fine.reset == 1 From 332fc2167006976f6c809a96114836b1b0be4829 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 14:13:49 -0700 Subject: [PATCH 03/39] fix(runs): address Codex review findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three findings the Bugbot pass did not cover. The other two Codex raised (environment version pinning, PRIME_RUN_ID re-init) were already fixed in 14bbac59, including its note that `setdefault` left descendants pointing at a stale ID when an explicit `id=` was supplied — `_announce` now assigns unconditionally. Non-idempotent retries. The client retried every method through ambiguous failures, so a lost response to POST /evaluations/ would create a second evaluation that the SDK never tracked, leaving an orphaned duplicate run. The module docstring even asserted this was safe "because run creation happens once per init()", which confuses the call site with the retry loop inside it. Retry safety is now decided per call. A failure is ambiguous when the request may already have been processed (gateway 502/504, read timeout, stream broken after the bytes went out); unambiguous when nothing reached the server (connect failure, 429 refused before any work). Unambiguous failures replay for every method. Ambiguous ones replay only when the caller declares `idempotent=True`, which defaults to `method != "POST"`. Run creation and sample appends keep the default; get-or-create, finalize and status writes declare themselves safe. Same classification prime-traces' client already uses. This also stops the samples sink duplicating rows on a lost response — duplicates silently skew every average on the dashboard, where a lost batch is at least recoverable. Uploader failures under on_error="raise". A sink fails on the uploader thread, where the raise went straight into the worker's own except and was discarded — so flush() and finish() returned success while records were being dropped, in exactly the mode documented as being for "tests and CI, where a silent upload failure is the bug". The failure is now held and re-raised at the next synchronization point the caller controls: flush(), or the very end of finish() so the run is still closed out properly first. The atexit and signal paths swallow it, since neither is a place to surface an exception. Signal handlers were never restored. `_restore_signal_handlers` compared `signal.getsignal(signum) is self._handle_signal`, but every access to a bound method builds a new object, so the identity check could never match. Handlers stayed installed for the life of the process: the finished Run was pinned, and the next run in that process saw a non-default handler and declined to install its own, leaving it unable to report signal termination. The bound method is created once in __init__ and compared against. 133 tests (up from 122). Co-Authored-By: Claude Opus 5 (1M context) --- packages/prime-runs/src/prime_runs/_http.py | 40 +++++++++-- .../src/prime_runs/backends/evals.py | 16 ++++- packages/prime-runs/src/prime_runs/run.py | 54 ++++++++++++-- .../src/prime_runs/sinks/samples.py | 11 ++- packages/prime-runs/tests/test_http.py | 70 +++++++++++++++++++ packages/prime-runs/tests/test_run.py | 62 ++++++++++++++++ 6 files changed, 237 insertions(+), 16 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/_http.py b/packages/prime-runs/src/prime_runs/_http.py index 07cb8edc1..027031e3a 100644 --- a/packages/prime-runs/src/prime_runs/_http.py +++ b/packages/prime-runs/src/prime_runs/_http.py @@ -9,10 +9,17 @@ - retries 429/502/503/504 and transport failures with exponential backoff, honouring ``Retry-After`` when the server sends one. -Retries are safe here because every call it makes is either idempotent (PUT, -GET) or create-shaped and guarded upstream: run creation happens exactly once -per ``init()``, and sample POSTs that get retried after a lost response are the -known duplicate-append case the traces sink exists to replace. +Retry safety is decided per call, not per client. A failure is *ambiguous* when +the request may already have been processed — a gateway 502/504, a read timeout, +a stream broken after the bytes went out. Replaying an ambiguous failure is fine +for a GET or a PUT and is not fine for ``POST /evaluations/``: if the platform +created the run and the response was lost, the retry creates a second one and +the SDK only ever knows about the second, leaving an orphaned duplicate. + +So callers declare intent with ``idempotent=``. Unambiguous failures — a +connection that was never established, a 429 refused before any work — are +replayed for every method, because there is nothing on the other side to +duplicate. """ import json @@ -37,6 +44,10 @@ # answering, so uploads get their own, much longer budget. UPLOAD_TIMEOUT = httpx.Timeout(300.0, connect=10.0) RETRY_STATUS = frozenset({429, 502, 503, 504}) +#: Refused before the server did any work, so replaying cannot duplicate +#: anything. 503 is deliberately *not* here: without a service error code it may +#: equally have come from an intermediary after the request was forwarded. +UNAMBIGUOUS_RETRY_STATUS = frozenset({429}) DEFAULT_MAX_ATTEMPTS = 5 MAX_BACKOFF_SECONDS = 16.0 @@ -133,15 +144,24 @@ def request( params: Optional[Mapping[str, Any]] = None, timeout: Union[httpx.Timeout, float, None] = None, max_attempts: Optional[int] = None, + idempotent: Optional[bool] = None, ) -> Dict[str, Any]: - """Send one request, retrying transient failures. Returns the JSON body.""" + """Send one request, retrying transient failures. Returns the JSON body. + + ``idempotent`` says whether replaying this request is safe when the + outcome is unknown. It defaults to ``method != "POST"``; a POST that is + in fact safe to replay (get-or-create, setting a terminal state) should + pass ``idempotent=True`` explicitly. + """ url = f"{self.api_prefix}{path}" body = content if content is not None else (encode_json(json_body) if json_body else None) headers = {"Content-Type": "application/json"} if body is not None else None attempts = max_attempts or self.max_attempts + replayable = idempotent if idempotent is not None else method.upper() != "POST" last_error: Optional[Exception] = None for attempt in range(1, attempts + 1): + ambiguous = True try: response = self._client.request( method, @@ -151,12 +171,17 @@ def request( params=dict(params) if params else None, timeout=timeout, ) + except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as exc: + # No connection was ever established, so the server saw nothing. + ambiguous = False + last_error = TransportError(f"{method} {path} failed to connect: {exc}") except httpx.TimeoutException as exc: last_error = TransportError(f"{method} {path} timed out: {exc}") except httpx.RequestError as exc: last_error = TransportError(f"{method} {path} failed: {type(exc).__name__}: {exc}") else: if response.status_code in RETRY_STATUS: + ambiguous = response.status_code not in UNAMBIGUOUS_RETRY_STATUS last_error = RetryableAPIError( _error_message(response), status_code=response.status_code, @@ -170,6 +195,11 @@ def request( if attempt == attempts: break + if ambiguous and not replayable: + # The request may already have been processed and replaying it + # would create a second resource. Surfacing the failure is the + # lesser harm: the caller can look, a duplicate cannot be undone. + break after = getattr(last_error, "retry_after", None) time.sleep(retry_delay(attempt, after)) diff --git a/packages/prime-runs/src/prime_runs/backends/evals.py b/packages/prime-runs/src/prime_runs/backends/evals.py index 0680506ad..a2bdaa585 100644 --- a/packages/prime-runs/src/prime_runs/backends/evals.py +++ b/packages/prime-runs/src/prime_runs/backends/evals.py @@ -81,6 +81,9 @@ def create(self, spec: RunSpec) -> RunHandle: _set_if(payload, "metrics", spec.summary or None) _set_if(payload, "team_id", spec.team_id or self._team_id) + # Not replayable: POST defaults to idempotent=False here on purpose. If + # the platform created the run and the response was lost, a retry would + # create a second one and only the second would be tracked. response = self._client.post("/evaluations/", json_body=payload) run_id = response.get("evaluation_id") if not run_id: @@ -158,7 +161,12 @@ def finalize( if status is RunStatus.COMPLETED: body: Dict[str, Any] = {} _set_if(body, "metrics", summary or None) - self._client.post(f"/evaluations/{run_id}/finalize", json_body=body or {"metrics": {}}) + # Setting a terminal state: replaying it lands on the same state. + self._client.post( + f"/evaluations/{run_id}/finalize", + json_body=body or {"metrics": {}}, + idempotent=True, + ) return self._report_failure(run_id, status=status, summary=summary, error=error, config=config) @@ -185,6 +193,7 @@ def _report_failure( f"/evaluations/{run_id}/status", json_body={"status": _PLATFORM_STATUS[status], "error": error}, max_attempts=1, + idempotent=True, ) return except NotFoundError: @@ -248,7 +257,10 @@ def _lookup_environment(self, ref: EnvironmentRef) -> str: body: Dict[str, Any] = {"name": ref.name} _set_if(body, "team_id", self._team_id) try: - response = self._client.post("/environmentshub/resolve", json_body=body) + # Get-or-create: a replay returns the same environment. + response = self._client.post( + "/environmentshub/resolve", json_body=body, idempotent=True + ) except RunAPIError as exc: raise EnvironmentResolutionError( f"Could not resolve environment {ref.name!r}: {exc}" diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 87ef09562..3d867b399 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -93,6 +93,11 @@ def __init__( self.config: Dict[str, Any] = dict(spec.config) self.summary: Dict[str, Any] = dict(spec.summary) self.errors: List[str] = [] + # Raised at the next synchronization point the caller controls. A sink + # fails on the uploader thread, where raising reaches nobody — so under + # on_error="raise" the exception is held and re-raised from flush() or + # finish(), which is where a test or a CI job is actually looking. + self._deferred_error: Optional[BaseException] = None self._summary_flush_seconds = summary_flush_seconds self._finish_timeout = finish_timeout @@ -122,6 +127,12 @@ def __init__( self._atexit_hook = self._on_process_exit atexit.register(self._atexit_hook) + # Bound once and kept. ``self._handle_signal`` builds a *new* bound + # method on every attribute access, so an ``is`` comparison against a + # freshly-made one is always False — which is how handlers end up + # installed forever, pinning a finished run and blocking the next run in + # the process from installing its own. + self._signal_handler = self._handle_signal self._previous_signal_handlers: Dict[int, Any] = {} # -------------------------------------------------------------- identity @@ -304,15 +315,23 @@ def finish( logger.warning( "Run %s finished with %d dropped record(s)", self.id, self._worker.dropped ) + # Last, so a run that failed to upload is still closed out properly + # before the failure reaches the caller. + self._raise_deferred() def fail(self, error: Union[str, BaseException]) -> None: """Close the run out as failed.""" self.finish(status=RunStatus.FAILED, error=_describe(error)) def flush(self, timeout: Optional[float] = 30.0) -> bool: - """Block until queued records have been written. Mostly for tests.""" + """Block until queued records have been written. + + Under ``on_error="raise"`` this is the first place an upload failure can + surface, since the failure itself happened on the uploader thread. + """ flushed = self._worker.flush(timeout=timeout) self._flush_summary() + self._raise_deferred() return flushed # -------------------------------------------------------- context manager @@ -352,7 +371,7 @@ def install_signal_handlers(self) -> None: if current not in (signal.SIG_DFL, signal.default_int_handler): continue try: - signal.signal(signum, self._handle_signal) + signal.signal(signum, self._signal_handler) except (ValueError, OSError): # pragma: no cover continue self._previous_signal_handlers[signum] = current @@ -368,7 +387,10 @@ def _handle_signal(self, signum: int, frame: Any) -> None: # CRASHED, not FAILED: the producer never said the run failed, it was # stopped from outside its own control flow. Same bucket as the # atexit path, and deliberately not the bucket a broken eval lands in. - self.finish(status=RunStatus.CRASHED, error=f"received {name}") + try: + self.finish(status=RunStatus.CRASHED, error=f"received {name}") + except Exception as exc: # noqa: BLE001 - the signal must still chain + logger.warning("Run %s: reporting %s failed: %s", self.id, name, exc) signal.signal(signum, previous) if callable(previous): previous(signum, frame) @@ -378,7 +400,7 @@ def _handle_signal(self, signum: int, frame: Any) -> None: def _restore_signal_handlers(self) -> None: for signum, previous in self._previous_signal_handlers.items(): try: - if signal.getsignal(signum) is self._handle_signal: + if signal.getsignal(signum) is self._signal_handler: signal.signal(signum, previous) except (ValueError, OSError): # pragma: no cover continue @@ -394,7 +416,12 @@ def _on_process_exit(self) -> None: if self._finished: return logger.warning("Run %s was never finished; reporting it as crashed", self.id) - self.finish(status=RunStatus.CRASHED, error="process exited without finishing the run") + # on_error="raise" must not turn interpreter shutdown into a traceback + # from atexit; the run is already being reported as crashed. + try: + self.finish(status=RunStatus.CRASHED, error="process exited without finishing the run") + except Exception as exc: # noqa: BLE001 + logger.warning("Run %s: reporting the crash failed: %s", self.id, exc) def _retract_run_id(self) -> None: """Stop advertising a finished run to processes started from here. @@ -439,7 +466,22 @@ def _flush_summary(self) -> None: ) def _record_sink_error(self, sink_name: str, exc: Exception) -> None: - self._report(f"writing to the {sink_name} sink", exc) + """Called on the uploader thread when a sink gives up.""" + message = f"writing to the {sink_name} sink failed: {type(exc).__name__}: {exc}" + self.errors.append(message) + if self._on_error == "raise": + if self._deferred_error is None: + self._deferred_error = exc + return + logger.warning("Run %s: %s", self._handle.id, message) + + def _raise_deferred(self) -> None: + """Re-raise the first upload failure, once.""" + exc = self._deferred_error + if exc is None: + return + self._deferred_error = None + raise exc def _report_guarded(self, what: str, call: Any) -> None: try: diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py index 55a1755f2..1be93ec96 100644 --- a/packages/prime-runs/src/prime_runs/sinks/samples.py +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -8,9 +8,11 @@ default sink list, with nothing to do in verifiers or prime-rl. Its known weakness is why traces is the primary: ``POST /samples`` *appends*, -so a retried request whose response was lost duplicates rows. Content-addressed -uploads do not have that problem, which is exactly the property the traces sink -was built on. +so a request whose response was lost cannot be safely replayed. The client +therefore does not retry it through an ambiguous failure — losing a batch is +recoverable, duplicated rows silently skew every average on the dashboard. +Content-addressed uploads have neither problem, which is exactly the property +the traces sink was built on. """ import logging @@ -61,6 +63,9 @@ def write( f"/evaluations/{self._run_id}/samples", content=encode_json({"samples": batch}), timeout=UPLOAD_TIMEOUT, + # Appends. Left non-replayable (the POST default) so a lost + # response cannot turn into duplicate rows. + idempotent=False, ) self.samples_written += len(batch) diff --git a/packages/prime-runs/tests/test_http.py b/packages/prime-runs/tests/test_http.py index 79e4b0118..399576324 100644 --- a/packages/prime-runs/tests/test_http.py +++ b/packages/prime-runs/tests/test_http.py @@ -108,3 +108,73 @@ def test_encoding_refuses_values_json_cannot_carry(): """Bare ``NaN`` is JavaScript, not JSON; it comes back as an opaque 400.""" with pytest.raises(ValueError): encode_json({"reward": float("nan")}) + + +def test_an_ambiguous_failure_does_not_replay_a_create(no_sleep): + """If the platform created the run and the response was lost, a retry makes + a second one and only the second is tracked — an orphaned duplicate run.""" + attempts = [] + + def handler(request): + attempts.append(request) + return httpx.Response(502, text="bad gateway") + + with pytest.raises(RetryableAPIError): + client_for(handler, max_attempts=5).post("/evaluations/", json_body={"name": "r"}) + + assert len(attempts) == 1 + assert no_sleep == [] + + +def test_a_refusal_is_replayed_even_for_a_create(no_sleep): + """429 is refused before the server does any work, so there is nothing on + the other side to duplicate.""" + responses = [httpx.Response(429), httpx.Response(201, json={"evaluation_id": "e1"})] + + client = client_for(lambda request: responses.pop(0)) + + assert client.post("/evaluations/", json_body={"name": "r"})["evaluation_id"] == "e1" + + +def test_a_connection_that_was_never_made_is_replayed_for_a_create(no_sleep): + """Nothing reached the server, so replaying cannot duplicate anything.""" + calls = [] + + def handler(request): + calls.append(request) + if len(calls) == 1: + raise httpx.ConnectError("refused", request=request) + return httpx.Response(201, json={"evaluation_id": "e1"}) + + assert client_for(handler).post("/evaluations/", json_body={"name": "r"}) + + +def test_a_read_timeout_does_not_replay_a_create(no_sleep): + """The bytes went out; the platform may have processed them.""" + calls = [] + + def handler(request): + calls.append(request) + raise httpx.ReadTimeout("slow", request=request) + + with pytest.raises(TransportError): + client_for(handler, max_attempts=5).post("/evaluations/", json_body={"name": "r"}) + + assert len(calls) == 1 + + +def test_a_post_declared_idempotent_still_retries(no_sleep): + """Get-or-create and terminal-state writes are safe to replay.""" + responses = [httpx.Response(502), httpx.Response(200, json={"data": {"id": "env-1"}})] + + client = client_for(lambda request: responses.pop(0)) + + assert client.post("/environmentshub/resolve", json_body={}, idempotent=True) + + +def test_idempotent_methods_still_replay_ambiguous_failures(no_sleep): + responses = [httpx.Response(504), httpx.Response(200, json={"ok": True})] + + client = client_for(lambda request: responses.pop(0)) + + assert client.put("/evaluations/x", json_body={"metrics": {}}) == {"ok": True} diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index a9e8cb80c..8e2b01054 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -340,3 +340,65 @@ def test_dropped_records_are_reported_on_the_handle(): assert run.dropped_records == 3 run.finish() + + +def test_an_upload_failure_reaches_the_caller_in_raise_mode(): + """A sink fails on the uploader thread, where raising reaches nobody. Under + on_error="raise" the failure has to surface where a test is looking.""" + run = make_run(sinks=[FakeSink("broken", fail_on_write=True)], on_error="raise") + + run.log_traces([{"id": "t1"}]) + with pytest.raises(RuntimeError, match="sink is broken"): + run.flush() + + run.finish() + + +def test_an_upload_failure_surfaces_from_finish_too(): + run = make_run(sinks=[FakeSink("broken", fail_on_write=True)], on_error="raise") + run.log_traces([{"id": "t1"}]) + + with pytest.raises(RuntimeError, match="sink is broken"): + run.finish() + + # Still closed out: the failure is reported after teardown, not instead of it. + assert run.finished + + +def test_an_upload_failure_is_reported_once(): + run = make_run(sinks=[FakeSink("broken", fail_on_write=True)], on_error="raise") + run.log_traces([{"id": "t1"}]) + + with pytest.raises(RuntimeError): + run.flush() + run.flush() # nothing left to raise + + run.finish() + + +def test_signal_handlers_are_restored_when_the_run_finishes(): + """`self._handle_signal` builds a new bound method on every access, so an + identity check against a fresh one never matches — leaving the handler + installed, pinning the finished run, and blocking the next run in the + process from installing its own.""" + original = signal.getsignal(signal.SIGTERM) + run = make_run() + run.install_signal_handlers() + assert signal.getsignal(signal.SIGTERM) is run._signal_handler + + run.finish() + + assert signal.getsignal(signal.SIGTERM) is original + + +def test_a_later_run_can_install_its_own_handlers(): + first = make_run() + first.install_signal_handlers() + first.finish() + + second = make_run() + second.install_signal_handlers() + + assert signal.getsignal(signal.SIGTERM) is second._signal_handler + second.finish() + assert signal.getsignal(signal.SIGTERM) is signal.SIG_DFL From fb831d24cb1fdd4416189746786ecdfe3dabeda7 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 14:22:17 -0700 Subject: [PATCH 04/39] fix(runs): unbuffer offline records, forgive transient sink failures MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two follow-up findings, both on the previous rounds' fixes. Offline records were buffered. reset_after_fork() dropped the inherited file handles, but on CPython the last reference going away closes them, and close() flushes — writing out the child's copy of the parent's buffer and duplicating every record still sitting in it. Dropping without closing is not expressible for a buffered writer, so the buffer is gone instead: records are written to an unbuffered append-mode handle, encoded here. Nothing is ever held in process memory, so a fork has nothing to copy and flush() has nothing to do. The earlier fork test passed on this path by luck. It forked before the uploader thread had opened the file, so no handle and no buffer existed in the child. Replaced with an assertion that records are readable through a separate handle with no flush and no close, which is what actually pins the property. Transient sample failures retired the sink. Making sample POSTs non-replayable (332fc216) was right on its own, but combined with the worker disabling a sink on any raise it meant a single 502 stopped every later batch — so one gateway blip could leave the rest of a run missing from the dashboard of exactly the accounts the v0 sample table exists to serve. That trade is worse than the duplicates it avoids. The worker now separates "this batch failed" from "this sink is finished". A permanent failure (gated account, rejected credential) will fail identically forever and still retires the sink immediately. A transient one gets three consecutive strikes, reset by any success, so a blip costs one batch and a sustained outage still stops the SDK re-attempting for hours. Dropped records are counted either way, so run.dropped_records reflects the loss. This applies to the traces sink too, which had the same all-or-nothing behaviour. 138 tests (up from 133). Co-Authored-By: Claude Opus 5 (1M context) --- .../prime-runs/src/prime_runs/exceptions.py | 21 +++++ .../src/prime_runs/sinks/offline.py | 41 ++++++--- packages/prime-runs/src/prime_runs/worker.py | 79 ++++++++++++---- packages/prime-runs/tests/test_init.py | 24 +++++ packages/prime-runs/tests/test_worker.py | 91 +++++++++++++++++++ 5 files changed, 228 insertions(+), 28 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/exceptions.py b/packages/prime-runs/src/prime_runs/exceptions.py index 412a85da4..e417e8c1c 100644 --- a/packages/prime-runs/src/prime_runs/exceptions.py +++ b/packages/prime-runs/src/prime_runs/exceptions.py @@ -67,6 +67,27 @@ class TransportError(RunAPIError): """The request failed below HTTP — connection refused, TLS failure, timeout.""" +def is_transient(exc: BaseException) -> bool: + """Whether a failure is about this moment rather than this run. + + The distinction decides whether a sink is retired. A gated account or a bad + credential will fail identically on every future batch, so the sink should + stop. A gateway blip or a dropped connection will not, and retiring a sink + for one of those means a single 502 empties the rest of the run's dashboard. + + Covers the traces service's exception family as well as this package's, + since both reach the uploader through the same path. + """ + if isinstance(exc, (RetryableAPIError, TransportError)): + return True + try: + from prime_traces.exceptions import RetryableAPIError as TracesRetryable + from prime_traces.exceptions import TransportError as TracesTransport + except ImportError: # pragma: no cover - dependency is declared + return False + return isinstance(exc, (TracesRetryable, TracesTransport)) + + class EnvironmentResolutionError(PrimeRunsError): """An environment named in ``init()`` could not be resolved to a hub ID. diff --git a/packages/prime-runs/src/prime_runs/sinks/offline.py b/packages/prime-runs/src/prime_runs/sinks/offline.py index e7d35f9d3..ebe0084dc 100644 --- a/packages/prime-runs/src/prime_runs/sinks/offline.py +++ b/packages/prime-runs/src/prime_runs/sinks/offline.py @@ -10,7 +10,7 @@ import json import logging from pathlib import Path -from typing import Any, Mapping, Optional, Sequence, TextIO, Union +from typing import Any, BinaryIO, Mapping, Optional, Sequence, Union from .. import _fork from .base import Sink, to_mapping @@ -29,16 +29,18 @@ def __init__(self, directory: Union[str, Path], *, stamp_run: bool = True) -> No self._stamp_run = stamp_run self._run_id: Optional[str] = None self._run_kind: Optional[str] = None - self._handles: dict[str, TextIO] = {} + self._handles: dict[str, BinaryIO] = {} self.records_written = 0 _fork.register(self) def reset_after_fork(self) -> None: """Abandon inherited file handles; ``_handle`` reopens on next write. - Dropped without flushing or closing: the inherited buffer holds records - the parent has not written yet and will write itself, so flushing it - here would put every one of them in the file twice. + Dropping a buffered file object would not be enough on its own: on + CPython the last reference going away closes it, and ``close()`` + *flushes* — writing out the parent's copied buffer and duplicating every + record in it. The handles are unbuffered (see ``_handle``) precisely so + that there is never anything in that buffer to duplicate. """ self._handles = {} @@ -69,22 +71,28 @@ def write( if self._run_kind: run["type"] = self._run_kind mapping["run"] = run - handle.write( - json.dumps(mapping, ensure_ascii=False, separators=(",", ":"), default=str) + "\n" - ) + line = json.dumps(mapping, ensure_ascii=False, separators=(",", ":"), default=str) + _write_all(handle, (line + "\n").encode("utf-8")) self.records_written += 1 - def _handle(self, name: str) -> TextIO: + def _handle(self, name: str) -> BinaryIO: + """An unbuffered append-mode handle for one line format. + + Unbuffered on purpose. A buffered writer keeps records in process memory + until it decides to flush, and a fork copies that buffer — after which + both processes eventually write it, putting every record in the file + twice. Writing straight through means the only copy of a record lives in + the file, and ``O_APPEND`` keeps concurrent writers from interleaving. + """ handle = self._handles.get(name) if handle is None: self._records_dir.mkdir(parents=True, exist_ok=True) - handle = (self._records_dir / f"{name}.jsonl").open("a", encoding="utf-8") + handle = open(self._records_dir / f"{name}.jsonl", "ab", buffering=0) self._handles[name] = handle return handle def flush(self) -> None: - for handle in self._handles.values(): - handle.flush() + """Nothing is held back — every write already went to the file.""" def close(self) -> None: for handle in self._handles.values(): @@ -95,6 +103,15 @@ def close(self) -> None: self._handles.clear() +def _write_all(handle: BinaryIO, data: bytes) -> None: + """Write every byte. A raw handle may report a short write.""" + while data: + written = handle.write(data) + if not written: # pragma: no cover - only on a non-blocking handle + raise OSError("offline record write made no progress") + data = data[written:] + + def _infer_format(record: Any) -> str: if isinstance(record, Mapping): return "episode" if "traces" in record else "trace" diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index 022ac3464..fa9c6f5b4 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -28,11 +28,16 @@ from typing import Any, Callable, List, Optional, Sequence from . import _fork +from .exceptions import is_transient logger = logging.getLogger(__name__) DEFAULT_QUEUE_SIZE = 256 DEFAULT_PUT_TIMEOUT = 5.0 +#: Consecutive transient failures before a sink is retired. One gateway blip +#: must not empty the rest of a run's dashboard; a sustained outage should still +#: stop the SDK from re-attempting every batch for hours. +TRANSIENT_FAILURE_LIMIT = 3 @dataclass @@ -84,6 +89,7 @@ def __init__( self._stopping = threading.Event() self._lock = threading.Lock() self.dropped = 0 + self._transient_failures: dict = {} self._pid = os.getpid() _fork.register(self) @@ -125,7 +131,9 @@ def _dispatch(self, item: WriteItem) -> None: try: sink.write(item.records, line_format=item.line_format, step=item.step) except Exception as exc: # noqa: BLE001 - one sink failing must not stop the others - self._fail_sink(sink, exc) + self._fail_sink(sink, exc, dropped=len(item.records)) + else: + self._transient_failures.pop(getattr(sink, "name", id(sink)), None) def _write_metrics(self, item: MetricItem) -> None: if self._metric_writer is None: @@ -151,23 +159,62 @@ def _flush_sinks(self) -> None: except Exception as exc: # noqa: BLE001 self._fail_sink(sink, exc) - def _fail_sink(self, sink: Any, exc: Exception) -> None: - """Disable a sink that raised, and report it exactly once. - - Not retried here: the transports already retry internally (traces on - content-addressed uploads, the platform client on 429/5xx), so an error - that reaches this point has already exhausted its budget. Continuing to - call a sink in that state produces one log line per batch for the rest - of the run and hides whatever failed first. + def _fail_sink(self, sink: Any, exc: Exception, *, dropped: int = 0) -> None: + """Handle a sink that raised, and report it. + + The batch is gone either way — the transports already retried internally + (traces on content-addressed uploads, the platform client on whatever it + can safely replay), so an error reaching this point has exhausted its + budget. What is decided here is whether the *sink* is finished: + + - A permanent failure — a gated account, a rejected credential — will + fail identically on every future batch, so the sink stops. Continuing + would produce one log line per batch for the rest of the run and bury + whatever failed first. + - A transient one gets ``TRANSIENT_FAILURE_LIMIT`` consecutive strikes, + reset by any success. Retiring a sink on a single gateway blip would + leave the rest of the run missing from the dashboard, which is a much + larger loss than the one batch that actually failed. """ name = getattr(sink, "name", type(sink).__name__) - sink.enabled = False - logger.warning("Sink %s disabled after an error: %s: %s", name, type(exc).__name__, exc) - if self._on_error is not None: - try: - self._on_error(name, exc) - except Exception: # noqa: BLE001 - the handler is the caller's problem - logger.debug("Error handler raised while reporting a sink failure", exc_info=True) + self.dropped += dropped + + if is_transient(exc): + strikes = self._transient_failures.get(name, 0) + 1 + self._transient_failures[name] = strikes + if strikes < TRANSIENT_FAILURE_LIMIT: + logger.warning( + "Sink %s dropped a batch of %d record(s) (%s: %s); " + "strike %d of %d, still enabled.", + name, + dropped, + type(exc).__name__, + exc, + strikes, + TRANSIENT_FAILURE_LIMIT, + ) + self._notify(name, exc) + return + sink.enabled = False + logger.warning( + "Sink %s disabled after %d consecutive transient failures: %s: %s", + name, + strikes, + type(exc).__name__, + exc, + ) + else: + sink.enabled = False + logger.warning("Sink %s disabled after an error: %s: %s", name, type(exc).__name__, exc) + self._notify(name, exc) + + def _notify(self, name: str, exc: Exception) -> None: + if self._on_error is None: + return + try: + self._on_error(name, exc) + except Exception: # noqa: BLE001 - the handler is the caller's problem + logger.debug("Error handler raised while reporting a sink failure", exc_info=True) # ------------------------------------------------------------------ queue diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index 0af0d074c..ac65f3d38 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -325,3 +325,27 @@ def test_a_forked_child_joins_the_run_without_duplicating_the_parents_records(tm ids = [json.loads(line)["id"] for line in lines] assert sorted(ids) == sorted([f"parent-{n}" for n in range(5)] + ["child-1"]) + + +def test_offline_records_are_on_disk_before_any_flush(tmp_path): + """Nothing may sit in a process-local write buffer. + + A buffered writer holds records in memory until it decides to flush, and a + fork copies that buffer — after which both processes eventually write it and + every buffered record lands in the file twice. Reading the file back through + a separate handle, with no flush and no close, is what proves the buffer is + not there to be copied. + """ + run = pr.init(mode="offline", dir=str(tmp_path)) + run.log_traces([{"id": "t1"}]) + run.flush() + + path = tmp_path / run.id / "records" / "trace.jsonl" + assert path.read_text().count('"t1"') == 1 + + run.log_traces([{"id": "t2"}]) + run.flush() + assert path.read_text().count('"t2"') == 1 + + run.finish() + assert len(path.read_text().splitlines()) == 2 diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index 83733ffc9..5b4c54e31 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -233,3 +233,94 @@ def reset_after_fork(self) -> None: _fork._reset_all() assert fine.reset == 1 + + +def test_a_transient_failure_drops_the_batch_but_keeps_the_sink(): + """One gateway blip must not empty the rest of the run's dashboard. The + batch is already lost; retiring the sink loses every batch after it too.""" + from prime_runs.exceptions import RetryableAPIError + + class BlipSink(FakeSink): + def write(self, records, *, line_format=None, step=None) -> None: + raise RetryableAPIError("bad gateway", status_code=502) + + sink = BlipSink("blippy") + worker = UploadWorker([sink]) + + worker.submit(WriteItem(records=[{"id": 1}, {"id": 2}])) + drain(worker) + + assert sink.enabled is True + assert worker.dropped == 2 + worker.close() + + +def test_a_sustained_outage_eventually_retires_the_sink(): + """A blip is forgiven; hours of re-attempting every batch is not useful.""" + from prime_runs.exceptions import TransportError + from prime_runs.worker import TRANSIENT_FAILURE_LIMIT + + class DeadSink(FakeSink): + def write(self, records, *, line_format=None, step=None) -> None: + raise TransportError("connection refused") + + sink = DeadSink("dead") + worker = UploadWorker([sink]) + + for _ in range(TRANSIENT_FAILURE_LIMIT): + worker.submit(WriteItem(records=[{"id": 1}])) + drain(worker) + + assert sink.enabled is False + worker.close() + + +def test_a_success_forgives_earlier_blips(): + """Strikes are consecutive: an intermittent gateway must never accumulate + its way to a retirement across an otherwise healthy run.""" + from prime_runs.exceptions import RetryableAPIError + from prime_runs.worker import TRANSIENT_FAILURE_LIMIT + + class FlakySink(FakeSink): + """Fails every other batch, forever.""" + + def __init__(self) -> None: + super().__init__("flaky") + self.calls = 0 + + def write(self, records, *, line_format=None, step=None) -> None: + self.calls += 1 + if self.calls % 2 == 1: + raise RetryableAPIError("bad gateway", status_code=502) + super().write(records, line_format=line_format, step=step) + + sink = FlakySink() + worker = UploadWorker([sink]) + + # Far more failures than the limit, but never two in a row. + for _ in range(TRANSIENT_FAILURE_LIMIT * 4): + worker.submit(WriteItem(records=[{"id": 1}])) + drain(worker) + + assert sink.calls == TRANSIENT_FAILURE_LIMIT * 4 + assert sink.enabled is True + assert len(sink.batches) == TRANSIENT_FAILURE_LIMIT * 2 + worker.close() + + +def test_a_permanent_failure_retires_the_sink_immediately(): + """A gated account or a rejected credential fails identically forever.""" + from prime_runs.exceptions import UnauthorizedError + + class DeniedSink(FakeSink): + def write(self, records, *, line_format=None, step=None) -> None: + raise UnauthorizedError("nope", status_code=401) + + sink = DeniedSink("denied") + worker = UploadWorker([sink]) + + worker.submit(WriteItem(records=[{"id": 1}])) + drain(worker) + + assert sink.enabled is False + worker.close() From 6282c24e83a95744eee2c3ca7e588242dd40ebbf Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 14:31:11 -0700 Subject: [PATCH 05/39] fix(runs): keep queue drops and per-sink write failures separate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous commit folded failed sink writes into `dropped`, which is documented and warned about as queue backpressure. Two things went wrong with that. A default online run writes to both the traces sink and the sample table, so one failed batch was counted twice. And a failure on one sink was counted at all even when the other sink stored the records, so `finish()` could warn about data missing from a run that has all of it. They are different losses and stay separate now. `dropped` counts records that reached no sink because the queue was full — the producer outran the uploader, and those records are stored nowhere. `failed_records` counts per sink, exposed as a mapping rather than a total, because summing it would recreate exactly the overstatement above. The finish warnings are phrased to match: one about records that reached nothing, one per sink about what that sink could not store. 139 tests. Co-Authored-By: Claude Opus 5 (1M context) --- packages/prime-runs/src/prime_runs/run.py | 29 ++++++++++++++++++-- packages/prime-runs/src/prime_runs/worker.py | 10 ++++++- packages/prime-runs/tests/test_worker.py | 27 +++++++++++++++++- 3 files changed, 62 insertions(+), 4 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 3d867b399..2a1b34a70 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -174,9 +174,25 @@ def finished(self) -> bool: @property def dropped_records(self) -> int: - """Records the uploader could not keep up with. Should be zero.""" + """Records that reached no sink because the queue was full. + + Backpressure only: the producer durably outran the uploader. A record + counted here was stored nowhere. Contrast ``failed_records``, which is + per-sink and usually means the record is still safe in another sink. + """ return self._worker.dropped + @property + def failed_records(self) -> Dict[str, int]: + """Records each sink could not store, by sink name. + + Not summed into one number and not merged into ``dropped_records``: with + traces and the sample table both enabled, the same batch failing on one + sink says nothing about whether the other stored it, so a single total + would report data missing that is not actually gone. + """ + return dict(self._worker.failed_records) + def __repr__(self) -> str: return ( f" None: larger loss than the one batch that actually failed. """ name = getattr(sink, "name", type(sink).__name__) - self.dropped += dropped + if dropped: + self.failed_records[name] = self.failed_records.get(name, 0) + dropped if is_transient(exc): strikes = self._transient_failures.get(name, 0) + 1 diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index 5b4c54e31..9477b1eec 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -251,7 +251,10 @@ def write(self, records, *, line_format=None, step=None) -> None: drain(worker) assert sink.enabled is True - assert worker.dropped == 2 + # Per sink, and not folded into `dropped`: the queue accepted these records + # fine, and with both sinks enabled the other one may well have stored them. + assert worker.failed_records == {"blippy": 2} + assert worker.dropped == 0 worker.close() @@ -324,3 +327,25 @@ def write(self, records, *, line_format=None, step=None) -> None: assert sink.enabled is False worker.close() + + +def test_a_failed_batch_is_counted_once_per_sink_not_once_per_run(): + """Default online runs write to two sinks. Adding both to one total would + report twice the loss, and would report loss at all when the other sink + stored the records.""" + from prime_runs.exceptions import RetryableAPIError + + class BlipSink(FakeSink): + def write(self, records, *, line_format=None, step=None) -> None: + raise RetryableAPIError("bad gateway", status_code=502) + + broken, healthy = BlipSink("broken"), FakeSink("healthy") + worker = UploadWorker([broken, healthy]) + + worker.submit(WriteItem(records=[{"id": 1}, {"id": 2}])) + drain(worker) + + assert worker.failed_records == {"broken": 2} + assert worker.dropped == 0 + assert healthy.batches, "the healthy sink stored them" + worker.close() From e0d0a4aafad3b2f9a88616bfcdb25c4bc005f3e0 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 15:24:59 -0700 Subject: [PATCH 06/39] chore(runs): align prime-runs with the other SDK packages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five points where prime-runs diverged from prime-traces, prime-evals, prime-sandboxes and prime-tunnel without meaning to: - Normalize an explicitly passed `base_url`. `Config` strips a trailing `/api/v1`; the `PlatformClient` constructor did not, so `pr.init( base_url=".../api/v1")` requested `/api/v1/api/v1` while the identical value in `PRIME_API_BASE_URL` worked. prime-traces carries the same helper for the same reason. - Map 403 to a typed `ForbiddenError`, matching `prime_traces.ForbiddenError` — which the traces sink already branches on for beta gating, while a 403 from the platform API collapsed into a generic `RunAPIError`. Behavior is unchanged (it was already classified permanent); callers can now branch. - Drop `pydantic` and `tenacity`. Neither was imported: this package models no response bodies and hand-rolls its retry loop. They do not belong in the dependency tree of a leaf package that lands inside verifiers. - Add the LICENSE file the other packages ship. - Declare 3.13, which CI has been testing all along. Also documents the three departures that are deliberate — private client instead of `core/`, dataclasses instead of pydantic models, no async client — and what to do about the three blocking calls when driving a run from async code, since every sibling SDK ships an async client and this one does not. Co-Authored-By: Claude Opus 5 (1M context) --- packages/prime-runs/LICENSE | 21 +++++++++ packages/prime-runs/README.md | 47 +++++++++++++++++++ packages/prime-runs/pyproject.toml | 8 +++- .../prime-runs/src/prime_runs/__init__.py | 2 + packages/prime-runs/src/prime_runs/_http.py | 17 ++++++- .../prime-runs/src/prime_runs/exceptions.py | 11 +++++ packages/prime-runs/tests/test_http.py | 35 +++++++++++++- uv.lock | 4 -- 8 files changed, 136 insertions(+), 9 deletions(-) create mode 100644 packages/prime-runs/LICENSE diff --git a/packages/prime-runs/LICENSE b/packages/prime-runs/LICENSE new file mode 100644 index 000000000..44a4d7b13 --- /dev/null +++ b/packages/prime-runs/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Prime Intellect + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index 66629359c..c465572d4 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -101,6 +101,26 @@ Set the mode explicitly, or through `$PRIME_RUNS_MODE`. - **Knows about ranks.** Rank 0 owns creation and finalization; other ranks join through `PRIME_RUN_ID` and upload their own records. +## Using it from async code + +There is no `AsyncRun`, unlike the async clients in `prime-traces`, +`prime-evals` and `prime-sandboxes`. The uploader thread is what replaces it: +`log()`, `log_traces()` and `update_config()` are queue puts, not requests, so +calling them straight from a coroutine does no network I/O on the event loop. + +Three calls do block, and all three are worth knowing about: + +| call | blocks on | when it matters | +| --- | --- | --- | +| `init()` | create + environment resolution | once, at startup | +| `finish()` | draining the queue, then finalize | once, at shutdown | +| `log_traces()` | up to `put_timeout` (5s) **only if the queue is full** | a producer durably outrunning the uploader | + +The first two are run boundaries — `await asyncio.to_thread(run.finish)` if a +stall there would matter. The third is the one to watch in a hot rollout loop: +the block is the backpressure, and past it the batch is dropped and counted in +`run.dropped_records`. Raise `queue_size=` before reaching for a thread. + ## Configuration Resolved from environment variables first, then `~/.prime/config.json`: @@ -148,6 +168,33 @@ them structurally, and no producer package is imported. This is a leaf package b design — the `prime` CLI depends on `verifiers`, so verifiers can never depend on `prime`. +## How this differs from the other prime SDKs + +`prime-sandboxes`, `prime-traces`, `prime-evals` and `prime-tunnel` are all built +the same way: a `core/` subpackage holding an `APIClient` and a `Config`, pydantic +models for the responses, and a sync/async client pair as the thing you import. +Three deliberate departures here, so the difference reads as a choice rather than +an oversight: + +- **The client is private.** `init()` is the surface, not a client object, so the + HTTP layer lives in `_http.py` rather than `core/client.py` and `PlatformClient` + is not exported. Config still is, and is the same class as everywhere else — + `~/.prime/config.json`, env wins, same variable names. +- **Responses are not modeled.** The platform returns more fields than any + producer reads; freezing them in pydantic would make every backend addition a + breaking SDK release. Backends take the two or three fields they need and + return a `RunHandle`. That is why the local types are dataclasses and why + pydantic is not a dependency. +- **No async client.** See [Using it from async code](#using-it-from-async-code) + — the background uploader covers the case an async client would exist for. + +## Related packages + +- [prime-traces](../prime-traces) — the traces service client this SDK streams + through, and the direct API for querying or exporting what a run produced. +- [prime](../prime) — the CLI and full SDK. Depends on this package's consumers, + never the other way around. + ## Status Eval runs are supported. Training runs (`kind="train"`, over diff --git a/packages/prime-runs/pyproject.toml b/packages/prime-runs/pyproject.toml index 058d90bdc..165788c58 100644 --- a/packages/prime-runs/pyproject.toml +++ b/packages/prime-runs/pyproject.toml @@ -13,10 +13,13 @@ authors = [ # verifiers can never depend on `prime` — and verifiers is a first-class # consumer of this SDK. Nothing here may pull in prime, verifiers, typer, # rich or textual, directly or transitively. +# +# Kept to what is actually imported. The other SDKs carry `pydantic` (they model +# response bodies) and `tenacity` (they retry through it); this package models +# nothing and hand-rolls its retry loop in `_http.py`, so neither belongs in the +# dependency tree of a package that lands inside verifiers and prime-rl. dependencies = [ "httpx>=0.25.0", - "pydantic>=2.0.0", - "tenacity>=9.1.2", "prime-traces>=0.0.2", ] keywords = ["evals", "evaluations", "training", "runs", "traces", "observability"] @@ -29,6 +32,7 @@ classifiers = [ "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Topic :: Software Development :: Libraries :: Python Modules", "Topic :: Scientific/Engineering :: Artificial Intelligence" ] diff --git a/packages/prime-runs/src/prime_runs/__init__.py b/packages/prime-runs/src/prime_runs/__init__.py index d295c9194..3d0b00830 100644 --- a/packages/prime-runs/src/prime_runs/__init__.py +++ b/packages/prime-runs/src/prime_runs/__init__.py @@ -30,6 +30,7 @@ from .exceptions import ( ConfigurationError, EnvironmentResolutionError, + ForbiddenError, NotFoundError, PaymentRequiredError, PrimeRunsError, @@ -88,6 +89,7 @@ "EnvironmentResolutionError", "RunAPIError", "RunFinishedError", + "ForbiddenError", "NotFoundError", "PaymentRequiredError", "RetryableAPIError", diff --git a/packages/prime-runs/src/prime_runs/_http.py b/packages/prime-runs/src/prime_runs/_http.py index 027031e3a..8db2c86ba 100644 --- a/packages/prime-runs/src/prime_runs/_http.py +++ b/packages/prime-runs/src/prime_runs/_http.py @@ -31,6 +31,7 @@ from . import _fork from .exceptions import ( + ForbiddenError, NotFoundError, PaymentRequiredError, RetryableAPIError, @@ -78,6 +79,18 @@ def _parse_retry_after(response: httpx.Response) -> Optional[float]: return None +def normalize_base_url(url: str) -> str: + """The same normalization ``Config`` applies to file/env URLs. + + ``PlatformClient`` appends ``/api/v1`` itself, and platform URLs are + commonly written with the suffix already on them. Without stripping it here + an explicit ``base_url=`` would request ``/api/v1/api/v1/...`` while the + identical value read from ``PRIME_API_BASE_URL`` worked — the config path + strips it and the constructor path did not. + """ + return url.rstrip("/").removesuffix("/api/v1") + + def encode_json(value: Any) -> bytes: """Compact UTF-8 JSON, matching the encoding used to size batches. @@ -102,7 +115,7 @@ def __init__( max_attempts: int = DEFAULT_MAX_ATTEMPTS, client: Optional[httpx.Client] = None, ) -> None: - self.base_url = base_url.rstrip("/") + self.base_url = normalize_base_url(base_url) self.api_prefix = f"{self.base_url}/api/v1" self.max_attempts = max(1, max_attempts) self._owns_client = client is None @@ -276,6 +289,8 @@ def _map_error(response: httpx.Response) -> RunAPIError: ) if status == 402: return PaymentRequiredError(message, status_code=status, code=code) + if status == 403: + return ForbiddenError(message, status_code=status, code=code) if status == 404: return NotFoundError(message, status_code=status, code=code) return RunAPIError(message, status_code=status, code=code) diff --git a/packages/prime-runs/src/prime_runs/exceptions.py b/packages/prime-runs/src/prime_runs/exceptions.py index e417e8c1c..1d89fe334 100644 --- a/packages/prime-runs/src/prime_runs/exceptions.py +++ b/packages/prime-runs/src/prime_runs/exceptions.py @@ -44,6 +44,17 @@ class PaymentRequiredError(RunAPIError): """402 — payment required. Check billing status.""" +class ForbiddenError(RunAPIError): + """403 — authenticated, but not allowed to do this. + + Distinct from 401 because the credential is fine and re-authenticating will + not help: the run belongs to another owner, the team header names a team the + key cannot act for, or the feature is gated to an allowlist. Named to match + ``prime_traces.ForbiddenError``, which the traces sink already branches on + to retire itself when an account is outside the closed beta. + """ + + class NotFoundError(RunAPIError): """404 — the run, environment or evaluation does not exist for this owner.""" diff --git a/packages/prime-runs/tests/test_http.py b/packages/prime-runs/tests/test_http.py index 399576324..d230aadd9 100644 --- a/packages/prime-runs/tests/test_http.py +++ b/packages/prime-runs/tests/test_http.py @@ -5,19 +5,21 @@ from prime_runs._http import PlatformClient, encode_json, retry_delay from prime_runs.exceptions import ( + ForbiddenError, NotFoundError, PaymentRequiredError, RetryableAPIError, RunAPIError, TransportError, UnauthorizedError, + is_transient, ) -def client_for(handler, **kwargs) -> PlatformClient: +def client_for(handler, *, base_url: str = "http://testserver", **kwargs) -> PlatformClient: return PlatformClient( api_key="test-key", - base_url="http://testserver", + base_url=base_url, client=httpx.Client(transport=httpx.MockTransport(handler)), **kwargs, ) @@ -28,6 +30,7 @@ def client_for(handler, **kwargs) -> PlatformClient: [ (401, UnauthorizedError), (402, PaymentRequiredError), + (403, ForbiddenError), (404, NotFoundError), (400, RunAPIError), (422, RunAPIError), @@ -43,6 +46,17 @@ def test_status_codes_map_to_types_callers_can_branch_on(status, expected): assert "nope" in str(caught.value) +def test_a_forbidden_response_is_permanent_so_a_sink_retires_on_it(): + """403 is the gated-account signal. Retrying it for the rest of a run would + log one failure per batch and never succeed.""" + client = client_for(lambda request: httpx.Response(403, json={"code": "service_not_enabled"})) + + with pytest.raises(ForbiddenError) as caught: + client.get("/evaluations/x") + + assert not is_transient(caught.value) + + def test_an_unauthorized_error_says_what_to_do_about_it(): client = client_for(lambda request: httpx.Response(401, json={"detail": "bad token"})) @@ -172,6 +186,23 @@ def test_a_post_declared_idempotent_still_retries(no_sleep): assert client.post("/environmentshub/resolve", json_body={}, idempotent=True) +@pytest.mark.parametrize( + "given", + [ + "http://testserver", + "http://testserver/", + "http://testserver/api/v1", + "http://testserver/api/v1/", + ], +) +def test_a_base_url_written_with_the_api_prefix_is_not_doubled(given): + """``Config`` strips the suffix, so an explicit ``base_url=`` that does not + would 404 on exactly the value that works through the environment.""" + assert client_for(lambda request: httpx.Response(200, json={}), base_url=given).api_prefix == ( + "http://testserver/api/v1" + ) + + def test_idempotent_methods_still_replay_ambiguous_failures(no_sleep): responses = [httpx.Response(504), httpx.Response(200, json={"ok": True})] diff --git a/uv.lock b/uv.lock index 6092621c5..e486c14a2 100644 --- a/uv.lock +++ b/uv.lock @@ -1453,8 +1453,6 @@ source = { editable = "packages/prime-runs" } dependencies = [ { name = "httpx" }, { name = "prime-traces" }, - { name = "pydantic" }, - { name = "tenacity" }, ] [package.optional-dependencies] @@ -1468,11 +1466,9 @@ dev = [ requires-dist = [ { name = "httpx", specifier = ">=0.25.0" }, { name = "prime-traces", editable = "packages/prime-traces" }, - { name = "pydantic", specifier = ">=2.0.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.23.0" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.13.1" }, - { name = "tenacity", specifier = ">=9.1.2" }, ] provides-extras = ["dev"] From 22c43bb30c28d23824a14d7e417a1db00a89d8c6 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 15:38:37 -0700 Subject: [PATCH 07/39] fix(runs): address SDK lifecycle and metric logging issues --- .../src/prime_runs/backends/evals.py | 5 + packages/prime-runs/src/prime_runs/run.py | 93 ++++++++++++++----- packages/prime-runs/src/prime_runs/worker.py | 37 +++++++- .../prime-runs/tests/test_evals_backend.py | 29 +++++- packages/prime-runs/tests/test_run.py | 47 ++++++++++ packages/prime-runs/tests/test_worker.py | 15 ++- 6 files changed, 202 insertions(+), 24 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/backends/evals.py b/packages/prime-runs/src/prime_runs/backends/evals.py index a2bdaa585..043e21f9c 100644 --- a/packages/prime-runs/src/prime_runs/backends/evals.py +++ b/packages/prime-runs/src/prime_runs/backends/evals.py @@ -25,6 +25,7 @@ EnvironmentResolutionError, NotFoundError, RunAPIError, + is_transient, ) from ..models import EnvironmentRef, RunHandle, RunSpec, RunStatus @@ -106,6 +107,10 @@ def attach(self, run_id: str) -> RunHandle: # Attach is a convenience — a resume or a non-primary rank joining. # Losing the run's name to a transient read is not worth failing on; # the ID is what everything downstream actually needs. + if not is_transient(exc) and not ( + exc.status_code is not None and exc.status_code >= 500 + ): + raise logger.debug("Could not read evaluation %s on attach: %s", run_id, exc) return RunHandle(id=run_id, url=self.url_for(run_id)) return RunHandle( diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 2a1b34a70..a0ed10be3 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -30,7 +30,7 @@ from .exceptions import ConfigurationError, RunFinishedError from .models import EnvironmentRef, Mode, OnError, RunHandle, RunKind, RunSpec, RunStatus from .sinks import EvalSamplesSink, OfflineSink, Sink, TracesSink -from .worker import MetricItem, UploadWorker, WriteItem +from .worker import MetricItem, RunUpdateItem, UploadWorker, WriteItem logger = logging.getLogger(__name__) @@ -104,6 +104,8 @@ def __init__( self._last_summary_flush = time.monotonic() self._summary_dirty = False self._config_dirty = False + self._pending_metrics: Dict[str, Any] = {} + self._pending_metric_step: Optional[int] = None self._finish_lock = threading.RLock() self._finished = False @@ -115,6 +117,7 @@ def __init__( sinks, on_error=self._record_sink_error, metric_writer=self._write_metrics if backend.supports_step_metrics else None, + update_writer=self._write_run_update, **worker_kwargs, ) context = _sink_context(spec, handle) @@ -227,9 +230,17 @@ def log( self.summary.update(cleaned) self._summary_dirty = True if not commit: + if self._backend.supports_step_metrics: + self._pending_metrics.update(cleaned) + if step is not None: + self._pending_metric_step = step return if self._backend.supports_step_metrics: - self._worker.submit(MetricItem(metrics=cleaned, step=step)) + committed = {**self._pending_metrics, **cleaned} + committed_step = step if step is not None else self._pending_metric_step + self._pending_metrics.clear() + self._pending_metric_step = None + self._worker.submit(MetricItem(metrics=committed, step=committed_step)) return self._maybe_flush_summary() @@ -284,13 +295,15 @@ def finish( with self._finish_lock: if self._finished: return + resolved = RunStatus(status) if not isinstance(status, RunStatus) else status self._finished = True - resolved = RunStatus(status) if not isinstance(status, RunStatus) else status if summary: self.summary.update(_clean_metrics(summary)) self._status = resolved + finish_error: Optional[BaseException] = None + # Order matters: records first, so a dashboard that reacts to the # terminal status never sees a finished run with samples still landing. if not self._worker.flush(timeout=self._finish_timeout): @@ -303,15 +316,16 @@ def finish( self._worker.close(timeout=self._finish_timeout) if self._owns_lifecycle: - self._report_guarded( + finish_error = self._finish_guarded( "updating the run", lambda: self._backend.update( self.id, config=self.config if (self._config_dirty or self.config) else None, summary=self.summary or None, ), + finish_error, ) - self._report_guarded( + finish_error = self._finish_guarded( "finalizing the run", lambda: self._backend.finalize( self.id, @@ -320,8 +334,11 @@ def finish( error=error or (self.errors[0] if self.errors else None), config=self.config or None, ), + finish_error, ) - self._report_guarded("closing the backend", self._backend.close) + finish_error = self._finish_guarded( + "closing the backend", self._backend.close, finish_error + ) atexit.unregister(self._atexit_hook) self._restore_signal_handlers() @@ -342,6 +359,8 @@ def finish( ) # Last, so a run that failed to upload is still closed out properly # before the failure reaches the caller. + if finish_error is not None: + raise finish_error self._raise_deferred() def fail(self, error: Union[str, BaseException]) -> None: @@ -354,8 +373,13 @@ def flush(self, timeout: Optional[float] = 30.0) -> bool: Under ``on_error="raise"`` this is the first place an upload failure can surface, since the failure itself happened on the uploader thread. """ + update_queued = self._queue_run_update() flushed = self._worker.flush(timeout=timeout) - self._flush_summary() + # A periodic update can lose a race for the last queue slot. Once the + # barrier drains that backlog, give the still-dirty snapshot one more + # chance so an explicit flush keeps its persistence guarantee. + if flushed and not update_queued and self._queue_run_update(): + flushed = self._worker.flush(timeout=timeout) self._raise_deferred() return flushed @@ -471,24 +495,40 @@ def _require_live(self, operation: str) -> None: def _write_metrics(self, metrics: Dict[str, Any], step: Optional[int]) -> None: self._backend.log_metrics(self.id, metrics, step) + def _write_run_update( + self, + config: Optional[Dict[str, Any]], + summary: Optional[Dict[str, Any]], + ) -> None: + self._backend.update(self.id, config=config, summary=summary) + def _maybe_flush_summary(self) -> None: now = time.monotonic() if now - self._last_summary_flush < self._summary_flush_seconds: return - self._flush_summary() + self._queue_run_update() - def _flush_summary(self) -> None: + def _queue_run_update(self) -> bool: if not (self._summary_dirty or self._config_dirty) or not self._owns_lifecycle: - return - config = self.config if self._config_dirty else None - summary = self.summary if self._summary_dirty else None - self._last_summary_flush = time.monotonic() - self._summary_dirty = False - self._config_dirty = False - self._report_guarded( - "flushing run metrics", - lambda: self._backend.update(self.id, config=config, summary=summary), + return True + config_dirty = self._config_dirty + summary_dirty = self._summary_dirty + item = RunUpdateItem( + config=dict(self.config) if config_dirty else None, + summary=dict(self.summary) if summary_dirty else None, ) + # Clear before the potentially blocking queue put. If another producer + # thread logs while this one waits, its new dirty bit must survive. + if config_dirty: + self._config_dirty = False + if summary_dirty: + self._summary_dirty = False + if not self._worker.submit(item): + self._config_dirty = self._config_dirty or config_dirty + self._summary_dirty = self._summary_dirty or summary_dirty + return False + self._last_summary_flush = time.monotonic() + return True def _record_sink_error(self, sink_name: str, exc: Exception) -> None: """Called on the uploader thread when a sink gives up.""" @@ -508,11 +548,22 @@ def _raise_deferred(self) -> None: self._deferred_error = None raise exc - def _report_guarded(self, what: str, call: Any) -> None: + def _finish_guarded( + self, + what: str, + call: Any, + first_error: Optional[BaseException], + ) -> Optional[BaseException]: + """Run one teardown step without letting it skip the steps after it.""" try: call() - except Exception as exc: # noqa: BLE001 - routed through the error policy - self._report(what, exc) + except Exception as exc: # noqa: BLE001 - teardown must continue + message = f"{what} failed: {type(exc).__name__}: {exc}" + self.errors.append(message) + if self._on_error == "raise": + return first_error or exc + logger.warning("Run %s: %s", self._handle.id, message) + return first_error def _report(self, what: str, exc: Exception) -> None: message = f"{what} failed: {type(exc).__name__}: {exc}" diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index 09009aef4..10bdb33fe 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -57,6 +57,14 @@ class MetricItem: step: Optional[int] = None +@dataclass +class RunUpdateItem: + """A config/summary snapshot destined for the run lifecycle backend.""" + + config: Optional[dict] = None + summary: Optional[dict] = None + + @dataclass class _Flush: """A barrier the caller waits on.""" @@ -75,6 +83,7 @@ def __init__( put_timeout: float = DEFAULT_PUT_TIMEOUT, on_error: Optional[Callable[[str, Exception], None]] = None, metric_writer: Optional[Callable[[dict, Optional[int]], None]] = None, + update_writer: Optional[Callable[[Optional[dict], Optional[dict]], None]] = None, ) -> None: self.sinks = sinks self.max_queue_size = max_queue_size @@ -84,6 +93,9 @@ def __init__( # same queue as records, so a per-step log() in a training loop costs a # queue put rather than an HTTP round trip. self._metric_writer = metric_writer + # Eval summaries have no time-series endpoint, but they still belong on + # this thread: periodic persistence must never block the producer loop. + self._update_writer = update_writer self._queue: "queue.Queue[Any]" = queue.Queue(maxsize=max_queue_size) self._thread: Optional[threading.Thread] = None self._stopping = threading.Event() @@ -125,6 +137,9 @@ def _run(self) -> None: if isinstance(item, MetricItem): self._write_metrics(item) continue + if isinstance(item, RunUpdateItem): + self._write_update(item) + continue self._dispatch(item) except Exception as exc: # noqa: BLE001 - the thread must outlive one bad batch logger.debug("Uploader iteration failed: %s", exc) @@ -157,6 +172,19 @@ def _write_metrics(self, item: MetricItem) -> None: except Exception: # noqa: BLE001 logger.debug("Error handler raised while reporting metrics", exc_info=True) + def _write_update(self, item: RunUpdateItem) -> None: + if self._update_writer is None: + return + try: + self._update_writer(item.config, item.summary) + except Exception as exc: # noqa: BLE001 - updates must not kill the uploader + logger.warning("Dropped a run metadata update: %s: %s", type(exc).__name__, exc) + if self._on_error is not None: + try: + self._on_error("run metadata", exc) + except Exception: # noqa: BLE001 + logger.debug("Error handler raised while reporting a run update", exc_info=True) + def _flush_sinks(self) -> None: for sink in self.sinks: if not getattr(sink, "enabled", True): @@ -242,7 +270,14 @@ def submit(self, item: Any) -> bool: self._queue.put(item, timeout=self.put_timeout) return True except queue.Full: - count = len(item.records) if isinstance(item, WriteItem) else 1 + if not isinstance(item, WriteItem): + logger.warning( + "Upload queue full after %.1fs; could not queue %s", + self.put_timeout, + type(item).__name__, + ) + return False + count = len(item.records) self.dropped += count logger.warning( "Upload queue full after %.1fs; dropped %d item(s) (%d total). " diff --git a/packages/prime-runs/tests/test_evals_backend.py b/packages/prime-runs/tests/test_evals_backend.py index fa3c68e87..bebea8e90 100644 --- a/packages/prime-runs/tests/test_evals_backend.py +++ b/packages/prime-runs/tests/test_evals_backend.py @@ -5,7 +5,13 @@ from conftest import RecordingHandler from prime_runs.backends import EvalsBackend -from prime_runs.exceptions import ConfigurationError, EnvironmentResolutionError +from prime_runs.exceptions import ( + ConfigurationError, + EnvironmentResolutionError, + ForbiddenError, + PaymentRequiredError, + UnauthorizedError, +) from prime_runs.models import EnvironmentRef, RunSpec, RunStatus @@ -169,6 +175,27 @@ def test_attach_survives_a_read_failure(make_platform_client, eval_routes): assert handle.url == "https://app.example/dashboard/evaluations/eval-abc" +@pytest.mark.parametrize( + ("status_code", "error_type"), + [ + (401, UnauthorizedError), + (402, PaymentRequiredError), + (403, ForbiddenError), + ], +) +def test_attach_propagates_permanent_access_failures( + make_platform_client, eval_routes, status_code, error_type +): + routes = dict(eval_routes) + routes["GET /api/v1/evaluations/eval-abc"] = lambda request: httpx.Response( + status_code, json={"detail": "denied"} + ) + backend, _ = make_backend(make_platform_client, routes) + + with pytest.raises(error_type): + backend.attach("eval-abc") + + def test_a_pinned_environment_version_reaches_the_api(make_platform_client, eval_routes): """The API's EnvironmentReference carries version_id. Dropping it attaches the run to whatever version the hub resolves today — the difference between diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index 8e2b01054..dd481dbbe 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -1,6 +1,7 @@ """The run handle: lifecycle, containment, ranks, terminal status.""" import signal +import threading from typing import Any, Dict, List, Optional import pytest @@ -112,6 +113,7 @@ def test_metrics_land_in_the_summary_when_the_backend_has_no_time_series(): run.log({"reward": 0.5}, step=1) run.log({"reward": 0.75}, step=2) + run.flush() assert run.summary["reward"] == 0.75 assert backend.points == [] @@ -142,6 +144,38 @@ def test_commit_false_stages_without_writing(): run.finish() +def test_commit_false_merges_staged_metrics_into_the_next_point(): + backend = FakeBackend(supports_step_metrics=True) + run = make_run(backend) + + run.log({"loss": 2.0}, step=7, commit=False) + run.log({"reward": 0.5}) + run.flush() + + assert backend.points == [({"loss": 2.0, "reward": 0.5}, 7)] + run.finish() + + +def test_periodic_summary_updates_run_on_the_uploader_thread(): + caller_thread = threading.get_ident() + update_threads = [] + + class ThreadRecordingBackend(FakeBackend): + def update(self, run_id, *, config=None, summary=None) -> None: + update_threads.append(threading.get_ident()) + super().update(run_id, config=config, summary=summary) + + backend = ThreadRecordingBackend(supports_step_metrics=False) + run = make_run(backend, summary_flush_seconds=0.0) + + run.log({"reward": 0.5}) + run.flush() + + assert update_threads + assert all(thread_id != caller_thread for thread_id in update_threads) + run.finish() + + def test_non_finite_metrics_are_dropped_rather_than_failing_the_request(): """A diverged loss serializes as bare ``NaN``, which strict JSON rejects — the whole request fails on a payload nobody can inspect.""" @@ -298,6 +332,19 @@ def test_on_error_raise_surfaces_the_failure_for_tests_and_ci(): with pytest.raises(RuntimeError, match="finalize exploded"): run.finish() + assert backend.closed is True + + +def test_update_failure_in_raise_mode_still_finalizes_and_closes(): + backend = FakeBackend(fail_on="update") + run = make_run(backend, on_error="raise") + + with pytest.raises(RuntimeError, match="update exploded"): + run.finish() + + assert len(backend.finalized) == 1 + assert backend.closed is True + def test_a_sink_error_is_recorded_on_the_run(): run = make_run(sinks=[FakeSink("broken", fail_on_write=True)]) diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index 9477b1eec..7c5db2540 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -5,7 +5,7 @@ from conftest import FakeSink -from prime_runs.worker import MetricItem, UploadWorker, WriteItem +from prime_runs.worker import MetricItem, RunUpdateItem, UploadWorker, WriteItem def drain(worker: UploadWorker) -> None: @@ -106,6 +106,19 @@ def test_metrics_ride_the_same_queue_when_the_backend_stores_a_time_series(): worker.close() +def test_run_updates_ride_the_uploader_queue(): + updates = [] + worker = UploadWorker( + [], update_writer=lambda config, summary: updates.append((config, summary)) + ) + + worker.submit(RunUpdateItem(config={"seed": 7}, summary={"reward": 0.5})) + drain(worker) + + assert updates == [({"seed": 7}, {"reward": 0.5})] + worker.close() + + def test_a_metric_write_that_raises_does_not_kill_the_uploader(): sink = FakeSink() From ab0fdfa48135536455646c823bcb772262ed2df9 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 16:14:35 -0700 Subject: [PATCH 08/39] fix(prime-runs): harden run lifecycle handling --- packages/prime-runs/src/prime_runs/_http.py | 19 ++++++++++------- packages/prime-runs/src/prime_runs/run.py | 20 ++++++++++++++++-- packages/prime-runs/tests/test_http.py | 23 +++++++++++++++++++++ packages/prime-runs/tests/test_init.py | 21 +++++++++++++++++++ packages/prime-runs/tests/test_run.py | 13 ++++++++++++ 5 files changed, 86 insertions(+), 10 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/_http.py b/packages/prime-runs/src/prime_runs/_http.py index 8db2c86ba..3fbda8517 100644 --- a/packages/prime-runs/src/prime_runs/_http.py +++ b/packages/prime-runs/src/prime_runs/_http.py @@ -176,14 +176,17 @@ def request( for attempt in range(1, attempts + 1): ambiguous = True try: - response = self._client.request( - method, - url, - content=body, - headers=headers, - params=dict(params) if params else None, - timeout=timeout, - ) + request_kwargs: Dict[str, Any] = { + "content": body, + "headers": headers, + "params": dict(params) if params else None, + } + # ``None`` disables httpx timeouts; it does not mean "use the + # client's default". Omit the override so ordinary lifecycle + # calls retain the timeout configured in ``_new_client``. + if timeout is not None: + request_kwargs["timeout"] = timeout + response = self._client.request(method, url, **request_kwargs) except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as exc: # No connection was ever established, so the server saw nothing. ambiguous = False diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index a0ed10be3..3fe936486 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -24,6 +24,7 @@ import time from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Union +from . import _fork from ._http import DEFAULT_TIMEOUT, UPLOAD_TIMEOUT, PlatformClient from .backends import Backend, EvalsBackend, OfflineBackend from .config import Config @@ -137,6 +138,7 @@ def __init__( # the process from installing its own. self._signal_handler = self._handle_signal self._previous_signal_handlers: Dict[int, Any] = {} + _fork.register(self) # -------------------------------------------------------------- identity @@ -169,7 +171,7 @@ def status(self) -> RunStatus: @property def is_primary(self) -> bool: """Whether this process owns the run's lifecycle (rank 0, or single-process).""" - return self._is_primary + return self._owns_lifecycle @property def finished(self) -> bool: @@ -455,6 +457,20 @@ def _restore_signal_handlers(self) -> None: continue self._previous_signal_handlers.clear() + def reset_after_fork(self) -> None: + """Make an inherited handle safe to use in a forked child. + + The child may keep using this handle, but the process that created the + run remains responsible for its lifecycle. In particular, the child's + inherited signal and atexit callbacks must never finalize the parent's + still-running run. The lock also has to be replaced because it may have + been owned at fork time by a thread that no longer exists. + """ + self._finish_lock = threading.RLock() + self._is_primary = False + self._owns_lifecycle = False + self._deferred_error = None + def _on_process_exit(self) -> None: """Last resort: the process is exiting and nobody called ``finish()``. @@ -834,7 +850,7 @@ def _resolve_mode( "recording this run offline instead." ) mode = "offline" - if mode == "online" and not is_primary and not run_id: + if mode in ("online", "offline") and not is_primary and not run_id: # A non-primary rank with no run to join would create a second run for # the same job. Recording nothing is better than that. logger.debug("Non-primary rank with no %s; disabling this run handle", RUN_ID_ENV) diff --git a/packages/prime-runs/tests/test_http.py b/packages/prime-runs/tests/test_http.py index d230aadd9..b4895b5bd 100644 --- a/packages/prime-runs/tests/test_http.py +++ b/packages/prime-runs/tests/test_http.py @@ -105,6 +105,29 @@ def handler(request): assert len(no_sleep) == 1 +def test_an_omitted_request_timeout_keeps_the_clients_default(): + seen = [] + + def handler(request): + seen.append(request.extensions["timeout"]) + return httpx.Response(200, json={"ok": True}) + + transport = httpx.MockTransport(handler) + http_client = httpx.Client( + transport=transport, + timeout=httpx.Timeout(17.0, connect=3.0), + ) + client = PlatformClient( + api_key="test-key", + base_url="http://testserver", + client=http_client, + ) + + assert client.get("/evaluations/x") == {"ok": True} + assert seen == [{"connect": 3.0, "read": 17.0, "write": 17.0, "pool": 17.0}] + http_client.close() + + def test_an_empty_body_is_a_valid_response(): client = client_for(lambda request: httpx.Response(204)) diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index ac65f3d38..286d26b62 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -136,6 +136,18 @@ def test_a_non_primary_rank_with_no_run_to_join_records_nothing(monkeypatch, tmp run.finish() +def test_a_non_primary_offline_rank_without_a_run_to_join_records_nothing(monkeypatch, tmp_path): + """An offline rank must not create a run it is forbidden to finalize.""" + monkeypatch.setenv("RANK", "3") + + run = pr.init(mode="offline", dir=str(tmp_path)) + + assert run.mode == "disabled" + assert run.is_primary is False + run.finish() + assert not list(tmp_path.iterdir()) + + def test_a_run_id_in_the_environment_is_joined_not_recreated(monkeypatch, tmp_path): monkeypatch.setenv("DP_RANK", "2") monkeypatch.setenv(RUN_ID_ENV, "offline-shared") @@ -264,6 +276,7 @@ def test_an_id_inherited_from_a_parent_process_is_joined(monkeypatch, tmp_path): run = pr.init(mode="offline", dir=str(tmp_path)) assert run.id == "offline-from-parent" + assert run.is_primary is False run.finish() @@ -310,8 +323,16 @@ def test_a_forked_child_joins_the_run_without_duplicating_the_parents_records(tm child = pr.init(mode="offline", dir=str(tmp_path), handle_signals=False) if child.id != run.id: code = 1 + if child.is_primary: + code = 3 child.log_traces([{"id": "child-1"}]) child.finish() + # The original handle is inherited too. Its atexit callback must be + # harmless in the child: the parent still owns this lifecycle. + run._on_process_exit() + state = json.loads((tmp_path / run.id / "run.json").read_text()) + if state["status"] != RunStatus.RUNNING.value: + code = 4 except BaseException: code = 2 finally: diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index dd481dbbe..5763706f7 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -315,6 +315,19 @@ def test_a_process_that_exits_without_finishing_reports_crashed(): assert backend.finalized[0]["status"] is RunStatus.CRASHED +def test_a_forked_handle_gets_a_fresh_lock_and_loses_lifecycle_ownership(): + backend = FakeBackend() + run = make_run(backend) + inherited_lock = run._finish_lock + + run.reset_after_fork() + + assert run._finish_lock is not inherited_lock + assert run.is_primary is False + run.finish() + assert backend.finalized == [] + + def test_a_backend_failure_does_not_escape_into_the_producer_by_default(): """Six hours of rollouts must not be lost to a 502 on a telemetry call.""" backend = FakeBackend(fail_on="finalize") From 3ca640fa76298c7133054046153b6081a5672530 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 16:33:00 -0700 Subject: [PATCH 09/39] fix(runs): handle trace samples and avoid duplicate finalization --- .../src/prime_runs/backends/evals.py | 5 ++-- .../src/prime_runs/sinks/samples.py | 24 +++++++++++-------- .../prime-runs/tests/test_evals_backend.py | 13 ++++++++++ .../prime-runs/tests/test_samples_sink.py | 13 ++++++++++ 4 files changed, 43 insertions(+), 12 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/backends/evals.py b/packages/prime-runs/src/prime_runs/backends/evals.py index 043e21f9c..4d6bb3e91 100644 --- a/packages/prime-runs/src/prime_runs/backends/evals.py +++ b/packages/prime-runs/src/prime_runs/backends/evals.py @@ -166,11 +166,12 @@ def finalize( if status is RunStatus.COMPLETED: body: Dict[str, Any] = {} _set_if(body, "metrics", summary or None) - # Setting a terminal state: replaying it lands on the same state. + # Finalization also enqueues the platform's asynchronous statistics + # task. A lost response leaves the outcome ambiguous, so replaying + # this POST can enqueue the work twice. self._client.post( f"/evaluations/{run_id}/finalize", json_body=body or {"metrics": {}}, - idempotent=True, ) return self._report_failure(run_id, status=status, summary=summary, error=error, config=config) diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py index 1be93ec96..ba7ec6e57 100644 --- a/packages/prime-runs/src/prime_runs/sinks/samples.py +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -19,7 +19,7 @@ from typing import Any, Dict, List, Mapping, Optional, Sequence from .._http import UPLOAD_TIMEOUT, PlatformClient, encode_json -from ..projection import batch_samples, build_samples +from ..projection import batch_samples, build_samples, trace_to_sample from .base import Sink logger = logging.getLogger(__name__) @@ -70,19 +70,19 @@ def write( self.samples_written += len(batch) def _to_samples(self, records: Sequence[Any]) -> List[Dict[str, Any]]: - """Split a batch into episodes to project and samples to pass through. + """Project native episodes/traces and pass through existing samples. A producer that already speaks the v0 sample format (a dict with ``sample_id``) sends it unchanged; anything with ``traces`` is a native - episode and gets projected. Anything else is skipped loudly rather than - posted as a malformed row the API would reject for the whole batch. + episode, and anything with ``branches`` is a native trace. Anything else + is skipped loudly rather than posted as a malformed row the API would + reject for the whole batch. """ - episodes: List[Any] = [] - passthrough: List[Dict[str, Any]] = [] + samples: List[Dict[str, Any]] = [] for record in records: if isinstance(record, Mapping): if "sample_id" in record: - passthrough.append(dict(record)) + samples.append(dict(record)) elif "traces" in record: logger.debug( "Skipping a pre-serialized episode: this sink projects native " @@ -92,10 +92,14 @@ def _to_samples(self, records: Sequence[Any]) -> List[Dict[str, Any]]: logger.debug("Skipping a record with no sample_id and no traces") continue if hasattr(record, "traces"): - episodes.append(record) + samples.extend(build_samples([record], self._rollout_numbers)) + elif hasattr(record, "branches"): + idx = record.task.data.idx + self._rollout_numbers[idx] = number = self._rollout_numbers.get(idx, 0) + 1 + samples.append(trace_to_sample(record, rollout_number=number)) else: - logger.debug("Skipping %s: not an episode", type(record).__name__) - return build_samples(episodes, self._rollout_numbers) + passthrough + logger.debug("Skipping %s: not an episode or trace", type(record).__name__) + return samples def flush(self) -> None: """Writes are synchronous; the uploader thread owns the asynchrony.""" diff --git a/packages/prime-runs/tests/test_evals_backend.py b/packages/prime-runs/tests/test_evals_backend.py index bebea8e90..a250ab2d4 100644 --- a/packages/prime-runs/tests/test_evals_backend.py +++ b/packages/prime-runs/tests/test_evals_backend.py @@ -10,6 +10,7 @@ EnvironmentResolutionError, ForbiddenError, PaymentRequiredError, + RetryableAPIError, UnauthorizedError, ) from prime_runs.models import EnvironmentRef, RunSpec, RunStatus @@ -107,6 +108,18 @@ def test_finalizing_a_completed_run_posts_its_metrics(make_platform_client, eval } +def test_an_ambiguous_finalize_failure_is_not_replayed(make_platform_client, eval_routes): + """Finalization enqueues asynchronous processing, so a retry can enqueue it twice.""" + routes = dict(eval_routes) + routes["POST /api/v1/evaluations/eval-abc/finalize"] = lambda request: httpx.Response(502) + backend, handler = make_backend(make_platform_client, routes) + + with pytest.raises(RetryableAPIError): + backend.finalize("eval-abc", status=RunStatus.COMPLETED) + + assert handler.paths().count("POST /api/v1/evaluations/eval-abc/finalize") == 1 + + def test_a_failed_run_falls_back_to_metadata_when_the_status_endpoint_is_missing( make_platform_client, eval_routes, caplog ): diff --git a/packages/prime-runs/tests/test_samples_sink.py b/packages/prime-runs/tests/test_samples_sink.py index 2e95737e4..1b4c5f0e1 100644 --- a/packages/prime-runs/tests/test_samples_sink.py +++ b/packages/prime-runs/tests/test_samples_sink.py @@ -34,6 +34,19 @@ def test_rollout_numbering_is_continuous_across_streamed_batches(make_platform_c assert [body["samples"][0]["rollout_number"] for body in posted] == [1, 2] +def test_individual_traces_are_projected_with_continuous_rollout_numbers( + make_platform_client, eval_routes +): + sink, handler = make_sink(make_platform_client, eval_routes) + + sink.write([make_trace(trace_id="trace-1", idx=0)]) + sink.write([make_trace(trace_id="trace-2", idx=0)]) + + posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples") + assert [body["samples"][0]["sample_id"] for body in posted] == ["trace-1", "trace-2"] + assert [body["samples"][0]["rollout_number"] for body in posted] == [1, 2] + + def test_a_producer_that_already_speaks_v0_is_passed_through(make_platform_client, eval_routes): sink, handler = make_sink(make_platform_client, eval_routes) From f5fb81be95efa1767edcd0a7dd168f5900bcedf4 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 17:17:51 -0700 Subject: [PATCH 10/39] fix(runs): harden finalization and release retries --- .github/workflows/release-runs.yml | 19 ++++++++++----- packages/prime-runs/src/prime_runs/run.py | 28 ++++++++++++++++++++--- packages/prime-runs/tests/test_run.py | 25 ++++++++++++++++++++ 3 files changed, 63 insertions(+), 9 deletions(-) diff --git a/.github/workflows/release-runs.yml b/.github/workflows/release-runs.yml index 432f4b57c..da263298a 100644 --- a/.github/workflows/release-runs.yml +++ b/.github/workflows/release-runs.yml @@ -46,6 +46,12 @@ jobs: echo "Tag $TAG does not exist" fi + # A forced release must build the tagged source, not whatever the + # dispatch ref currently points at. + - name: Check out existing tag + if: steps.check_tag.outputs.exists == 'true' && inputs.force_release == 'true' + run: git checkout --detach "prime-runs-v${{ steps.version.outputs.version }}" + - name: Set up Python if: steps.check_tag.outputs.exists != 'true' || inputs.force_release == 'true' uses: actions/setup-python@v6 @@ -62,6 +68,13 @@ jobs: run: | uv build --out-dir dist + - name: Publish to PyPI + if: steps.check_tag.outputs.exists != 'true' || inputs.force_release == 'true' + uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # v1.14.2 + with: + packages-dir: packages/prime-runs/dist + skip-existing: ${{ inputs.force_release == 'true' }} + - name: Create tag if: steps.check_tag.outputs.exists != 'true' run: | @@ -70,9 +83,3 @@ jobs: git config user.email 'github-actions[bot]@users.noreply.github.com' git tag -a "$TAG" -m "Release $TAG" git push origin "$TAG" - - - name: Publish to PyPI - if: steps.check_tag.outputs.exists != 'true' || inputs.force_release == 'true' - uses: pypa/gh-action-pypi-publish@dc37677b2e1c63e2034f94d8a5b11f265b73ba33 # v1.14.2 - with: - packages-dir: packages/prime-runs/dist diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 3fe936486..2649ff47f 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -290,6 +290,7 @@ def finish( ) -> None: """Flush everything and close the run out. Idempotent. + ``status`` must be one of the terminal :class:`RunStatus` values. Safe to call from ``__exit__``, an atexit hook and a signal handler at once — whichever gets there first reports the status, and the rest return. @@ -298,6 +299,8 @@ def finish( if self._finished: return resolved = RunStatus(status) if not isinstance(status, RunStatus) else status + if not resolved.is_terminal(): + raise ValueError(f"finish() requires a terminal status, got {resolved.value!r}") self._finished = True if summary: @@ -393,13 +396,32 @@ def __enter__(self) -> "Run": def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: if exc_type is None: self.finish() - elif isinstance(exc, KeyboardInterrupt): + return False + + if isinstance(exc, KeyboardInterrupt): # An interrupt is a decision, not a fault, so it must not land in # the same bucket as broken ones. Matches the SIGINT handler, which # normally gets there first when signal handling is on. - self.finish(status=RunStatus.CRASHED, error="interrupted") + status = RunStatus.CRASHED + error = "interrupted" else: - self.finish(status=RunStatus.FAILED, error=_describe(exc)) + status = RunStatus.FAILED + error = _describe(exc) + + try: + self.finish(status=status, error=error) + except BaseException as finish_error: + # The producer exception is the reason this context is unwinding. + # A telemetry teardown error must not replace it, even in strict + # mode; finish() has already recorded the failure on the run. + logger.warning( + "Run %s: finishing after %s also failed: %s: %s", + self.id, + exc_type.__name__, + type(finish_error).__name__, + finish_error, + exc_info=True, + ) return False # ------------------------------------------------------------- internals diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index 5763706f7..c0307f066 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -221,6 +221,21 @@ def test_finish_is_idempotent(): assert backend.finalized[0]["summary"] == {"avg_reward": 1.0} +@pytest.mark.parametrize("status", [RunStatus.RUNNING, "running"]) +def test_finish_rejects_a_nonterminal_status_without_closing_the_run(status): + backend = FakeBackend() + run = make_run(backend) + + with pytest.raises(ValueError, match="requires a terminal status"): + run.finish(status=status) + + assert not run.finished + assert backend.finalized == [] + + run.finish() + assert backend.finalized[0]["status"] is RunStatus.COMPLETED + + def test_logging_after_finish_is_a_producer_bug(): run = make_run() run.finish() @@ -250,6 +265,16 @@ def test_an_exception_inside_the_block_fails_the_run_and_still_propagates(): assert "rollout blew up" in backend.finalized[0]["error"] +def test_a_finish_failure_does_not_mask_the_context_exception(): + backend = FakeBackend(fail_on="finalize") + + with pytest.raises(ValueError, match="rollout blew up"): + with make_run(backend, on_error="raise"): + raise ValueError("rollout blew up") + + assert backend.closed is True + + def test_an_interrupt_is_recorded_as_a_decision_not_a_fault(): """Ctrl-C must not land in the same bucket as a broken eval — and it must agree with the SIGINT handler, which normally gets there first.""" From e8d9f6c988aca54de1191908cffea78344204a16 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 17:21:55 -0700 Subject: [PATCH 11/39] fix(runs): drain full upload queues during shutdown --- packages/prime-runs/src/prime_runs/worker.py | 33 +++++++++-- packages/prime-runs/tests/test_worker.py | 62 ++++++++++++++++---- 2 files changed, 77 insertions(+), 18 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index 10bdb33fe..a709a7652 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -24,6 +24,7 @@ import os import queue import threading +import time from dataclasses import dataclass, field from typing import Any, Callable, List, Optional, Sequence @@ -72,6 +73,18 @@ class _Flush: event: threading.Event = field(default_factory=threading.Event) +def _deadline(timeout: Optional[float]) -> Optional[float]: + if timeout is None: + return None + return time.monotonic() + max(0.0, timeout) + + +def _remaining(deadline: Optional[float]) -> Optional[float]: + if deadline is None: + return None + return max(0.0, deadline - time.monotonic()) + + class UploadWorker: """Drains a bounded queue into a list of sinks on one daemon thread.""" @@ -293,24 +306,32 @@ def flush(self, timeout: Optional[float] = None) -> bool: if self._thread is None or not self._thread.is_alive(): self._flush_sinks() return True + deadline = _deadline(timeout) barrier = _Flush() try: - self._queue.put(barrier, timeout=self.put_timeout) + # Synchronization belongs to the caller's drain budget, not the + # short producer backpressure budget. A full queue is precisely + # when finish() most needs to wait for room for this barrier. + self._queue.put(barrier, timeout=_remaining(deadline)) except queue.Full: - logger.warning("Could not enqueue a flush barrier; the queue is saturated") + logger.warning("Could not enqueue a flush barrier before the drain timeout") return False - return barrier.event.wait(timeout) + return barrier.event.wait(_remaining(deadline)) def close(self, timeout: Optional[float] = 30.0) -> None: """Drain, stop the thread, and close every sink.""" self._stopping.set() thread = self._thread if thread is not None and thread.is_alive(): + deadline = _deadline(timeout) try: - self._queue.put(None, timeout=self.put_timeout) + # The sentinel sits behind every accepted item. Give it the + # close budget so a temporarily full queue can make room and + # then drain in FIFO order before the thread exits. + self._queue.put(None, timeout=_remaining(deadline)) except queue.Full: - logger.warning("Upload queue saturated at close; some records may be lost") - thread.join(timeout) + logger.warning("Upload queue remained saturated through the close timeout") + thread.join(_remaining(deadline)) if thread.is_alive(): # Closing the sinks now would pull an httpx client, or a file # handle, out from under a request that is still running on that diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index 7c5db2540..a2a6f18c9 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -8,6 +8,18 @@ from prime_runs.worker import MetricItem, RunUpdateItem, UploadWorker, WriteItem +class BlockingSink(FakeSink): + def __init__(self) -> None: + super().__init__("blocking") + self.entered = threading.Event() + self.released = threading.Event() + + def write(self, records, *, line_format=None, step=None) -> None: + self.entered.set() + self.released.wait(5.0) + super().write(records, line_format=line_format, step=step) + + def drain(worker: UploadWorker) -> None: assert worker.flush(timeout=5.0) @@ -66,18 +78,6 @@ def test_a_failed_sink_is_not_called_again(): def test_a_full_queue_drops_rather_than_blocking_the_producer(): """Stalling a training run to protect telemetry is the wrong trade.""" - - class BlockingSink(FakeSink): - def __init__(self) -> None: - super().__init__("blocking") - self.entered = threading.Event() - self.released = threading.Event() - - def write(self, records, *, line_format=None, step=None) -> None: - self.entered.set() - self.released.wait(5.0) - super().write(records, line_format=line_format, step=step) - sink = BlockingSink() worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.05) @@ -95,6 +95,44 @@ def write(self, records, *, line_format=None, step=None) -> None: worker.close() +def test_flush_uses_the_drain_budget_to_get_behind_a_full_queue(): + sink = BlockingSink() + worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.01) + assert worker.submit(WriteItem(records=[{"id": 0}])) + assert sink.entered.wait(1.0) + assert worker.submit(WriteItem(records=[{"id": 1}])) + + release = threading.Timer(0.05, sink.released.set) + release.start() + try: + assert worker.flush(timeout=0.5) is True + finally: + sink.released.set() + release.join() + worker.close(timeout=1.0) + + assert [batch[0][0]["id"] for batch in sink.batches] == [0, 1] + + +def test_close_uses_its_budget_to_queue_the_stop_behind_pending_records(): + sink = BlockingSink() + worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.01) + assert worker.submit(WriteItem(records=[{"id": 0}])) + assert sink.entered.wait(1.0) + assert worker.submit(WriteItem(records=[{"id": 1}])) + + release = threading.Timer(0.05, sink.released.set) + release.start() + try: + worker.close(timeout=0.5) + finally: + sink.released.set() + release.join() + + assert [batch[0][0]["id"] for batch in sink.batches] == [0, 1] + assert sink.closed is True + + def test_metrics_ride_the_same_queue_when_the_backend_stores_a_time_series(): points = [] worker = UploadWorker([], metric_writer=lambda metrics, step: points.append((metrics, step))) From 86c6e1207daba82db2f2a995da259fee9a2c78af Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 17:44:25 -0700 Subject: [PATCH 12/39] fix(prime-runs): harden lifecycle and trace persistence --- .../src/prime_runs/backends/evals.py | 18 +- .../prime-runs/src/prime_runs/projection.py | 279 +++++++++++++++++- packages/prime-runs/src/prime_runs/run.py | 47 ++- .../src/prime_runs/sinks/offline.py | 8 +- .../src/prime_runs/sinks/samples.py | 22 +- .../prime-runs/src/prime_runs/sinks/traces.py | 24 +- .../prime-runs/tests/test_evals_backend.py | 28 ++ packages/prime-runs/tests/test_init.py | 17 ++ packages/prime-runs/tests/test_run.py | 53 ++++ .../prime-runs/tests/test_samples_sink.py | 56 +++- packages/prime-runs/tests/test_traces_sink.py | 11 + 11 files changed, 522 insertions(+), 41 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/backends/evals.py b/packages/prime-runs/src/prime_runs/backends/evals.py index 4d6bb3e91..8e393e76b 100644 --- a/packages/prime-runs/src/prime_runs/backends/evals.py +++ b/packages/prime-runs/src/prime_runs/backends/evals.py @@ -198,7 +198,6 @@ def _report_failure( self._client.post( f"/evaluations/{run_id}/status", json_body={"status": _PLATFORM_STATUS[status], "error": error}, - max_attempts=1, idempotent=True, ) return @@ -211,10 +210,17 @@ def _report_failure( "recording terminal state in metadata instead" ) except RunAPIError as exc: - if exc.status_code not in (405, 422): + if is_transient(exc) or (exc.status_code is not None and exc.status_code >= 500): + logger.warning( + "Status endpoint remained unavailable after retries (%s); " + "recording terminal state in metadata instead", + exc, + ) + elif exc.status_code not in (405, 422): raise - self._status_endpoint_missing = True - logger.debug("Status endpoint rejected the request (%s); using metadata", exc) + else: + self._status_endpoint_missing = True + logger.debug("Status endpoint rejected the request (%s); using metadata", exc) # Fallback: the run cannot be moved out of RUNNING, but the failure is # at least recorded where an operator and the dashboard can both read it. @@ -227,8 +233,8 @@ def _report_failure( summary=summary, ) logger.warning( - "Run %s %s, but the platform has no way to mark an evaluation failed; " - "it will keep showing as running. Recorded the failure in metadata.prime_runs.", + "Run %s %s, but its evaluation status could not be updated; it will keep " + "showing as running. Recorded the failure in metadata.prime_runs.", run_id, status.value, ) diff --git a/packages/prime-runs/src/prime_runs/projection.py b/packages/prime-runs/src/prime_runs/projection.py index de50348f4..4a4af5e47 100644 --- a/packages/prime-runs/src/prime_runs/projection.py +++ b/packages/prime-runs/src/prime_runs/projection.py @@ -17,7 +17,7 @@ """ import logging -from typing import Any, Dict, Iterable, List, Optional, Sequence +from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence from ._http import encode_json @@ -113,6 +113,283 @@ def trace_to_sample( return sample +def trace_record_to_sample( + trace: Mapping[str, Any], rollout_number: int = 1, episode_id: Optional[str] = None +) -> Dict[str, Any]: + """Project a serialized trace without importing its producer package. + + Verifiers persists its message graph as ``nodes``/``calls`` rather than the + derived ``branches`` property used by :func:`trace_to_sample`. Older records + may carry branches directly, so both representations are accepted. Sparse + trace mappings still produce a visible row as long as they have an ID; + fields the legacy viewer cannot derive remain empty instead of making the + entire record disappear. + """ + trace_id = trace.get("id") + if not trace_id: + raise TypeError("serialized trace records must contain a non-empty 'id'") + + task_container = _as_mapping(trace.get("task")) + task = dict(_as_mapping(task_container.get("data"))) + agent = _as_mapping(trace.get("agent")) + branches = _serialized_branches(trace) + rewards = _as_mapping(trace.get("rewards")) + errors = trace.get("errors") + last_error = trace.get("last_error") + if last_error is None and isinstance(errors, list) and errors: + last_error = errors[-1] + stop_condition = trace.get("stop_condition") + + sample: Dict[str, Any] = { + "sample_id": trace_id, + "example_id": task.get("idx"), + "rollout_number": rollout_number, + "episode_id": episode_id, + "agent": agent.get("name"), + "trainable": agent.get("trainable", True), + "task": task, + "prompt": [], + "completion": branches[-1]["messages"] if branches else [], + "answer": task.get("answer"), + "tool_defs": _mapping_list(trace.get("tools")) or None, + "reward": trace["reward"] if "reward" in trace else _total_reward(rewards), + "timing": dict(_as_mapping(trace.get("timing"))) or None, + "is_completed": trace.get("is_completed", False), + "is_truncated": trace.get("is_truncated", _is_truncated(stop_condition, trace)), + "metrics": dict(_as_mapping(trace.get("metrics"))), + "error": dict(last_error) if isinstance(last_error, Mapping) else last_error, + "stop_condition": stop_condition, + "trajectory": branches, + "token_usage": dict(_as_mapping(trace.get("usage"))) or _aggregate_usage(trace), + "info": dict(_as_mapping(trace.get("info"))) or None, + } + for name, reward in rewards.items(): + if reward is None: + continue + score = reward.get("score") if isinstance(reward, Mapping) else reward + sample.setdefault(name, score) + return sample + + +def record_to_samples( + record: Mapping[str, Any], rollout_numbers: Optional[Dict[Any, int]] = None +) -> List[Dict[str, Any]]: + """Project one serialized trace or episode to legacy viewer samples.""" + counts = rollout_numbers if rollout_numbers is not None else {} + if "traces" not in record: + task = _as_mapping(_as_mapping(record.get("task")).get("data")) + idx = task.get("idx") + rollout_key = idx if idx is not None else record.get("id") + counts[rollout_key] = number = counts.get(rollout_key, 0) + 1 + return [trace_record_to_sample(record, rollout_number=number)] + + episode_id = record.get("id") + if not episode_id: + raise TypeError("serialized episode records must contain a non-empty 'id'") + raw_traces = record.get("traces") + if not isinstance(raw_traces, list): + raise TypeError("serialized episode 'traces' must be a list") + traces: List[Mapping[str, Any]] = [] + for trace in raw_traces: + if not isinstance(trace, Mapping): + raise TypeError("serialized episode traces must be mappings") + traces.append(trace) + if not traces: + return [] + + summary_index = next( + ( + index + for index, trace in enumerate(traces) + if _as_mapping(trace.get("agent")).get("trainable", True) + ), + 0, + ) + summary_task = _as_mapping(_as_mapping(traces[summary_index].get("task")).get("data")) + idx = summary_task.get("idx") + rollout_key = idx if idx is not None else episode_id + counts[rollout_key] = number = counts.get(rollout_key, 0) + 1 + sample = trace_record_to_sample(traces[summary_index], number, str(episode_id)) + sample["sample_id"] = episode_id + sample["info"] = { + **(sample["info"] or {}), + "native_wrapper": dict(record), + "native_trace_index": summary_index, + } + if ENVELOPE_BYTES + json_bytes(sample) <= MAX_SAMPLES_PAYLOAD_BYTES: + return [sample] + + logger.warning( + "Episode %s exceeds the platform sample limit; uploading projected traces", + episode_id, + ) + return [trace_record_to_sample(trace, number, str(episode_id)) for trace in traces] + + +def _as_mapping(value: Any) -> Mapping[str, Any]: + return value if isinstance(value, Mapping) else {} + + +def _mapping_list(value: Any) -> List[Dict[str, Any]]: + if not isinstance(value, list): + return [] + return [dict(item) for item in value if isinstance(item, Mapping)] + + +def _serialized_branches(trace: Mapping[str, Any]) -> List[Dict[str, Any]]: + raw_branches = trace.get("branches") + if isinstance(raw_branches, list): + return [ + { + "messages": _mapping_list(_as_mapping(branch).get("messages")), + "num_input_tokens": _as_mapping(branch).get("num_input_tokens", 0), + "num_output_tokens": _as_mapping(branch).get("num_output_tokens", 0), + } + for branch in raw_branches + if isinstance(branch, Mapping) + ] + + raw_nodes = trace.get("nodes") + if not isinstance(raw_nodes, list): + return [] + nodes = [_as_mapping(node) for node in raw_nodes] + parents = {node.get("parent") for node in nodes if isinstance(node.get("parent"), int)} + leaves = [index for index in range(len(nodes)) if index not in parents] + calls = trace.get("calls") + calls_by_node = ( + { + call.get("node"): call + for call in calls + if isinstance(call, Mapping) and isinstance(call.get("node"), int) + } + if isinstance(calls, list) + else {} + ) + + branches: List[Dict[str, Any]] = [] + for leaf in leaves: + path: List[int] = [] + seen = set() + node_index: Any = leaf + while ( + isinstance(node_index, int) and 0 <= node_index < len(nodes) and node_index not in seen + ): + seen.add(node_index) + path.append(node_index) + node_index = nodes[node_index].get("parent") + path.reverse() + branch_calls = [calls_by_node[index] for index in path if index in calls_by_node] + input_tokens, output_tokens = _branch_token_counts(branch_calls) + branches.append( + { + "messages": [ + dict(message) + for index in path + if isinstance((message := nodes[index].get("message")), Mapping) + ], + "num_input_tokens": input_tokens, + "num_output_tokens": output_tokens, + } + ) + return branches + + +def _branch_token_counts(calls: Sequence[Mapping[str, Any]]) -> tuple[int, int]: + input_tokens = 0 + output_tokens = 0 + previous_total = 0 + for call in calls: + usage = _as_mapping(call.get("usage")) + current_input, current_output = _usage_counts(usage) + input_tokens += max(0, current_input - previous_total) + output_tokens += current_output + previous_total = int(usage.get("total_tokens", current_input + current_output) or 0) + return input_tokens, output_tokens + + +def _aggregate_usage(trace: Mapping[str, Any]) -> Optional[Dict[str, Any]]: + calls = trace.get("calls") + if not isinstance(calls, list): + return None + usages = [usage for call in calls if (usage := _as_mapping(_as_mapping(call).get("usage")))] + if not usages: + return None + if any("prompt_tokens" in usage or "completion_tokens" in usage for usage in usages): + result: Dict[str, Any] = { + "prompt_tokens": sum(int(usage.get("prompt_tokens", 0) or 0) for usage in usages), + "completion_tokens": sum( + int(usage.get("completion_tokens", 0) or 0) for usage in usages + ), + } + for key in ("cached_input_tokens", "reasoning_tokens", "cost"): + values = [usage[key] for usage in usages if usage.get(key) is not None] + if values: + result[key] = sum(values) + return result + + input_tokens = 0 + output_tokens = 0 + for usage in usages: + current_input, current_output = _usage_counts(usage) + input_tokens += current_input + output_tokens += current_output + if not input_tokens and not output_tokens: + return None + return { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens, + } + + +def _usage_counts(usage: Mapping[str, Any]) -> tuple[int, int]: + if "input_tokens" in usage: + input_tokens = int(usage.get("input_tokens", 0) or 0) + else: + input_tokens = int(usage.get("prompt_tokens", 0) or 0) + int( + usage.get("cached_input_tokens", 0) or 0 + ) + output_tokens = int(usage.get("output_tokens", usage.get("completion_tokens", 0)) or 0) + return input_tokens, output_tokens + + +def _total_reward(rewards: Mapping[str, Any]) -> float: + total = 0.0 + for reward in rewards.values(): + if reward is None: + continue + if isinstance(reward, Mapping): + total += float(reward.get("score", 0.0) or 0.0) * float( + reward.get("weight", 1.0) or 0.0 + ) + else: + total += float(reward) + return total + + +def _is_truncated(stop_condition: Any, trace: Mapping[str, Any]) -> bool: + if stop_condition in { + "max_turns", + "max_input_tokens", + "max_output_tokens", + "max_total_tokens", + "context_length", + }: + return True + calls = trace.get("calls") + if not isinstance(calls, list): + return False + last_successful = next( + ( + call + for call in reversed(calls) + if isinstance(call, Mapping) and call.get("error") is None + ), + None, + ) + return bool(last_successful and last_successful.get("finish_reason") == "length") + + def episode_to_samples(episode: Any, rollout_number: int) -> List[Dict[str, Any]]: """One episode -> the sample rows the platform should store for it. diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 2649ff47f..63ab15053 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -138,6 +138,13 @@ def __init__( # the process from installing its own. self._signal_handler = self._handle_signal self._previous_signal_handlers: Dict[int, Any] = {} + # ``signal.signal`` can only run on the main thread. If finish() runs in + # an executor, or this object is inherited across a fork, its handler + # may remain as the process disposition until the main thread gets a + # chance to replace it. Marking forked handlers as relinquishable lets + # a child run take ownership without mistaking the inherited callback + # for an application-installed handler. + self._signal_handler_stale = False _fork.register(self) # -------------------------------------------------------------- identity @@ -410,10 +417,12 @@ def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: try: self.finish(status=status, error=error) - except BaseException as finish_error: + except Exception as finish_error: # The producer exception is the reason this context is unwinding. # A telemetry teardown error must not replace it, even in strict # mode; finish() has already recorded the failure on the run. + # Control-flow exceptions such as KeyboardInterrupt and SystemExit + # deliberately bypass this handler so teardown cannot swallow them. logger.warning( "Run %s: finishing after %s also failed: %s: %s", self.id, @@ -429,10 +438,11 @@ def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: def install_signal_handlers(self) -> None: """Report a terminal status when the process is killed. - Only installed on the main thread, and only over a *default* handler: - replacing a handler the application chose would be worse than missing a - status. The previous handler is always called afterwards, so SIGINT - still raises ``KeyboardInterrupt`` and SIGTERM still terminates. + Only installed on the main thread, and only over a *default* handler or + one relinquished by a finished/forked ``Run``. Replacing a handler the + application chose would be worse than missing a status. The previous + handler is always called afterwards, so SIGINT still raises + ``KeyboardInterrupt`` and SIGTERM still terminates. """ if threading.current_thread() is not threading.main_thread(): return @@ -441,13 +451,27 @@ def install_signal_handlers(self) -> None: current = signal.getsignal(signum) except (ValueError, OSError): # pragma: no cover - platform dependent continue + previous = current + relinquished_owner: Optional[Run] = None if current not in (signal.SIG_DFL, signal.default_int_handler): - continue + owner = getattr(current, "__self__", None) + if ( + isinstance(owner, Run) + and current is owner._signal_handler + and (owner._finished or owner._signal_handler_stale) + and signum in owner._previous_signal_handlers + ): + previous = owner._previous_signal_handlers[signum] + relinquished_owner = owner + else: + continue try: signal.signal(signum, self._signal_handler) except (ValueError, OSError): # pragma: no cover continue - self._previous_signal_handlers[signum] = current + self._previous_signal_handlers[signum] = previous + if relinquished_owner is not None: + relinquished_owner._previous_signal_handlers.pop(signum, None) def _handle_signal(self, signum: int, frame: Any) -> None: name = signal.Signals(signum).name @@ -471,13 +495,17 @@ def _handle_signal(self, signum: int, frame: Any) -> None: os.kill(os.getpid(), signum) def _restore_signal_handlers(self) -> None: + remaining: Dict[int, Any] = {} for signum, previous in self._previous_signal_handlers.items(): try: if signal.getsignal(signum) is self._signal_handler: signal.signal(signum, previous) except (ValueError, OSError): # pragma: no cover - continue - self._previous_signal_handlers.clear() + # Most commonly finish() was deliberately run in an executor. + # Keep the displaced handler so the main thread can restore it + # from the signal callback or hand it to the next Run. + remaining[signum] = previous + self._previous_signal_handlers = remaining def reset_after_fork(self) -> None: """Make an inherited handle safe to use in a forked child. @@ -492,6 +520,7 @@ def reset_after_fork(self) -> None: self._is_primary = False self._owns_lifecycle = False self._deferred_error = None + self._signal_handler_stale = True def _on_process_exit(self) -> None: """Last resort: the process is exiting and nobody called ``finish()``. diff --git a/packages/prime-runs/src/prime_runs/sinks/offline.py b/packages/prime-runs/src/prime_runs/sinks/offline.py index ebe0084dc..015a3a5cc 100644 --- a/packages/prime-runs/src/prime_runs/sinks/offline.py +++ b/packages/prime-runs/src/prime_runs/sinks/offline.py @@ -7,12 +7,12 @@ and does not change on sync. """ -import json import logging from pathlib import Path from typing import Any, BinaryIO, Mapping, Optional, Sequence, Union from .. import _fork +from .._http import encode_json from .base import Sink, to_mapping logger = logging.getLogger(__name__) @@ -71,8 +71,10 @@ def write( if self._run_kind: run["type"] = self._run_kind mapping["run"] = run - line = json.dumps(mapping, ensure_ascii=False, separators=(",", ":"), default=str) - _write_all(handle, (line + "\n").encode("utf-8")) + # Match the online JSON encoder exactly. In particular, rejecting + # NaN/Infinity here prevents creating an archive that exists but + # cannot later be uploaded by Prime Traces' strict parser. + _write_all(handle, encode_json(mapping) + b"\n") self.records_written += 1 def _handle(self, name: str) -> BinaryIO: diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py index ba7ec6e57..81cdbbd61 100644 --- a/packages/prime-runs/src/prime_runs/sinks/samples.py +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -19,7 +19,7 @@ from typing import Any, Dict, List, Mapping, Optional, Sequence from .._http import UPLOAD_TIMEOUT, PlatformClient, encode_json -from ..projection import batch_samples, build_samples, trace_to_sample +from ..projection import batch_samples, build_samples, record_to_samples, trace_to_sample from .base import Sink logger = logging.getLogger(__name__) @@ -73,23 +73,18 @@ def _to_samples(self, records: Sequence[Any]) -> List[Dict[str, Any]]: """Project native episodes/traces and pass through existing samples. A producer that already speaks the v0 sample format (a dict with - ``sample_id``) sends it unchanged; anything with ``traces`` is a native - episode, and anything with ``branches`` is a native trace. Anything else - is skipped loudly rather than posted as a malformed row the API would - reject for the whole batch. + ``sample_id``) sends it unchanged. Serialized trace and episode records + are projected alongside their native object forms; unsupported mappings + fail explicitly so a gated traces sink cannot turn data loss into a + successful-looking empty run. """ samples: List[Dict[str, Any]] = [] for record in records: if isinstance(record, Mapping): if "sample_id" in record: samples.append(dict(record)) - elif "traces" in record: - logger.debug( - "Skipping a pre-serialized episode: this sink projects native " - "episode objects, not their JSON records" - ) else: - logger.debug("Skipping a record with no sample_id and no traces") + samples.extend(record_to_samples(record, self._rollout_numbers)) continue if hasattr(record, "traces"): samples.extend(build_samples([record], self._rollout_numbers)) @@ -98,7 +93,10 @@ def _to_samples(self, records: Sequence[Any]) -> List[Dict[str, Any]]: self._rollout_numbers[idx] = number = self._rollout_numbers.get(idx, 0) + 1 samples.append(trace_to_sample(record, rollout_number=number)) else: - logger.debug("Skipping %s: not an episode or trace", type(record).__name__) + raise TypeError( + f"EvalSamplesSink cannot project {type(record).__name__}; expected a " + "mapping, trace, or episode" + ) return samples def flush(self) -> None: diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py index 8e02d24df..c2a3f2dcf 100644 --- a/packages/prime-runs/src/prime_runs/sinks/traces.py +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -23,6 +23,8 @@ logger = logging.getLogger(__name__) +DEFAULT_RECEIPT_HISTORY_SIZE = 100 + class TracesSink(Sink): """Uploads records through the Prime Traces service.""" @@ -38,7 +40,10 @@ def __init__( team_id: Optional[str] = None, stamp_run: bool = True, compress: bool = True, + receipt_history_size: int = DEFAULT_RECEIPT_HISTORY_SIZE, ) -> None: + if receipt_history_size < 0: + raise ValueError("receipt_history_size must be non-negative") self.enabled = True self._client = client self._injected_client = client is not None @@ -59,6 +64,8 @@ def __init__( self._run_kind: Optional[str] = None self._context: Dict[str, str] = {} self.receipts: list = [] + self.receipts_received = 0 + self._receipt_history_size = receipt_history_size _fork.register(self) # ------------------------------------------------------------------ setup @@ -119,11 +126,13 @@ def write( payload = [self._prepare(record) for record in records] try: - receipts = self._client.upload_records( - payload, - line_format=resolved, - context=context or None, - compress=self._compress, + receipts = list( + self._client.upload_records( + payload, + line_format=resolved, + context=context or None, + compress=self._compress, + ) ) except Exception as exc: # noqa: BLE001 - classified below if self._is_gated(exc): @@ -133,7 +142,10 @@ def write( ) return raise - self.receipts.extend(receipts) + self.receipts_received += len(receipts) + if self._receipt_history_size: + self.receipts.extend(receipts) + del self.receipts[: -self._receipt_history_size] def _prepare(self, record: Any) -> Any: """Stamp the run onto plain mappings that do not already carry one. diff --git a/packages/prime-runs/tests/test_evals_backend.py b/packages/prime-runs/tests/test_evals_backend.py index a250ab2d4..9cb3edece 100644 --- a/packages/prime-runs/tests/test_evals_backend.py +++ b/packages/prime-runs/tests/test_evals_backend.py @@ -166,6 +166,34 @@ def test_a_status_endpoint_that_exists_is_used_instead_of_the_fallback( assert "PUT /api/v1/evaluations/eval-abc" not in handler.paths() +def test_a_transient_status_failure_is_retried(make_platform_client, eval_routes): + routes = dict(eval_routes) + responses = [httpx.Response(503), httpx.Response(200, json={"evaluation_id": "eval-abc"})] + routes["POST /api/v1/evaluations/eval-abc/status"] = lambda request: responses.pop(0) + backend, handler = make_backend(make_platform_client, routes) + + backend.finalize("eval-abc", status=RunStatus.FAILED, error="boom") + + assert handler.paths().count("POST /api/v1/evaluations/eval-abc/status") == 2 + assert "PUT /api/v1/evaluations/eval-abc" not in handler.paths() + + +def test_exhausted_status_retries_fall_back_to_metadata(make_platform_client, eval_routes, caplog): + routes = dict(eval_routes) + routes["POST /api/v1/evaluations/eval-abc/status"] = lambda request: httpx.Response(503) + handler = RecordingHandler(routes) + client = make_platform_client(handler, max_attempts=2) + backend = EvalsBackend(client, frontend_url="https://app.example") + + with caplog.at_level("WARNING"): + backend.finalize("eval-abc", status=RunStatus.FAILED, error="boom") + + assert handler.paths().count("POST /api/v1/evaluations/eval-abc/status") == 2 + terminal = handler.bodies_for("/api/v1/evaluations/eval-abc")[0]["metadata"]["prime_runs"] + assert terminal["status"] == "failed" + assert "remained unavailable after retries" in caplog.text + + def test_update_sends_nothing_when_there_is_nothing_to_send(make_platform_client, eval_routes): backend, handler = make_backend(make_platform_client, eval_routes) diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index 286d26b62..ceea3bbb9 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -370,3 +370,20 @@ def test_offline_records_are_on_disk_before_any_flush(tmp_path): run.finish() assert len(path.read_text().splitlines()) == 2 + + +def test_offline_records_reject_nonfinite_json_instead_of_writing_invalid_jsonl(tmp_path): + run = pr.init( + mode="offline", + dir=str(tmp_path), + handle_signals=False, + on_error="raise", + ) + run.log_traces([{"id": "bad", "reward": float("nan")}]) + + with pytest.raises(ValueError, match="Out of range float values"): + run.flush() + + path = tmp_path / run.id / "records" / "trace.jsonl" + assert path.read_text() == "" + run.finish() diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index c0307f066..b7bd1151c 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -2,6 +2,7 @@ import signal import threading +from concurrent.futures import ThreadPoolExecutor from typing import Any, Dict, List, Optional import pytest @@ -275,6 +276,24 @@ def test_a_finish_failure_does_not_mask_the_context_exception(): assert backend.closed is True +@pytest.mark.parametrize("teardown_error", [KeyboardInterrupt, SystemExit]) +def test_teardown_does_not_swallow_control_flow_exceptions(teardown_error): + run = make_run() + original_finish = run.finish + + def interrupt_finish(*args, **kwargs): + raise teardown_error() + + run.finish = interrupt_finish + try: + with pytest.raises(teardown_error): + with run: + raise ValueError("rollout blew up") + finally: + run.finish = original_finish + run.finish() + + def test_an_interrupt_is_recorded_as_a_decision_not_a_fault(): """Ctrl-C must not land in the same bucket as a broken eval — and it must agree with the SIGINT handler, which normally gets there first.""" @@ -487,3 +506,37 @@ def test_a_later_run_can_install_its_own_handlers(): assert signal.getsignal(signal.SIGTERM) is second._signal_handler second.finish() assert signal.getsignal(signal.SIGTERM) is signal.SIG_DFL + + +def test_a_run_finished_in_an_executor_relinquishes_handlers_to_the_next_run(): + original = signal.getsignal(signal.SIGTERM) + first = make_run() + first.install_signal_handlers() + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(first.finish).result() + + # Python forbids signal.signal() off the main thread, so restoration is + # deferred rather than forgetting which handler was displaced. + assert signal.getsignal(signal.SIGTERM) is first._signal_handler + second = make_run() + second.install_signal_handlers() + assert signal.getsignal(signal.SIGTERM) is second._signal_handler + + second.finish() + assert signal.getsignal(signal.SIGTERM) is original + + +def test_a_child_run_can_replace_an_inherited_signal_handler(): + original = signal.getsignal(signal.SIGTERM) + inherited = make_run() + inherited.install_signal_handlers() + inherited.reset_after_fork() + + child = make_run() + child.install_signal_handlers() + + assert signal.getsignal(signal.SIGTERM) is child._signal_handler + child.finish() + inherited.finish() + assert signal.getsignal(signal.SIGTERM) is original diff --git a/packages/prime-runs/tests/test_samples_sink.py b/packages/prime-runs/tests/test_samples_sink.py index 1b4c5f0e1..379afde1a 100644 --- a/packages/prime-runs/tests/test_samples_sink.py +++ b/packages/prime-runs/tests/test_samples_sink.py @@ -1,5 +1,6 @@ """The legacy sample sink that keeps today's viewer working.""" +import pytest from _fakes import make_episode, make_trace from conftest import RecordingHandler @@ -56,12 +57,59 @@ def test_a_producer_that_already_speaks_v0_is_passed_through(make_platform_clien assert body["samples"] == [{"sample_id": "s1", "reward": 1.0}] -def test_records_this_sink_cannot_project_are_skipped_not_posted(make_platform_client, eval_routes): - """A malformed row would be rejected for the whole batch, taking the valid - rows with it.""" +def test_serialized_trace_records_are_projected(make_platform_client, eval_routes): sink, handler = make_sink(make_platform_client, eval_routes) - sink.write([{"unrelated": True}]) + sink.write([make_trace(trace_id="serialized-trace", reward=0.75).to_record()]) + + body = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")[0] + assert body["samples"][0]["sample_id"] == "serialized-trace" + assert body["samples"][0]["reward"] == 0.75 + + +def test_serialized_episode_records_keep_the_native_wrapper(make_platform_client, eval_routes): + sink, handler = make_sink(make_platform_client, eval_routes) + record = make_episode("serialized-episode", [make_trace()]).to_record() + + sink.write([record]) + + sample = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")[0]["samples"][0] + assert sample["sample_id"] == "serialized-episode" + assert sample["info"]["native_wrapper"] == record + + +def test_serialized_message_graphs_recover_the_viewer_completion(make_platform_client, eval_routes): + sink, handler = make_sink(make_platform_client, eval_routes) + record = { + "id": "graph-trace", + "task": {"data": {"idx": 7, "answer": "42"}}, + "agent": {"name": "solver", "trainable": True}, + "nodes": [ + {"parent": None, "message": {"role": "user", "content": "6 * 7?"}}, + {"parent": 0, "message": {"role": "assistant", "content": "42"}}, + ], + "calls": [ + { + "node": 1, + "usage": {"prompt_tokens": 4, "completion_tokens": 1}, + } + ], + "rewards": {"correct": {"score": 1.0, "weight": 1.0}}, + } + + sink.write([record]) + + sample = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")[0]["samples"][0] + assert sample["example_id"] == 7 + assert sample["completion"][-1] == {"role": "assistant", "content": "42"} + assert sample["reward"] == 1.0 + + +def test_records_this_sink_cannot_project_fail_explicitly(make_platform_client, eval_routes): + sink, handler = make_sink(make_platform_client, eval_routes) + + with pytest.raises(TypeError, match="non-empty 'id'"): + sink.write([{"unrelated": True}]) assert handler.requests == [] diff --git a/packages/prime-runs/tests/test_traces_sink.py b/packages/prime-runs/tests/test_traces_sink.py index 612e99988..fbbb75b62 100644 --- a/packages/prime-runs/tests/test_traces_sink.py +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -128,6 +128,17 @@ def test_a_disabled_sink_stops_calling_the_service(): assert client.calls == [] +def test_receipt_history_is_bounded_while_the_total_is_retained(): + client = FakeTracesClient() + sink = make_sink(client, receipt_history_size=2) + + for index in range(5): + sink.write([{"id": f"t{index}"}]) + + assert sink.receipts_received == 5 + assert len(sink.receipts) == 2 + + def test_closing_the_sink_closes_the_client(): client = FakeTracesClient() sink = make_sink(client) From 5cf17d8adc53733552c4e15d5eca2a80a3222df9 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 18:03:28 -0700 Subject: [PATCH 13/39] fix(runs): harden shutdown and upload error handling --- packages/prime-runs/src/prime_runs/run.py | 79 +++++++++++-- .../src/prime_runs/sinks/samples.py | 6 +- .../prime-runs/src/prime_runs/sinks/traces.py | 6 +- packages/prime-runs/tests/test_init.py | 56 +++++++++- packages/prime-runs/tests/test_run.py | 105 ++++++++++++++++++ packages/prime-runs/tests/test_traces_sink.py | 5 +- 6 files changed, 241 insertions(+), 16 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 63ab15053..2a8973bd4 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -108,7 +108,15 @@ def __init__( self._pending_metrics: Dict[str, Any] = {} self._pending_metric_step: Optional[int] = None self._finish_lock = threading.RLock() + self._finish_condition = threading.Condition(self._finish_lock) + self._finishing = False + self._finishing_thread_id: Optional[int] = None self._finished = False + # A Python signal handler can interrupt finish() on the same thread. + # Re-entering teardown would duplicate finalization, while chaining the + # signal immediately would kill the process before teardown completes. + # Keep the first such signal and deliver it once the run is closed. + self._pending_signal: Optional[tuple[int, Any, Any]] = None sinks = sinks or [] worker_kwargs: Dict[str, Any] = {} @@ -299,33 +307,68 @@ def finish( ``status`` must be one of the terminal :class:`RunStatus` values. Safe to call from ``__exit__``, an atexit hook and a signal handler at - once — whichever gets there first reports the status, and the rest - return. + once — whichever gets there first reports the status, and the rest wait + for that teardown to complete. """ - with self._finish_lock: + thread_id = threading.get_ident() + with self._finish_condition: + while self._finishing and not self._finished: + # A signal handler can interrupt this very finish() call. It + # cannot wait for itself, so _handle_signal defers chaining the + # signal and this nested call simply yields to the active one. + if self._finishing_thread_id == thread_id: + return + self._finish_condition.wait() if self._finished: return resolved = RunStatus(status) if not isinstance(status, RunStatus) else status if not resolved.is_terminal(): raise ValueError(f"finish() requires a terminal status, got {resolved.value!r}") - self._finished = True + self._finishing = True + self._finishing_thread_id = thread_id + + try: + self._finish_once(summary, resolved, error) + finally: + with self._finish_condition: + self._finishing = False + self._finishing_thread_id = None + self._finished = True + pending_signal = self._pending_signal + self._pending_signal = None + self._finish_condition.notify_all() + + if pending_signal is not None: + self._chain_signal(*pending_signal) + + def _finish_once( + self, + summary: Optional[Mapping[str, Any]], + resolved: RunStatus, + error: Optional[str], + ) -> None: + """Perform the single teardown owned by the first ``finish()`` caller.""" if summary: self.summary.update(_clean_metrics(summary)) self._status = resolved finish_error: Optional[BaseException] = None + deadline = time.monotonic() + max(0.0, self._finish_timeout) + + def remaining_finish_time() -> float: + return max(0.0, deadline - time.monotonic()) # Order matters: records first, so a dashboard that reacts to the # terminal status never sees a finished run with samples still landing. - if not self._worker.flush(timeout=self._finish_timeout): + if not self._worker.flush(timeout=remaining_finish_time()): logger.warning( "Run %s: uploads did not drain within %ss; finalizing anyway. " "Some records may be missing from this run.", self.id, self._finish_timeout, ) - self._worker.close(timeout=self._finish_timeout) + self._worker.close(timeout=remaining_finish_time()) if self._owns_lifecycle: finish_error = self._finish_guarded( @@ -480,7 +523,13 @@ def _handle_signal(self, signum: int, frame: Any) -> None: # SIG_DFL — which re-raises the signal at its default disposition and # kills the process instead of running the handler the app installed. previous = self._previous_signal_handlers.get(signum, signal.SIG_DFL) - if not self._finished: + with self._finish_condition: + if self._finishing and self._finishing_thread_id == threading.get_ident(): + if self._pending_signal is None: + self._pending_signal = (signum, frame, previous) + return + finished = self._finished + if not finished: # CRASHED, not FAILED: the producer never said the run failed, it was # stopped from outside its own control flow. Same bucket as the # atexit path, and deliberately not the bucket a broken eval lands in. @@ -488,6 +537,11 @@ def _handle_signal(self, signum: int, frame: Any) -> None: self.finish(status=RunStatus.CRASHED, error=f"received {name}") except Exception as exc: # noqa: BLE001 - the signal must still chain logger.warning("Run %s: reporting %s failed: %s", self.id, name, exc) + self._chain_signal(signum, frame, previous) + + @staticmethod + def _chain_signal(signum: int, frame: Any, previous: Any) -> None: + """Restore and invoke the handler displaced by this run.""" signal.signal(signum, previous) if callable(previous): previous(signum, frame) @@ -517,6 +571,10 @@ def reset_after_fork(self) -> None: been owned at fork time by a thread that no longer exists. """ self._finish_lock = threading.RLock() + self._finish_condition = threading.Condition(self._finish_lock) + self._finishing = False + self._finishing_thread_id = None + self._pending_signal = None self._is_primary = False self._owns_lifecycle = False self._deferred_error = None @@ -553,7 +611,7 @@ def _retract_run_id(self) -> None: _exported_run_ids.pop(self.id, None) def _require_live(self, operation: str) -> None: - if self._finished: + if self._finishing or self._finished: raise RunFinishedError( f"{operation}() was called on run {self.id}, which is already finished. " "The platform has closed this run out; start a new one." @@ -761,7 +819,10 @@ def init( # Both transports run during the transition: traces is the system of # record, the sample table is what today's viewer reads, and Prime # Traces is still gated to an account allowlist. - run_sinks.append(EvalSamplesSink(client)) + samples_client = PlatformClient( + api_key=api_key, base_url=base_url, timeout=DEFAULT_TIMEOUT + ) + run_sinks.append(EvalSamplesSink(samples_client, close_client=True)) else: run_sinks = list(sinks) diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py index 81cdbbd61..b77f7f2f7 100644 --- a/packages/prime-runs/src/prime_runs/sinks/samples.py +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -30,9 +30,10 @@ class EvalSamplesSink(Sink): name = "eval_samples" - def __init__(self, client: PlatformClient) -> None: + def __init__(self, client: PlatformClient, *, close_client: bool = False) -> None: self.enabled = True self._client = client + self._close_client = close_client self._run_id: Optional[str] = None # Carried across calls so a streaming producer numbers rollouts the same # way a one-shot upload does: the Nth episode for an example is rollout N, @@ -103,4 +104,5 @@ def flush(self) -> None: """Writes are synchronous; the uploader thread owns the asynchrony.""" def close(self) -> None: - """The platform client is shared with the backend, which closes it.""" + if self._close_client: + self._client.close() diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py index c2a3f2dcf..a3bf7b453 100644 --- a/packages/prime-runs/src/prime_runs/sinks/traces.py +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -140,7 +140,11 @@ def write( f"Prime Traces is not enabled for this account ({exc}); " "falling back to the remaining sinks" ) - return + # The worker contains this error in the default warn mode and + # continues with the remaining sinks. Raising is still required + # so strict callers see the failed batch and loss accounting is + # updated instead of reporting a successful traces-only run. + raise raise self.receipts_received += len(receipts) if self._receipt_history_size: diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index ceea3bbb9..a042f8578 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -177,8 +177,9 @@ def online(monkeypatch, make_platform_client, eval_routes): def _init(routes=None, **kwargs): handler = RecordingHandler(routes or eval_routes) - client = make_platform_client(handler) - monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: client) + monkeypatch.setattr( + "prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler) + ) run = pr.init( name="test-run", environments=["gsm8k"], @@ -193,6 +194,57 @@ def _init(routes=None, **kwargs): return _init +def test_samples_use_a_separate_client_from_run_finalization( + monkeypatch, make_platform_client, eval_routes +): + handler = RecordingHandler(eval_routes) + + class TrackingClient: + def __init__(self): + self._delegate = make_platform_client(handler) + self.closed = False + + def get(self, *args, **kwargs): + return self._delegate.get(*args, **kwargs) + + def post(self, *args, **kwargs): + return self._delegate.post(*args, **kwargs) + + def put(self, *args, **kwargs): + return self._delegate.put(*args, **kwargs) + + def close(self): + self.closed = True + + clients = [] + + def make_client(**kwargs): + client = TrackingClient() + clients.append(client) + return client + + monkeypatch.setattr("prime_runs.run.PlatformClient", make_client) + run = pr.init( + name="test-run", + environments=["gsm8k"], + api_key="test-key", + traces=False, + handle_signals=False, + ) + assert len(clients) == 2 + + # Model UploadWorker.close() timing out: it intentionally leaves its sink + # open, while lifecycle finalization still closes the backend transport. + run._worker.close = lambda timeout=None: None + run.finish() + + backend_client, samples_client = clients + assert backend_client.closed is True + assert samples_client.closed is False + run._worker.sinks[0].close() + assert samples_client.closed is True + + def test_an_online_run_returns_the_platforms_id_and_viewer_url(online): run, _ = online() diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index b7bd1151c..dd3d5fc7e 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -2,15 +2,18 @@ import signal import threading +import time from concurrent.futures import ThreadPoolExecutor from typing import Any, Dict, List, Optional import pytest from conftest import FakeSink +from prime_traces.exceptions import ForbiddenError from prime_runs.exceptions import RunFinishedError from prime_runs.models import RunHandle, RunSpec, RunStatus from prime_runs.run import Run +from prime_runs.sinks import TracesSink class FakeBackend: @@ -222,6 +225,45 @@ def test_finish_is_idempotent(): assert backend.finalized[0]["summary"] == {"avg_reward": 1.0} +def test_concurrent_finish_waits_for_the_first_teardown_to_complete(): + class BlockingBackend(FakeBackend): + def __init__(self) -> None: + super().__init__() + self.finalize_started = threading.Event() + self.release_finalize = threading.Event() + + def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> None: + self.finalize_started.set() + assert self.release_finalize.wait(2.0) + super().finalize( + run_id, status=status, summary=summary, error=error, config=config + ) + + backend = BlockingBackend() + run = make_run(backend) + second_started = threading.Event() + + def finish_second() -> None: + second_started.set() + run.finish(status=RunStatus.FAILED) + + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit(run.finish) + assert backend.finalize_started.wait(2.0) + assert run.finished is False + + second = executor.submit(finish_second) + assert second_started.wait(2.0) + assert second.done() is False + + backend.release_finalize.set() + first.result(timeout=2.0) + second.result(timeout=2.0) + + assert run.finished is True + assert [entry["status"] for entry in backend.finalized] == [RunStatus.COMPLETED] + + @pytest.mark.parametrize("status", [RunStatus.RUNNING, "running"]) def test_finish_rejects_a_nonterminal_status_without_closing_the_run(status): backend = FakeBackend() @@ -348,6 +390,27 @@ def test_finish_warns_when_uploads_do_not_drain(caplog): assert "did not drain" in caplog.text +def test_finish_uses_one_timeout_budget_for_flush_and_close(): + run = make_run() + run._finish_timeout = 0.1 + observed = {} + + def slow_flush(timeout=None): + observed["flush"] = timeout + time.sleep(0.02) + return False + + def record_close(timeout=None): + observed["close"] = timeout + + run._worker.flush = slow_flush + run._worker.close = record_close + + run.finish() + + assert 0.0 <= observed["close"] < observed["flush"] <= run._finish_timeout + + def test_a_process_that_exits_without_finishing_reports_crashed(): """The producer never said the run failed — it stopped existing. The distinction tells an operator where to look.""" @@ -480,6 +543,48 @@ def test_an_upload_failure_is_reported_once(): run.finish() +def test_a_gated_trace_upload_reaches_strict_callers_and_loss_accounting(): + class GatedClient: + def upload_records(self, records, **kwargs): + raise ForbiddenError( + "not in beta", status_code=403, code="service_not_enabled" + ) + + def close(self) -> None: + pass + + run = make_run(sinks=[TracesSink(client=GatedClient())], on_error="raise") + run.log_traces([{"id": "t1"}]) + + with pytest.raises(ForbiddenError, match="not in beta"): + run.finish() + + assert run.failed_records == {"traces": 1} + + +def test_a_signal_interrupting_finish_is_chained_after_teardown(monkeypatch): + events = [] + + class SignallingBackend(FakeBackend): + def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> None: + events.append("finalize-start") + run._handle_signal(signal.SIGTERM, None) + events.append("finalize-end") + super().finalize( + run_id, status=status, summary=summary, error=error, config=config + ) + + backend = SignallingBackend() + run = make_run(backend) + run._previous_signal_handlers[signal.SIGTERM] = lambda *_: events.append("signal") + monkeypatch.setattr("prime_runs.run.signal.signal", lambda *_: None) + + run.finish() + + assert events == ["finalize-start", "finalize-end", "signal"] + assert backend.closed is True + + def test_signal_handlers_are_restored_when_the_run_finishes(): """`self._handle_signal` builds a new bound method on every access, so an identity check against a fresh one never matches — leaving the handler diff --git a/packages/prime-runs/tests/test_traces_sink.py b/packages/prime-runs/tests/test_traces_sink.py index fbbb75b62..b31b9cb75 100644 --- a/packages/prime-runs/tests/test_traces_sink.py +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -93,7 +93,7 @@ def test_producer_objects_are_passed_through_untouched(): assert client.calls[0][0][0] is trace -def test_a_gated_account_disables_the_sink_instead_of_failing_the_run(caplog): +def test_a_gated_account_disables_the_sink_and_reports_the_failed_batch(caplog): """Prime Traces is in closed beta; no runtime action fixes a 403, so retrying it for the rest of the run only produces noise.""" client = FakeTracesClient( @@ -102,7 +102,8 @@ def test_a_gated_account_disables_the_sink_instead_of_failing_the_run(caplog): sink = make_sink(client) with caplog.at_level("WARNING"): - sink.write([{"id": "t1"}]) + with pytest.raises(ForbiddenError, match="not in beta"): + sink.write([{"id": "t1"}]) assert sink.enabled is False assert "not enabled" in caplog.text From 7288b41e3719e313931d2e41044dc5f0ef0022ad Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 18:20:20 -0700 Subject: [PATCH 14/39] fix: preserve resumed state and support environment slugs --- .../src/prime_runs/backends/evals.py | 16 +++++- packages/prime-runs/src/prime_runs/models.py | 15 ++++-- packages/prime-runs/src/prime_runs/run.py | 52 +++++++++++++++++-- .../prime-runs/tests/test_evals_backend.py | 17 ++++++ packages/prime-runs/tests/test_init.py | 38 ++++++++++++++ 5 files changed, 130 insertions(+), 8 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/backends/evals.py b/packages/prime-runs/src/prime_runs/backends/evals.py index 8e393e76b..bd2259787 100644 --- a/packages/prime-runs/src/prime_runs/backends/evals.py +++ b/packages/prime-runs/src/prime_runs/backends/evals.py @@ -265,7 +265,21 @@ def _resolve_environments(self, refs: List[EnvironmentRef]) -> List[Dict[str, An return resolved def _lookup_environment(self, ref: EnvironmentRef) -> str: - """Resolve one environment name to a hub ID (get-or-create).""" + """Resolve one environment reference to a hub ID.""" + if ref.slug: + owner_slug, name = ref.slug.split("/", 1) + try: + response = self._client.get(f"/environmentshub/{owner_slug}/{name}/@latest") + except RunAPIError as exc: + raise EnvironmentResolutionError( + f"Could not resolve environment {ref.slug!r}: {exc}" + ) from exc + details = response.get("data") or response + environment_id = details.get("id") + if not environment_id: + raise EnvironmentResolutionError(f"Hub returned no id for environment {ref.slug!r}") + return str(environment_id) + body: Dict[str, Any] = {"name": ref.name} _set_if(body, "team_id", self._team_id) try: diff --git a/packages/prime-runs/src/prime_runs/models.py b/packages/prime-runs/src/prime_runs/models.py index e783f140a..ea724b0b5 100644 --- a/packages/prime-runs/src/prime_runs/models.py +++ b/packages/prime-runs/src/prime_runs/models.py @@ -46,22 +46,25 @@ class EnvironmentRef: """An environment as a producer names it, before hub resolution. ``id`` short-circuits resolution; ``name`` goes through the hub's - get-or-create so a local run uploads without a prior ``prime env push``. + get-or-create so a local run uploads without a prior ``prime env push``; + ``slug`` looks up an already-published ``owner/name`` environment. """ name: Optional[str] = None id: Optional[str] = None version_id: Optional[str] = None + slug: Optional[str] = None @classmethod def coerce(cls, value: Any) -> "EnvironmentRef": if isinstance(value, EnvironmentRef): return value if isinstance(value, str): - return cls(name=value) + return cls(slug=value) if "/" in value else cls(name=value) if isinstance(value, dict): return cls( name=value.get("name"), + slug=value.get("slug"), id=value.get("id"), version_id=value.get("version_id"), ) @@ -71,8 +74,12 @@ def coerce(cls, value: Any) -> "EnvironmentRef": ) def __post_init__(self) -> None: - if not self.name and not self.id: - raise ValueError("EnvironmentRef needs a name or an id") + if not self.name and not self.slug and not self.id: + raise ValueError("EnvironmentRef needs a name, slug or id") + if self.slug: + owner, name = self.slug.split("/", 1) if "/" in self.slug else ("", "") + if not owner or not name: + raise ValueError("EnvironmentRef slug must use owner/name format") @dataclass diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 2a8973bd4..7e4a99b32 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -91,8 +91,9 @@ def __init__( self._owns_lifecycle = owns_lifecycle and is_primary self._status = RunStatus.RUNNING - self.config: Dict[str, Any] = dict(spec.config) - self.summary: Dict[str, Any] = dict(spec.summary) + attached_config, attached_summary = _attached_state(handle) + self.config: Dict[str, Any] = {**attached_config, **spec.config} + self.summary: Dict[str, Any] = {**attached_summary, **spec.summary} self.errors: List[str] = [] # Raised at the next synchronization point the caller controls. A sink # fails on the uploader thread, where raising reaches nobody — so under @@ -726,6 +727,7 @@ def init( on_error: OnError = "warn", handle_signals: bool = True, queue_size: Optional[int] = None, + finish_timeout: float = DEFAULT_FINISH_TIMEOUT, ) -> Run: """Start a run and return a handle to it. @@ -741,6 +743,9 @@ def init( ``id`` attaches to an existing run instead of creating one, for resuming after a crash and for non-primary ranks joining a run rank 0 created. + + ``finish_timeout`` is the total number of seconds ``finish()`` gives queued + uploads to drain and close before finalizing the run anyway. """ resolved_config = Config() api_key = api_key if api_key is not None else resolved_config.api_key @@ -775,7 +780,16 @@ def init( backend: Backend = _DisabledBackend() handle = RunHandle(id=inherited_id or _local_id(), name=name) return _build( - spec, backend, handle, [], resolved_mode, on_error, is_primary, False, queue_size + spec, + backend, + handle, + [], + resolved_mode, + on_error, + is_primary, + False, + queue_size, + finish_timeout, ) if resolved_mode == "offline": @@ -792,6 +806,7 @@ def init( is_primary, owns_lifecycle, queue_size, + finish_timeout, ) _announce(run, handle_signals) return run @@ -836,6 +851,7 @@ def init( is_primary, owns_lifecycle, queue_size, + finish_timeout, ) _announce(run, handle_signals) return run @@ -851,6 +867,7 @@ def _build( is_primary: bool, owns_lifecycle: bool, queue_size: Optional[int], + finish_timeout: float, ) -> Run: return Run( backend=backend, @@ -862,6 +879,7 @@ def _build( is_primary=is_primary, owns_lifecycle=owns_lifecycle, queue_size=queue_size, + finish_timeout=finish_timeout, ) @@ -985,6 +1003,34 @@ def _sink_context(spec: RunSpec, handle: RunHandle) -> Dict[str, str]: return context +def _attached_state(handle: RunHandle) -> tuple[Dict[str, Any], Dict[str, Any]]: + """Recover config and summary from an attached backend response. + + Online evaluations expose them as ``metadata`` and ``metrics``. Offline + archives keep their initial values under ``spec`` and subsequent values at + the top level. Supporting both shapes keeps resume lossless for either + backend; values supplied to the new ``init()`` call are merged afterwards + and therefore win. + """ + config: Dict[str, Any] = {} + summary: Dict[str, Any] = {} + raw = handle.raw + nested_spec = raw.get("spec") + if isinstance(nested_spec, Mapping): + _merge_mapping(config, nested_spec.get("config")) + _merge_mapping(summary, nested_spec.get("summary")) + _merge_mapping(config, raw.get("config")) + _merge_mapping(config, raw.get("metadata")) + _merge_mapping(summary, raw.get("summary")) + _merge_mapping(summary, raw.get("metrics")) + return config, summary + + +def _merge_mapping(target: Dict[str, Any], value: Any) -> None: + if isinstance(value, Mapping): + target.update(value) + + def _describe(error: Union[str, BaseException]) -> str: if isinstance(error, BaseException): return f"{type(error).__name__}: {error}" diff --git a/packages/prime-runs/tests/test_evals_backend.py b/packages/prime-runs/tests/test_evals_backend.py index 9cb3edece..fb696d9b8 100644 --- a/packages/prime-runs/tests/test_evals_backend.py +++ b/packages/prime-runs/tests/test_evals_backend.py @@ -54,6 +54,23 @@ def test_an_explicit_environment_id_skips_the_hub(make_platform_client, eval_rou assert handler.bodies_for("/api/v1/evaluations/")[0]["environments"] == [{"id": "env-999"}] +@pytest.mark.parametrize("environment", ["alice/gsm8k", {"slug": "alice/gsm8k"}]) +def test_a_published_environment_slug_uses_owner_aware_lookup( + make_platform_client, eval_routes, environment +): + routes = dict(eval_routes) + routes["GET /api/v1/environmentshub/alice/gsm8k/@latest"] = {"data": {"id": "env-published"}} + backend, handler = make_backend(make_platform_client, routes) + + backend.create(RunSpec(name="r", environments=[EnvironmentRef.coerce(environment)])) + + assert handler.paths()[0] == "GET /api/v1/environmentshub/alice/gsm8k/@latest" + assert "POST /api/v1/environmentshub/resolve" not in handler.paths() + assert handler.bodies_for("/api/v1/evaluations/")[0]["environments"] == [ + {"id": "env-published"} + ] + + def test_an_unresolvable_environment_fails_the_run_rather_than_being_skipped( make_platform_client, eval_routes ): diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index a042f8578..89ef3d9f0 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -342,6 +342,44 @@ def test_an_explicit_id_is_a_resume_and_still_finalizes(online): assert "POST /api/v1/evaluations/eval-abc/finalize" in handler.paths() +def test_resuming_preserves_existing_config_and_summary(online, eval_routes): + routes = dict(eval_routes) + routes["GET /api/v1/evaluations/eval-abc"] = { + **routes["GET /api/v1/evaluations/eval-abc"], + "metadata": {"before_crash": True, "overridden": "old"}, + "metrics": {"old_reward": 0.5, "overridden": "old"}, + } + + run, handler = online( + routes=routes, + id="eval-abc", + config={"overridden": "new"}, + summary={"overridden": "new"}, + ) + run.update_config({"after_resume": True}) + run.log({"new_reward": 1.0}) + run.finish() + + update = handler.bodies_for("/api/v1/evaluations/eval-abc")[0] + assert update["metadata"] == { + "before_crash": True, + "overridden": "new", + "after_resume": True, + } + assert update["metrics"] == { + "old_reward": 0.5, + "overridden": "new", + "new_reward": 1.0, + } + + +def test_init_forwards_the_finish_timeout(tmp_path): + run = pr.init(mode="offline", dir=str(tmp_path), finish_timeout=0.25) + + assert run._finish_timeout == 0.25 + run.finish() + + def test_an_id_inherited_from_the_environment_does_not_finalize(monkeypatch, online): monkeypatch.setenv(RUN_ID_ENV, "eval-abc") From 4531655544072837b070f97b79e5acfcef93a0c3 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 23:28:59 -0700 Subject: [PATCH 15/39] fix: harden startup and environment metadata --- .../src/prime_runs/backends/evals.py | 5 + packages/prime-runs/src/prime_runs/run.py | 105 ++++++++++++------ .../prime-runs/src/prime_runs/sinks/traces.py | 10 +- .../prime-runs/tests/test_evals_backend.py | 16 +++ packages/prime-runs/tests/test_init.py | 30 ++++- packages/prime-runs/tests/test_run.py | 23 ++++ packages/prime-runs/tests/test_traces_sink.py | 9 +- 7 files changed, 159 insertions(+), 39 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/backends/evals.py b/packages/prime-runs/src/prime_runs/backends/evals.py index bd2259787..1a0c031b4 100644 --- a/packages/prime-runs/src/prime_runs/backends/evals.py +++ b/packages/prime-runs/src/prime_runs/backends/evals.py @@ -306,6 +306,11 @@ def _first_environment_name(spec: RunSpec) -> Optional[str]: for ref in spec.environments: if ref.name: return ref.name + if ref.slug: + # Published environments are addressed as ``owner/name``, but the + # name portion is still the dataset/run-name fallback used for + # ordinary environment references. + return ref.slug.split("/", 1)[1] return None diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 7e4a99b32..deab24085 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -27,6 +27,7 @@ from . import _fork from ._http import DEFAULT_TIMEOUT, UPLOAD_TIMEOUT, PlatformClient from .backends import Backend, EvalsBackend, OfflineBackend +from .backends.offline import DEFAULT_DIR_ENV from .config import Config from .exceptions import ConfigurationError, RunFinishedError from .models import EnvironmentRef, Mode, OnError, RunHandle, RunKind, RunSpec, RunStatus @@ -118,6 +119,29 @@ def __init__( # signal immediately would kill the process before teardown completes. # Keep the first such signal and deliver it once the run is closed. self._pending_signal: Optional[tuple[int, Any, Any]] = None + # ``_announce`` publishes enough context for child processes to join + # this exact run. Keep the previous values so finishing a nested or + # sequential run restores the caller's environment instead of leaking + # our resolved mode and offline directory into the next run. + self._published_join_env: Dict[str, str] = {} + self._previous_join_env: Dict[str, Optional[str]] = {} + + self._atexit_hook = self._on_process_exit + # Bound once and kept. ``self._handle_signal`` builds a *new* bound + # method on every attribute access, so an ``is`` comparison against a + # freshly-made one is always False — which is how handlers end up + # installed forever, pinning a finished run and blocking the next run in + # the process from installing its own. + self._signal_handler = self._handle_signal + self._previous_signal_handlers: Dict[int, Any] = {} + # ``signal.signal`` can only run on the main thread. If finish() runs in + # an executor, or this object is inherited across a fork, its handler + # may remain as the process disposition until the main thread gets a + # chance to replace it. Marking forked handlers as relinquishable lets + # a child run take ownership without mistaking the inherited callback + # for an application-installed handler. + self._signal_handler_stale = False + _fork.register(self) sinks = sinks or [] worker_kwargs: Dict[str, Any] = {} @@ -136,25 +160,25 @@ def __init__( sink.start(handle.id, context) except Exception as exc: # noqa: BLE001 - a bad sink is not a bad run sink.enabled = False - self._report(f"starting sink {getattr(sink, 'name', sink)}", exc) + try: + self._report(f"starting sink {getattr(sink, 'name', sink)}", exc) + except Exception: + # The backend may already have created a remote run. In + # strict mode the caller never receives this handle, so + # close it out before propagating the startup failure. + try: + self.finish(status=RunStatus.FAILED, error=_describe(exc)) + except Exception as cleanup_exc: # noqa: BLE001 - preserve the cause + logger.warning( + "Run %s: cleanup after sink startup failure also failed: %s: %s", + self.id, + type(cleanup_exc).__name__, + cleanup_exc, + exc_info=True, + ) + raise - self._atexit_hook = self._on_process_exit atexit.register(self._atexit_hook) - # Bound once and kept. ``self._handle_signal`` builds a *new* bound - # method on every attribute access, so an ``is`` comparison against a - # freshly-made one is always False — which is how handlers end up - # installed forever, pinning a finished run and blocking the next run in - # the process from installing its own. - self._signal_handler = self._handle_signal - self._previous_signal_handlers: Dict[int, Any] = {} - # ``signal.signal`` can only run on the main thread. If finish() runs in - # an executor, or this object is inherited across a fork, its handler - # may remain as the process disposition until the main thread gets a - # chance to replace it. Marking forked handlers as relinquishable lets - # a child run take ownership without mistaking the inherited callback - # for an application-installed handler. - self._signal_handler_stale = False - _fork.register(self) # -------------------------------------------------------------- identity @@ -398,7 +422,7 @@ def remaining_finish_time() -> float: atexit.unregister(self._atexit_hook) self._restore_signal_handlers() - self._retract_run_id() + self._retract_join_context() if self._worker.dropped: logger.warning( @@ -598,17 +622,39 @@ def _on_process_exit(self) -> None: except Exception as exc: # noqa: BLE001 logger.warning("Run %s: reporting the crash failed: %s", self.id, exc) - def _retract_run_id(self) -> None: + def _publish_join_context(self) -> None: + """Publish the resolved context a child needs to join this run.""" + values = { + RUN_ID_ENV: self.id, + MODE_ENV: self.mode, + } + if self.mode == "offline" and isinstance(self._backend, OfflineBackend): + # Resolve the path because a subprocess may use a different cwd. + values[DEFAULT_DIR_ENV] = str(self._backend.directory.resolve()) + + for name, value in values.items(): + self._previous_join_env.setdefault(name, os.environ.get(name)) + self._published_join_env[name] = value + os.environ[name] = value + _exported_run_ids[self.id] = os.getpid() + + def _retract_join_context(self) -> None: """Stop advertising a finished run to processes started from here. - Only retracts what this run published: if the value now points somewhere - else, another run owns it and clearing it would orphan that one's + Only restores values this run still owns: if one now points somewhere + else, another run owns it and changing it would orphan that one's children. """ if not self._owns_lifecycle: return - if os.environ.get(RUN_ID_ENV) == self.id: - os.environ.pop(RUN_ID_ENV, None) + for name, published in self._published_join_env.items(): + if os.environ.get(name) != published: + continue + previous = self._previous_join_env.get(name) + if previous is None: + os.environ.pop(name, None) + else: + os.environ[name] = previous _exported_run_ids.pop(self.id, None) def _require_live(self, operation: str) -> None: @@ -884,16 +930,13 @@ def _build( def _announce(run: Run, handle_signals: bool) -> None: - """Publish the run ID to child processes and arm crash reporting. + """Publish the run's join context to child processes and arm crash reporting. - Exporting ``PRIME_RUN_ID`` is how forked workers and subprocess launchers - join the run their parent created instead of each opening their own — the - same trick prime-rl's monitor used with ``RUN_ID``, generalized so every - producer gets it. The PID is recorded alongside so that *this* process does - not later mistake its own export for a parent's. + The ID says which run to attach to; the resolved mode and offline directory + say where it lives. The PID is recorded alongside so that *this* process + does not later mistake its own export for a parent's. """ - os.environ[RUN_ID_ENV] = run.id - _exported_run_ids[run.id] = os.getpid() + run._publish_join_context() if handle_signals: run.install_signal_handlers() if run.url: diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py index a3bf7b453..c5ef3acfe 100644 --- a/packages/prime-runs/src/prime_runs/sinks/traces.py +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -84,17 +84,19 @@ def _ensure_client(self) -> bool: if self._injected_client: # The caller handed us a client and a fork took it away. Rebuilding # would silently swap their transport for a default one. - return False + exc = RuntimeError("an injected traces client cannot be reused after a fork") + self._disable(str(exc)) + raise exc try: from prime_traces import TracesClient except ImportError as exc: # pragma: no cover - dependency is declared self._disable(f"prime-traces is not installed ({exc})") - return False + raise try: self._client = TracesClient(**self._client_kwargs) - except Exception as exc: # noqa: BLE001 - construction must not kill a run + except Exception as exc: # noqa: BLE001 - the run applies its error policy self._disable(f"could not construct the traces client ({exc})") - return False + raise return True def reset_after_fork(self) -> None: diff --git a/packages/prime-runs/tests/test_evals_backend.py b/packages/prime-runs/tests/test_evals_backend.py index fb696d9b8..09e1d735c 100644 --- a/packages/prime-runs/tests/test_evals_backend.py +++ b/packages/prime-runs/tests/test_evals_backend.py @@ -71,6 +71,22 @@ def test_a_published_environment_slug_uses_owner_aware_lookup( ] +def test_a_published_environment_slug_supplies_dataset_and_default_name( + make_platform_client, eval_routes +): + routes = dict(eval_routes) + routes["GET /api/v1/environmentshub/alice/gsm8k/@latest"] = { + "data": {"id": "env-published"} + } + backend, handler = make_backend(make_platform_client, routes) + + backend.create(RunSpec(environments=[EnvironmentRef.coerce("alice/gsm8k")])) + + created = handler.bodies_for("/api/v1/evaluations/")[0] + assert created["dataset"] == "gsm8k" + assert created["name"].startswith("gsm8k-") + + def test_an_unresolvable_environment_fails_the_run_rather_than_being_skipped( make_platform_client, eval_routes ): diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index 89ef3d9f0..554f612ff 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -8,9 +8,10 @@ from conftest import RecordingHandler import prime_runs as pr +from prime_runs.backends.offline import DEFAULT_DIR_ENV from prime_runs.exceptions import ConfigurationError from prime_runs.models import RunStatus -from prime_runs.run import RUN_ID_ENV +from prime_runs.run import MODE_ENV, RUN_ID_ENV, _exported_run_ids # ------------------------------------------------------------------- offline @@ -168,6 +169,33 @@ def test_init_publishes_the_run_id_for_child_processes(tmp_path): run.finish() +def test_a_child_inherits_the_resolved_offline_mode_and_directory(monkeypatch, tmp_path): + """An API key must not make the child switch an explicit offline parent online.""" + monkeypatch.setenv("PRIME_API_KEY", "test-key") + monkeypatch.setattr( + "prime_runs.run.PlatformClient", + lambda **_: pytest.fail("the inherited child unexpectedly selected online mode"), + ) + parent = pr.init(mode="offline", dir=str(tmp_path)) + + assert os.environ[MODE_ENV] == "offline" + assert os.environ[DEFAULT_DIR_ENV] == str(tmp_path.resolve()) + + # Emulate the PID distinction a real child process inherits. + _exported_run_ids[parent.id] = os.getpid() - 1 + child = pr.init(handle_signals=False) + child.log_traces([{"id": "from-child"}]) + child.finish() + parent.finish() + + assert child.id == parent.id + assert child.mode == "offline" + assert (tmp_path / parent.id / "records" / "trace.jsonl").exists() + assert RUN_ID_ENV not in os.environ + assert MODE_ENV not in os.environ + assert DEFAULT_DIR_ENV not in os.environ + + # ------------------------------------------------------------------- online diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index dd3d5fc7e..b43d4b24a 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -455,6 +455,29 @@ def test_on_error_raise_surfaces_the_failure_for_tests_and_ci(): assert backend.closed is True +def test_strict_sink_start_failure_closes_the_created_run(): + class StartFailingSink(FakeSink): + def start(self, run_id, context) -> None: + raise RuntimeError("could not start sink") + + backend = FakeBackend() + sink = StartFailingSink() + + with pytest.raises(RuntimeError, match="could not start sink"): + make_run(backend, sinks=[sink], on_error="raise") + + assert backend.finalized == [ + { + "status": RunStatus.FAILED, + "summary": None, + "error": "RuntimeError: could not start sink", + "config": None, + } + ] + assert backend.closed is True + assert sink.closed is True + + def test_update_failure_in_raise_mode_still_finalizes_and_closes(): backend = FakeBackend(fail_on="update") run = make_run(backend, on_error="raise") diff --git a/packages/prime-runs/tests/test_traces_sink.py b/packages/prime-runs/tests/test_traces_sink.py index b31b9cb75..f5948d7a8 100644 --- a/packages/prime-runs/tests/test_traces_sink.py +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -149,8 +149,10 @@ def test_closing_the_sink_closes_the_client(): assert client.closed is True -def test_a_missing_traces_client_disables_the_sink_rather_than_raising(monkeypatch, caplog): - """Construction failures must not take down a run that has not started.""" +def test_a_missing_traces_client_disables_the_sink_and_reports_the_failure( + monkeypatch, caplog +): + """The run applies warn/raise policy, so the sink must surface this failure.""" def explode(**kwargs): raise RuntimeError("no credentials") @@ -159,6 +161,7 @@ def explode(**kwargs): sink = TracesSink() with caplog.at_level("WARNING"): - sink.start("run-1", {"run_kind": "eval"}) + with pytest.raises(RuntimeError, match="no credentials"): + sink.start("run-1", {"run_kind": "eval"}) assert sink.enabled is False From d2a97a6d7291f0151a6df9383a39df8fa3a907e8 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 20 Aug 2026 10:58:02 -0700 Subject: [PATCH 16/39] feat(runs): record the config a run was actually configured with MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The platform's Config tab is unusable for both run kinds, and neither cause is the platform's: eval runs store a v0 blob of four keys, training runs store a fully-resolved dump where three chosen values sit under hundreds of defaults nobody picked. Both producers are now launched from one user-authored file — `uv run eval @ eval.toml`, `uv run rl @ train.toml` — and nothing captured it. Two changes, matching the two failures: `config_source=` takes the path to that file and stores it byte for byte, comments and section grouping intact. It rides inside the run's config under a reserved key, so every write that already carries the config carries the source too — create, the periodic update, finalize, the failure fallback, and the offline archive — with no extra plumbing and no chance of one path forgetting it. A str or Path is always a path, never inline text: guessing between them would turn a mistyped filename into a run whose config tab displays the filename. `config=` now accepts a pydantic model and dumps it with `exclude_unset=True`, so only fields somebody actually set are recorded. A mapping is still stored exactly as given, and a caller who wants every resolved default can still pass `cfg.model_dump()`. The asymmetry is deliberate — the shorter call should give the more useful answer. Note this is only the upload half. Both Config tabs are derived projections today (the eval tab filters through a 31-key allowlist, the training tab reconstructs TOML from stored fields), so rendering `metadata.config_source` verbatim is a separate frontend change. Co-Authored-By: Claude Opus 5 (1M context) --- packages/prime-runs/README.md | 39 +++ .../prime-runs/src/prime_runs/__init__.py | 4 + packages/prime-runs/src/prime_runs/models.py | 134 +++++++++- packages/prime-runs/src/prime_runs/run.py | 71 +++++- .../prime-runs/tests/test_config_source.py | 232 ++++++++++++++++++ 5 files changed, 476 insertions(+), 4 deletions(-) create mode 100644 packages/prime-runs/tests/test_config_source.py diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index c465572d4..ae839239b 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -59,6 +59,45 @@ than a scan over upload metadata. `init()` also exports `PRIME_RUN_ID`, so forked workers and subprocess launchers join the run their parent opened instead of each opening their own. +## Config + +A run records its configuration two ways, because they answer different questions. + +```python +run = pr.init( + ..., + config=cfg, # structured — queryable + config_source="eval.toml", # verbatim — readable +) +``` + +**`config_source`** is the file the run was launched from — `uv run eval @ +eval.toml`, `uv run rl @ train.toml` — stored byte for byte, comments and section +grouping intact. That file *is* the run's real configuration, and it is the thing +worth putting in front of a person. Pass a path (a `str` or `Path` is always a +path, never inline text); pass `pr.ConfigSource(text=..., format=...)` if you +already hold it in memory. It lands under the reserved `config_source` key inside +the run's config, so every write that carries the config carries it too — +including the offline archive. + +**`config`** is the structured form. A mapping is stored exactly as given. A +pydantic model is dumped with `exclude_unset=True`, so only fields somebody +actually set are recorded: + +| you pass | you get | +| --- | --- | +| `config={"n": 4}` | `{"n": 4}` | +| `config=cfg` (a model) | only the fields set on `cfg` | +| `config=cfg.model_dump()` | every field, defaults included | + +A resolved dump of a deep config tree is hundreds of lines nobody chose, and a +reader scrolling it cannot tell which three values were the experiment. The +shorter call gives the more useful answer; the full dump is still available to +anyone who explicitly asks for it. + +Nothing is redacted. A config file holding a secret puts that secret on the run's +page — keep credentials in the environment, not in the file. + ## Modes | mode | what happens | diff --git a/packages/prime-runs/src/prime_runs/__init__.py b/packages/prime-runs/src/prime_runs/__init__.py index 3d0b00830..acd2d844a 100644 --- a/packages/prime-runs/src/prime_runs/__init__.py +++ b/packages/prime-runs/src/prime_runs/__init__.py @@ -41,6 +41,8 @@ UnauthorizedError, ) from .models import ( + CONFIG_SOURCE_KEY, + ConfigSource, EnvironmentRef, Mode, OnError, @@ -63,6 +65,8 @@ "projection", # Types "Config", + "ConfigSource", + "CONFIG_SOURCE_KEY", "EnvironmentRef", "Mode", "OnError", diff --git a/packages/prime-runs/src/prime_runs/models.py b/packages/prime-runs/src/prime_runs/models.py index ea724b0b5..4ec361a23 100644 --- a/packages/prime-runs/src/prime_runs/models.py +++ b/packages/prime-runs/src/prime_runs/models.py @@ -7,9 +7,13 @@ back a ``RunHandle``. """ +import os from dataclasses import dataclass, field from enum import Enum -from typing import Any, Dict, List, Literal, Optional +from pathlib import Path +from typing import Any, Dict, List, Literal, Mapping, Optional, Union + +from .exceptions import ConfigurationError RunKind = Literal["eval", "train"] """Which run system owns the lifecycle. Selects the backend.""" @@ -82,6 +86,134 @@ def __post_init__(self) -> None: raise ValueError("EnvironmentRef slug must use owner/name format") +CONFIG_SOURCE_KEY = "config_source" +"""Reserved key inside a run's config holding :class:`ConfigSource` as a dict. + +It rides *inside* the config rather than beside it so that every path which +already carries the config carries the source too — create, the periodic update, +finalize, the failure fallback, and the offline archive — with no extra +plumbing and no chance of one of them forgetting it. +""" + +MAX_CONFIG_SOURCE_BYTES = 256 * 1024 +"""Ceiling on a stored config file. A hand-written run config is single-digit +kilobytes; anything past this is a dataset or a log that would bloat the run's +metadata document, so it is refused loudly at ``init()`` rather than truncated.""" + +_CONFIG_SOURCE_FORMATS = { + ".toml": "toml", + ".json": "json", + ".yaml": "yaml", + ".yml": "yaml", +} + + +@dataclass +class ConfigSource: + """The config file a run was started from, kept byte-for-byte. + + Both producers are now launched from one user-authored file — ``uv run eval + @ eval.toml``, ``uv run rl @ train.toml`` — and that file *is* the run's + real configuration. A resolved model dump is a different artifact: it + answers "what did every knob end up as", not "what did someone write", and + it loses comments, key order and section grouping on the way through. + + So this is stored verbatim, next to (not instead of) the structured config. + The structured form stays queryable; this form stays readable. + + Nothing here is redacted. A config file that carries a secret will carry it + onto the run's page, the same way it already reaches anyone who can read the + repository it lives in — keep credentials in the environment, not the file. + """ + + text: str + format: str = "toml" + filename: Optional[str] = None + + def to_dict(self) -> Dict[str, Any]: + data: Dict[str, Any] = {"format": self.format, "text": self.text} + if self.filename: + data["filename"] = self.filename + return data + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> Optional["ConfigSource"]: + """Rebuild from stored metadata. ``None`` if the mapping is not one of ours.""" + text = value.get("text") + if not isinstance(text, str): + return None + raw_format = value.get("format") + raw_filename = value.get("filename") + return cls( + text=text, + format=str(raw_format) if raw_format else "toml", + filename=str(raw_filename) if raw_filename else None, + ) + + @classmethod + def from_file(cls, path: Union[str, "os.PathLike[str]"]) -> "ConfigSource": + """Read a config file, inferring its format from the suffix.""" + resolved = Path(path) + try: + raw = resolved.read_bytes() + except FileNotFoundError as exc: + raise ConfigurationError( + f"config_source={str(path)!r} does not exist. Pass the path to the file the " + "run was started from, or a ConfigSource(text=...) if it is already in memory." + ) from exc + except OSError as exc: + raise ConfigurationError( + f"config_source={str(path)!r} could not be read: {exc}" + ) from exc + if len(raw) > MAX_CONFIG_SOURCE_BYTES: + raise ConfigurationError( + f"config_source={str(path)!r} is {len(raw)} bytes, over the " + f"{MAX_CONFIG_SOURCE_BYTES}-byte limit for a stored run config." + ) + try: + text = raw.decode("utf-8") + except UnicodeDecodeError as exc: + raise ConfigurationError( + f"config_source={str(path)!r} is not UTF-8 text; a run config must be readable." + ) from exc + return cls( + text=text, + format=_CONFIG_SOURCE_FORMATS.get(resolved.suffix.lower(), "text"), + filename=resolved.name, + ) + + @classmethod + def coerce(cls, value: Any) -> Optional["ConfigSource"]: + """Normalize whatever ``init(config_source=...)`` was given. + + A ``str`` or ``PathLike`` is a *path*, never inline text — that is how + every caller will reach for it, and guessing between the two would turn + a mistyped filename into a run whose config tab shows the filename. + Inline text goes through ``ConfigSource(text=...)`` explicitly. + """ + if value is None or isinstance(value, cls): + return value + if isinstance(value, Mapping): + source = cls.from_mapping(value) + if source is None: + raise ValueError("config_source mapping must contain a 'text' string") + return source + if isinstance(value, (str, os.PathLike)): + return cls.from_file(value) + raise TypeError( + f"config_source must be a path, a ConfigSource or a mapping, got {type(value).__name__}" + ) + + def __post_init__(self) -> None: + if not isinstance(self.text, str): + raise TypeError("ConfigSource.text must be a string") + if len(self.text.encode("utf-8")) > MAX_CONFIG_SOURCE_BYTES: + raise ConfigurationError( + f"config_source is over the {MAX_CONFIG_SOURCE_BYTES}-byte limit " + "for a stored run config." + ) + + @dataclass class RunSpec: """Everything a backend needs to open a run, in producer vocabulary. diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index deab24085..27a34e9c9 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -30,7 +30,17 @@ from .backends.offline import DEFAULT_DIR_ENV from .config import Config from .exceptions import ConfigurationError, RunFinishedError -from .models import EnvironmentRef, Mode, OnError, RunHandle, RunKind, RunSpec, RunStatus +from .models import ( + CONFIG_SOURCE_KEY, + ConfigSource, + EnvironmentRef, + Mode, + OnError, + RunHandle, + RunKind, + RunSpec, + RunStatus, +) from .sinks import EvalSamplesSink, OfflineSink, Sink, TracesSink from .worker import MetricItem, RunUpdateItem, UploadWorker, WriteItem @@ -200,6 +210,16 @@ def url(self) -> Optional[str]: def kind(self) -> RunKind: return self._spec.kind + @property + def config_source(self) -> Optional[ConfigSource]: + """The config file this run was launched from, if one was given. + + Read back out of ``config`` rather than cached, so a resumed run reports + the source recovered from the platform and not an empty one. + """ + raw = self.config.get(CONFIG_SOURCE_KEY) + return ConfigSource.from_mapping(raw) if isinstance(raw, Mapping) else None + @property def mode(self) -> Mode: return self._mode @@ -758,7 +778,8 @@ def init( dataset: Optional[str] = None, description: Optional[str] = None, tags: Optional[Sequence[str]] = None, - config: Optional[Mapping[str, Any]] = None, + config: Optional[Any] = None, + config_source: Optional[Any] = None, summary: Optional[Mapping[str, Any]] = None, id: Optional[str] = None, mode: Optional[Mode] = None, @@ -790,6 +811,15 @@ def init( ``id`` attaches to an existing run instead of creating one, for resuming after a crash and for non-primary ranks joining a run rank 0 created. + ``config`` is what the run was configured *with*. Pass a plain mapping and it + is stored as given; pass a pydantic model and only the fields someone + actually set are stored — see :func:`_normalize_config` for why that default + is the useful one. + + ``config_source`` is the path to the file the run was launched from + (``eval.toml``, ``train.toml``), stored verbatim alongside the structured + config so a reader sees what was written rather than what it expanded into. + ``finish_timeout`` is the total number of seconds ``finish()`` gives queued uploads to drain and close before finalizing the run anyway. """ @@ -798,6 +828,11 @@ def init( base_url = base_url or resolved_config.base_url team_id = team_id if team_id is not None else resolved_config.team_id + run_config = _normalize_config(config) + source = ConfigSource.coerce(config_source) + if source is not None: + run_config[CONFIG_SOURCE_KEY] = source.to_dict() + spec = RunSpec( name=name, kind=kind, @@ -808,7 +843,7 @@ def init( description=description, tags=list(tags or []), team_id=team_id, - config=dict(config or {}), + config=run_config, summary=dict(summary or {}), ) @@ -1074,6 +1109,36 @@ def _merge_mapping(target: Dict[str, Any], value: Any) -> None: target.update(value) +def _normalize_config(value: Any) -> Dict[str, Any]: + """A producer's config as a plain dict, preferring what was actually set. + + A mapping is taken as given — the caller already decided what it wanted to + say. A pydantic model is dumped with ``exclude_unset=True``, which is the + whole point of accepting one: a resolved dump of a deep config tree is + hundreds of lines of defaults nobody chose, and a reader scrolling it cannot + tell which three values were the experiment. ``exclude_unset`` leaves + exactly the fields someone typed. + + A caller who genuinely wants every resolved default can still pass + ``cfg.model_dump()`` explicitly. That asymmetry is deliberate: the shorter + call should give the more useful answer. + """ + if value is None: + return {} + if isinstance(value, Mapping): + return dict(value) + dump = getattr(value, "model_dump", None) + if callable(dump): + try: + dumped = dump(mode="json", exclude_unset=True) + except TypeError: # pragma: no cover - a model_dump with a different signature + dumped = dump() + if isinstance(dumped, Mapping): + return dict(dumped) + raise TypeError(f"{type(value).__name__}.model_dump() did not return a mapping") + raise TypeError(f"config must be a mapping or a pydantic model, got {type(value).__name__}") + + def _describe(error: Union[str, BaseException]) -> str: if isinstance(error, BaseException): return f"{type(error).__name__}: {error}" diff --git a/packages/prime-runs/tests/test_config_source.py b/packages/prime-runs/tests/test_config_source.py new file mode 100644 index 000000000..603c24b6c --- /dev/null +++ b/packages/prime-runs/tests/test_config_source.py @@ -0,0 +1,232 @@ +"""The config a run is actually configured with. + +Two failures this covers, both of which produced a useless Config tab on the +platform: a resolved model dump that buries three chosen values under hundreds +of defaults, and a structured projection that cannot show the file someone +actually wrote. +""" + +import json + +import pytest +from conftest import RecordingHandler + +import prime_runs as pr +from prime_runs.exceptions import ConfigurationError +from prime_runs.models import ( + CONFIG_SOURCE_KEY, + MAX_CONFIG_SOURCE_BYTES, + ConfigSource, + RunSpec, +) +from prime_runs.run import _normalize_config + +EVAL_TOML = """\ +# the environment we are measuring +environment = "primeintellect/terminal-bench-2" +model = "deepseek/deepseek-v4-flash" + +[env] +num_examples = 1 +""" + + +# ------------------------------------------------------------------ coercion + + +def test_a_path_is_read_verbatim(tmp_path): + """Comments, ordering and section grouping are the point — a dict loses all three.""" + path = tmp_path / "eval.toml" + path.write_text(EVAL_TOML) + + source = ConfigSource.coerce(str(path)) + + assert source.text == EVAL_TOML + assert "# the environment we are measuring" in source.text + assert source.format == "toml" + assert source.filename == "eval.toml" + + +def test_the_format_is_inferred_from_the_suffix(tmp_path): + for name, expected in [ + ("train.toml", "toml"), + ("eval.json", "json"), + ("run.yaml", "yaml"), + ("run.yml", "yaml"), + ("config", "text"), + ]: + path = tmp_path / name + path.write_text("x = 1") + assert ConfigSource.coerce(path).format == expected + + +def test_inline_text_has_to_be_explicit(tmp_path): + """A bare string is a path. Guessing would turn a typo'd filename into a run + whose config tab proudly displays the filename.""" + with pytest.raises(ConfigurationError, match="does not exist"): + ConfigSource.coerce("environment = 'gsm8k'") + + source = ConfigSource.coerce(ConfigSource(text="environment = 'gsm8k'")) + assert source.text == "environment = 'gsm8k'" + + +def test_an_oversized_file_is_refused_at_init_not_truncated(tmp_path): + """Silently storing half a config is worse than not storing one.""" + path = tmp_path / "eval.toml" + path.write_text("x = 1\n" * MAX_CONFIG_SOURCE_BYTES) + + with pytest.raises(ConfigurationError, match="over the"): + ConfigSource.coerce(path) + + +def test_a_binary_file_is_refused(tmp_path): + path = tmp_path / "eval.toml" + path.write_bytes(b"\xff\xfe\x00binary") + + with pytest.raises(ConfigurationError, match="not UTF-8"): + ConfigSource.coerce(path) + + +def test_an_unusable_type_says_what_was_expected(): + with pytest.raises(TypeError, match="config_source must be"): + ConfigSource.coerce(object()) + + +def test_a_mapping_round_trips(): + source = ConfigSource(text="a = 1", format="toml", filename="eval.toml") + + assert ConfigSource.from_mapping(source.to_dict()) == source + + +def test_a_mapping_without_text_is_not_a_config_source(): + assert ConfigSource.from_mapping({"format": "toml"}) is None + with pytest.raises(ValueError, match="must contain a 'text' string"): + ConfigSource.coerce({"format": "toml"}) + + +# ------------------------------------------------------- config normalization + + +class FakeModel: + """Duck-types the pydantic v2 surface ``_normalize_config`` looks for.""" + + def __init__(self, set_fields, all_fields): + self._set = set_fields + self._all = all_fields + + def model_dump(self, mode=None, exclude_unset=False): + return dict(self._set if exclude_unset else self._all) + + +def test_a_model_contributes_only_the_fields_someone_set(): + """The training Config tab's actual bug: ``exclude_none`` keeps every default, + so three chosen values arrive buried in a hundred lines nobody picked.""" + model = FakeModel( + set_fields={"model": "Qwen/Qwen3-8B", "max_steps": 1000}, + all_fields={"model": "Qwen/Qwen3-8B", "max_steps": 1000, "seed": 0, "log_level": "info"}, + ) + + assert _normalize_config(model) == {"model": "Qwen/Qwen3-8B", "max_steps": 1000} + + +def test_a_mapping_is_taken_exactly_as_given(): + """The caller already decided what to say; second-guessing it would be worse.""" + assert _normalize_config({"a": 1, "b": None}) == {"a": 1, "b": None} + + +def test_a_mapping_is_copied_not_aliased(): + original = {"a": 1} + assert _normalize_config(original) is not original + + +def test_no_config_is_an_empty_config(): + assert _normalize_config(None) == {} + + +def test_an_unusable_config_says_what_was_expected(): + with pytest.raises(TypeError, match="config must be a mapping"): + _normalize_config(object()) + + +# -------------------------------------------------------------- through init + + +def test_an_offline_run_stores_the_source_next_to_the_config(tmp_path): + path = tmp_path / "eval.toml" + path.write_text(EVAL_TOML) + + run = pr.init( + name="tb2", + environments=["gsm8k"], + mode="offline", + dir=str(tmp_path / "runs"), + config={"num_examples": 1}, + config_source=path, + ) + run.finish() + + state = json.loads((tmp_path / "runs" / run.id / "run.json").read_text()) + stored = state["spec"]["config"] + # Both, not either: the structured form stays queryable, the source stays readable. + assert stored["num_examples"] == 1 + assert stored[CONFIG_SOURCE_KEY]["text"] == EVAL_TOML + assert stored[CONFIG_SOURCE_KEY]["filename"] == "eval.toml" + + +def test_the_run_reports_its_own_source(tmp_path): + path = tmp_path / "train.toml" + path.write_text(EVAL_TOML) + + run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path), config_source=path) + + assert run.config_source.filename == "train.toml" + assert run.config_source.text == EVAL_TOML + run.finish() + + +def test_a_run_without_a_source_reports_none(tmp_path): + run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path)) + + assert run.config_source is None + run.finish() + + +def test_an_online_run_sends_the_source_in_create_metadata( + tmp_path, monkeypatch, make_platform_client, eval_routes +): + path = tmp_path / "eval.toml" + path.write_text(EVAL_TOML) + handler = RecordingHandler(eval_routes) + monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler)) + + run = pr.init( + name="tb2", + environments=["gsm8k"], + api_key="test-key", + config={"num_examples": 1}, + config_source=path, + sinks=[], + ) + run.finish() + + create = next(r for r in handler.requests if r.url.path == "/api/v1/evaluations/") + metadata = json.loads(create.content)["metadata"] + assert metadata[CONFIG_SOURCE_KEY]["text"] == EVAL_TOML + assert metadata["num_examples"] == 1 + + +def test_the_source_survives_the_failure_fallback(tmp_path): + """The fallback rewrites metadata to record a terminal state. It merges into + the whole config, so the source must still be there afterwards.""" + path = tmp_path / "eval.toml" + path.write_text(EVAL_TOML) + + run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path), config_source=path) + run.fail("something broke") + + state = json.loads((tmp_path / run.id / "run.json").read_text()) + assert state["config"][CONFIG_SOURCE_KEY]["text"] == EVAL_TOML + + +def test_a_spec_defaults_to_no_source(): + assert RunSpec().config == {} From ba1b00673fe96cc436eb2af08dba303ece1a6f49 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 20 Aug 2026 11:05:43 -0700 Subject: [PATCH 17/39] refactor(runs): one config parameter, not two MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit config_source= was a second way to say the same thing. Collapse it into config=, which now takes whichever form the caller has — the path to the file the run was launched from, a mapping, or a pydantic model — the same polymorphism environments= already has for slugs, dicts and EnvironmentRefs. The forms are distinguished by type, never by inspecting keys, so a config that happens to carry a `text` field is not mistaken for a launch file. Storage is unchanged: a path still lands under the reserved config_source key inside the run's config, which stays the single key a config-tab renderer has to know about. A run launched from a file that also wants a derived value adds it with run.update_config({...}) rather than the SDK carrying a second argument for the uncommon case. Co-Authored-By: Claude Opus 5 (1M context) --- packages/prime-runs/README.md | 55 +++++++++-------- packages/prime-runs/src/prime_runs/models.py | 37 ++++++------ packages/prime-runs/src/prime_runs/run.py | 60 +++++++++++-------- .../prime-runs/tests/test_config_source.py | 59 ++++++++++++++---- 4 files changed, 126 insertions(+), 85 deletions(-) diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index ae839239b..748a67585 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -61,39 +61,38 @@ join the run their parent opened instead of each opening their own. ## Config -A run records its configuration two ways, because they answer different questions. +One parameter, in whatever form you have it — the same shape `environments=` +already has. ```python -run = pr.init( - ..., - config=cfg, # structured — queryable - config_source="eval.toml", # verbatim — readable -) +pr.init(config="eval.toml") # the file the run was launched from +pr.init(config={"n": 4}) # a mapping +pr.init(config=cfg) # a pydantic model ``` -**`config_source`** is the file the run was launched from — `uv run eval @ -eval.toml`, `uv run rl @ train.toml` — stored byte for byte, comments and section -grouping intact. That file *is* the run's real configuration, and it is the thing -worth putting in front of a person. Pass a path (a `str` or `Path` is always a -path, never inline text); pass `pr.ConfigSource(text=..., format=...)` if you -already hold it in memory. It lands under the reserved `config_source` key inside -the run's config, so every write that carries the config carries it too — -including the offline archive. - -**`config`** is the structured form. A mapping is stored exactly as given. A -pydantic model is dumped with `exclude_unset=True`, so only fields somebody -actually set are recorded: - -| you pass | you get | +| you pass | what is stored | | --- | --- | -| `config={"n": 4}` | `{"n": 4}` | -| `config=cfg` (a model) | only the fields set on `cfg` | -| `config=cfg.model_dump()` | every field, defaults included | - -A resolved dump of a deep config tree is hundreds of lines nobody chose, and a -reader scrolling it cannot tell which three values were the experiment. The -shorter call gives the more useful answer; the full dump is still available to -anyone who explicitly asks for it. +| a path | the file, byte for byte, under `config_source` | +| a mapping | exactly as given | +| a pydantic model | only the fields somebody set (`exclude_unset=True`) | +| `cfg.model_dump()` | every field, defaults included | + +**A path** is the common case: both producers are launched from one user-authored +file (`uv run eval @ eval.toml`, `uv run rl @ train.toml`), and that file *is* the +run's configuration. It is kept verbatim — comments, key order and section +grouping included, none of which survive a dict round-trip. A `str` or `Path` is +always a path, never inline text; use `pr.ConfigSource(text=...)` if you already +hold the contents. It is stored, not parsed: parsing would buy a second +representation of something the platform can already read, at the cost of a TOML +dependency in a package that deliberately has two. + +**A pydantic model** is dumped with `exclude_unset=True` because a resolved dump +of a deep config tree is hundreds of lines nobody chose, and a reader scrolling it +cannot tell which three values were the experiment. The full dump is still one +explicit `cfg.model_dump()` away — the shorter call gives the more useful answer. + +Values land in the run's config either way, so a run launched from a file that +also wants a derived value adds it with `run.update_config({...})`. Nothing is redacted. A config file holding a secret puts that secret on the run's page — keep credentials in the environment, not in the file. diff --git a/packages/prime-runs/src/prime_runs/models.py b/packages/prime-runs/src/prime_runs/models.py index 4ec361a23..a8f22154f 100644 --- a/packages/prime-runs/src/prime_runs/models.py +++ b/packages/prime-runs/src/prime_runs/models.py @@ -87,12 +87,13 @@ def __post_init__(self) -> None: CONFIG_SOURCE_KEY = "config_source" -"""Reserved key inside a run's config holding :class:`ConfigSource` as a dict. +"""Where a config file lands inside a run's config, as a :class:`ConfigSource` dict. -It rides *inside* the config rather than beside it so that every path which -already carries the config carries the source too — create, the periodic update, -finalize, the failure fallback, and the offline archive — with no extra -plumbing and no chance of one of them forgetting it. +Inside the config rather than beside it, so every path that already carries the +config carries the file too — create, the periodic update, finalize, the failure +fallback, and the offline archive — with no extra plumbing and no chance of one +of them forgetting it. It is also the one key a config-tab renderer has to know +about: present means "show this verbatim", absent means "show the structure". """ MAX_CONFIG_SOURCE_BYTES = 256 * 1024 @@ -118,8 +119,9 @@ class ConfigSource: answers "what did every knob end up as", not "what did someone write", and it loses comments, key order and section grouping on the way through. - So this is stored verbatim, next to (not instead of) the structured config. - The structured form stays queryable; this form stays readable. + So when ``init(config=...)`` is given a path, this is what gets stored: the + bytes, not a parse of them. Callers without a file pass a mapping or a model + to the same parameter and get the structured form instead. Nothing here is redacted. A config file that carries a secret will carry it onto the run's page, the same way it already reaches anyone who can read the @@ -158,23 +160,21 @@ def from_file(cls, path: Union[str, "os.PathLike[str]"]) -> "ConfigSource": raw = resolved.read_bytes() except FileNotFoundError as exc: raise ConfigurationError( - f"config_source={str(path)!r} does not exist. Pass the path to the file the " - "run was started from, or a ConfigSource(text=...) if it is already in memory." + f"config={str(path)!r} does not exist. Pass the path to the file the run was " + "started from, or a ConfigSource(text=...) if it is already in memory." ) from exc except OSError as exc: - raise ConfigurationError( - f"config_source={str(path)!r} could not be read: {exc}" - ) from exc + raise ConfigurationError(f"config={str(path)!r} could not be read: {exc}") from exc if len(raw) > MAX_CONFIG_SOURCE_BYTES: raise ConfigurationError( - f"config_source={str(path)!r} is {len(raw)} bytes, over the " + f"config={str(path)!r} is {len(raw)} bytes, over the " f"{MAX_CONFIG_SOURCE_BYTES}-byte limit for a stored run config." ) try: text = raw.decode("utf-8") except UnicodeDecodeError as exc: raise ConfigurationError( - f"config_source={str(path)!r} is not UTF-8 text; a run config must be readable." + f"config={str(path)!r} is not UTF-8 text; a run config must be readable." ) from exc return cls( text=text, @@ -184,7 +184,7 @@ def from_file(cls, path: Union[str, "os.PathLike[str]"]) -> "ConfigSource": @classmethod def coerce(cls, value: Any) -> Optional["ConfigSource"]: - """Normalize whatever ``init(config_source=...)`` was given. + """Normalize the config-file form of ``init(config=...)``. A ``str`` or ``PathLike`` is a *path*, never inline text — that is how every caller will reach for it, and guessing between the two would turn @@ -196,12 +196,13 @@ def coerce(cls, value: Any) -> Optional["ConfigSource"]: if isinstance(value, Mapping): source = cls.from_mapping(value) if source is None: - raise ValueError("config_source mapping must contain a 'text' string") + raise ValueError("a config-source mapping must contain a 'text' string") return source if isinstance(value, (str, os.PathLike)): return cls.from_file(value) raise TypeError( - f"config_source must be a path, a ConfigSource or a mapping, got {type(value).__name__}" + "a config source must be a path, a ConfigSource or a mapping, " + f"got {type(value).__name__}" ) def __post_init__(self) -> None: @@ -209,7 +210,7 @@ def __post_init__(self) -> None: raise TypeError("ConfigSource.text must be a string") if len(self.text.encode("utf-8")) > MAX_CONFIG_SOURCE_BYTES: raise ConfigurationError( - f"config_source is over the {MAX_CONFIG_SOURCE_BYTES}-byte limit " + f"the config source is over the {MAX_CONFIG_SOURCE_BYTES}-byte limit " "for a stored run config." ) diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 27a34e9c9..0dc11da1c 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -779,7 +779,6 @@ def init( description: Optional[str] = None, tags: Optional[Sequence[str]] = None, config: Optional[Any] = None, - config_source: Optional[Any] = None, summary: Optional[Mapping[str, Any]] = None, id: Optional[str] = None, mode: Optional[Mode] = None, @@ -811,14 +810,9 @@ def init( ``id`` attaches to an existing run instead of creating one, for resuming after a crash and for non-primary ranks joining a run rank 0 created. - ``config`` is what the run was configured *with*. Pass a plain mapping and it - is stored as given; pass a pydantic model and only the fields someone - actually set are stored — see :func:`_normalize_config` for why that default - is the useful one. - - ``config_source`` is the path to the file the run was launched from - (``eval.toml``, ``train.toml``), stored verbatim alongside the structured - config so a reader sees what was written rather than what it expanded into. + ``config`` is what the run was configured *with*, in whatever form you have + it — the path to the file it was launched from, a mapping, or a pydantic + model. See :func:`_normalize_config`. ``finish_timeout`` is the total number of seconds ``finish()`` gives queued uploads to drain and close before finalizing the run anyway. @@ -828,11 +822,6 @@ def init( base_url = base_url or resolved_config.base_url team_id = team_id if team_id is not None else resolved_config.team_id - run_config = _normalize_config(config) - source = ConfigSource.coerce(config_source) - if source is not None: - run_config[CONFIG_SOURCE_KEY] = source.to_dict() - spec = RunSpec( name=name, kind=kind, @@ -843,7 +832,7 @@ def init( description=description, tags=list(tags or []), team_id=team_id, - config=run_config, + config=_normalize_config(config), summary=dict(summary or {}), ) @@ -1110,23 +1099,39 @@ def _merge_mapping(target: Dict[str, Any], value: Any) -> None: def _normalize_config(value: Any) -> Dict[str, Any]: - """A producer's config as a plain dict, preferring what was actually set. - - A mapping is taken as given — the caller already decided what it wanted to - say. A pydantic model is dumped with ``exclude_unset=True``, which is the - whole point of accepting one: a resolved dump of a deep config tree is - hundreds of lines of defaults nobody chose, and a reader scrolling it cannot - tell which three values were the experiment. ``exclude_unset`` leaves - exactly the fields someone typed. + """A producer's config as a plain dict, in the most faithful form available. + + One parameter takes whatever form the caller has it in, the same way + ``environments=`` accepts a slug, a dict or an ``EnvironmentRef``. What + changes per form is only how much fidelity there is to preserve: + + - **A path** is the file the run was launched from (``uv run eval @ + eval.toml``). That file *is* the run's configuration, so it is kept byte + for byte under ``config_source`` — comments, key order and section + grouping included. Nothing else can reproduce those. + - **A mapping** is taken exactly as given. The caller already decided what + it wanted to say and second-guessing it would be worse. + - **A pydantic model** is dumped with ``exclude_unset=True``. A resolved + dump of a deep config tree is hundreds of lines of defaults nobody chose, + and a reader scrolling it cannot tell which three values were the + experiment. ``exclude_unset`` leaves exactly the fields someone typed. A caller who genuinely wants every resolved default can still pass - ``cfg.model_dump()`` explicitly. That asymmetry is deliberate: the shorter - call should give the more useful answer. + ``cfg.model_dump()``. That asymmetry is deliberate: the shorter call should + give the more useful answer. + + The file is stored, not parsed. Parsing would buy a second representation of + something the platform can already read, at the cost of a TOML dependency in + a package that deliberately has two. """ if value is None: return {} if isinstance(value, Mapping): return dict(value) + if isinstance(value, (str, os.PathLike, ConfigSource)): + source = ConfigSource.coerce(value) + assert source is not None # coerce only returns None for None + return {CONFIG_SOURCE_KEY: source.to_dict()} dump = getattr(value, "model_dump", None) if callable(dump): try: @@ -1136,7 +1141,10 @@ def _normalize_config(value: Any) -> Dict[str, Any]: if isinstance(dumped, Mapping): return dict(dumped) raise TypeError(f"{type(value).__name__}.model_dump() did not return a mapping") - raise TypeError(f"config must be a mapping or a pydantic model, got {type(value).__name__}") + raise TypeError( + "config must be a path to the run's config file, a mapping or a pydantic model, " + f"got {type(value).__name__}" + ) def _describe(error: Union[str, BaseException]) -> str: diff --git a/packages/prime-runs/tests/test_config_source.py b/packages/prime-runs/tests/test_config_source.py index 603c24b6c..dd891fda1 100644 --- a/packages/prime-runs/tests/test_config_source.py +++ b/packages/prime-runs/tests/test_config_source.py @@ -31,7 +31,7 @@ """ -# ------------------------------------------------------------------ coercion +# ------------------------------------------------------ the config-file form def test_a_path_is_read_verbatim(tmp_path): @@ -88,7 +88,7 @@ def test_a_binary_file_is_refused(tmp_path): def test_an_unusable_type_says_what_was_expected(): - with pytest.raises(TypeError, match="config_source must be"): + with pytest.raises(TypeError, match="config source must be"): ConfigSource.coerce(object()) @@ -144,14 +144,35 @@ def test_no_config_is_an_empty_config(): def test_an_unusable_config_says_what_was_expected(): - with pytest.raises(TypeError, match="config must be a mapping"): + with pytest.raises(TypeError, match="config must be a path"): _normalize_config(object()) +def test_a_path_becomes_a_stored_source(tmp_path): + """One parameter, three forms — the same shape ``environments=`` already has.""" + path = tmp_path / "eval.toml" + path.write_text(EVAL_TOML) + + assert _normalize_config(path) == { + CONFIG_SOURCE_KEY: {"format": "toml", "text": EVAL_TOML, "filename": "eval.toml"} + } + assert _normalize_config(str(path))[CONFIG_SOURCE_KEY]["text"] == EVAL_TOML + assert _normalize_config(ConfigSource(text="a = 1"))[CONFIG_SOURCE_KEY]["text"] == "a = 1" + + +def test_a_mapping_that_looks_like_a_source_is_still_just_a_mapping(): + """Form is decided by type, never by inspecting keys — so a config that + happens to have a ``text`` field is not mistaken for a launch file.""" + assert _normalize_config({"text": "hello", "format": "toml"}) == { + "text": "hello", + "format": "toml", + } + + # -------------------------------------------------------------- through init -def test_an_offline_run_stores_the_source_next_to_the_config(tmp_path): +def test_an_offline_run_stores_the_launch_file(tmp_path): path = tmp_path / "eval.toml" path.write_text(EVAL_TOML) @@ -160,24 +181,37 @@ def test_an_offline_run_stores_the_source_next_to_the_config(tmp_path): environments=["gsm8k"], mode="offline", dir=str(tmp_path / "runs"), - config={"num_examples": 1}, - config_source=path, + config=path, ) run.finish() state = json.loads((tmp_path / "runs" / run.id / "run.json").read_text()) stored = state["spec"]["config"] - # Both, not either: the structured form stays queryable, the source stays readable. - assert stored["num_examples"] == 1 assert stored[CONFIG_SOURCE_KEY]["text"] == EVAL_TOML assert stored[CONFIG_SOURCE_KEY]["filename"] == "eval.toml" +def test_extra_values_can_be_merged_onto_a_launch_file(tmp_path): + """One parameter takes one form. A run launched from a file that also wants + a derived value adds it explicitly, rather than the SDK growing a second + config argument for a case that is not the common one.""" + path = tmp_path / "eval.toml" + path.write_text(EVAL_TOML) + + run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path), config=path) + run.update_config({"resolved_model": "deepseek/deepseek-v4-flash"}) + run.finish() + + state = json.loads((tmp_path / run.id / "run.json").read_text()) + assert state["config"]["resolved_model"] == "deepseek/deepseek-v4-flash" + assert state["config"][CONFIG_SOURCE_KEY]["text"] == EVAL_TOML + + def test_the_run_reports_its_own_source(tmp_path): path = tmp_path / "train.toml" path.write_text(EVAL_TOML) - run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path), config_source=path) + run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path), config=path) assert run.config_source.filename == "train.toml" assert run.config_source.text == EVAL_TOML @@ -203,8 +237,7 @@ def test_an_online_run_sends_the_source_in_create_metadata( name="tb2", environments=["gsm8k"], api_key="test-key", - config={"num_examples": 1}, - config_source=path, + config=path, sinks=[], ) run.finish() @@ -212,7 +245,7 @@ def test_an_online_run_sends_the_source_in_create_metadata( create = next(r for r in handler.requests if r.url.path == "/api/v1/evaluations/") metadata = json.loads(create.content)["metadata"] assert metadata[CONFIG_SOURCE_KEY]["text"] == EVAL_TOML - assert metadata["num_examples"] == 1 + assert metadata[CONFIG_SOURCE_KEY]["format"] == "toml" def test_the_source_survives_the_failure_fallback(tmp_path): @@ -221,7 +254,7 @@ def test_the_source_survives_the_failure_fallback(tmp_path): path = tmp_path / "eval.toml" path.write_text(EVAL_TOML) - run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path), config_source=path) + run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path), config=path) run.fail("something broke") state = json.loads((tmp_path / run.id / "run.json").read_text()) From 454c9116b1b3fc51b9d0fcc631bfc9a81853e2bc Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 20 Aug 2026 11:14:19 -0700 Subject: [PATCH 18/39] refactor(runs): drop four parameters nothing needed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four bits of surface that were each a second way to say something the SDK already had a first way to say. dataset= is always the environment under a different name, and no UI reads the column. The API field is still populated, derived from the first environment — producers just no longer repeat themselves to fill it. summary= on init() asked for a run's *outputs* at the moment it opens, before it has any. finish(summary=) and run.summary already cover the real case, so RunSpec loses the field too and create() stops sending an empty metrics blob. sinks= was a third way to configure transport next to traces= and samples=, with no caller outside one of my own tests. Run(sinks=) stays, so tests still inject fakes and a future custom sink has somewhere to land. Re-adding a keyword argument is non-breaking; removing one is not, which is the argument for cutting it now rather than later. log_episodes was a third name for log_traces. log_samples stays — it is the method name prime-rl's Monitor ABC already uses, so the adapter it exists for is real. Co-Authored-By: Claude Opus 5 (1M context) --- packages/prime-runs/README.md | 3 +- .../src/prime_runs/backends/evals.py | 6 ++-- packages/prime-runs/src/prime_runs/models.py | 8 ++--- packages/prime-runs/src/prime_runs/run.py | 35 ++++++------------- .../prime-runs/tests/test_config_source.py | 3 +- packages/prime-runs/tests/test_init.py | 10 ++---- 6 files changed, 24 insertions(+), 41 deletions(-) diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index 748a67585..94cc666c1 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -189,8 +189,7 @@ a traces-only client would leave non-allowlisted accounts with an empty dashboard. When the Viewer API reads traces natively, the default sink list drops one entry — and no producer changes. -Turn either off with `pr.init(traces=False)` / `pr.init(samples=False)`, or pass -`sinks=[...]` to supply your own. +Turn either off with `pr.init(traces=False)` / `pr.init(samples=False)`. ## Also here diff --git a/packages/prime-runs/src/prime_runs/backends/evals.py b/packages/prime-runs/src/prime_runs/backends/evals.py index 1a0c031b4..5a64aadf4 100644 --- a/packages/prime-runs/src/prime_runs/backends/evals.py +++ b/packages/prime-runs/src/prime_runs/backends/evals.py @@ -75,11 +75,13 @@ def create(self, spec: RunSpec) -> RunHandle: "tags": list(spec.tags), } _set_if(payload, "model_name", spec.model) - _set_if(payload, "dataset", spec.dataset or _first_environment_name(spec)) + # Derived, not asked for: the API has the column and it is always the + # environment under a different name, so making producers repeat it + # bought nothing. + _set_if(payload, "dataset", _first_environment_name(spec)) _set_if(payload, "framework", spec.framework) _set_if(payload, "description", spec.description) _set_if(payload, "metadata", spec.config or None) - _set_if(payload, "metrics", spec.summary or None) _set_if(payload, "team_id", spec.team_id or self._team_id) # Not replayable: POST defaults to idempotent=False here on purpose. If diff --git a/packages/prime-runs/src/prime_runs/models.py b/packages/prime-runs/src/prime_runs/models.py index a8f22154f..a37112268 100644 --- a/packages/prime-runs/src/prime_runs/models.py +++ b/packages/prime-runs/src/prime_runs/models.py @@ -229,15 +229,13 @@ class RunSpec: environments: List[EnvironmentRef] = field(default_factory=list) model: Optional[str] = None framework: Optional[str] = None - dataset: Optional[str] = None description: Optional[str] = None tags: List[str] = field(default_factory=list) team_id: Optional[str] = None - # W&B's split, which maps cleanly onto the platform's existing columns: - # `config` is what you set going in (-> metadata), `summary` is what came - # out (-> metrics). + # Only the inputs. A run's outputs (-> the platform's `metrics`) are not + # here because a run being opened does not have any yet; they accumulate on + # the handle and are written by update() and finalize(). config: Dict[str, Any] = field(default_factory=dict) - summary: Dict[str, Any] = field(default_factory=dict) @dataclass diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 0dc11da1c..2e93fbcc7 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -104,7 +104,7 @@ def __init__( attached_config, attached_summary = _attached_state(handle) self.config: Dict[str, Any] = {**attached_config, **spec.config} - self.summary: Dict[str, Any] = {**attached_summary, **spec.summary} + self.summary: Dict[str, Any] = dict(attached_summary) self.errors: List[str] = [] # Raised at the next synchronization point the caller controls. A sink # fails on the uploader thread, where raising reaches nobody — so under @@ -326,9 +326,6 @@ def log_traces( return self._worker.submit(WriteItem(records=batch, line_format=line_format, step=step)) - # Producers that think in episodes rather than traces; same path. - log_episodes = log_traces - def log_samples(self, records: Iterable[Any], *, step: Optional[int] = None) -> None: """Alias for :meth:`log_traces`, matching prime-rl's ``Monitor`` vocabulary.""" self.log_traces(records, step=step) @@ -775,11 +772,9 @@ def init( environments: Optional[Sequence[Any]] = None, model: Optional[str] = None, framework: Optional[str] = None, - dataset: Optional[str] = None, description: Optional[str] = None, tags: Optional[Sequence[str]] = None, config: Optional[Any] = None, - summary: Optional[Mapping[str, Any]] = None, id: Optional[str] = None, mode: Optional[Mode] = None, dir: Optional[str] = None, @@ -789,7 +784,6 @@ def init( traces_url: Optional[str] = None, traces: bool = True, samples: bool = True, - sinks: Optional[List[Sink]] = None, on_error: OnError = "warn", handle_signals: bool = True, queue_size: Optional[int] = None, @@ -828,12 +822,10 @@ def init( environments=[EnvironmentRef.coerce(entry) for entry in (environments or [])], model=model, framework=framework, - dataset=dataset, description=description, tags=list(tags or []), team_id=team_id, config=_normalize_config(config), - summary=dict(summary or {}), ) is_primary = _is_primary_rank() @@ -865,7 +857,7 @@ def init( if resolved_mode == "offline": offline = OfflineBackend(dir) handle = offline.attach(inherited_id) if inherited_id else offline.create(spec) - run_sinks = sinks if sinks is not None else [OfflineSink(offline.directory)] + run_sinks: List[Sink] = [OfflineSink(offline.directory)] run = _build( spec, offline, @@ -896,20 +888,15 @@ def init( backend = EvalsBackend(client, frontend_url=resolved_config.frontend_url, team_id=team_id) handle = backend.attach(inherited_id) if inherited_id else backend.create(spec) - if sinks is None: - run_sinks = [] - if traces: - run_sinks.append(TracesSink(api_key=api_key, traces_url=traces_url, team_id=team_id)) - if samples: - # Both transports run during the transition: traces is the system of - # record, the sample table is what today's viewer reads, and Prime - # Traces is still gated to an account allowlist. - samples_client = PlatformClient( - api_key=api_key, base_url=base_url, timeout=DEFAULT_TIMEOUT - ) - run_sinks.append(EvalSamplesSink(samples_client, close_client=True)) - else: - run_sinks = list(sinks) + run_sinks = [] + if traces: + run_sinks.append(TracesSink(api_key=api_key, traces_url=traces_url, team_id=team_id)) + if samples: + # Both transports run during the transition: traces is the system of + # record, the sample table is what today's viewer reads, and Prime + # Traces is still gated to an account allowlist. + samples_client = PlatformClient(api_key=api_key, base_url=base_url, timeout=DEFAULT_TIMEOUT) + run_sinks.append(EvalSamplesSink(samples_client, close_client=True)) run = _build( spec, diff --git a/packages/prime-runs/tests/test_config_source.py b/packages/prime-runs/tests/test_config_source.py index dd891fda1..c838f419e 100644 --- a/packages/prime-runs/tests/test_config_source.py +++ b/packages/prime-runs/tests/test_config_source.py @@ -238,7 +238,8 @@ def test_an_online_run_sends_the_source_in_create_metadata( environments=["gsm8k"], api_key="test-key", config=path, - sinks=[], + traces=False, + samples=False, ) run.finish() diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index 554f612ff..e94314727 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -378,14 +378,9 @@ def test_resuming_preserves_existing_config_and_summary(online, eval_routes): "metrics": {"old_reward": 0.5, "overridden": "old"}, } - run, handler = online( - routes=routes, - id="eval-abc", - config={"overridden": "new"}, - summary={"overridden": "new"}, - ) + run, handler = online(routes=routes, id="eval-abc", config={"overridden": "new"}) run.update_config({"after_resume": True}) - run.log({"new_reward": 1.0}) + run.log({"overridden": "new", "new_reward": 1.0}) run.finish() update = handler.bodies_for("/api/v1/evaluations/eval-abc")[0] @@ -394,6 +389,7 @@ def test_resuming_preserves_existing_config_and_summary(online, eval_routes): "overridden": "new", "after_resume": True, } + # Recovered metrics survive; anything logged after the resume wins. assert update["metrics"] == { "old_reward": 0.5, "overridden": "new", From 4a367dc08f5c56101fc16a458ac4a7895e879bd5 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 20 Aug 2026 11:20:12 -0700 Subject: [PATCH 19/39] fix(runs): ask the signature instead of catching TypeError MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _normalize_config guessed why model_dump() failed: any TypeError was read as "this callable has a different signature" and retried bare. Inferring a cause from an exception type is the problem, because the only recovery available — dumping every field, defaults included — is exactly the outcome passing a model was meant to avoid. A wrong guess therefore degrades silently, in the one direction that matters. Now the signature is inspected before the call, so a failure during serialization surfaces as itself, and the fallback fires only when the keywords genuinely are not accepted. When it does fire it logs, because silently recording a hundred defaults nobody chose is the behaviour this path exists to prevent. Note for the record: Bugbot reported this as PydanticSerializationError being caught by `except TypeError`. That specific claim is wrong — PydanticSerializationError subclasses ValueError, and pydantic wraps even a serializer's own TypeError into it, so the handler could only ever catch a signature mismatch. The underlying concern about guessing was still worth acting on. Co-Authored-By: Claude Opus 5 (1M context) --- packages/prime-runs/src/prime_runs/run.py | 31 ++++++++++++++- .../prime-runs/tests/test_config_source.py | 39 +++++++++++++++++++ 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 2e93fbcc7..dd91570e6 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -16,6 +16,7 @@ """ import atexit +import inspect import logging import math import os @@ -1085,6 +1086,25 @@ def _merge_mapping(target: Dict[str, Any], value: Any) -> None: target.update(value) +def _accepts_exclude_unset(dump: Any) -> bool: + """Whether a ``model_dump`` takes the keywords we want to hand it. + + Asked of the signature rather than discovered by catching ``TypeError``. + Catching would mean inferring *why* a call failed from its exception type, + and the only available recovery — dumping every field — is precisely the + outcome this whole path exists to avoid. So a wrong inference degrades + silently, in the one direction that matters. A real serialization failure + should surface as itself. + """ + try: + params = inspect.signature(dump).parameters + except (TypeError, ValueError): # pragma: no cover - uninspectable callables + return False + if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): + return True + return {"mode", "exclude_unset"} <= params.keys() + + def _normalize_config(value: Any) -> Dict[str, Any]: """A producer's config as a plain dict, in the most faithful form available. @@ -1121,9 +1141,16 @@ def _normalize_config(value: Any) -> Dict[str, Any]: return {CONFIG_SOURCE_KEY: source.to_dict()} dump = getattr(value, "model_dump", None) if callable(dump): - try: + if _accepts_exclude_unset(dump): dumped = dump(mode="json", exclude_unset=True) - except TypeError: # pragma: no cover - a model_dump with a different signature + else: + # Not silent: the fallback records every default, which is the exact + # outcome the caller was trying to avoid by handing us a model. + logger.warning( + "%s.model_dump() does not accept exclude_unset; recording the fully " + "resolved config instead, defaults included.", + type(value).__name__, + ) dumped = dump() if isinstance(dumped, Mapping): return dict(dumped) diff --git a/packages/prime-runs/tests/test_config_source.py b/packages/prime-runs/tests/test_config_source.py index c838f419e..5d5e0aafb 100644 --- a/packages/prime-runs/tests/test_config_source.py +++ b/packages/prime-runs/tests/test_config_source.py @@ -129,6 +129,45 @@ def test_a_model_contributes_only_the_fields_someone_set(): assert _normalize_config(model) == {"model": "Qwen/Qwen3-8B", "max_steps": 1000} +def test_a_dump_that_cannot_serialize_says_so_instead_of_dumping_everything(): + """The recovery on offer — dump every field — is the exact outcome passing a + model was meant to avoid, so it must never be reached by guessing at why a + call failed. A broken serializer is the caller's bug and surfaces as itself.""" + + class Broken: + def model_dump(self, mode=None, exclude_unset=False): + raise TypeError("serializer blew up") + + with pytest.raises(TypeError, match="serializer blew up"): + _normalize_config(Broken()) + + +def test_a_dump_without_exclude_unset_falls_back_loudly(caplog): + """Degrading to the full config is allowed, going quiet about it is not.""" + + class Old: + def model_dump(self): + return {"model": "Qwen/Qwen3-8B", "seed": 0} + + with caplog.at_level("WARNING"): + assert _normalize_config(Old()) == {"model": "Qwen/Qwen3-8B", "seed": 0} + + assert "exclude_unset" in caplog.text + assert "defaults included" in caplog.text + + +def test_a_dump_taking_kwargs_is_given_the_keywords(): + seen = {} + + class Flexible: + def model_dump(self, **kwargs): + seen.update(kwargs) + return {"a": 1} + + assert _normalize_config(Flexible()) == {"a": 1} + assert seen == {"mode": "json", "exclude_unset": True} + + def test_a_mapping_is_taken_exactly_as_given(): """The caller already decided what to say; second-guessing it would be worse.""" assert _normalize_config({"a": 1, "b": None}) == {"a": 1, "b": None} From 7d3cb5dd1aa0a158d1c4cd17e6c179a000a15009 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 20 Aug 2026 21:39:22 -0700 Subject: [PATCH 20/39] refactor(runs): cut the SDK to what its consumer uses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Design review found ~40% of the package serving scenarios with no backend and no consumer: training runs, multi-rank lifecycle ownership, PRIME_RUN_ID join, signal handling, a probe for a /status endpoint the platform does not have, and a dict-based reimplementation of the v0 projection. verifiers — the one producer — uses seven init() arguments, log_traces() and finish(). Removed: - training / multi-rank / join machinery: RANK_ENV_VARS, is_primary, PRIME_RUN_ID publish/retract, id= resume, attach(), log()/log_samples()/ update_config()/commit=, MetricItem/RunUpdateItem, supports_step_metrics, the summary timer, RunKind/kind, RunHandle.raw - signal handling (install/relinquish/chain protocol, _pending_signal); atexit and the context manager still report CRASHED/FAILED - the /evaluations/{id}/status probe; finalize() writes metadata.prime_runs directly for non-COMPLETED runs - projection.py's serialized-mapping path and the bare-Trace branch of EvalSamplesSink; the sink takes Episode objects or v0 sample dicts - init() params kind, id, traces_url, traces, samples, handle_signals, queue_size, finish_timeout, and the pydantic-model form of config= - _build(); the second PlatformClient (backend and samples sink share one) - 20 names from __all__ (Backend, Sink, RunSpec, RunHandle, sink and backend classes, RUN_ID_ENV) Consolidated: is_episode/stamp_run live once in sinks/base.py; line_format and step no longer thread through log_traces -> WriteItem -> Sink.write. Docstrings trimmed to contracts; README rewritten for the reduced surface. Kept offline mode, the fork hooks and the prime_runs-local HTTP/config/ exception plumbing — each is a separate discussion. Source 4,105 -> 2,524 lines; tests 3,201 -> 2,492 (165 passing); init() 21 -> 13 parameters. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01W4S4DNVRX5BNtsLva7TZz1 --- packages/prime-runs/README.md | 225 ++--- .../prime-runs/src/prime_runs/__init__.py | 57 +- packages/prime-runs/src/prime_runs/_fork.py | 31 +- packages/prime-runs/src/prime_runs/_http.py | 82 +- .../src/prime_runs/backends/__init__.py | 2 +- .../src/prime_runs/backends/base.py | 49 +- .../src/prime_runs/backends/evals.py | 176 +--- .../src/prime_runs/backends/offline.py | 43 +- packages/prime-runs/src/prime_runs/config.py | 14 +- .../prime-runs/src/prime_runs/exceptions.py | 55 +- packages/prime-runs/src/prime_runs/metrics.py | 14 +- packages/prime-runs/src/prime_runs/models.py | 79 +- .../prime-runs/src/prime_runs/projection.py | 302 +----- packages/prime-runs/src/prime_runs/run.py | 951 +++--------------- .../src/prime_runs/sinks/__init__.py | 6 +- .../prime-runs/src/prime_runs/sinks/base.py | 55 +- .../src/prime_runs/sinks/offline.py | 72 +- .../src/prime_runs/sinks/samples.py | 72 +- .../prime-runs/src/prime_runs/sinks/traces.py | 135 +-- packages/prime-runs/src/prime_runs/worker.py | 184 +--- packages/prime-runs/tests/conftest.py | 16 +- .../prime-runs/tests/test_config_source.py | 95 +- .../prime-runs/tests/test_evals_backend.py | 110 +- packages/prime-runs/tests/test_init.py | 301 +----- packages/prime-runs/tests/test_projection.py | 7 +- packages/prime-runs/tests/test_run.py | 254 +---- .../prime-runs/tests/test_samples_sink.py | 65 +- packages/prime-runs/tests/test_traces_sink.py | 18 +- packages/prime-runs/tests/test_worker.py | 105 +- 29 files changed, 591 insertions(+), 2984 deletions(-) diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index 94cc666c1..d07462720 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -1,6 +1,6 @@ # Prime Runs SDK -Track eval and training runs on the Prime Intellect platform. +Track eval runs on the Prime Intellect platform. ```bash pip install prime-runs @@ -16,148 +16,92 @@ run = pr.init( environments=["gsm8k"], model="Qwen/Qwen3-8B", framework="verifiers", - config={"num_rollouts": 4, "max_tokens": 2048}, + config="eval.toml", # the file the run was launched from ) print(run.url) # https://app.primeintellect.ai/dashboard/evaluations/eval-... for episode in rollouts: # episodes carry run.id — see "Identity" below run.log_traces([episode]) - run.log({"reward": episode.reward}) run.finish(summary=pr.metrics.from_episodes(episodes)) ``` `init()` opens the run and returns a handle carrying its ID and dashboard URL. -Records stream out on a background thread while the run proceeds — the dashboard -fills in as rollouts land, rather than all at once at the end. `finish()` closes -the run out. - -Prefer a `with` block, which also reports a terminal status on the paths where -you never reach `finish()`: - -```python -with pr.init(name="gsm8k-qwen3-8b", environments=["gsm8k"]) as run: - ... -``` +Records stream out on a background thread while the run proceeds, so the +dashboard fills in as rollouts land. `finish()` closes the run out; a `with` +block does that for you, including when the body raises. ## Identity `init()` is called **before** the first rollout, and the ID it returns is *the* -run ID everywhere — including inside every trace document you write, and -including the local archive. Stamp it once at rollout time: +run ID everywhere — including inside every trace document you write: ```python run = pr.init(...) trace.record_run(EvalRunInfo(id=run.id)) # verifiers ``` -Nothing is re-stamped afterwards and no record of yours is rewritten. The -ingestion service extracts `run.id` from the trace body into an indexed column -with a delete-by-run path, so "every trace for this run" is a fast query rather -than a scan over upload metadata. - -`init()` also exports `PRIME_RUN_ID`, so forked workers and subprocess launchers -join the run their parent opened instead of each opening their own. +The ingestion service indexes `run.id` from the trace body, so "every trace for +this run" is a fast query. Bare dicts without a `run` key are stamped for you; +producer objects are passed through untouched. ## Config -One parameter, in whatever form you have it — the same shape `environments=` -already has. - -```python -pr.init(config="eval.toml") # the file the run was launched from -pr.init(config={"n": 4}) # a mapping -pr.init(config=cfg) # a pydantic model -``` +`config=` takes the path to the file the run was launched from, or a mapping. | you pass | what is stored | | --- | --- | | a path | the file, byte for byte, under `config_source` | | a mapping | exactly as given | -| a pydantic model | only the fields somebody set (`exclude_unset=True`) | -| `cfg.model_dump()` | every field, defaults included | - -**A path** is the common case: both producers are launched from one user-authored -file (`uv run eval @ eval.toml`, `uv run rl @ train.toml`), and that file *is* the -run's configuration. It is kept verbatim — comments, key order and section -grouping included, none of which survive a dict round-trip. A `str` or `Path` is -always a path, never inline text; use `pr.ConfigSource(text=...)` if you already -hold the contents. It is stored, not parsed: parsing would buy a second -representation of something the platform can already read, at the cost of a TOML -dependency in a package that deliberately has two. -**A pydantic model** is dumped with `exclude_unset=True` because a resolved dump -of a deep config tree is hundreds of lines nobody chose, and a reader scrolling it -cannot tell which three values were the experiment. The full dump is still one -explicit `cfg.model_dump()` away — the shorter call gives the more useful answer. +The file is the run's real configuration — comments, key order and section +grouping included — so it is stored verbatim, not parsed. A `str` or `Path` is +always a path; use `pr.ConfigSource(text=...)` for contents already in memory. +To send structured values *and* the file, put the file under +`pr.CONFIG_SOURCE_KEY` in the mapping: -Values land in the run's config either way, so a run launched from a file that -also wants a derived value adds it with `run.update_config({...})`. +```python +config = {**cfg.model_dump(exclude_unset=True), + pr.CONFIG_SOURCE_KEY: pr.ConfigSource.from_file("eval.toml").to_dict()} +``` -Nothing is redacted. A config file holding a secret puts that secret on the run's -page — keep credentials in the environment, not in the file. +Nothing is redacted — keep credentials in the environment, not in the file. ## Modes | mode | what happens | | --- | --- | | `online` | the run lives on the platform (default when an API key is present) | -| `offline` | the run lives in a local directory, ready to sync later | +| `offline` | the run lives in a local directory (`dir=` / `PRIME_RUNS_DIR`) | | `disabled` | every call is a no-op, with the same object shape | -An offline run is a real run: a real ID, a status, a config, a summary, a metrics -stream, and records written in the JSONL wire format the traces service accepts. -That is why producers do not need a `--no-push` branch — the call sites are -identical either way, and a missing API key degrades to offline rather than -skipping the run. - -Set the mode explicitly, or through `$PRIME_RUNS_MODE`. +Set the mode explicitly, or through `$PRIME_RUNS_MODE`. A missing API key +degrades to offline with a warning. ## What the run handle does for you -- **Streams instead of buffering.** Records go out as they are produced, so a - run with a hundred thousand episodes does not hold them all in memory. +- **Streams instead of buffering.** Records go out as they are produced. - **Contains its own errors.** With the default `on_error="warn"`, nothing the - platform raises escapes into your loop. Use `on_error="raise"` in tests and CI, - where a silent upload failure is the bug. -- **Applies backpressure.** The upload queue is bounded; if a producer durably - outruns the uploader, records are dropped and counted (`run.dropped_records`) - rather than stalling the run. -- **Waits for its own uploads.** `finish()` gives queued records the same budget - a single upload gets (300s, `finish_timeout=`) and says so in a warning if - they do not drain, rather than finalizing over records still in flight. -- **Survives forks.** A forked child gets a fresh uploader, a fresh connection - pool and fresh file handles instead of writing the parent's — which would - interleave two processes into one HTTP stream and flush the parent's buffered - records a second time. It also joins the parent's run rather than opening - its own. -- **Reports a terminal status.** Context manager, `atexit` and signal handlers - all route to the same idempotent `finish()`. A run the producer decided had - failed is `failed`; one stopped from outside its control flow — Ctrl-C, - SIGTERM, an exit that never reached `finish()` — is `crashed`. Neither is - left running forever. -- **Knows about ranks.** Rank 0 owns creation and finalization; other ranks join - through `PRIME_RUN_ID` and upload their own records. - -## Using it from async code - -There is no `AsyncRun`, unlike the async clients in `prime-traces`, -`prime-evals` and `prime-sandboxes`. The uploader thread is what replaces it: -`log()`, `log_traces()` and `update_config()` are queue puts, not requests, so -calling them straight from a coroutine does no network I/O on the event loop. - -Three calls do block, and all three are worth knowing about: - -| call | blocks on | when it matters | -| --- | --- | --- | -| `init()` | create + environment resolution | once, at startup | -| `finish()` | draining the queue, then finalize | once, at shutdown | -| `log_traces()` | up to `put_timeout` (5s) **only if the queue is full** | a producer durably outrunning the uploader | - -The first two are run boundaries — `await asyncio.to_thread(run.finish)` if a -stall there would matter. The third is the one to watch in a hot rollout loop: -the block is the backpressure, and past it the batch is dropped and counted in -`run.dropped_records`. Raise `queue_size=` before reaching for a thread. + platform raises escapes into your loop. Use `on_error="raise"` in tests and + CI, where a silent upload failure is the bug; failures surface from `flush()` + and `finish()`. +- **Applies backpressure.** The upload queue is bounded; a producer that + durably outruns the uploader has records dropped and counted + (`run.dropped_records`) rather than stalled. Per-sink losses are in + `run.failed_records`. +- **Waits for its own uploads.** `finish()` gives queued records the same + budget a single upload gets (300s) and warns if they do not drain. +- **Reports a terminal status.** The context manager and an `atexit` hook both + route to the same idempotent `finish()`. A run you said failed is `failed`; + one that stopped without saying — Ctrl-C inside a `with` block, an exit that + never reached `finish()` — is `crashed`. + +## From async code + +`log_traces()` is a queue put, not a request, so it is safe to call from a +coroutine; it blocks only if the queue is full (up to 5s), which is the +backpressure. `init()` and `finish()` do network I/O — wrap them in +`asyncio.to_thread` if a stall there would matter. ## Configuration @@ -172,74 +116,27 @@ Resolved from environment variables first, then `~/.prime/config.json`: | traces service | `PRIME_TRACES_URL` | resolved by `prime-traces` | | offline runs | `PRIME_RUNS_DIR` | `./prime-runs` | -Or pass them to `init()` directly (`api_key=`, `base_url=`, `team_id=`, `dir=`). - -## Backends and sinks - -Two independent axes: - -- **Backends** own run *lifecycle* — `EvalsBackend` (`/api/v1/evaluations/*`), - `OfflineBackend` (a local directory). Selected by `kind`. -- **Sinks** own sample *transport* — `TracesSink` (primary; streaming, - episode-aware, content-addressed and therefore idempotent on retry) and - `EvalSamplesSink` (the flat v0 sample table today's viewer reads). - -Both sinks run during the transition, because Prime Traces is in closed beta and -a traces-only client would leave non-allowlisted accounts with an empty -dashboard. When the Viewer API reads traces natively, the default sink list drops -one entry — and no producer changes. - -Turn either off with `pr.init(traces=False)` / `pr.init(samples=False)`. - -## Also here - -`prime_runs.projection` holds `trace_to_sample` / `build_samples`, the projection -from native traces onto the platform's v0 eval-sample format. It lives here -because it is knowledge about a platform wire format, not about any one eval -framework. `prime_runs.metrics.from_episodes` is the run-level aggregation the -eval dashboard reads — opt-in, because what a run's headline number means is a -judgement that belongs next to the producer. - -Both are duck-typed: verifiers `Trace`/`Episode` and prime-rl `Rollout` satisfy -them structurally, and no producer package is imported. This is a leaf package by -design — the `prime` CLI depends on `verifiers`, so verifiers can never depend on -`prime`. - -## How this differs from the other prime SDKs - -`prime-sandboxes`, `prime-traces`, `prime-evals` and `prime-tunnel` are all built -the same way: a `core/` subpackage holding an `APIClient` and a `Config`, pydantic -models for the responses, and a sync/async client pair as the thing you import. -Three deliberate departures here, so the difference reads as a choice rather than -an oversight: +Or pass `api_key=`, `base_url=`, `team_id=`, `dir=` to `init()`. -- **The client is private.** `init()` is the surface, not a client object, so the - HTTP layer lives in `_http.py` rather than `core/client.py` and `PlatformClient` - is not exported. Config still is, and is the same class as everywhere else — - `~/.prime/config.json`, env wins, same variable names. -- **Responses are not modeled.** The platform returns more fields than any - producer reads; freezing them in pydantic would make every backend addition a - breaking SDK release. Backends take the two or three fields they need and - return a `RunHandle`. That is why the local types are dataclasses and why - pydantic is not a dependency. -- **No async client.** See [Using it from async code](#using-it-from-async-code) - — the background uploader covers the case an async client would exist for. +## Transports -## Related packages +An online run writes every record to two sinks: Prime Traces (the system of +record — streaming, episode-aware, content-addressed and therefore idempotent +on retry) and the flat v0 sample table today's viewer reads. Both run because +Prime Traces is gated to an allowlist; when the viewer reads traces natively the +default sink list drops one entry and no producer changes. -- [prime-traces](../prime-traces) — the traces service client this SDK streams - through, and the direct API for querying or exporting what a run produced. -- [prime](../prime) — the CLI and full SDK. Depends on this package's consumers, - never the other way around. +`prime_runs.projection` holds `trace_to_sample` / `build_samples`, the v0 +projection moved here from verifiers. `prime_runs.metrics.from_episodes` is the +run-level aggregation the eval dashboard reads. Both are duck-typed; no producer +package is imported — this is a leaf package, because the `prime` CLI depends on +`verifiers` and verifiers depends on this. ## Status -Eval runs are supported. Training runs (`kind="train"`, over -`/api/v1/rft/external-runs`) are next; `pr.init(kind="train")` raises a clear -error until then. +Eval runs only. Training runs will arrive with a backend over +`/api/v1/rft/external-runs`, designed against prime-rl's actual needs. -One platform gap is worth knowing about: there is currently no producer-facing -way to mark an evaluation **failed**. The SDK calls the status endpoint it needs, -treats its absence as expected, and records the terminal state in the run's -metadata as a fallback — a failed run will keep showing as running on the -dashboard, and the SDK says so in a warning. +There is currently no producer-facing way to mark an evaluation **failed**. The +SDK records the terminal state under `metadata.prime_runs` and warns that the +run will keep showing as running on the dashboard. diff --git a/packages/prime-runs/src/prime_runs/__init__.py b/packages/prime-runs/src/prime_runs/__init__.py index acd2d844a..5ce95149e 100644 --- a/packages/prime-runs/src/prime_runs/__init__.py +++ b/packages/prime-runs/src/prime_runs/__init__.py @@ -1,6 +1,6 @@ """Prime Intellect Runs SDK. -Track eval and training runs on the Prime platform:: +Track eval runs on the Prime platform:: import prime_runs as pr @@ -9,24 +9,15 @@ for episode in rollouts: # stamp run.id onto the traces run.log_traces([episode]) - run.log({"reward": episode.reward}, step=step) run.finish(summary=pr.metrics.from_episodes(episodes)) -``init()`` opens the run and returns a handle carrying its ID and dashboard -URL; records stream out on a background thread as the run proceeds; ``finish()`` -closes it out. A ``with`` block does the last part for you, including on the -paths where the producer never gets to it. - -This is a leaf package on purpose. The ``prime`` CLI depends on ``verifiers``, -so verifiers can never depend on ``prime`` — and verifiers is one of the two -producers this SDK exists to serve. Nothing here imports a producer package; -records are duck-typed through ``to_record()``. +This is a leaf package: the ``prime`` CLI depends on ``verifiers``, so verifiers +can never depend on ``prime``. Nothing here imports a producer package; records +are duck-typed through ``to_record()``. """ from . import metrics, projection -from .backends import Backend, EvalsBackend, OfflineBackend -from .config import Config from .exceptions import ( ConfigurationError, EnvironmentResolutionError, @@ -40,54 +31,24 @@ TransportError, UnauthorizedError, ) -from .models import ( - CONFIG_SOURCE_KEY, - ConfigSource, - EnvironmentRef, - Mode, - OnError, - RunHandle, - RunKind, - RunSpec, - RunStatus, -) +from .models import CONFIG_SOURCE_KEY, ConfigSource, EnvironmentRef, RunStatus from .projection import build_samples, trace_to_sample -from .run import MODE_ENV, RUN_ID_ENV, Run, init -from .sinks import EvalSamplesSink, OfflineSink, Sink, TracesSink +from .run import MODE_ENV, Run, init __version__ = "0.1.0" __all__ = [ - # The surface almost every caller needs "init", "Run", - "metrics", - "projection", - # Types - "Config", + "RunStatus", "ConfigSource", "CONFIG_SOURCE_KEY", "EnvironmentRef", - "Mode", - "OnError", - "RunHandle", - "RunKind", - "RunSpec", - "RunStatus", "MODE_ENV", - "RUN_ID_ENV", - # Backends & sinks, for callers assembling their own - "Backend", - "EvalsBackend", - "OfflineBackend", - "Sink", - "EvalSamplesSink", - "OfflineSink", - "TracesSink", - # The v0 sample projection, moved here from verifiers + "metrics", + "projection", "build_samples", "trace_to_sample", - # Exceptions "PrimeRunsError", "ConfigurationError", "EnvironmentResolutionError", diff --git a/packages/prime-runs/src/prime_runs/_fork.py b/packages/prime-runs/src/prime_runs/_fork.py index 77b44cf36..ff2eefb2f 100644 --- a/packages/prime-runs/src/prime_runs/_fork.py +++ b/packages/prime-runs/src/prime_runs/_fork.py @@ -1,26 +1,15 @@ """One process-wide ``os.register_at_fork`` hook, shared by everything stateful. -Hosted evals fork after the SDK is initialized, and a forked child inherits far -more than the queue: it gets copies of every open socket and every buffered file -handle. Using those is not merely untidy — two processes writing the same TCP -connection interleave bytes into one HTTP stream, and a duplicated write buffer -gets flushed twice, once from each side. - -So anything holding a connection or a file registers here and gets told to start -over in the child. Two rules for a ``reset_after_fork`` implementation: - -- **Drop, do not close.** Closing an inherited transport can send bytes — a TLS - ``close_notify``, an HTTP ``Connection: close`` — down a socket the parent is - still using. Release the reference and let the child's copies of the - descriptors go when it exits. -- **Do not flush.** A buffer inherited from the parent holds records the parent - has not written yet and will write itself. Flushing it in the child writes - them a second time. - -Registration is weak and the hook is installed once. A per-object -``register_at_fork`` call cannot be undone, so registering per instance would -pin every run the process ever opened in memory and re-run hooks for runs that -finished hours ago. +A forked child inherits copies of every open socket and buffered file handle. +Anything holding a connection or a file registers here and gets told to start +over in the child. Two rules for ``reset_after_fork``: + +- **Drop, do not close.** Closing an inherited transport can send bytes down a + socket the parent is still using. +- **Do not flush.** An inherited buffer holds records the parent will write. + +Registration is weak and the hook is installed once: ``register_at_fork`` +cannot be undone, so per-instance registration would pin every run forever. """ import logging diff --git a/packages/prime-runs/src/prime_runs/_http.py b/packages/prime-runs/src/prime_runs/_http.py index 3fbda8517..2dd31743f 100644 --- a/packages/prime-runs/src/prime_runs/_http.py +++ b/packages/prime-runs/src/prime_runs/_http.py @@ -1,25 +1,13 @@ """Shared HTTP client for the platform run APIs. -One client for all backends and the legacy samples sink, because they share a -host, a credential and a retry policy. Two things it does that a bare -``httpx.Client`` does not: - -- maps status codes onto :mod:`prime_runs.exceptions` so callers branch on a - type rather than on a message; -- retries 429/502/503/504 and transport failures with exponential backoff, - honouring ``Retry-After`` when the server sends one. - -Retry safety is decided per call, not per client. A failure is *ambiguous* when -the request may already have been processed — a gateway 502/504, a read timeout, -a stream broken after the bytes went out. Replaying an ambiguous failure is fine -for a GET or a PUT and is not fine for ``POST /evaluations/``: if the platform -created the run and the response was lost, the retry creates a second one and -the SDK only ever knows about the second, leaving an orphaned duplicate. - -So callers declare intent with ``idempotent=``. Unambiguous failures — a -connection that was never established, a 429 refused before any work — are -replayed for every method, because there is nothing on the other side to -duplicate. +Maps status codes onto :mod:`prime_runs.exceptions` and retries 429/502/503/504 +and transport failures with backoff, honouring ``Retry-After``. + +Retry safety is decided per call. A failure is *ambiguous* when the request +may already have been processed (a 502/504, a read timeout). Replaying one is +fine for a GET or PUT and not for ``POST /evaluations/``, which would create a +second run. Callers declare intent with ``idempotent=``; unambiguous failures +(connect errors, 429) are replayed for every method. """ import json @@ -41,13 +29,11 @@ ) DEFAULT_TIMEOUT = httpx.Timeout(60.0, connect=10.0) -# Sample batches are megabytes and the platform fans them out to storage before -# answering, so uploads get their own, much longer budget. +# Sample batches are megabytes; uploads get a longer budget. UPLOAD_TIMEOUT = httpx.Timeout(300.0, connect=10.0) RETRY_STATUS = frozenset({429, 502, 503, 504}) -#: Refused before the server did any work, so replaying cannot duplicate -#: anything. 503 is deliberately *not* here: without a service error code it may -#: equally have come from an intermediary after the request was forwarded. +#: Refused before any work was done, so replaying cannot duplicate anything. +#: 503 is excluded: it may come from an intermediary after forwarding. UNAMBIGUOUS_RETRY_STATUS = frozenset({429}) DEFAULT_MAX_ATTEMPTS = 5 MAX_BACKOFF_SECONDS = 16.0 @@ -74,30 +60,18 @@ def _parse_retry_after(response: httpx.Response) -> Optional[float]: try: return float(raw) except ValueError: - # HTTP-date form. Not worth parsing for a backoff hint — fall back to - # the exponential schedule rather than guessing a clock skew. + # HTTP-date form; fall back to the exponential schedule. return None def normalize_base_url(url: str) -> str: - """The same normalization ``Config`` applies to file/env URLs. - - ``PlatformClient`` appends ``/api/v1`` itself, and platform URLs are - commonly written with the suffix already on them. Without stripping it here - an explicit ``base_url=`` would request ``/api/v1/api/v1/...`` while the - identical value read from ``PRIME_API_BASE_URL`` worked — the config path - strips it and the constructor path did not. - """ + """Strip a trailing ``/api/v1``; the client appends it itself.""" return url.rstrip("/").removesuffix("/api/v1") def encode_json(value: Any) -> bytes: - """Compact UTF-8 JSON, matching the encoding used to size batches. - - ``allow_nan=False`` matters: a NaN reward serialized as JavaScript's bare - ``NaN`` is rejected by strict JSON parsers server-side, and the failure - surfaces as an opaque 400 on a payload the producer cannot inspect. - """ + """Compact UTF-8 JSON. ``allow_nan=False``: strict parsers server-side + reject bare ``NaN`` and the failure is an opaque 400.""" return json.dumps(value, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode( "utf-8" ) @@ -126,8 +100,7 @@ def __init__( self._timeout = timeout self._client = client or self._new_client() if self._owns_client: - # An injected client belongs to the caller (tests, the CLI); only a - # pool we opened ourselves is ours to rebuild after a fork. + # Only a pool we opened ourselves is ours to rebuild after a fork. _fork.register(self) def _new_client(self) -> httpx.Client: @@ -138,13 +111,8 @@ def _new_client(self) -> httpx.Client: ) def reset_after_fork(self) -> None: - """Rebuild the connection pool in a forked child. - - The inherited pool's sockets are the parent's: writing them would - interleave two processes' requests into one HTTP stream. The old client - is dropped rather than closed, because closing can send ``close_notify`` - on a connection the parent is still reading. - """ + """Rebuild the pool in a forked child. The old one is dropped, not + closed: closing could send ``close_notify`` on the parent's socket.""" self._client = self._new_client() def request( @@ -161,10 +129,8 @@ def request( ) -> Dict[str, Any]: """Send one request, retrying transient failures. Returns the JSON body. - ``idempotent`` says whether replaying this request is safe when the - outcome is unknown. It defaults to ``method != "POST"``; a POST that is - in fact safe to replay (get-or-create, setting a terminal state) should - pass ``idempotent=True`` explicitly. + ``idempotent`` defaults to ``method != "POST"``; a POST that is safe to + replay (get-or-create) passes ``idempotent=True`` explicitly. """ url = f"{self.api_prefix}{path}" body = content if content is not None else (encode_json(json_body) if json_body else None) @@ -181,9 +147,7 @@ def request( "headers": headers, "params": dict(params) if params else None, } - # ``None`` disables httpx timeouts; it does not mean "use the - # client's default". Omit the override so ordinary lifecycle - # calls retain the timeout configured in ``_new_client``. + # ``None`` would disable httpx timeouts, not restore the default. if timeout is not None: request_kwargs["timeout"] = timeout response = self._client.request(method, url, **request_kwargs) @@ -212,9 +176,7 @@ def request( if attempt == attempts: break if ambiguous and not replayable: - # The request may already have been processed and replaying it - # would create a second resource. Surfacing the failure is the - # lesser harm: the caller can look, a duplicate cannot be undone. + # Possibly processed already; a duplicate cannot be undone. break after = getattr(last_error, "retry_after", None) time.sleep(retry_delay(attempt, after)) diff --git a/packages/prime-runs/src/prime_runs/backends/__init__.py b/packages/prime-runs/src/prime_runs/backends/__init__.py index 1b8693ade..f94eb5343 100644 --- a/packages/prime-runs/src/prime_runs/backends/__init__.py +++ b/packages/prime-runs/src/prime_runs/backends/__init__.py @@ -1,4 +1,4 @@ -"""Run lifecycle backends, one per platform run system.""" +"""Run lifecycle backends.""" from .base import Backend from .evals import EvalsBackend diff --git a/packages/prime-runs/src/prime_runs/backends/base.py b/packages/prime-runs/src/prime_runs/backends/base.py index f8fc7809d..47f705faf 100644 --- a/packages/prime-runs/src/prime_runs/backends/base.py +++ b/packages/prime-runs/src/prime_runs/backends/base.py @@ -1,12 +1,8 @@ """The contract a run backend implements. -A backend owns one thing: the *lifecycle* of a run — bringing it into -existence, updating what is known about it, and closing it out with a terminal -status. It does not move samples; that is a sink's job (see -:mod:`prime_runs.sinks`). Keeping the two axes independent is what lets the -eval and training APIs — which agree on almost nothing at the wire level — -share a single ``Run`` handle, and what lets the sample transport change -underneath without touching either. +A backend owns the *lifecycle* of a run — creating it, updating what is known +about it, closing it out with a terminal status. It does not move records; +that is a sink's job (see :mod:`prime_runs.sinks`). """ from typing import Any, Dict, Optional, Protocol, runtime_checkable @@ -16,26 +12,8 @@ @runtime_checkable class Backend(Protocol): - """Lifecycle operations for one family of runs.""" - - kind: str - """The ``RunKind`` this backend serves.""" - - supports_step_metrics: bool - """Whether ``log_metrics`` records a point per step. - - ``False`` means the API has no time series and the run keeps a last-value - summary instead. The ``Run`` handle reads this to decide whether - ``log(..., step=)`` is a real write or a summary merge, so producers get - the same call either way. - """ - def create(self, spec: RunSpec) -> RunHandle: - """Open a new run and return its platform identity.""" - ... - - def attach(self, run_id: str) -> RunHandle: - """Re-acquire an existing run, for resume and for non-primary ranks.""" + """Open a new run and return its identity.""" ... def update( @@ -45,18 +23,13 @@ def update( config: Optional[Dict[str, Any]] = None, summary: Optional[Dict[str, Any]] = None, ) -> None: - """Persist config (inputs) and/or summary (outputs) mid-run. + """Persist config (inputs) and/or summary (outputs). - ``config`` is the run's *whole* config, not a patch. The evaluations API - stores metadata with a document-level ``$set``, so a partial write - replaces whatever was there — every caller must send the full picture. + ``config`` is the run's *whole* config, not a patch: the evaluations API + replaces the stored metadata document. """ ... - def log_metrics(self, run_id: str, metrics: Dict[str, Any], step: Optional[int] = None) -> None: - """Append one point to the run's time series. No-op when unsupported.""" - ... - def finalize( self, run_id: str, @@ -66,12 +39,8 @@ def finalize( error: Optional[str] = None, config: Optional[Dict[str, Any]] = None, ) -> None: - """Close the run out. Called exactly once per run. - - ``config`` is passed so a backend that has to record the terminal state - *inside* metadata can merge it into the full config rather than - replacing the document with one key. - """ + """Close the run out. Called exactly once per run. ``config`` is passed + so a backend recording terminal state inside metadata can merge it.""" ... def close(self) -> None: diff --git a/packages/prime-runs/src/prime_runs/backends/evals.py b/packages/prime-runs/src/prime_runs/backends/evals.py index 5a64aadf4..81398dc7c 100644 --- a/packages/prime-runs/src/prime_runs/backends/evals.py +++ b/packages/prime-runs/src/prime_runs/backends/evals.py @@ -1,17 +1,10 @@ -"""Eval runs over ``/api/v1/evaluations/*``. +"""Eval runs over ``/api/v1/evaluations/*``, plus environment resolution +through the hub's get-or-create so a local run uploads without ``prime env push``. -Wraps the endpoints verifiers previously called inline, plus environment -resolution through the hub's get-or-create so a local run uploads without a -prior ``prime env push``. - -One gap is worth stating plainly, because it shapes the code below: **the eval -API has no producer-facing way to mark a run failed.** ``finalize`` moves a run -PROCESSING -> COMPLETED, ``UpdateEvaluationRequest`` carries no ``status``, and -FAILED is written only when an internal Cloud Task trigger fails. So a crashed -run stays RUNNING forever. ``_report_failure`` calls the status endpoint this -SDK needs, treats its absence as expected, and falls back to recording the -terminal state in ``metadata`` so the failure is at least visible and machine -readable. When the endpoint lands, the fallback stops firing on its own. +The eval API has no producer-facing way to mark a run failed: ``finalize`` +moves a run to COMPLETED and ``UpdateEvaluationRequest`` carries no status. A +failed or crashed run therefore keeps showing as running; the terminal state is +recorded in ``metadata.prime_runs`` so it is at least visible. """ import logging @@ -20,32 +13,15 @@ from typing import Any, Dict, List, Optional from .._http import PlatformClient -from ..exceptions import ( - ConfigurationError, - EnvironmentResolutionError, - NotFoundError, - RunAPIError, - is_transient, -) +from ..exceptions import ConfigurationError, EnvironmentResolutionError, RunAPIError from ..models import EnvironmentRef, RunHandle, RunSpec, RunStatus logger = logging.getLogger(__name__) -# Statuses the platform's EvaluationStatus enum uses, keyed by ours. -_PLATFORM_STATUS = { - RunStatus.COMPLETED: "COMPLETED", - RunStatus.FAILED: "FAILED", - RunStatus.CRASHED: "FAILED", -} - class EvalsBackend: """Lifecycle for evaluation runs.""" - kind = "eval" - # The evaluations API stores a single metrics blob, not a time series. - supports_step_metrics = False - def __init__( self, client: PlatformClient, @@ -56,7 +32,6 @@ def __init__( self._client = client self._frontend_url = frontend_url.rstrip("/") self._team_id = team_id - self._status_endpoint_missing = False # ------------------------------------------------------------------ create @@ -75,18 +50,15 @@ def create(self, spec: RunSpec) -> RunHandle: "tags": list(spec.tags), } _set_if(payload, "model_name", spec.model) - # Derived, not asked for: the API has the column and it is always the - # environment under a different name, so making producers repeat it - # bought nothing. + # The API's `dataset` column is always the environment under another name. _set_if(payload, "dataset", _first_environment_name(spec)) _set_if(payload, "framework", spec.framework) _set_if(payload, "description", spec.description) _set_if(payload, "metadata", spec.config or None) _set_if(payload, "team_id", spec.team_id or self._team_id) - # Not replayable: POST defaults to idempotent=False here on purpose. If - # the platform created the run and the response was lost, a retry would - # create a second one and only the second would be tracked. + # Not replayable (the POST default): a retry after a lost response would + # create a second run and only the second would be tracked. response = self._client.post("/evaluations/", json_body=payload) run_id = response.get("evaluation_id") if not run_id: @@ -97,29 +69,6 @@ def create(self, spec: RunSpec) -> RunHandle: id=run_id, name=str(response.get("name") or run_name), url=response.get("viewer_url") or self.url_for(run_id), - raw=response, - ) - - def attach(self, run_id: str) -> RunHandle: - try: - response = self._client.get(f"/evaluations/{run_id}") - except NotFoundError: - raise - except RunAPIError as exc: - # Attach is a convenience — a resume or a non-primary rank joining. - # Losing the run's name to a transient read is not worth failing on; - # the ID is what everything downstream actually needs. - if not is_transient(exc) and not ( - exc.status_code is not None and exc.status_code >= 500 - ): - raise - logger.debug("Could not read evaluation %s on attach: %s", run_id, exc) - return RunHandle(id=run_id, url=self.url_for(run_id)) - return RunHandle( - id=run_id, - name=response.get("name"), - url=response.get("viewer_url") or self.url_for(run_id), - raw=response, ) def url_for(self, run_id: str) -> str: @@ -134,12 +83,8 @@ def update( config: Optional[Dict[str, Any]] = None, summary: Optional[Dict[str, Any]] = None, ) -> None: - """Persist config and/or summary. - - ``config`` must be the run's whole config: the service writes metadata - with ``{"$set": {"metadata": ...}}``, which replaces the stored document - rather than merging into it. - """ + """``config`` must be the whole config: the service writes metadata with + a document-level ``$set``.""" payload: Dict[str, Any] = {} _set_if(payload, "metadata", config or None) _set_if(payload, "metrics", summary or None) @@ -147,13 +92,6 @@ def update( return self._client.put(f"/evaluations/{run_id}", json_body=payload) - def log_metrics(self, run_id: str, metrics: Dict[str, Any], step: Optional[int] = None) -> None: - """No-op: see ``supports_step_metrics``. - - The run keeps a last-value summary and flushes it through ``update``, - which is the only shape this API can store. - """ - # ---------------------------------------------------------------- finalize def finalize( @@ -168,74 +106,25 @@ def finalize( if status is RunStatus.COMPLETED: body: Dict[str, Any] = {} _set_if(body, "metrics", summary or None) - # Finalization also enqueues the platform's asynchronous statistics - # task. A lost response leaves the outcome ambiguous, so replaying - # this POST can enqueue the work twice. + # Finalization enqueues the platform's statistics task; replaying an + # ambiguous failure could enqueue it twice, so this stays non-idempotent. self._client.post( f"/evaluations/{run_id}/finalize", json_body=body or {"metrics": {}}, ) return - self._report_failure(run_id, status=status, summary=summary, error=error, config=config) - def _report_failure( - self, - run_id: str, - *, - status: RunStatus, - summary: Optional[Dict[str, Any]], - error: Optional[str], - config: Optional[Dict[str, Any]] = None, - ) -> None: - """Mark a run failed, or record why we could not.""" + # No status endpoint exists yet. Record the terminal state in metadata, + # merged into the full config because this PUT replaces the document. terminal = { "status": status.value, "finished_at": datetime.now(timezone.utc).isoformat(), } if error: terminal["error"] = error - - if not self._status_endpoint_missing: - try: - self._client.post( - f"/evaluations/{run_id}/status", - json_body={"status": _PLATFORM_STATUS[status], "error": error}, - idempotent=True, - ) - return - except NotFoundError: - # Expected until the status endpoint ships. Latch so a run that - # fails repeatedly does not pay for the probe every time. - self._status_endpoint_missing = True - logger.debug( - "Platform has no /evaluations/{id}/status endpoint; " - "recording terminal state in metadata instead" - ) - except RunAPIError as exc: - if is_transient(exc) or (exc.status_code is not None and exc.status_code >= 500): - logger.warning( - "Status endpoint remained unavailable after retries (%s); " - "recording terminal state in metadata instead", - exc, - ) - elif exc.status_code not in (405, 422): - raise - else: - self._status_endpoint_missing = True - logger.debug("Status endpoint rejected the request (%s); using metadata", exc) - - # Fallback: the run cannot be moved out of RUNNING, but the failure is - # at least recorded where an operator and the dashboard can both read it. - # The terminal block is merged into the full config because this PUT - # replaces the stored metadata document — sending it alone would erase - # everything finish() just wrote. - self.update( - run_id, - config={**(config or {}), "prime_runs": terminal}, - summary=summary, - ) + self.update(run_id, config={**(config or {}), "prime_runs": terminal}, summary=summary) logger.warning( - "Run %s %s, but its evaluation status could not be updated; it will keep " + "Run %s %s, but the evaluations API cannot record that; it will keep " "showing as running. Recorded the failure in metadata.prime_runs.", run_id, status.value, @@ -247,18 +136,10 @@ def close(self) -> None: # ----------------------------------------------------------- environments def _resolve_environments(self, refs: List[EnvironmentRef]) -> List[Dict[str, Any]]: - """Hub references as the API's ``EnvironmentReference`` objects. - - ``version_id`` is carried through when the producer pinned one — the - API accepts it, and dropping it would silently attach the run to - whatever version the hub resolves today, which is the difference - between a reproducible eval and one that quietly moved. - - Unlike the old client, a reference that cannot be resolved raises - instead of being skipped: dropping one silently produces a run attached - to the wrong environments, which looks like a successful upload and is - found much later. - """ + """Hub references as the API's ``EnvironmentReference`` objects, carrying + a pinned ``version_id`` through. A reference that cannot be resolved + raises rather than being skipped: a run silently attached to the wrong + environments looks like a successful upload.""" resolved: List[Dict[str, Any]] = [] for ref in refs: entry: Dict[str, Any] = {"id": ref.id or self._lookup_environment(ref)} @@ -267,7 +148,6 @@ def _resolve_environments(self, refs: List[EnvironmentRef]) -> List[Dict[str, An return resolved def _lookup_environment(self, ref: EnvironmentRef) -> str: - """Resolve one environment reference to a hub ID.""" if ref.slug: owner_slug, name = ref.slug.split("/", 1) try: @@ -309,19 +189,11 @@ def _first_environment_name(spec: RunSpec) -> Optional[str]: if ref.name: return ref.name if ref.slug: - # Published environments are addressed as ``owner/name``, but the - # name portion is still the dataset/run-name fallback used for - # ordinary environment references. return ref.slug.split("/", 1)[1] return None def _default_name(spec: RunSpec) -> str: - """A name for producers that did not pick one. - - The API requires a name, so the alternative to generating one is a 422 at - the worst possible moment. Leads with the environment so runs sort together - in the dashboard list. - """ - stem = _first_environment_name(spec) or spec.framework or spec.kind + """The API requires a name; lead with the environment so runs sort together.""" + stem = _first_environment_name(spec) or spec.framework or "eval" return f"{stem}-{uuid.uuid4().hex[:8]}" diff --git a/packages/prime-runs/src/prime_runs/backends/offline.py b/packages/prime-runs/src/prime_runs/backends/offline.py index e0c508dde..50703f5c7 100644 --- a/packages/prime-runs/src/prime_runs/backends/offline.py +++ b/packages/prime-runs/src/prime_runs/backends/offline.py @@ -1,17 +1,9 @@ -"""Offline runs: a local directory that looks exactly like a platform run. - -This is not a debugging affordance, it is the reason producers can delete their -``--no-push`` branching. A run that never reaches the network still has an ID, -a status, a config, a summary and a metrics stream, so the call sites above it -are identical whether or not anyone is logged in. The locally issued ID is used -as the run ID everywhere — including inside the trace documents — so a later -sync attaches the archive to a platform run without rewriting a single record. +"""Offline runs: a local directory with the same lifecycle as a platform run. Layout, one directory per run:: - //run.json spec + status + timestamps - //metrics.jsonl one JSON object per log() call - //records/ whatever the offline sink wrote + //run.json spec + status + timestamps + //records/ whatever the offline sink wrote """ import json @@ -43,9 +35,6 @@ def new_run_id() -> str: class OfflineBackend: """Run lifecycle recorded on the local filesystem.""" - kind = "offline" - supports_step_metrics = True - def __init__(self, directory: Union[str, Path, None] = None) -> None: self.directory = Path(directory) if directory is not None else default_dir() @@ -60,24 +49,12 @@ def create(self, spec: RunSpec) -> RunHandle: state: Dict[str, Any] = { "id": run_id, "name": run_name, - "kind": spec.kind, "status": RunStatus.RUNNING.value, "created_at": _now(), "spec": _spec_to_json(spec), } self._write_state(run_id, state) - return RunHandle(id=run_id, name=run_name, url=str(path.resolve()), raw=state) - - def attach(self, run_id: str) -> RunHandle: - path = self.run_dir(run_id) - path.mkdir(parents=True, exist_ok=True) - state = self._read_state(run_id) - return RunHandle( - id=run_id, - name=str(state.get("name") or run_id), - url=str(path.resolve()), - raw=state, - ) + return RunHandle(id=run_id, name=run_name, url=str(path.resolve())) def update( self, @@ -94,13 +71,6 @@ def update( state["updated_at"] = _now() self._write_state(run_id, state) - def log_metrics(self, run_id: str, metrics: Dict[str, Any], step: Optional[int] = None) -> None: - line = {"step": step, "timestamp": _now(), **metrics} - path = self.run_dir(run_id) / "metrics.jsonl" - path.parent.mkdir(parents=True, exist_ok=True) - with path.open("a", encoding="utf-8") as handle: - handle.write(json.dumps(line, ensure_ascii=False, default=str) + "\n") - def finalize( self, run_id: str, @@ -122,7 +92,7 @@ def finalize( self._write_state(run_id, state) def close(self) -> None: - """Nothing to release — every write is already flushed to disk.""" + """Every write is already on disk.""" # ------------------------------------------------------------------ state @@ -143,8 +113,7 @@ def _read_state(self, run_id: str) -> Dict[str, Any]: def _write_state(self, run_id: str, state: Dict[str, Any]) -> None: path = self._state_path(run_id) path.parent.mkdir(parents=True, exist_ok=True) - # Write-then-rename: a crash mid-write must not leave the run's own - # record truncated, since it is the only description of what ran. + # Write-then-rename so a crash mid-write never truncates the record. temp = path.with_suffix(".json.tmp") temp.write_text(json.dumps(state, indent=2, ensure_ascii=False, default=str), "utf-8") temp.replace(path) diff --git a/packages/prime-runs/src/prime_runs/config.py b/packages/prime-runs/src/prime_runs/config.py index e84ec2985..2326c8208 100644 --- a/packages/prime-runs/src/prime_runs/config.py +++ b/packages/prime-runs/src/prime_runs/config.py @@ -1,9 +1,5 @@ -"""Lightweight configuration for the Prime Runs SDK. - -Same shape as the other prime SDK packages: reads ``~/.prime/config.json`` -plus environment variables, env taking precedence. Adds ``frontend_url`` -(where run URLs point) on top of the shared platform settings. -""" +"""Configuration: ``~/.prime/config.json`` plus environment variables, env +taking precedence. Same shape as the other prime SDKs, plus ``frontend_url``.""" import json import os @@ -64,11 +60,7 @@ def base_url(self) -> str: @property def frontend_url(self) -> str: - """Dashboard base URL, used to build the run URL a producer prints. - - The platform returns a ``viewer_url`` on create; this is the fallback - for responses that omit it and for offline/legacy paths. - """ + """Dashboard base URL; fallback when a create response omits ``viewer_url``.""" env_val = os.getenv("PRIME_FRONTEND_URL") if env_val: return env_val.rstrip("/") diff --git a/packages/prime-runs/src/prime_runs/exceptions.py b/packages/prime-runs/src/prime_runs/exceptions.py index 1d89fe334..c17ab94c2 100644 --- a/packages/prime-runs/src/prime_runs/exceptions.py +++ b/packages/prime-runs/src/prime_runs/exceptions.py @@ -1,10 +1,6 @@ -"""Exceptions for the Prime Runs SDK. - -Producers run for hours; the default posture is that nothing here escapes into -a training loop (``on_error="warn"``). These types exist so that callers who -opt into ``on_error="raise"`` — tests, CI, hosted workers — can branch on what -actually failed instead of matching log strings. -""" +"""Exceptions for the Prime Runs SDK. Nothing here escapes into a producer +loop by default (``on_error="warn"``); callers opting into ``on_error="raise"`` +can branch on these types.""" from typing import Optional @@ -14,11 +10,8 @@ class PrimeRunsError(Exception): class ConfigurationError(PrimeRunsError): - """The SDK was asked to do something the local configuration cannot support. - - Missing API key, an unknown ``kind``, ``mode="online"`` with no way to reach - the platform. Raised before any request is made. - """ + """Missing API key, unreadable config file, unknown mode. Raised before + any request is made.""" class RunAPIError(PrimeRunsError): @@ -45,14 +38,9 @@ class PaymentRequiredError(RunAPIError): class ForbiddenError(RunAPIError): - """403 — authenticated, but not allowed to do this. - - Distinct from 401 because the credential is fine and re-authenticating will - not help: the run belongs to another owner, the team header names a team the - key cannot act for, or the feature is gated to an allowlist. Named to match - ``prime_traces.ForbiddenError``, which the traces sink already branches on - to retire itself when an account is outside the closed beta. - """ + """403 — authenticated, but not allowed: another owner's run, a team the + key cannot act for, or a feature gated to an allowlist. Named to match + ``prime_traces.ForbiddenError``.""" class NotFoundError(RunAPIError): @@ -79,16 +67,9 @@ class TransportError(RunAPIError): def is_transient(exc: BaseException) -> bool: - """Whether a failure is about this moment rather than this run. - - The distinction decides whether a sink is retired. A gated account or a bad - credential will fail identically on every future batch, so the sink should - stop. A gateway blip or a dropped connection will not, and retiring a sink - for one of those means a single 502 empties the rest of the run's dashboard. - - Covers the traces service's exception family as well as this package's, - since both reach the uploader through the same path. - """ + """Whether a failure is about this moment (retry later) rather than this + run (stop). Decides whether a sink is retired. Covers the traces service's + exception family too, since both reach the uploader through one path.""" if isinstance(exc, (RetryableAPIError, TransportError)): return True try: @@ -100,17 +81,9 @@ def is_transient(exc: BaseException) -> bool: class EnvironmentResolutionError(PrimeRunsError): - """An environment named in ``init()`` could not be resolved to a hub ID. - - Distinct from a generic API error because it is usually a typo or a - permissions problem on the environment, not an outage, and because an eval - run cannot be created without at least one resolved environment. - """ + """An environment named in ``init()`` could not be resolved to a hub ID — + usually a typo or a permissions problem, not an outage.""" class RunFinishedError(PrimeRunsError): - """A finished run was written to again. - - Terminal status is reported once. Logging after ``finish()`` is a producer - bug — the data would land on a run the platform has already closed out. - """ + """A finished run was written to again: a producer bug.""" diff --git a/packages/prime-runs/src/prime_runs/metrics.py b/packages/prime-runs/src/prime_runs/metrics.py index 25ec6d888..7b31d9928 100644 --- a/packages/prime-runs/src/prime_runs/metrics.py +++ b/packages/prime-runs/src/prime_runs/metrics.py @@ -1,14 +1,6 @@ -"""Run-level aggregates over native episodes. - -Opt-in, not automatic. The SDK does not decide what a run's headline number is -— that judgement belongs next to the producer, which knows which agents are -being scored and what counts as an error. This module ships the aggregation -verifiers already used, so callers migrating off ``verifiers.v1.utils.platform`` -keep byte-identical dashboard numbers, and anyone else can pass their own dict -to ``run.finish(summary=...)``. - -Duck-typed like :mod:`prime_runs.projection`: no producer package is imported. -""" +"""Run-level aggregates over native episodes — the aggregation verifiers used, +so migrated runs keep identical dashboard numbers. Opt-in: pass the result to +``run.finish(summary=...)``. Duck-typed; no producer package is imported.""" from typing import Any, Dict, Optional, Sequence diff --git a/packages/prime-runs/src/prime_runs/models.py b/packages/prime-runs/src/prime_runs/models.py index a37112268..15a5d7cb4 100644 --- a/packages/prime-runs/src/prime_runs/models.py +++ b/packages/prime-runs/src/prime_runs/models.py @@ -1,10 +1,7 @@ """Types shared across backends, sinks and the ``Run`` handle. -Only the values that cross a module boundary live here. Response bodies are -deliberately *not* modeled: the platform returns more fields than any producer -reads, and freezing them in pydantic here would make every backend addition a -breaking SDK release. Backends pull the two or three fields they need and hand -back a ``RunHandle``. +Response bodies are deliberately not modeled: backends pull the two or three +fields they need and hand back a ``RunHandle``. """ import os @@ -15,25 +12,21 @@ from .exceptions import ConfigurationError -RunKind = Literal["eval", "train"] -"""Which run system owns the lifecycle. Selects the backend.""" - Mode = Literal["online", "offline", "disabled"] -"""``online`` talks to the platform, ``offline`` writes a local run directory -that can be synced later, ``disabled`` makes every call a no-op while keeping -the same object shape so producer code needs no branching.""" +"""``online`` talks to the platform, ``offline`` writes a local run directory, +``disabled`` makes every call a no-op with the same object shape.""" OnError = Literal["warn", "raise"] +RUN_KIND = "eval" +"""Stamped as ``run.type`` on records and sent as upload provenance.""" + class RunStatus(str, Enum): """Terminal state a producer can report. - ``crashed`` is distinct from ``failed``: ``failed`` means the producer - decided the run failed, ``crashed`` means the process exited without ever - saying. Only the second one is inferred by the SDK (atexit / signal), and - the distinction is what tells an operator whether to look at the run's own - error or at the machine it ran on. + ``failed`` means the producer said the run failed; ``crashed`` means the + process exited without saying (only the SDK's atexit hook reports it). """ RUNNING = "running" @@ -50,8 +43,7 @@ class EnvironmentRef: """An environment as a producer names it, before hub resolution. ``id`` short-circuits resolution; ``name`` goes through the hub's - get-or-create so a local run uploads without a prior ``prime env push``; - ``slug`` looks up an already-published ``owner/name`` environment. + get-or-create; ``slug`` looks up a published ``owner/name`` environment. """ name: Optional[str] = None @@ -87,19 +79,12 @@ def __post_init__(self) -> None: CONFIG_SOURCE_KEY = "config_source" -"""Where a config file lands inside a run's config, as a :class:`ConfigSource` dict. - -Inside the config rather than beside it, so every path that already carries the -config carries the file too — create, the periodic update, finalize, the failure -fallback, and the offline archive — with no extra plumbing and no chance of one -of them forgetting it. It is also the one key a config-tab renderer has to know -about: present means "show this verbatim", absent means "show the structure". -""" +"""Where a config file lands inside a run's config, as a :class:`ConfigSource` +dict. Present means "render this verbatim"; absent means "render the structure".""" MAX_CONFIG_SOURCE_BYTES = 256 * 1024 -"""Ceiling on a stored config file. A hand-written run config is single-digit -kilobytes; anything past this is a dataset or a log that would bloat the run's -metadata document, so it is refused loudly at ``init()`` rather than truncated.""" +"""A hand-written run config is kilobytes; anything past this is refused at +``init()`` rather than truncated.""" _CONFIG_SOURCE_FORMATS = { ".toml": "toml", @@ -113,19 +98,9 @@ def __post_init__(self) -> None: class ConfigSource: """The config file a run was started from, kept byte-for-byte. - Both producers are now launched from one user-authored file — ``uv run eval - @ eval.toml``, ``uv run rl @ train.toml`` — and that file *is* the run's - real configuration. A resolved model dump is a different artifact: it - answers "what did every knob end up as", not "what did someone write", and - it loses comments, key order and section grouping on the way through. - - So when ``init(config=...)`` is given a path, this is what gets stored: the - bytes, not a parse of them. Callers without a file pass a mapping or a model - to the same parameter and get the structured form instead. - - Nothing here is redacted. A config file that carries a secret will carry it - onto the run's page, the same way it already reaches anyone who can read the - repository it lives in — keep credentials in the environment, not the file. + That file *is* the run's configuration — comments, key order and section + grouping included — where a resolved model dump is a different artifact. + Nothing here is redacted: keep credentials in the environment, not the file. """ text: str @@ -186,10 +161,8 @@ def from_file(cls, path: Union[str, "os.PathLike[str]"]) -> "ConfigSource": def coerce(cls, value: Any) -> Optional["ConfigSource"]: """Normalize the config-file form of ``init(config=...)``. - A ``str`` or ``PathLike`` is a *path*, never inline text — that is how - every caller will reach for it, and guessing between the two would turn - a mistyped filename into a run whose config tab shows the filename. - Inline text goes through ``ConfigSource(text=...)`` explicitly. + A ``str`` or ``PathLike`` is a *path*, never inline text; inline text + goes through ``ConfigSource(text=...)`` explicitly. """ if value is None or isinstance(value, cls): return value @@ -217,24 +190,17 @@ def __post_init__(self) -> None: @dataclass class RunSpec: - """Everything a backend needs to open a run, in producer vocabulary. - - This is the argument surface of ``init()`` after normalization — backends - translate it into whatever their API family calls these things, which is - the whole reason eval and training runs can share one handle. - """ + """Everything a backend needs to open a run: ``init()``'s arguments after + normalization. ``config`` is the run's inputs; outputs accumulate on the + handle as ``summary``.""" name: Optional[str] = None - kind: RunKind = "eval" environments: List[EnvironmentRef] = field(default_factory=list) model: Optional[str] = None framework: Optional[str] = None description: Optional[str] = None tags: List[str] = field(default_factory=list) team_id: Optional[str] = None - # Only the inputs. A run's outputs (-> the platform's `metrics`) are not - # here because a run being opened does not have any yet; they accumulate on - # the handle and are written by update() and finalize(). config: Dict[str, Any] = field(default_factory=dict) @@ -245,4 +211,3 @@ class RunHandle: id: str name: Optional[str] = None url: Optional[str] = None - raw: Dict[str, Any] = field(default_factory=dict) diff --git a/packages/prime-runs/src/prime_runs/projection.py b/packages/prime-runs/src/prime_runs/projection.py index 4a4af5e47..1474bb3ae 100644 --- a/packages/prime-runs/src/prime_runs/projection.py +++ b/packages/prime-runs/src/prime_runs/projection.py @@ -1,23 +1,15 @@ """Projection of native traces onto the platform's v0 eval-sample format. -Moved here from ``verifiers.v1.utils.platform``. This is knowledge about a -*platform wire format*, so it belongs in the client for that wire — not in an -eval framework that prime-rl then has to reach across a repo boundary to import -(``from verifiers.v1.push import trace_to_sample``, which had already drifted -from the module's real path). +Moved here from ``verifiers.v1.utils.platform``: this is knowledge about a +platform wire format, so it lives in the client for that wire. Duck-typed — +verifiers ``Trace``/``Episode`` satisfy it structurally and are not imported. -Everything is duck-typed. Verifiers ``Trace``/``Episode`` and prime-rl -``Rollout`` satisfy it structurally, and none of them is imported: the leaf -package that both producers depend on cannot depend back on either. - -The projection exists for the *current* viewer, which reads the flat sample -table. Once the Viewer API reads traces natively this module stops being on -the default path — which is why it is a standalone function rather than -something woven through the run lifecycle. +The projection serves the *current* viewer, which reads the flat sample table; +it leaves the default path once the Viewer API reads traces natively. """ import logging -from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence +from typing import Any, Dict, Iterable, List, Optional, Sequence from ._http import encode_json @@ -38,11 +30,6 @@ def _dump(items: Iterable[Any]) -> List[Dict[str, Any]]: return [item.model_dump(mode="json", exclude_none=True) for item in items] -def is_episode(record: Any) -> bool: - """Whether a record is an episode (a group of traces) rather than a trace.""" - return hasattr(record, "traces") and not hasattr(record, "branches") - - def summary_trace_index(episode: Any) -> int: """Index of the trace whose flat projection represents the episode. @@ -113,283 +100,6 @@ def trace_to_sample( return sample -def trace_record_to_sample( - trace: Mapping[str, Any], rollout_number: int = 1, episode_id: Optional[str] = None -) -> Dict[str, Any]: - """Project a serialized trace without importing its producer package. - - Verifiers persists its message graph as ``nodes``/``calls`` rather than the - derived ``branches`` property used by :func:`trace_to_sample`. Older records - may carry branches directly, so both representations are accepted. Sparse - trace mappings still produce a visible row as long as they have an ID; - fields the legacy viewer cannot derive remain empty instead of making the - entire record disappear. - """ - trace_id = trace.get("id") - if not trace_id: - raise TypeError("serialized trace records must contain a non-empty 'id'") - - task_container = _as_mapping(trace.get("task")) - task = dict(_as_mapping(task_container.get("data"))) - agent = _as_mapping(trace.get("agent")) - branches = _serialized_branches(trace) - rewards = _as_mapping(trace.get("rewards")) - errors = trace.get("errors") - last_error = trace.get("last_error") - if last_error is None and isinstance(errors, list) and errors: - last_error = errors[-1] - stop_condition = trace.get("stop_condition") - - sample: Dict[str, Any] = { - "sample_id": trace_id, - "example_id": task.get("idx"), - "rollout_number": rollout_number, - "episode_id": episode_id, - "agent": agent.get("name"), - "trainable": agent.get("trainable", True), - "task": task, - "prompt": [], - "completion": branches[-1]["messages"] if branches else [], - "answer": task.get("answer"), - "tool_defs": _mapping_list(trace.get("tools")) or None, - "reward": trace["reward"] if "reward" in trace else _total_reward(rewards), - "timing": dict(_as_mapping(trace.get("timing"))) or None, - "is_completed": trace.get("is_completed", False), - "is_truncated": trace.get("is_truncated", _is_truncated(stop_condition, trace)), - "metrics": dict(_as_mapping(trace.get("metrics"))), - "error": dict(last_error) if isinstance(last_error, Mapping) else last_error, - "stop_condition": stop_condition, - "trajectory": branches, - "token_usage": dict(_as_mapping(trace.get("usage"))) or _aggregate_usage(trace), - "info": dict(_as_mapping(trace.get("info"))) or None, - } - for name, reward in rewards.items(): - if reward is None: - continue - score = reward.get("score") if isinstance(reward, Mapping) else reward - sample.setdefault(name, score) - return sample - - -def record_to_samples( - record: Mapping[str, Any], rollout_numbers: Optional[Dict[Any, int]] = None -) -> List[Dict[str, Any]]: - """Project one serialized trace or episode to legacy viewer samples.""" - counts = rollout_numbers if rollout_numbers is not None else {} - if "traces" not in record: - task = _as_mapping(_as_mapping(record.get("task")).get("data")) - idx = task.get("idx") - rollout_key = idx if idx is not None else record.get("id") - counts[rollout_key] = number = counts.get(rollout_key, 0) + 1 - return [trace_record_to_sample(record, rollout_number=number)] - - episode_id = record.get("id") - if not episode_id: - raise TypeError("serialized episode records must contain a non-empty 'id'") - raw_traces = record.get("traces") - if not isinstance(raw_traces, list): - raise TypeError("serialized episode 'traces' must be a list") - traces: List[Mapping[str, Any]] = [] - for trace in raw_traces: - if not isinstance(trace, Mapping): - raise TypeError("serialized episode traces must be mappings") - traces.append(trace) - if not traces: - return [] - - summary_index = next( - ( - index - for index, trace in enumerate(traces) - if _as_mapping(trace.get("agent")).get("trainable", True) - ), - 0, - ) - summary_task = _as_mapping(_as_mapping(traces[summary_index].get("task")).get("data")) - idx = summary_task.get("idx") - rollout_key = idx if idx is not None else episode_id - counts[rollout_key] = number = counts.get(rollout_key, 0) + 1 - sample = trace_record_to_sample(traces[summary_index], number, str(episode_id)) - sample["sample_id"] = episode_id - sample["info"] = { - **(sample["info"] or {}), - "native_wrapper": dict(record), - "native_trace_index": summary_index, - } - if ENVELOPE_BYTES + json_bytes(sample) <= MAX_SAMPLES_PAYLOAD_BYTES: - return [sample] - - logger.warning( - "Episode %s exceeds the platform sample limit; uploading projected traces", - episode_id, - ) - return [trace_record_to_sample(trace, number, str(episode_id)) for trace in traces] - - -def _as_mapping(value: Any) -> Mapping[str, Any]: - return value if isinstance(value, Mapping) else {} - - -def _mapping_list(value: Any) -> List[Dict[str, Any]]: - if not isinstance(value, list): - return [] - return [dict(item) for item in value if isinstance(item, Mapping)] - - -def _serialized_branches(trace: Mapping[str, Any]) -> List[Dict[str, Any]]: - raw_branches = trace.get("branches") - if isinstance(raw_branches, list): - return [ - { - "messages": _mapping_list(_as_mapping(branch).get("messages")), - "num_input_tokens": _as_mapping(branch).get("num_input_tokens", 0), - "num_output_tokens": _as_mapping(branch).get("num_output_tokens", 0), - } - for branch in raw_branches - if isinstance(branch, Mapping) - ] - - raw_nodes = trace.get("nodes") - if not isinstance(raw_nodes, list): - return [] - nodes = [_as_mapping(node) for node in raw_nodes] - parents = {node.get("parent") for node in nodes if isinstance(node.get("parent"), int)} - leaves = [index for index in range(len(nodes)) if index not in parents] - calls = trace.get("calls") - calls_by_node = ( - { - call.get("node"): call - for call in calls - if isinstance(call, Mapping) and isinstance(call.get("node"), int) - } - if isinstance(calls, list) - else {} - ) - - branches: List[Dict[str, Any]] = [] - for leaf in leaves: - path: List[int] = [] - seen = set() - node_index: Any = leaf - while ( - isinstance(node_index, int) and 0 <= node_index < len(nodes) and node_index not in seen - ): - seen.add(node_index) - path.append(node_index) - node_index = nodes[node_index].get("parent") - path.reverse() - branch_calls = [calls_by_node[index] for index in path if index in calls_by_node] - input_tokens, output_tokens = _branch_token_counts(branch_calls) - branches.append( - { - "messages": [ - dict(message) - for index in path - if isinstance((message := nodes[index].get("message")), Mapping) - ], - "num_input_tokens": input_tokens, - "num_output_tokens": output_tokens, - } - ) - return branches - - -def _branch_token_counts(calls: Sequence[Mapping[str, Any]]) -> tuple[int, int]: - input_tokens = 0 - output_tokens = 0 - previous_total = 0 - for call in calls: - usage = _as_mapping(call.get("usage")) - current_input, current_output = _usage_counts(usage) - input_tokens += max(0, current_input - previous_total) - output_tokens += current_output - previous_total = int(usage.get("total_tokens", current_input + current_output) or 0) - return input_tokens, output_tokens - - -def _aggregate_usage(trace: Mapping[str, Any]) -> Optional[Dict[str, Any]]: - calls = trace.get("calls") - if not isinstance(calls, list): - return None - usages = [usage for call in calls if (usage := _as_mapping(_as_mapping(call).get("usage")))] - if not usages: - return None - if any("prompt_tokens" in usage or "completion_tokens" in usage for usage in usages): - result: Dict[str, Any] = { - "prompt_tokens": sum(int(usage.get("prompt_tokens", 0) or 0) for usage in usages), - "completion_tokens": sum( - int(usage.get("completion_tokens", 0) or 0) for usage in usages - ), - } - for key in ("cached_input_tokens", "reasoning_tokens", "cost"): - values = [usage[key] for usage in usages if usage.get(key) is not None] - if values: - result[key] = sum(values) - return result - - input_tokens = 0 - output_tokens = 0 - for usage in usages: - current_input, current_output = _usage_counts(usage) - input_tokens += current_input - output_tokens += current_output - if not input_tokens and not output_tokens: - return None - return { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - "total_tokens": input_tokens + output_tokens, - } - - -def _usage_counts(usage: Mapping[str, Any]) -> tuple[int, int]: - if "input_tokens" in usage: - input_tokens = int(usage.get("input_tokens", 0) or 0) - else: - input_tokens = int(usage.get("prompt_tokens", 0) or 0) + int( - usage.get("cached_input_tokens", 0) or 0 - ) - output_tokens = int(usage.get("output_tokens", usage.get("completion_tokens", 0)) or 0) - return input_tokens, output_tokens - - -def _total_reward(rewards: Mapping[str, Any]) -> float: - total = 0.0 - for reward in rewards.values(): - if reward is None: - continue - if isinstance(reward, Mapping): - total += float(reward.get("score", 0.0) or 0.0) * float( - reward.get("weight", 1.0) or 0.0 - ) - else: - total += float(reward) - return total - - -def _is_truncated(stop_condition: Any, trace: Mapping[str, Any]) -> bool: - if stop_condition in { - "max_turns", - "max_input_tokens", - "max_output_tokens", - "max_total_tokens", - "context_length", - }: - return True - calls = trace.get("calls") - if not isinstance(calls, list): - return False - last_successful = next( - ( - call - for call in reversed(calls) - if isinstance(call, Mapping) and call.get("error") is None - ), - None, - ) - return bool(last_successful and last_successful.get("finish_reason") == "length") - - def episode_to_samples(episode: Any, rollout_number: int) -> List[Dict[str, Any]]: """One episode -> the sample rows the platform should store for it. diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index dd91570e6..2eb93db9e 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -1,79 +1,57 @@ """The run handle, and ``init()`` that produces one. -A run is a long-lived thing with a status, so it is an object, not three -stateless calls. That single change is what lets the SDK take on the work every -producer was doing privately: streaming instead of buffering, containing its -own errors, reporting a terminal status when the process dies, and behaving the -same on rank 3 of a training job as on a laptop. - -The identity rule matters most, so it is worth stating once. ``init()`` is -called *before* rollouts start, and the ID it returns is *the* run ID -everywhere — including inside every trace document the producer writes, and -including the local archive. Nothing is re-stamped afterwards and no producer -record is rewritten. Verifiers already stamps ``EvalRunInfo(id=config.run.id)`` -at rollout time; the only change is where that ID comes from. Offline runs get -a locally issued ID through the same path, so there is one code path, not two. +A run is a long-lived thing with a status, so it is an object rather than three +stateless calls: records stream instead of buffering, errors are contained, and +a process that dies still reports a terminal status. + +``init()`` is called *before* rollouts start, and the ID it returns is *the* run +ID everywhere — including inside every trace document the producer writes. +Nothing is re-stamped afterwards. Offline runs get a locally issued ID through +the same path. """ import atexit -import inspect import logging import math import os -import signal import threading import time from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Union from . import _fork from ._http import DEFAULT_TIMEOUT, UPLOAD_TIMEOUT, PlatformClient -from .backends import Backend, EvalsBackend, OfflineBackend -from .backends.offline import DEFAULT_DIR_ENV +from .backends import Backend, EvalsBackend, OfflineBackend, new_run_id from .config import Config from .exceptions import ConfigurationError, RunFinishedError from .models import ( CONFIG_SOURCE_KEY, + RUN_KIND, ConfigSource, EnvironmentRef, Mode, OnError, RunHandle, - RunKind, RunSpec, RunStatus, ) from .sinks import EvalSamplesSink, OfflineSink, Sink, TracesSink -from .worker import MetricItem, RunUpdateItem, UploadWorker, WriteItem +from .worker import UploadWorker logger = logging.getLogger(__name__) -RUN_ID_ENV = "PRIME_RUN_ID" MODE_ENV = "PRIME_RUNS_MODE" -#: Rank variables, in the order prime-rl sets them. Rank 0 owns the lifecycle. -RANK_ENV_VARS = ("RANK", "DP_RANK", "LOCAL_RANK") -DEFAULT_SUMMARY_FLUSH_SECONDS = 10.0 -#: How long ``finish()`` waits for queued records. Derived from the upload -#: timeout rather than picked: a single in-flight sample POST is allowed 300s, -#: so a shorter budget here would routinely abandon an upload that was about to -#: succeed and then finalize the run without it. +#: How long ``finish()`` gives queued uploads to drain before finalizing anyway. +#: Derived from the upload timeout: a single in-flight sample POST may take +#: this long, and a shorter budget would abandon an upload about to succeed. DEFAULT_FINISH_TIMEOUT = float(UPLOAD_TIMEOUT.read or 300.0) -#: Run IDs this process exported into ``PRIME_RUN_ID``, mapped to the PID that -#: exported them. The PID is the whole point: it is what distinguishes "my -#: parent opened this run and I should join it" from "I opened this run a moment -#: ago and the variable is still lying around". A forked child sees a different -#: PID and correctly treats the entry as inherited. -_exported_run_ids: Dict[str, int] = {} - class Run: - """A live run: an ID, a URL, somewhere to put metrics, somewhere to put traces. + """A live run: an ID, a URL, somewhere to put traces, a summary. - Every method that touches the network is contained. With the default - ``on_error="warn"`` nothing raised by the platform escapes into a producer's - loop — a run that has been going for six hours does not get killed by a 502 - on a telemetry call. ``on_error="raise"`` inverts that for tests and CI, - where a silent upload failure is the bug. + With the default ``on_error="warn"`` nothing raised by the platform escapes + into a producer's loop; ``on_error="raise"`` surfaces the first failure + from :meth:`flush` or :meth:`finish`, for tests and CI. """ def __init__( @@ -85,109 +63,45 @@ def __init__( sinks: Optional[List[Sink]] = None, mode: Mode = "online", on_error: OnError = "warn", - is_primary: bool = True, - owns_lifecycle: bool = True, - summary_flush_seconds: float = DEFAULT_SUMMARY_FLUSH_SECONDS, - finish_timeout: float = DEFAULT_FINISH_TIMEOUT, - queue_size: Optional[int] = None, ) -> None: self._backend = backend self._handle = handle self._spec = spec self._mode: Mode = mode self._on_error: OnError = on_error - self._is_primary = is_primary - # A non-primary rank shares the run but must not create or close it: - # eight ranks racing to finalize produce seven confusing failures and - # one winner. - self._owns_lifecycle = owns_lifecycle and is_primary self._status = RunStatus.RUNNING + # A forked child inherits this handle but must not close the parent's run. + self._owns_lifecycle = True - attached_config, attached_summary = _attached_state(handle) - self.config: Dict[str, Any] = {**attached_config, **spec.config} - self.summary: Dict[str, Any] = dict(attached_summary) + self.config: Dict[str, Any] = dict(spec.config) + self.summary: Dict[str, Any] = {} self.errors: List[str] = [] - # Raised at the next synchronization point the caller controls. A sink - # fails on the uploader thread, where raising reaches nobody — so under - # on_error="raise" the exception is held and re-raised from flush() or - # finish(), which is where a test or a CI job is actually looking. + # Under on_error="raise", a failure on the uploader thread is held here + # and re-raised from flush() or finish(), where the caller is looking. self._deferred_error: Optional[BaseException] = None - self._summary_flush_seconds = summary_flush_seconds - self._finish_timeout = finish_timeout - self._last_summary_flush = time.monotonic() - self._summary_dirty = False - self._config_dirty = False - self._pending_metrics: Dict[str, Any] = {} - self._pending_metric_step: Optional[int] = None - self._finish_lock = threading.RLock() - self._finish_condition = threading.Condition(self._finish_lock) + self._finish_timeout = DEFAULT_FINISH_TIMEOUT + self._finish_lock = threading.Lock() self._finishing = False - self._finishing_thread_id: Optional[int] = None self._finished = False - # A Python signal handler can interrupt finish() on the same thread. - # Re-entering teardown would duplicate finalization, while chaining the - # signal immediately would kill the process before teardown completes. - # Keep the first such signal and deliver it once the run is closed. - self._pending_signal: Optional[tuple[int, Any, Any]] = None - # ``_announce`` publishes enough context for child processes to join - # this exact run. Keep the previous values so finishing a nested or - # sequential run restores the caller's environment instead of leaking - # our resolved mode and offline directory into the next run. - self._published_join_env: Dict[str, str] = {} - self._previous_join_env: Dict[str, Optional[str]] = {} - self._atexit_hook = self._on_process_exit - # Bound once and kept. ``self._handle_signal`` builds a *new* bound - # method on every attribute access, so an ``is`` comparison against a - # freshly-made one is always False — which is how handlers end up - # installed forever, pinning a finished run and blocking the next run in - # the process from installing its own. - self._signal_handler = self._handle_signal - self._previous_signal_handlers: Dict[int, Any] = {} - # ``signal.signal`` can only run on the main thread. If finish() runs in - # an executor, or this object is inherited across a fork, its handler - # may remain as the process disposition until the main thread gets a - # chance to replace it. Marking forked handlers as relinquishable lets - # a child run take ownership without mistaking the inherited callback - # for an application-installed handler. - self._signal_handler_stale = False _fork.register(self) sinks = sinks or [] - worker_kwargs: Dict[str, Any] = {} - if queue_size is not None: - worker_kwargs["max_queue_size"] = queue_size - self._worker = UploadWorker( - sinks, - on_error=self._record_sink_error, - metric_writer=self._write_metrics if backend.supports_step_metrics else None, - update_writer=self._write_run_update, - **worker_kwargs, - ) - context = _sink_context(spec, handle) + self._worker = UploadWorker(sinks, on_error=self._record_sink_error) + context = _sink_context(spec) for sink in sinks: try: sink.start(handle.id, context) except Exception as exc: # noqa: BLE001 - a bad sink is not a bad run sink.enabled = False - try: - self._report(f"starting sink {getattr(sink, 'name', sink)}", exc) - except Exception: - # The backend may already have created a remote run. In - # strict mode the caller never receives this handle, so - # close it out before propagating the startup failure. - try: - self.finish(status=RunStatus.FAILED, error=_describe(exc)) - except Exception as cleanup_exc: # noqa: BLE001 - preserve the cause - logger.warning( - "Run %s: cleanup after sink startup failure also failed: %s: %s", - self.id, - type(cleanup_exc).__name__, - cleanup_exc, - exc_info=True, - ) - raise + self._note(f"starting sink {getattr(sink, 'name', sink)}", exc) + if self._on_error == "raise": + # The backend may already have created a remote run; close + # it out before the failure reaches the caller. finish() + # re-raises the error noted above. + self.finish(status=RunStatus.FAILED, error=_describe(exc)) + raise exc atexit.register(self._atexit_hook) @@ -207,17 +121,9 @@ def url(self) -> Optional[str]: """Where to open this run — a dashboard URL, or a local path offline.""" return self._handle.url - @property - def kind(self) -> RunKind: - return self._spec.kind - @property def config_source(self) -> Optional[ConfigSource]: - """The config file this run was launched from, if one was given. - - Read back out of ``config`` rather than cached, so a resumed run reports - the source recovered from the platform and not an empty one. - """ + """The config file this run was launched from, if one was given.""" raw = self.config.get(CONFIG_SOURCE_KEY) return ConfigSource.from_mapping(raw) if isinstance(raw, Mapping) else None @@ -229,113 +135,43 @@ def mode(self) -> Mode: def status(self) -> RunStatus: return self._status - @property - def is_primary(self) -> bool: - """Whether this process owns the run's lifecycle (rank 0, or single-process).""" - return self._owns_lifecycle - @property def finished(self) -> bool: return self._finished @property def dropped_records(self) -> int: - """Records that reached no sink because the queue was full. - - Backpressure only: the producer durably outran the uploader. A record - counted here was stored nowhere. Contrast ``failed_records``, which is - per-sink and usually means the record is still safe in another sink. - """ + """Records that reached no sink because the queue was full.""" return self._worker.dropped @property def failed_records(self) -> Dict[str, int]: - """Records each sink could not store, by sink name. - - Not summed into one number and not merged into ``dropped_records``: with - traces and the sample table both enabled, the same batch failing on one - sink says nothing about whether the other stored it, so a single total - would report data missing that is not actually gone. - """ + """Records each sink could not store, by sink name. Per sink because + another sink may still hold them.""" return dict(self._worker.failed_records) def __repr__(self) -> str: - return ( - f"" - ) + return f"" # ------------------------------------------------------------------- log - def log( - self, - metrics: Mapping[str, Any], - *, - step: Optional[int] = None, - commit: bool = True, - ) -> None: - """Record scalar metrics. - - Values always land in ``summary`` last-value-wins, the way W&B's - implicit summary works. Whether they *also* become a time series - depends on the backend: the training API stores one, the evaluations - API stores a single metrics blob, and rather than making producers care, - a backend without a time series simply keeps the summary — flushed on a - timer so a tight loop does not turn into one PUT per step. - - ``commit=False`` stages values without scheduling a write, for callers - assembling a step from several places. - """ - self._require_live("log") - cleaned = _clean_metrics(metrics) - if not cleaned: - return - self.summary.update(cleaned) - self._summary_dirty = True - if not commit: - if self._backend.supports_step_metrics: - self._pending_metrics.update(cleaned) - if step is not None: - self._pending_metric_step = step - return - if self._backend.supports_step_metrics: - committed = {**self._pending_metrics, **cleaned} - committed_step = step if step is not None else self._pending_metric_step - self._pending_metrics.clear() - self._pending_metric_step = None - self._worker.submit(MetricItem(metrics=committed, step=committed_step)) - return - self._maybe_flush_summary() - - def log_traces( - self, - records: Iterable[Any], - *, - line_format: Optional[str] = None, - step: Optional[int] = None, - ) -> None: + def log_traces(self, records: Iterable[Any]) -> None: """Hand traces or episodes to the sinks. Returns immediately. - Accepts whatever the producer already has: verifiers ``Trace`` / - ``Episode``, prime-rl ``Rollout``, or plain JSON mappings. Nothing is - buffered until the end of the run — call this as rollouts complete and - the dashboard fills in while the run is still going. + Accepts verifiers ``Trace``/``Episode`` objects or plain JSON mappings. + Call this as rollouts complete; nothing is buffered until the end. """ self._require_live("log_traces") batch = list(records) - if not batch: - return - self._worker.submit(WriteItem(records=batch, line_format=line_format, step=step)) - - def log_samples(self, records: Iterable[Any], *, step: Optional[int] = None) -> None: - """Alias for :meth:`log_traces`, matching prime-rl's ``Monitor`` vocabulary.""" - self.log_traces(records, step=step) + if batch: + self._worker.submit(batch) - def update_config(self, values: Mapping[str, Any]) -> None: - """Merge into the run's config (its inputs). Flushed with the summary.""" - self._require_live("update_config") - self.config.update(values) - self._config_dirty = True + def flush(self, timeout: Optional[float] = 30.0) -> bool: + """Block until queued records have been written. Under + ``on_error="raise"`` this is the first place an upload failure surfaces.""" + flushed = self._worker.flush(timeout=timeout) + self._raise_deferred() + return flushed # ---------------------------------------------------------------- finish @@ -346,43 +182,20 @@ def finish( status: Union[RunStatus, str] = RunStatus.COMPLETED, error: Optional[str] = None, ) -> None: - """Flush everything and close the run out. Idempotent. - - ``status`` must be one of the terminal :class:`RunStatus` values. - Safe to call from ``__exit__``, an atexit hook and a signal handler at - once — whichever gets there first reports the status, and the rest wait - for that teardown to complete. - """ - thread_id = threading.get_ident() - with self._finish_condition: - while self._finishing and not self._finished: - # A signal handler can interrupt this very finish() call. It - # cannot wait for itself, so _handle_signal defers chaining the - # signal and this nested call simply yields to the active one. - if self._finishing_thread_id == thread_id: - return - self._finish_condition.wait() + """Flush everything and close the run out. Idempotent: the first caller + reports the status, concurrent callers wait for that teardown.""" + with self._finish_lock: if self._finished: return resolved = RunStatus(status) if not isinstance(status, RunStatus) else status if not resolved.is_terminal(): raise ValueError(f"finish() requires a terminal status, got {resolved.value!r}") self._finishing = True - self._finishing_thread_id = thread_id - - try: - self._finish_once(summary, resolved, error) - finally: - with self._finish_condition: + try: + self._finish_once(summary, resolved, error) + finally: self._finishing = False - self._finishing_thread_id = None self._finished = True - pending_signal = self._pending_signal - self._pending_signal = None - self._finish_condition.notify_all() - - if pending_signal is not None: - self._chain_signal(*pending_signal) def _finish_once( self, @@ -390,40 +203,34 @@ def _finish_once( resolved: RunStatus, error: Optional[str], ) -> None: - """Perform the single teardown owned by the first ``finish()`` caller.""" - if summary: self.summary.update(_clean_metrics(summary)) self._status = resolved - finish_error: Optional[BaseException] = None deadline = time.monotonic() + max(0.0, self._finish_timeout) - def remaining_finish_time() -> float: + def remaining() -> float: return max(0.0, deadline - time.monotonic()) - # Order matters: records first, so a dashboard that reacts to the - # terminal status never sees a finished run with samples still landing. - if not self._worker.flush(timeout=remaining_finish_time()): + # Records first, so a dashboard reacting to the terminal status never + # sees a finished run with samples still landing. + if not self._worker.flush(timeout=remaining()): logger.warning( "Run %s: uploads did not drain within %ss; finalizing anyway. " "Some records may be missing from this run.", self.id, self._finish_timeout, ) - self._worker.close(timeout=remaining_finish_time()) + self._worker.close(timeout=remaining()) if self._owns_lifecycle: - finish_error = self._finish_guarded( + self._teardown_step( "updating the run", lambda: self._backend.update( - self.id, - config=self.config if (self._config_dirty or self.config) else None, - summary=self.summary or None, + self.id, config=self.config or None, summary=self.summary or None ), - finish_error, ) - finish_error = self._finish_guarded( + self._teardown_step( "finalizing the run", lambda: self._backend.finalize( self.id, @@ -432,15 +239,9 @@ def remaining_finish_time() -> float: error=error or (self.errors[0] if self.errors else None), config=self.config or None, ), - finish_error, ) - finish_error = self._finish_guarded( - "closing the backend", self._backend.close, finish_error - ) - + self._teardown_step("closing the backend", self._backend.close) atexit.unregister(self._atexit_hook) - self._restore_signal_handlers() - self._retract_join_context() if self._worker.dropped: logger.warning( @@ -450,37 +251,16 @@ def remaining_finish_time() -> float: self._worker.dropped, ) for sink_name, count in self._worker.failed_records.items(): - # Deliberately phrased per sink: another sink may hold these records, - # so claiming they are missing from the run would overstate the loss. logger.warning( "Run %s: the %s sink could not store %d record(s)", self.id, sink_name, count ) - # Last, so a run that failed to upload is still closed out properly - # before the failure reaches the caller. - if finish_error is not None: - raise finish_error + # Last, so a run whose teardown failed is still closed out first. self._raise_deferred() def fail(self, error: Union[str, BaseException]) -> None: """Close the run out as failed.""" self.finish(status=RunStatus.FAILED, error=_describe(error)) - def flush(self, timeout: Optional[float] = 30.0) -> bool: - """Block until queued records have been written. - - Under ``on_error="raise"`` this is the first place an upload failure can - surface, since the failure itself happened on the uploader thread. - """ - update_queued = self._queue_run_update() - flushed = self._worker.flush(timeout=timeout) - # A periodic update can lose a race for the last queue slot. Once the - # barrier drains that backlog, give the still-dirty snapshot one more - # chance so an explicit flush keeps its persistence guarantee. - if flushed and not update_queued and self._queue_run_update(): - flushed = self._worker.flush(timeout=timeout) - self._raise_deferred() - return flushed - # -------------------------------------------------------- context manager def __enter__(self) -> "Run": @@ -492,23 +272,17 @@ def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: return False if isinstance(exc, KeyboardInterrupt): - # An interrupt is a decision, not a fault, so it must not land in - # the same bucket as broken ones. Matches the SIGINT handler, which - # normally gets there first when signal handling is on. - status = RunStatus.CRASHED - error = "interrupted" + # An interrupt is a decision, not a fault. + status, error = RunStatus.CRASHED, "interrupted" else: - status = RunStatus.FAILED - error = _describe(exc) + status, error = RunStatus.FAILED, _describe(exc) try: self.finish(status=status, error=error) except Exception as finish_error: - # The producer exception is the reason this context is unwinding. - # A telemetry teardown error must not replace it, even in strict - # mode; finish() has already recorded the failure on the run. - # Control-flow exceptions such as KeyboardInterrupt and SystemExit - # deliberately bypass this handler so teardown cannot swallow them. + # The producer's exception is why this block is unwinding; a + # telemetry teardown error must not replace it. Control-flow + # exceptions (KeyboardInterrupt, SystemExit) deliberately pass. logger.warning( "Run %s: finishing after %s also failed: %s: %s", self.id, @@ -521,160 +295,26 @@ def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: # ------------------------------------------------------------- internals - def install_signal_handlers(self) -> None: - """Report a terminal status when the process is killed. - - Only installed on the main thread, and only over a *default* handler or - one relinquished by a finished/forked ``Run``. Replacing a handler the - application chose would be worse than missing a status. The previous - handler is always called afterwards, so SIGINT still raises - ``KeyboardInterrupt`` and SIGTERM still terminates. - """ - if threading.current_thread() is not threading.main_thread(): - return - for signum in (signal.SIGINT, signal.SIGTERM): - try: - current = signal.getsignal(signum) - except (ValueError, OSError): # pragma: no cover - platform dependent - continue - previous = current - relinquished_owner: Optional[Run] = None - if current not in (signal.SIG_DFL, signal.default_int_handler): - owner = getattr(current, "__self__", None) - if ( - isinstance(owner, Run) - and current is owner._signal_handler - and (owner._finished or owner._signal_handler_stale) - and signum in owner._previous_signal_handlers - ): - previous = owner._previous_signal_handlers[signum] - relinquished_owner = owner - else: - continue - try: - signal.signal(signum, self._signal_handler) - except (ValueError, OSError): # pragma: no cover - continue - self._previous_signal_handlers[signum] = previous - if relinquished_owner is not None: - relinquished_owner._previous_signal_handlers.pop(signum, None) - - def _handle_signal(self, signum: int, frame: Any) -> None: - name = signal.Signals(signum).name - # Read the displaced handler *before* finishing: finish() restores and - # then clears this table, so looking it up afterwards always yields - # SIG_DFL — which re-raises the signal at its default disposition and - # kills the process instead of running the handler the app installed. - previous = self._previous_signal_handlers.get(signum, signal.SIG_DFL) - with self._finish_condition: - if self._finishing and self._finishing_thread_id == threading.get_ident(): - if self._pending_signal is None: - self._pending_signal = (signum, frame, previous) - return - finished = self._finished - if not finished: - # CRASHED, not FAILED: the producer never said the run failed, it was - # stopped from outside its own control flow. Same bucket as the - # atexit path, and deliberately not the bucket a broken eval lands in. - try: - self.finish(status=RunStatus.CRASHED, error=f"received {name}") - except Exception as exc: # noqa: BLE001 - the signal must still chain - logger.warning("Run %s: reporting %s failed: %s", self.id, name, exc) - self._chain_signal(signum, frame, previous) - - @staticmethod - def _chain_signal(signum: int, frame: Any, previous: Any) -> None: - """Restore and invoke the handler displaced by this run.""" - signal.signal(signum, previous) - if callable(previous): - previous(signum, frame) - else: - os.kill(os.getpid(), signum) - - def _restore_signal_handlers(self) -> None: - remaining: Dict[int, Any] = {} - for signum, previous in self._previous_signal_handlers.items(): - try: - if signal.getsignal(signum) is self._signal_handler: - signal.signal(signum, previous) - except (ValueError, OSError): # pragma: no cover - # Most commonly finish() was deliberately run in an executor. - # Keep the displaced handler so the main thread can restore it - # from the signal callback or hand it to the next Run. - remaining[signum] = previous - self._previous_signal_handlers = remaining - def reset_after_fork(self) -> None: - """Make an inherited handle safe to use in a forked child. - - The child may keep using this handle, but the process that created the - run remains responsible for its lifecycle. In particular, the child's - inherited signal and atexit callbacks must never finalize the parent's - still-running run. The lock also has to be replaced because it may have - been owned at fork time by a thread that no longer exists. - """ - self._finish_lock = threading.RLock() - self._finish_condition = threading.Condition(self._finish_lock) + """Make an inherited handle safe in a forked child: fresh lock, and the + parent keeps ownership of the lifecycle so the child's atexit hook + cannot finalize a still-running run.""" + self._finish_lock = threading.Lock() self._finishing = False - self._finishing_thread_id = None - self._pending_signal = None - self._is_primary = False self._owns_lifecycle = False self._deferred_error = None - self._signal_handler_stale = True def _on_process_exit(self) -> None: - """Last resort: the process is exiting and nobody called ``finish()``. - - Reported as CRASHED rather than FAILED — the producer never said the run - failed, it just stopped existing, and the distinction tells an operator - whether to read the run's error or go look at the machine. - """ + """The process is exiting and nobody called ``finish()``: report CRASHED + rather than FAILED, since the producer never said the run failed.""" if self._finished: return logger.warning("Run %s was never finished; reporting it as crashed", self.id) - # on_error="raise" must not turn interpreter shutdown into a traceback - # from atexit; the run is already being reported as crashed. try: self.finish(status=RunStatus.CRASHED, error="process exited without finishing the run") - except Exception as exc: # noqa: BLE001 + except Exception as exc: # noqa: BLE001 - never a traceback from atexit logger.warning("Run %s: reporting the crash failed: %s", self.id, exc) - def _publish_join_context(self) -> None: - """Publish the resolved context a child needs to join this run.""" - values = { - RUN_ID_ENV: self.id, - MODE_ENV: self.mode, - } - if self.mode == "offline" and isinstance(self._backend, OfflineBackend): - # Resolve the path because a subprocess may use a different cwd. - values[DEFAULT_DIR_ENV] = str(self._backend.directory.resolve()) - - for name, value in values.items(): - self._previous_join_env.setdefault(name, os.environ.get(name)) - self._published_join_env[name] = value - os.environ[name] = value - _exported_run_ids[self.id] = os.getpid() - - def _retract_join_context(self) -> None: - """Stop advertising a finished run to processes started from here. - - Only restores values this run still owns: if one now points somewhere - else, another run owns it and changing it would orphan that one's - children. - """ - if not self._owns_lifecycle: - return - for name, published in self._published_join_env.items(): - if os.environ.get(name) != published: - continue - previous = self._previous_join_env.get(name) - if previous is None: - os.environ.pop(name, None) - else: - os.environ[name] = previous - _exported_run_ids.pop(self.id, None) - def _require_live(self, operation: str) -> None: if self._finishing or self._finished: raise RunFinishedError( @@ -682,86 +322,35 @@ def _require_live(self, operation: str) -> None: "The platform has closed this run out; start a new one." ) - def _write_metrics(self, metrics: Dict[str, Any], step: Optional[int]) -> None: - self._backend.log_metrics(self.id, metrics, step) - - def _write_run_update( - self, - config: Optional[Dict[str, Any]], - summary: Optional[Dict[str, Any]], - ) -> None: - self._backend.update(self.id, config=config, summary=summary) + def _record_sink_error(self, sink_name: str, exc: Exception) -> None: + """Called on the uploader thread when a sink gives up on a batch.""" + self._note(f"writing to the {sink_name} sink", exc) - def _maybe_flush_summary(self) -> None: - now = time.monotonic() - if now - self._last_summary_flush < self._summary_flush_seconds: - return - self._queue_run_update() - - def _queue_run_update(self) -> bool: - if not (self._summary_dirty or self._config_dirty) or not self._owns_lifecycle: - return True - config_dirty = self._config_dirty - summary_dirty = self._summary_dirty - item = RunUpdateItem( - config=dict(self.config) if config_dirty else None, - summary=dict(self.summary) if summary_dirty else None, - ) - # Clear before the potentially blocking queue put. If another producer - # thread logs while this one waits, its new dirty bit must survive. - if config_dirty: - self._config_dirty = False - if summary_dirty: - self._summary_dirty = False - if not self._worker.submit(item): - self._config_dirty = self._config_dirty or config_dirty - self._summary_dirty = self._summary_dirty or summary_dirty - return False - self._last_summary_flush = time.monotonic() - return True + def _teardown_step(self, what: str, call: Any) -> None: + """Run one teardown step without letting it skip the steps after it.""" + try: + call() + except Exception as exc: # noqa: BLE001 - teardown must continue + self._note(what, exc) - def _record_sink_error(self, sink_name: str, exc: Exception) -> None: - """Called on the uploader thread when a sink gives up.""" - message = f"writing to the {sink_name} sink failed: {type(exc).__name__}: {exc}" + def _note(self, what: str, exc: BaseException) -> None: + """Record a contained failure: warn, or hold it for the next sync point.""" + message = f"{what} failed: {type(exc).__name__}: {exc}" self.errors.append(message) if self._on_error == "raise": if self._deferred_error is None: self._deferred_error = exc - return - logger.warning("Run %s: %s", self._handle.id, message) + else: + logger.warning("Run %s: %s", self.id, message) def _raise_deferred(self) -> None: - """Re-raise the first upload failure, once.""" + """Re-raise the first held failure, once.""" exc = self._deferred_error if exc is None: return self._deferred_error = None raise exc - def _finish_guarded( - self, - what: str, - call: Any, - first_error: Optional[BaseException], - ) -> Optional[BaseException]: - """Run one teardown step without letting it skip the steps after it.""" - try: - call() - except Exception as exc: # noqa: BLE001 - teardown must continue - message = f"{what} failed: {type(exc).__name__}: {exc}" - self.errors.append(message) - if self._on_error == "raise": - return first_error or exc - logger.warning("Run %s: %s", self._handle.id, message) - return first_error - - def _report(self, what: str, exc: Exception) -> None: - message = f"{what} failed: {type(exc).__name__}: {exc}" - self.errors.append(message) - if self._on_error == "raise": - raise exc - logger.warning("Run %s: %s", self._handle.id, message) - # --------------------------------------------------------------------- init @@ -769,57 +358,36 @@ def _report(self, what: str, exc: Exception) -> None: def init( *, name: Optional[str] = None, - kind: RunKind = "eval", environments: Optional[Sequence[Any]] = None, model: Optional[str] = None, framework: Optional[str] = None, description: Optional[str] = None, tags: Optional[Sequence[str]] = None, config: Optional[Any] = None, - id: Optional[str] = None, mode: Optional[Mode] = None, dir: Optional[str] = None, team_id: Optional[str] = None, api_key: Optional[str] = None, base_url: Optional[str] = None, - traces_url: Optional[str] = None, - traces: bool = True, - samples: bool = True, on_error: OnError = "warn", - handle_signals: bool = True, - queue_size: Optional[int] = None, - finish_timeout: float = DEFAULT_FINISH_TIMEOUT, ) -> Run: """Start a run and return a handle to it. Call this *before* the first rollout: the ID it returns is what every trace - in the run should carry, and the URL it returns is what a producer prints so - someone can watch the run fill in. - - ``mode`` selects where the run lives. Left unset it is read from - ``$PRIME_RUNS_MODE``, and failing that inferred: online when there is an API - key, offline when there is not. Offline is a real run with a real ID and a - real directory, which is why producers no longer need a ``--no-push`` - branch — the call sites are identical either way. - - ``id`` attaches to an existing run instead of creating one, for resuming - after a crash and for non-primary ranks joining a run rank 0 created. - - ``config`` is what the run was configured *with*, in whatever form you have - it — the path to the file it was launched from, a mapping, or a pydantic - model. See :func:`_normalize_config`. + in the run should carry, and the URL is what a producer prints. - ``finish_timeout`` is the total number of seconds ``finish()`` gives queued - uploads to drain and close before finalizing the run anyway. + ``mode`` defaults to ``$PRIME_RUNS_MODE``, else online when there is an API + key and offline when there is not. ``config`` is what the run was configured + with: the path to the file it was launched from (stored byte for byte under + ``config_source``), or a mapping taken as given. """ - resolved_config = Config() - api_key = api_key if api_key is not None else resolved_config.api_key - base_url = base_url or resolved_config.base_url - team_id = team_id if team_id is not None else resolved_config.team_id + settings = Config() + api_key = api_key if api_key is not None else settings.api_key + base_url = base_url or settings.base_url + team_id = team_id if team_id is not None else settings.team_id spec = RunSpec( name=name, - kind=kind, environments=[EnvironmentRef.coerce(entry) for entry in (environments or [])], model=model, framework=framework, @@ -828,167 +396,54 @@ def init( team_id=team_id, config=_normalize_config(config), ) + resolved_mode = _resolve_mode(mode, api_key=api_key) - is_primary = _is_primary_rank() - joined_id = _inherited_run_id() - inherited_id = id or joined_id - # Owning the lifecycle means "this call is responsible for creating and - # finalizing the run". An explicit `id=` is a deliberate resume, so it owns. - # An ID picked up from the environment belongs to whoever exported it, so it - # does not. (A non-primary rank never owns either way — see Run.__init__.) - owns_lifecycle = id is not None or joined_id is None - resolved_mode = _resolve_mode(mode, api_key=api_key, is_primary=is_primary, run_id=inherited_id) - + backend: Backend + sinks: List[Sink] if resolved_mode == "disabled": - backend: Backend = _DisabledBackend() - handle = RunHandle(id=inherited_id or _local_id(), name=name) - return _build( - spec, - backend, - handle, - [], - resolved_mode, - on_error, - is_primary, - False, - queue_size, - finish_timeout, - ) - - if resolved_mode == "offline": + backend = _DisabledBackend() + handle = RunHandle(id=new_run_id(), name=name) + sinks = [] + elif resolved_mode == "offline": offline = OfflineBackend(dir) - handle = offline.attach(inherited_id) if inherited_id else offline.create(spec) - run_sinks: List[Sink] = [OfflineSink(offline.directory)] - run = _build( - spec, - offline, - handle, - run_sinks, - resolved_mode, - on_error, - is_primary, - owns_lifecycle, - queue_size, - finish_timeout, - ) - _announce(run, handle_signals) - return run - - if kind != "eval": - raise ConfigurationError( - f"kind={kind!r} is not supported yet — training runs arrive with the RFT backend. " - 'Use kind="eval", or mode="offline" to record the run locally.' - ) - if not api_key: - raise ConfigurationError( - 'mode="online" needs an API key. Set PRIME_API_KEY, run `prime login`, ' - 'or pass mode="offline".' - ) - - client = PlatformClient(api_key=api_key, base_url=base_url, timeout=DEFAULT_TIMEOUT) - backend = EvalsBackend(client, frontend_url=resolved_config.frontend_url, team_id=team_id) - handle = backend.attach(inherited_id) if inherited_id else backend.create(spec) - - run_sinks = [] - if traces: - run_sinks.append(TracesSink(api_key=api_key, traces_url=traces_url, team_id=team_id)) - if samples: + backend = offline + handle = offline.create(spec) + sinks = [OfflineSink(offline.directory)] + else: + if not api_key: + raise ConfigurationError( + 'mode="online" needs an API key. Set PRIME_API_KEY, run `prime login`, ' + 'or pass mode="offline".' + ) + client = PlatformClient(api_key=api_key, base_url=base_url, timeout=DEFAULT_TIMEOUT) + backend = EvalsBackend(client, frontend_url=settings.frontend_url, team_id=team_id) + handle = backend.create(spec) # Both transports run during the transition: traces is the system of - # record, the sample table is what today's viewer reads, and Prime - # Traces is still gated to an account allowlist. - samples_client = PlatformClient(api_key=api_key, base_url=base_url, timeout=DEFAULT_TIMEOUT) - run_sinks.append(EvalSamplesSink(samples_client, close_client=True)) - - run = _build( - spec, - backend, - handle, - run_sinks, - resolved_mode, - on_error, - is_primary, - owns_lifecycle, - queue_size, - finish_timeout, - ) - _announce(run, handle_signals) - return run + # record, the sample table is what today's viewer reads. + sinks = [TracesSink(api_key=api_key, team_id=team_id), EvalSamplesSink(client)] - -def _build( - spec: RunSpec, - backend: Backend, - handle: RunHandle, - sinks: List[Sink], - mode: Mode, - on_error: OnError, - is_primary: bool, - owns_lifecycle: bool, - queue_size: Optional[int], - finish_timeout: float, -) -> Run: - return Run( + run = Run( backend=backend, handle=handle, spec=spec, sinks=sinks, - mode=mode, + mode=resolved_mode, on_error=on_error, - is_primary=is_primary, - owns_lifecycle=owns_lifecycle, - queue_size=queue_size, - finish_timeout=finish_timeout, ) - - -def _announce(run: Run, handle_signals: bool) -> None: - """Publish the run's join context to child processes and arm crash reporting. - - The ID says which run to attach to; the resolved mode and offline directory - say where it lives. The PID is recorded alongside so that *this* process - does not later mistake its own export for a parent's. - """ - run._publish_join_context() - if handle_signals: - run.install_signal_handlers() if run.url: logger.info("Run %s: %s", run.id, run.url) - - -def _inherited_run_id() -> Optional[str]: - """A run ID this process should join, or ``None`` to open a fresh run. - - ``PRIME_RUN_ID`` set by an ancestor means "join that run". The same variable - set by an earlier ``init()`` *in this process* means nothing of the sort — - without this check, a second eval in one process would silently attach to - the first, and would never create or finalize a run of its own. - """ - value = os.getenv(RUN_ID_ENV) - if not value: - return None - if _exported_run_ids.get(value) == os.getpid(): - return None - return value + return run class _DisabledBackend: """No-op lifecycle, so ``mode="disabled"`` needs no branching upstream.""" - kind = "disabled" - supports_step_metrics = False - def create(self, spec: RunSpec) -> RunHandle: - return RunHandle(id=_local_id()) - - def attach(self, run_id: str) -> RunHandle: - return RunHandle(id=run_id) + return RunHandle(id=new_run_id()) def update(self, run_id: str, **kwargs: Any) -> None: return None - def log_metrics(self, run_id: str, metrics: Dict[str, Any], step: Optional[int] = None) -> None: - return None - def finalize(self, run_id: str, **kwargs: Any) -> None: return None @@ -996,30 +451,7 @@ def close(self) -> None: return None -def _local_id() -> str: - from .backends.offline import new_run_id - - return new_run_id() - - -def _is_primary_rank() -> bool: - """Whether this process should own the run's lifecycle. - - Any rank variable set to something other than 0 means a peer process is - rank 0 and owns creation and finalization. Non-primary ranks still upload - their own records — the point is that eight processes contribute to one run - rather than creating eight. - """ - for name in RANK_ENV_VARS: - value = os.getenv(name) - if value and value.strip() not in ("0", ""): - return False - return True - - -def _resolve_mode( - mode: Optional[Mode], *, api_key: str, is_primary: bool, run_id: Optional[str] -) -> Mode: +def _resolve_mode(mode: Optional[Mode], *, api_key: str) -> Mode: if mode is None: env_mode = os.getenv(MODE_ENV) if env_mode: @@ -1035,22 +467,13 @@ def _resolve_mode( "recording this run offline instead." ) mode = "offline" - if mode in ("online", "offline") and not is_primary and not run_id: - # A non-primary rank with no run to join would create a second run for - # the same job. Recording nothing is better than that. - logger.debug("Non-primary rank with no %s; disabling this run handle", RUN_ID_ENV) - return "disabled" return mode # type: ignore[return-value] -def _sink_context(spec: RunSpec, handle: RunHandle) -> Dict[str, str]: - """Upload-scoped provenance. - - Not the join key — that is ``run.id`` inside the trace document, which the - ingestion service extracts into an indexed column. What goes here is what - you would want when looking at an upload and asking where it came from. - """ - context = {"source": "prime-runs", "run_kind": spec.kind} +def _sink_context(spec: RunSpec) -> Dict[str, str]: + """Upload-scoped provenance. Not the join key — that is ``run.id`` inside + the trace document.""" + context = {"source": "prime-runs", "run_kind": RUN_KIND} if spec.framework: context["framework"] = spec.framework if spec.model: @@ -1058,78 +481,12 @@ def _sink_context(spec: RunSpec, handle: RunHandle) -> Dict[str, str]: return context -def _attached_state(handle: RunHandle) -> tuple[Dict[str, Any], Dict[str, Any]]: - """Recover config and summary from an attached backend response. - - Online evaluations expose them as ``metadata`` and ``metrics``. Offline - archives keep their initial values under ``spec`` and subsequent values at - the top level. Supporting both shapes keeps resume lossless for either - backend; values supplied to the new ``init()`` call are merged afterwards - and therefore win. - """ - config: Dict[str, Any] = {} - summary: Dict[str, Any] = {} - raw = handle.raw - nested_spec = raw.get("spec") - if isinstance(nested_spec, Mapping): - _merge_mapping(config, nested_spec.get("config")) - _merge_mapping(summary, nested_spec.get("summary")) - _merge_mapping(config, raw.get("config")) - _merge_mapping(config, raw.get("metadata")) - _merge_mapping(summary, raw.get("summary")) - _merge_mapping(summary, raw.get("metrics")) - return config, summary - - -def _merge_mapping(target: Dict[str, Any], value: Any) -> None: - if isinstance(value, Mapping): - target.update(value) - - -def _accepts_exclude_unset(dump: Any) -> bool: - """Whether a ``model_dump`` takes the keywords we want to hand it. - - Asked of the signature rather than discovered by catching ``TypeError``. - Catching would mean inferring *why* a call failed from its exception type, - and the only available recovery — dumping every field — is precisely the - outcome this whole path exists to avoid. So a wrong inference degrades - silently, in the one direction that matters. A real serialization failure - should surface as itself. - """ - try: - params = inspect.signature(dump).parameters - except (TypeError, ValueError): # pragma: no cover - uninspectable callables - return False - if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): - return True - return {"mode", "exclude_unset"} <= params.keys() - - def _normalize_config(value: Any) -> Dict[str, Any]: - """A producer's config as a plain dict, in the most faithful form available. - - One parameter takes whatever form the caller has it in, the same way - ``environments=`` accepts a slug, a dict or an ``EnvironmentRef``. What - changes per form is only how much fidelity there is to preserve: - - - **A path** is the file the run was launched from (``uv run eval @ - eval.toml``). That file *is* the run's configuration, so it is kept byte - for byte under ``config_source`` — comments, key order and section - grouping included. Nothing else can reproduce those. - - **A mapping** is taken exactly as given. The caller already decided what - it wanted to say and second-guessing it would be worse. - - **A pydantic model** is dumped with ``exclude_unset=True``. A resolved - dump of a deep config tree is hundreds of lines of defaults nobody chose, - and a reader scrolling it cannot tell which three values were the - experiment. ``exclude_unset`` leaves exactly the fields someone typed. - - A caller who genuinely wants every resolved default can still pass - ``cfg.model_dump()``. That asymmetry is deliberate: the shorter call should - give the more useful answer. - - The file is stored, not parsed. Parsing would buy a second representation of - something the platform can already read, at the cost of a TOML dependency in - a package that deliberately has two. + """A producer's config as a plain dict. + + A path is the file the run was launched from, kept byte for byte under + ``config_source`` (stored, not parsed — the platform can read TOML). A + mapping is taken exactly as given. """ if value is None: return {} @@ -1139,25 +496,8 @@ def _normalize_config(value: Any) -> Dict[str, Any]: source = ConfigSource.coerce(value) assert source is not None # coerce only returns None for None return {CONFIG_SOURCE_KEY: source.to_dict()} - dump = getattr(value, "model_dump", None) - if callable(dump): - if _accepts_exclude_unset(dump): - dumped = dump(mode="json", exclude_unset=True) - else: - # Not silent: the fallback records every default, which is the exact - # outcome the caller was trying to avoid by handing us a model. - logger.warning( - "%s.model_dump() does not accept exclude_unset; recording the fully " - "resolved config instead, defaults included.", - type(value).__name__, - ) - dumped = dump() - if isinstance(dumped, Mapping): - return dict(dumped) - raise TypeError(f"{type(value).__name__}.model_dump() did not return a mapping") raise TypeError( - "config must be a path to the run's config file, a mapping or a pydantic model, " - f"got {type(value).__name__}" + f"config must be a path to the run's config file or a mapping, got {type(value).__name__}" ) @@ -1168,13 +508,8 @@ def _describe(error: Union[str, BaseException]) -> str: def _clean_metrics(metrics: Mapping[str, Any]) -> Dict[str, Any]: - """Drop values JSON cannot carry. - - NaN and infinity are the ones that matter: a diverged loss serializes as - JavaScript's bare ``NaN``, which strict JSON rejects, and the failure - surfaces as an opaque 400 on a payload nobody can inspect. Dropping the key - loses one point; sending it loses the request. - """ + """Drop NaN/infinity: strict JSON rejects them, and the failure would + surface as an opaque 400 on the whole request.""" cleaned: Dict[str, Any] = {} for key, value in metrics.items(): if isinstance(value, float) and not math.isfinite(value): @@ -1189,4 +524,4 @@ def _clean_metrics(metrics: Mapping[str, Any]) -> Dict[str, Any]: return cleaned -__all__ = ["Run", "init", "RUN_ID_ENV", "MODE_ENV"] +__all__ = ["Run", "init", "MODE_ENV"] diff --git a/packages/prime-runs/src/prime_runs/sinks/__init__.py b/packages/prime-runs/src/prime_runs/sinks/__init__.py index 115536d25..d75fdb283 100644 --- a/packages/prime-runs/src/prime_runs/sinks/__init__.py +++ b/packages/prime-runs/src/prime_runs/sinks/__init__.py @@ -1,12 +1,14 @@ -"""Sample transports. Independent of backends, and of each other.""" +"""Record transports. Independent of backends, and of each other.""" -from .base import Sink, to_mapping +from .base import Sink, is_episode, stamp_run, to_mapping from .offline import OfflineSink from .samples import EvalSamplesSink from .traces import TracesSink __all__ = [ "Sink", + "is_episode", + "stamp_run", "to_mapping", "EvalSamplesSink", "OfflineSink", diff --git a/packages/prime-runs/src/prime_runs/sinks/base.py b/packages/prime-runs/src/prime_runs/sinks/base.py index e6a4771e1..e8ec3edc1 100644 --- a/packages/prime-runs/src/prime_runs/sinks/base.py +++ b/packages/prime-runs/src/prime_runs/sinks/base.py @@ -1,20 +1,17 @@ -"""The contract a sample sink implements. +"""The contract a record sink implements, plus the record helpers sinks share. -A sink moves *records* — traces, episodes, rollouts — to wherever they are -stored. It knows nothing about run lifecycle; a backend closing a run and a -sink flushing its last batch are separate events on purpose. +A sink moves records — traces, episodes — to wherever they are stored. It knows +nothing about run lifecycle. Sinks are independent of each other: during the +transition the traces sink and the legacy eval-samples sink both run, and +retiring one is a change to the default sink list, not to any producer. -Sinks are independent of backends and of each other. During the transition both -the traces sink and the legacy eval-samples sink run at once, so the dashboard -keeps working for accounts outside the traces beta while traces becomes the -system of record. When the Viewer API reads traces natively, the default sink -list drops one entry — and no producer changes. - -Every sink must be *degradable*: a sink that cannot write sets ``enabled = -False`` and says why, once. A run whose traces are gated is still a valid run. +Every sink must be degradable: one that cannot write sets ``enabled = False`` +and says why, once. """ -from typing import Any, Mapping, Optional, Protocol, Sequence, runtime_checkable +from typing import Any, Dict, Mapping, Protocol, Sequence, runtime_checkable + +from ..models import RUN_KIND @runtime_checkable @@ -28,13 +25,7 @@ def start(self, run_id: str, context: Mapping[str, str]) -> None: """Bind the sink to a run before the first write.""" ... - def write( - self, - records: Sequence[Any], - *, - line_format: Optional[str] = None, - step: Optional[int] = None, - ) -> None: + def write(self, records: Sequence[Any]) -> None: """Send one batch. Called from the uploader thread, never inline.""" ... @@ -48,12 +39,8 @@ def close(self) -> None: def to_mapping(record: Any) -> Mapping[str, Any]: - """The JSON mapping for a record, whatever shape the producer handed us. - - Mirrors ``prime_traces.SupportsToRecord``: verifiers ``Trace``/``Episode`` - and prime-rl ``Rollout`` all implement ``to_record()``, and plain dicts pass - straight through. - """ + """The JSON mapping for a record: a dict passes through, anything else must + implement ``to_record()`` (verifiers ``Trace``/``Episode`` do).""" if isinstance(record, Mapping): return record to_record = getattr(record, "to_record", None) @@ -63,3 +50,19 @@ def to_mapping(record: Any) -> Mapping[str, Any]: return value raise TypeError(f"{type(record).__name__}.to_record() must return a mapping") raise TypeError(f"{type(record).__name__} is not a mapping and has no to_record()") + + +def is_episode(record: Any) -> bool: + """Whether a record (object or mapping) is an episode rather than a trace.""" + if isinstance(record, Mapping): + return "traces" in record + return hasattr(record, "traces") + + +def stamp_run(mapping: Mapping[str, Any], run_id: str) -> Dict[str, Any]: + """A copy of ``mapping`` carrying ``run`` if it did not already. Producer + objects are never stamped — they carry their own ``run`` — but a bare dict + with no ``run.id`` is an orphaned, unqueryable upload.""" + if mapping.get("run"): + return dict(mapping) + return {**mapping, "run": {"id": run_id, "type": RUN_KIND}} diff --git a/packages/prime-runs/src/prime_runs/sinks/offline.py b/packages/prime-runs/src/prime_runs/sinks/offline.py index 015a3a5cc..d63a72dc2 100644 --- a/packages/prime-runs/src/prime_runs/sinks/offline.py +++ b/packages/prime-runs/src/prime_runs/sinks/offline.py @@ -1,11 +1,5 @@ -"""Local JSONL sink, written in the wire format Prime Traces accepts. - -Deliberately not a debug dump: the files this writes are valid trace/episode -JSONL, so ``prime_traces.TracesClient.upload_file`` can send them later -untouched. That is what makes an offline run a *deferred* run rather than a -different one — the run ID stamped into the records was issued at ``init()`` -and does not change on sync. -""" +"""Local JSONL sink, written in the wire format Prime Traces accepts, so the +files can later be sent by ``prime_traces.TracesClient.upload_file`` untouched.""" import logging from pathlib import Path @@ -13,79 +7,53 @@ from .. import _fork from .._http import encode_json -from .base import Sink, to_mapping +from .base import Sink, is_episode, stamp_run, to_mapping logger = logging.getLogger(__name__) class OfflineSink(Sink): - """Appends records to ``//records/.jsonl``.""" + """Appends records to ``//records/{trace,episode}.jsonl``.""" name = "offline" - def __init__(self, directory: Union[str, Path], *, stamp_run: bool = True) -> None: + def __init__(self, directory: Union[str, Path]) -> None: self.enabled = True self.directory = Path(directory) - self._stamp_run = stamp_run self._run_id: Optional[str] = None - self._run_kind: Optional[str] = None self._handles: dict[str, BinaryIO] = {} self.records_written = 0 _fork.register(self) def reset_after_fork(self) -> None: - """Abandon inherited file handles; ``_handle`` reopens on next write. - - Dropping a buffered file object would not be enough on its own: on - CPython the last reference going away closes it, and ``close()`` - *flushes* — writing out the parent's copied buffer and duplicating every - record in it. The handles are unbuffered (see ``_handle``) precisely so - that there is never anything in that buffer to duplicate. - """ + """Abandon inherited file handles. They are unbuffered (see ``_handle``), + so dropping them cannot flush a copy of the parent's buffer.""" self._handles = {} def start(self, run_id: str, context: Mapping[str, str]) -> None: self._run_id = run_id - self._run_kind = context.get("run_kind") self._records_dir.mkdir(parents=True, exist_ok=True) @property def _records_dir(self) -> Path: return self.directory / (self._run_id or "unknown") / "records" - def write( - self, - records: Sequence[Any], - *, - line_format: Optional[str] = None, - step: Optional[int] = None, - ) -> None: + def write(self, records: Sequence[Any]) -> None: if not self.enabled or not records: return - name = str(getattr(line_format, "value", line_format) or _infer_format(records[0])) - handle = self._handle(name) + handle = self._handle("episode" if is_episode(records[0]) else "trace") for record in records: - mapping = dict(to_mapping(record)) - if self._stamp_run and not mapping.get("run") and self._run_id: - run: dict[str, Any] = {"id": self._run_id} - if self._run_kind: - run["type"] = self._run_kind - mapping["run"] = run - # Match the online JSON encoder exactly. In particular, rejecting - # NaN/Infinity here prevents creating an archive that exists but - # cannot later be uploaded by Prime Traces' strict parser. + mapping = to_mapping(record) + if self._run_id: + mapping = stamp_run(mapping, self._run_id) + # Same strict encoder as the online path: an archive that holds + # NaN cannot later be uploaded. _write_all(handle, encode_json(mapping) + b"\n") self.records_written += 1 def _handle(self, name: str) -> BinaryIO: - """An unbuffered append-mode handle for one line format. - - Unbuffered on purpose. A buffered writer keeps records in process memory - until it decides to flush, and a fork copies that buffer — after which - both processes eventually write it, putting every record in the file - twice. Writing straight through means the only copy of a record lives in - the file, and ``O_APPEND`` keeps concurrent writers from interleaving. - """ + """An unbuffered append-mode handle per line format. Unbuffered so a + fork never copies pending records; ``O_APPEND`` keeps writers whole.""" handle = self._handles.get(name) if handle is None: self._records_dir.mkdir(parents=True, exist_ok=True) @@ -94,7 +62,7 @@ def _handle(self, name: str) -> BinaryIO: return handle def flush(self) -> None: - """Nothing is held back — every write already went to the file.""" + """Every write already went to the file.""" def close(self) -> None: for handle in self._handles.values(): @@ -112,9 +80,3 @@ def _write_all(handle: BinaryIO, data: bytes) -> None: if not written: # pragma: no cover - only on a non-blocking handle raise OSError("offline record write made no progress") data = data[written:] - - -def _infer_format(record: Any) -> str: - if isinstance(record, Mapping): - return "episode" if "traces" in record else "trace" - return "episode" if hasattr(record, "traces") else "trace" diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py index b77f7f2f7..d99b0d719 100644 --- a/packages/prime-runs/src/prime_runs/sinks/samples.py +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -1,26 +1,19 @@ """Legacy sink: the flat eval-sample table behind today's viewer. -This exists so the migration is a refactor rather than a regression. The viewer -reads the v0 sample table; Prime Traces is in closed beta on an account -allowlist. Shipping traces-only would leave every non-allowlisted account -staring at an empty dashboard — so both sinks run, and this one retires when -the Viewer API reads traces natively. Retiring it is a one-line change to the -default sink list, with nothing to do in verifiers or prime-rl. +Prime Traces is gated to an allowlist and the viewer reads the v0 sample table, +so both sinks run until the viewer reads traces natively; retiring this one is +a change to the default sink list. -Its known weakness is why traces is the primary: ``POST /samples`` *appends*, -so a request whose response was lost cannot be safely replayed. The client -therefore does not retry it through an ambiguous failure — losing a batch is -recoverable, duplicated rows silently skew every average on the dashboard. -Content-addressed uploads have neither problem, which is exactly the property -the traces sink was built on. +``POST /samples`` appends, so a request whose response was lost is not +replayed: a lost batch is recoverable, duplicated rows skew every average. """ import logging from typing import Any, Dict, List, Mapping, Optional, Sequence from .._http import UPLOAD_TIMEOUT, PlatformClient, encode_json -from ..projection import batch_samples, build_samples, record_to_samples, trace_to_sample -from .base import Sink +from ..projection import batch_samples, build_samples +from .base import Sink, is_episode logger = logging.getLogger(__name__) @@ -30,27 +23,19 @@ class EvalSamplesSink(Sink): name = "eval_samples" - def __init__(self, client: PlatformClient, *, close_client: bool = False) -> None: + def __init__(self, client: PlatformClient) -> None: self.enabled = True self._client = client - self._close_client = close_client self._run_id: Optional[str] = None - # Carried across calls so a streaming producer numbers rollouts the same - # way a one-shot upload does: the Nth episode for an example is rollout N, - # whether it arrived alone or in a batch of five hundred. + # Carried across calls so a streaming producer numbers rollouts the + # same way a one-shot upload did. self._rollout_numbers: Dict[Any, int] = {} self.samples_written = 0 def start(self, run_id: str, context: Mapping[str, str]) -> None: self._run_id = run_id - def write( - self, - records: Sequence[Any], - *, - line_format: Optional[str] = None, - step: Optional[int] = None, - ) -> None: + def write(self, records: Sequence[Any]) -> None: if not self.enabled or not records: return if self._run_id is None: @@ -64,39 +49,23 @@ def write( f"/evaluations/{self._run_id}/samples", content=encode_json({"samples": batch}), timeout=UPLOAD_TIMEOUT, - # Appends. Left non-replayable (the POST default) so a lost - # response cannot turn into duplicate rows. - idempotent=False, + idempotent=False, # appends; a lost response must not duplicate rows ) self.samples_written += len(batch) def _to_samples(self, records: Sequence[Any]) -> List[Dict[str, Any]]: - """Project native episodes/traces and pass through existing samples. - - A producer that already speaks the v0 sample format (a dict with - ``sample_id``) sends it unchanged. Serialized trace and episode records - are projected alongside their native object forms; unsupported mappings - fail explicitly so a gated traces sink cannot turn data loss into a - successful-looking empty run. - """ + """Episode objects are projected; v0 sample dicts (``sample_id``) pass + through. Anything else fails loudly rather than vanishing.""" samples: List[Dict[str, Any]] = [] for record in records: - if isinstance(record, Mapping): - if "sample_id" in record: - samples.append(dict(record)) - else: - samples.extend(record_to_samples(record, self._rollout_numbers)) - continue - if hasattr(record, "traces"): + if isinstance(record, Mapping) and "sample_id" in record: + samples.append(dict(record)) + elif not isinstance(record, Mapping) and is_episode(record): samples.extend(build_samples([record], self._rollout_numbers)) - elif hasattr(record, "branches"): - idx = record.task.data.idx - self._rollout_numbers[idx] = number = self._rollout_numbers.get(idx, 0) + 1 - samples.append(trace_to_sample(record, rollout_number=number)) else: raise TypeError( - f"EvalSamplesSink cannot project {type(record).__name__}; expected a " - "mapping, trace, or episode" + f"EvalSamplesSink cannot project {type(record).__name__}; expected an " + "episode object or a v0 sample dict" ) return samples @@ -104,5 +73,4 @@ def flush(self) -> None: """Writes are synchronous; the uploader thread owns the asynchrony.""" def close(self) -> None: - if self._close_client: - self._client.close() + """The client is shared with the backend, which closes it.""" diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py index c5ef3acfe..61a6f5412 100644 --- a/packages/prime-runs/src/prime_runs/sinks/traces.py +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -1,25 +1,17 @@ -"""Streaming sink over Prime Traces — the primary sample transport. - -Records go out as they are produced, in content-addressed JSONL batches. Two -properties of that transport are why the run handle can promise what it does: -uploads are idempotent (the same bytes resolve to the same upload ID, so a -retry after a lost response replays rather than duplicates), and they are -episode-aware, so a multi-trace rollout keeps its grouping instead of being -flattened into one summary row. - -**The join key is ``run.id`` inside the trace document, not an upload context -key.** The ingestion service extracts ``run.id`` into an indexed column with a -delete-by-run path; ``context`` is an upload-scoped map that answers a -different question. Producers already stamp the run onto their traces, and -``init()`` returns the ID they stamp — so this sink adds nothing to the join -and uses ``context`` only for provenance. +"""Streaming sink over Prime Traces — the primary record transport. + +Uploads are content-addressed, so a retry after a lost response replays rather +than duplicates, and episode-aware, so a multi-trace rollout keeps its grouping. + +The join key is ``run.id`` inside the trace document, which the ingestion +service indexes; ``context`` is upload-scoped provenance only. """ import logging from typing import Any, Dict, Mapping, Optional, Sequence from .. import _fork -from .base import Sink +from .base import Sink, is_episode, stamp_run logger = logging.getLogger(__name__) @@ -36,9 +28,7 @@ def __init__( *, client: Optional[Any] = None, api_key: Optional[str] = None, - traces_url: Optional[str] = None, team_id: Optional[str] = None, - stamp_run: bool = True, compress: bool = True, receipt_history_size: int = DEFAULT_RECEIPT_HISTORY_SIZE, ) -> None: @@ -47,21 +37,15 @@ def __init__( self.enabled = True self._client = client self._injected_client = client is not None - # Left unset, prime-traces resolves its own endpoint. That matters: - # the service has its own URL (PRIME_TRACES_URL / config `traces_url`) - # which is not necessarily the platform API's, and passing the - # platform base URL through here would quietly override it. + # The traces service has its own URL (PRIME_TRACES_URL / `traces_url`), + # which prime-traces resolves itself. self._client_kwargs: Dict[str, Any] = {} if api_key is not None: self._client_kwargs["api_key"] = api_key - if traces_url is not None: - self._client_kwargs["base_url"] = traces_url if team_id is not None: self._client_kwargs["team_id"] = team_id - self._stamp_run = stamp_run self._compress = compress self._run_id: Optional[str] = None - self._run_kind: Optional[str] = None self._context: Dict[str, str] = {} self.receipts: list = [] self.receipts_received = 0 @@ -72,18 +56,14 @@ def __init__( def start(self, run_id: str, context: Mapping[str, str]) -> None: self._run_id = run_id - self._run_kind = context.get("run_kind") self._context = {key: str(value) for key, value in context.items() if value is not None} self._ensure_client() def _ensure_client(self) -> bool: - """Build the traces client if we do not have one. Lazy so that a fork - reset — which drops the inherited client — is repaired on next write.""" + """Build the traces client lazily, so a fork reset is repaired on next write.""" if self._client is not None: return True if self._injected_client: - # The caller handed us a client and a fork took it away. Rebuilding - # would silently swap their transport for a default one. exc = RuntimeError("an injected traces client cannot be reused after a fork") self._disable(str(exc)) raise exc @@ -100,53 +80,38 @@ def _ensure_client(self) -> bool: return True def reset_after_fork(self) -> None: - """Drop the inherited traces client; the next write builds a fresh one. - - Not closed: the child's copy of the socket is the parent's connection, - and shutting it down here would cut the parent off mid-upload. - """ + """Drop (not close) the inherited client; its socket is the parent's.""" self._client = None # ------------------------------------------------------------------ write - def write( - self, - records: Sequence[Any], - *, - line_format: Optional[str] = None, - step: Optional[int] = None, - ) -> None: + def write(self, records: Sequence[Any]) -> None: if not self.enabled or not records or not self._ensure_client(): return from prime_traces import LineFormat - resolved = _resolve_line_format(line_format, records, LineFormat) - context = dict(self._context) - if step is not None: - context["step"] = str(step) - + # The same bytes under a different format are rejected as a conflict, + # so infer from the first record, which is stable within a batch. + line_format = LineFormat.EPISODE if is_episode(records[0]) else LineFormat.TRACE payload = [self._prepare(record) for record in records] try: receipts = list( self._client.upload_records( payload, - line_format=resolved, - context=context or None, + line_format=line_format, + context=dict(self._context) or None, compress=self._compress, ) ) - except Exception as exc: # noqa: BLE001 - classified below + except Exception as exc: if self._is_gated(exc): self._disable( f"Prime Traces is not enabled for this account ({exc}); " "falling back to the remaining sinks" ) - # The worker contains this error in the default warn mode and - # continues with the remaining sinks. Raising is still required - # so strict callers see the failed batch and loss accounting is - # updated instead of reporting a successful traces-only run. - raise + # Re-raised either way so the worker's loss accounting and strict + # callers see the failed batch. raise self.receipts_received += len(receipts) if self._receipt_history_size: @@ -154,30 +119,13 @@ def write( del self.receipts[: -self._receipt_history_size] def _prepare(self, record: Any) -> Any: - """Stamp the run onto plain mappings that do not already carry one. - - Producer objects are passed through untouched — verifiers and prime-rl - both stamp the run themselves at rollout time, and rewriting a caller's - object to add something it already has is how two sources of truth for - the run ID appear. A bare dict has no such convention, so filling in - the indexed field is the difference between a queryable run and an - orphaned upload. - """ - if not self._stamp_run or not isinstance(record, Mapping): - return record - if record.get("run"): + """Stamp the run onto bare mappings; producer objects pass through.""" + if not isinstance(record, Mapping) or self._run_id is None: return record - run: Dict[str, Any] = {"id": self._run_id} - if self._run_kind: - run["type"] = self._run_kind - return {**record, "run": run} + return stamp_run(record, self._run_id) def flush(self) -> None: - """Uploads are synchronous, so nothing is held back here. - - Batching happens inside ``upload_records``; the asynchrony a producer - cares about lives one level up, in the uploader thread. - """ + """Uploads are synchronous; nothing is held back here.""" def close(self) -> None: client = self._client @@ -197,39 +145,10 @@ def _disable(self, reason: str) -> None: @staticmethod def _is_gated(exc: Exception) -> bool: - """Whether this failure means "not allowed", not "try again". - - Prime Traces is in closed beta: a non-allowlisted account gets 403 - ``service_not_enabled``, and a write-only hosted-eval token gets 403 - ``forbidden`` on anything it may not do. Neither is fixable at runtime, - so the sink turns itself off instead of retrying for the rest of the run. - """ + """A 403 (``service_not_enabled``, or a write-only token) is not + fixable at runtime, so the sink turns itself off instead of retrying.""" try: from prime_traces.exceptions import ForbiddenError except ImportError: # pragma: no cover - dependency is declared return False return isinstance(exc, ForbiddenError) - - -def _resolve_line_format(line_format: Optional[str], records: Sequence[Any], enum: Any) -> Any: - """Pick the wire format, preferring what the caller said. - - The default is inferred from the records themselves: anything carrying - ``traces`` is an episode. Guessing wrong is not cosmetic — the same bytes - submitted under a different format are rejected as a conflict — so the - inference only ever looks at the first record's shape, which is stable - within a batch a producer handed over as a unit. - """ - if line_format is not None: - return enum(line_format) if not isinstance(line_format, enum) else line_format - first = records[0] - mapping = _try_mapping(first) - if mapping is not None: - return enum.EPISODE if "traces" in mapping else enum.TRACE - return enum.EPISODE if hasattr(first, "traces") else enum.TRACE - - -def _try_mapping(record: Any) -> Optional[Mapping[str, Any]]: - if isinstance(record, Mapping): - return record - return None diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index a709a7652..2a352f836 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -1,27 +1,16 @@ -"""Background uploader: the thread that keeps the network off the rollout loop. +"""Background uploader: one daemon thread draining a bounded queue into sinks. -Three things a producer should never have to think about, handled once here: +Backpressure: the queue is bounded, so a producer that outruns the uploader +blocks briefly and then drops (counted) rather than stalling the run. -**Backpressure.** The queue is bounded. Verifiers' uploader held every episode -of a run in memory and posted them all at the end, which is fine at a hundred -episodes and is an OOM at a hundred thousand. A bounded queue trades that for a -short block, and — past the block — a counted drop, because stalling a training -run to protect telemetry is the wrong trade in the other direction. +Fork safety: a forked child inherits the queue's memory but not the thread. +The child starts over empty; the queued records belong to the parent. -**Fork safety.** Hosted evals fork after the SDK is initialized. A forked child -inherits the queue's *memory* but not the thread that drains it, so anything -already queued would sit there forever and any lock held mid-write stays held. -The child therefore starts over with an empty queue and a fresh thread, and -drops what it inherited: those records belong to the parent, which is still -running and will upload them itself. - -**Containment.** A sink that raises is retried once, then disabled for the rest -of the run with the error reported through the run's error handler. The upload -thread never propagates into the producer, and never dies quietly either. +Containment: a sink that raises is given a few consecutive transient strikes, +then disabled for the rest of the run. The thread never dies on one bad batch. """ import logging -import os import queue import threading import time @@ -35,37 +24,10 @@ DEFAULT_QUEUE_SIZE = 256 DEFAULT_PUT_TIMEOUT = 5.0 -#: Consecutive transient failures before a sink is retired. One gateway blip -#: must not empty the rest of a run's dashboard; a sustained outage should still -#: stop the SDK from re-attempting every batch for hours. +#: Consecutive transient failures before a sink is retired. TRANSIENT_FAILURE_LIMIT = 3 -@dataclass -class WriteItem: - """One batch of records destined for every enabled sink.""" - - records: Sequence[Any] - line_format: Optional[str] = None - step: Optional[int] = None - - -@dataclass -class MetricItem: - """One ``log()`` call destined for a backend that stores a time series.""" - - metrics: dict - step: Optional[int] = None - - -@dataclass -class RunUpdateItem: - """A config/summary snapshot destined for the run lifecycle backend.""" - - config: Optional[dict] = None - summary: Optional[dict] = None - - @dataclass class _Flush: """A barrier the caller waits on.""" @@ -86,7 +48,7 @@ def _remaining(deadline: Optional[float]) -> Optional[float]: class UploadWorker: - """Drains a bounded queue into a list of sinks on one daemon thread.""" + """Drains a bounded queue of record batches into a list of sinks.""" def __init__( self, @@ -95,34 +57,21 @@ def __init__( max_queue_size: int = DEFAULT_QUEUE_SIZE, put_timeout: float = DEFAULT_PUT_TIMEOUT, on_error: Optional[Callable[[str, Exception], None]] = None, - metric_writer: Optional[Callable[[dict, Optional[int]], None]] = None, - update_writer: Optional[Callable[[Optional[dict], Optional[dict]], None]] = None, ) -> None: self.sinks = sinks self.max_queue_size = max_queue_size self.put_timeout = put_timeout self._on_error = on_error - # Set when the backend stores a real time series. Metrics then ride the - # same queue as records, so a per-step log() in a training loop costs a - # queue put rather than an HTTP round trip. - self._metric_writer = metric_writer - # Eval summaries have no time-series endpoint, but they still belong on - # this thread: periodic persistence must never block the producer loop. - self._update_writer = update_writer self._queue: "queue.Queue[Any]" = queue.Queue(maxsize=max_queue_size) self._thread: Optional[threading.Thread] = None self._stopping = threading.Event() self._lock = threading.Lock() - # Two different losses, deliberately not merged. `dropped` is records - # never handed to any sink because the queue was full — the producer - # outran the uploader. `failed_records` is records a *particular* sink - # could not store, which says nothing about the others: with traces and - # the sample table both enabled, one sink failing usually means the - # records are still safe in the other. + #: Records never handed to any sink because the queue was full. self.dropped = 0 + #: Records a particular sink could not store, by sink name. Kept apart + #: from ``dropped``: another sink may well have stored them. self.failed_records: dict = {} self._transient_failures: dict = {} - self._pid = os.getpid() _fork.register(self) # ----------------------------------------------------------------- thread @@ -146,58 +95,24 @@ def _run(self) -> None: if isinstance(item, _Flush): self._flush_sinks() item.event.set() - continue - if isinstance(item, MetricItem): - self._write_metrics(item) - continue - if isinstance(item, RunUpdateItem): - self._write_update(item) - continue - self._dispatch(item) + else: + self._dispatch(item) except Exception as exc: # noqa: BLE001 - the thread must outlive one bad batch logger.debug("Uploader iteration failed: %s", exc) finally: self._queue.task_done() - def _dispatch(self, item: WriteItem) -> None: + def _dispatch(self, records: Sequence[Any]) -> None: for sink in self.sinks: if not getattr(sink, "enabled", True): continue try: - sink.write(item.records, line_format=item.line_format, step=item.step) + sink.write(records) except Exception as exc: # noqa: BLE001 - one sink failing must not stop the others - self._fail_sink(sink, exc, dropped=len(item.records)) + self._fail_sink(sink, exc, dropped=len(records)) else: self._transient_failures.pop(getattr(sink, "name", id(sink)), None) - def _write_metrics(self, item: MetricItem) -> None: - if self._metric_writer is None: - return - try: - self._metric_writer(item.metrics, item.step) - except Exception as exc: # noqa: BLE001 - metrics must not kill the uploader - logger.warning( - "Dropped metrics for step %s: %s: %s", item.step, type(exc).__name__, exc - ) - if self._on_error is not None: - try: - self._on_error("metrics", exc) - except Exception: # noqa: BLE001 - logger.debug("Error handler raised while reporting metrics", exc_info=True) - - def _write_update(self, item: RunUpdateItem) -> None: - if self._update_writer is None: - return - try: - self._update_writer(item.config, item.summary) - except Exception as exc: # noqa: BLE001 - updates must not kill the uploader - logger.warning("Dropped a run metadata update: %s: %s", type(exc).__name__, exc) - if self._on_error is not None: - try: - self._on_error("run metadata", exc) - except Exception: # noqa: BLE001 - logger.debug("Error handler raised while reporting a run update", exc_info=True) - def _flush_sinks(self) -> None: for sink in self.sinks: if not getattr(sink, "enabled", True): @@ -208,22 +123,9 @@ def _flush_sinks(self) -> None: self._fail_sink(sink, exc) def _fail_sink(self, sink: Any, exc: Exception, *, dropped: int = 0) -> None: - """Handle a sink that raised, and report it. - - The batch is gone either way — the transports already retried internally - (traces on content-addressed uploads, the platform client on whatever it - can safely replay), so an error reaching this point has exhausted its - budget. What is decided here is whether the *sink* is finished: - - - A permanent failure — a gated account, a rejected credential — will - fail identically on every future batch, so the sink stops. Continuing - would produce one log line per batch for the rest of the run and bury - whatever failed first. - - A transient one gets ``TRANSIENT_FAILURE_LIMIT`` consecutive strikes, - reset by any success. Retiring a sink on a single gateway blip would - leave the rest of the run missing from the dashboard, which is a much - larger loss than the one batch that actually failed. - """ + """The batch is gone (the transports already retried). Decide whether + the sink is too: permanent failures retire it at once, transient ones + after ``TRANSIENT_FAILURE_LIMIT`` consecutive strikes.""" name = getattr(sink, "name", type(sink).__name__) if dropped: self.failed_records[name] = self.failed_records.get(name, 0) + dropped @@ -267,30 +169,19 @@ def _notify(self, name: str, exc: Exception) -> None: # ------------------------------------------------------------------ queue - def submit(self, item: Any) -> bool: - """Hand a batch or a metric point to the uploader. - - ``False`` means it was dropped: the queue stayed full for the whole - timeout, so the producer is durably outrunning the uploader. Blocking - further would turn a telemetry backlog into a stalled training run. - """ + def submit(self, records: Sequence[Any]) -> bool: + """Hand a batch to the uploader. ``False`` means it was dropped: the + queue stayed full for the whole put timeout.""" if self._stopping.is_set(): return False thread = self._thread if thread is None or not thread.is_alive(): self.start() try: - self._queue.put(item, timeout=self.put_timeout) + self._queue.put(records, timeout=self.put_timeout) return True except queue.Full: - if not isinstance(item, WriteItem): - logger.warning( - "Upload queue full after %.1fs; could not queue %s", - self.put_timeout, - type(item).__name__, - ) - return False - count = len(item.records) + count = len(records) self.dropped += count logger.warning( "Upload queue full after %.1fs; dropped %d item(s) (%d total). " @@ -309,9 +200,8 @@ def flush(self, timeout: Optional[float] = None) -> bool: deadline = _deadline(timeout) barrier = _Flush() try: - # Synchronization belongs to the caller's drain budget, not the - # short producer backpressure budget. A full queue is precisely - # when finish() most needs to wait for room for this barrier. + # The barrier gets the caller's drain budget, not the short + # producer put timeout: a full queue is when finish() most needs it. self._queue.put(barrier, timeout=_remaining(deadline)) except queue.Full: logger.warning("Could not enqueue a flush barrier before the drain timeout") @@ -325,18 +215,13 @@ def close(self, timeout: Optional[float] = 30.0) -> None: if thread is not None and thread.is_alive(): deadline = _deadline(timeout) try: - # The sentinel sits behind every accepted item. Give it the - # close budget so a temporarily full queue can make room and - # then drain in FIFO order before the thread exits. self._queue.put(None, timeout=_remaining(deadline)) except queue.Full: logger.warning("Upload queue remained saturated through the close timeout") thread.join(_remaining(deadline)) if thread.is_alive(): - # Closing the sinks now would pull an httpx client, or a file - # handle, out from under a request that is still running on that - # thread — turning a slow upload into a crash inside a daemon - # thread nobody is watching. Leave them to the interpreter. + # Closing the sinks now would pull a client or file handle out + # from under a request still running on that thread. logger.warning( "Uploader still running after %ss; leaving it and its sinks open. " "Records still in flight may not finish before the process exits.", @@ -353,16 +238,7 @@ def close(self, timeout: Optional[float] = 30.0) -> None: # ------------------------------------------------------------------- fork def reset_after_fork(self) -> None: - """Give the child a clean uploader. - - Everything queued at fork time belongs to the parent, which still has a - live thread and will send it. Inheriting that queue would upload each - record twice; inheriting the lock could deadlock the child on its first - write. The sinks reset themselves through the same hook — their sockets - and file buffers are the parent's too, and using those from two - processes interleaves one HTTP stream or writes one buffer twice. - """ - self._pid = os.getpid() + """Give the child a clean uploader; what was queued belongs to the parent.""" self._queue = queue.Queue(maxsize=self.max_queue_size) self._thread = None self._stopping = threading.Event() diff --git a/packages/prime-runs/tests/conftest.py b/packages/prime-runs/tests/conftest.py index 96e4b995e..c266341c9 100644 --- a/packages/prime-runs/tests/conftest.py +++ b/packages/prime-runs/tests/conftest.py @@ -15,18 +15,14 @@ "PRIME_BASE_URL", "PRIME_TRACES_URL", "PRIME_FRONTEND_URL", - "PRIME_RUN_ID", "PRIME_RUNS_MODE", "PRIME_RUNS_DIR", - "RANK", - "DP_RANK", - "LOCAL_RANK", ) @pytest.fixture(autouse=True) def isolated_prime_config(monkeypatch, tmp_path): - """Never read the developer's real ~/.prime, env vars or rank.""" + """Never read the developer's real ~/.prime or env vars.""" monkeypatch.setattr(Path, "home", lambda: tmp_path) for name in _PRIME_ENV_VARS: monkeypatch.delenv(name, raising=False) @@ -107,12 +103,6 @@ def eval_routes() -> Dict[str, Any]: "evaluation_id": "eval-abc", "status": "PROCESSING", }, - "GET /api/v1/evaluations/eval-abc": { - "evaluation_id": "eval-abc", - "name": "test-run", - "status": "RUNNING", - "viewer_url": "https://app.example/dashboard/evaluations/eval-abc", - }, "PUT /api/v1/evaluations/eval-abc": { "evaluation_id": "eval-abc", "name": "test-run", @@ -137,10 +127,10 @@ def __init__(self, name: str = "fake", fail_on_write: bool = False) -> None: def start(self, run_id: str, context: Dict[str, str]) -> None: self.started.append((run_id, dict(context))) - def write(self, records, *, line_format=None, step=None) -> None: + def write(self, records) -> None: if self.fail_on_write: raise RuntimeError("sink is broken") - self.batches.append((list(records), line_format, step)) + self.batches.append(list(records)) def flush(self) -> None: self.flushes += 1 diff --git a/packages/prime-runs/tests/test_config_source.py b/packages/prime-runs/tests/test_config_source.py index 5d5e0aafb..8ef867fa7 100644 --- a/packages/prime-runs/tests/test_config_source.py +++ b/packages/prime-runs/tests/test_config_source.py @@ -1,10 +1,5 @@ -"""The config a run is actually configured with. - -Two failures this covers, both of which produced a useless Config tab on the -platform: a resolved model dump that buries three chosen values under hundreds -of defaults, and a structured projection that cannot show the file someone -actually wrote. -""" +"""The config a run is actually configured with: the file someone wrote, kept +byte for byte, rather than a projection that cannot show it.""" import json @@ -107,67 +102,6 @@ def test_a_mapping_without_text_is_not_a_config_source(): # ------------------------------------------------------- config normalization -class FakeModel: - """Duck-types the pydantic v2 surface ``_normalize_config`` looks for.""" - - def __init__(self, set_fields, all_fields): - self._set = set_fields - self._all = all_fields - - def model_dump(self, mode=None, exclude_unset=False): - return dict(self._set if exclude_unset else self._all) - - -def test_a_model_contributes_only_the_fields_someone_set(): - """The training Config tab's actual bug: ``exclude_none`` keeps every default, - so three chosen values arrive buried in a hundred lines nobody picked.""" - model = FakeModel( - set_fields={"model": "Qwen/Qwen3-8B", "max_steps": 1000}, - all_fields={"model": "Qwen/Qwen3-8B", "max_steps": 1000, "seed": 0, "log_level": "info"}, - ) - - assert _normalize_config(model) == {"model": "Qwen/Qwen3-8B", "max_steps": 1000} - - -def test_a_dump_that_cannot_serialize_says_so_instead_of_dumping_everything(): - """The recovery on offer — dump every field — is the exact outcome passing a - model was meant to avoid, so it must never be reached by guessing at why a - call failed. A broken serializer is the caller's bug and surfaces as itself.""" - - class Broken: - def model_dump(self, mode=None, exclude_unset=False): - raise TypeError("serializer blew up") - - with pytest.raises(TypeError, match="serializer blew up"): - _normalize_config(Broken()) - - -def test_a_dump_without_exclude_unset_falls_back_loudly(caplog): - """Degrading to the full config is allowed, going quiet about it is not.""" - - class Old: - def model_dump(self): - return {"model": "Qwen/Qwen3-8B", "seed": 0} - - with caplog.at_level("WARNING"): - assert _normalize_config(Old()) == {"model": "Qwen/Qwen3-8B", "seed": 0} - - assert "exclude_unset" in caplog.text - assert "defaults included" in caplog.text - - -def test_a_dump_taking_kwargs_is_given_the_keywords(): - seen = {} - - class Flexible: - def model_dump(self, **kwargs): - seen.update(kwargs) - return {"a": 1} - - assert _normalize_config(Flexible()) == {"a": 1} - assert seen == {"mode": "json", "exclude_unset": True} - - def test_a_mapping_is_taken_exactly_as_given(): """The caller already decided what to say; second-guessing it would be worse.""" assert _normalize_config({"a": 1, "b": None}) == {"a": 1, "b": None} @@ -231,19 +165,22 @@ def test_an_offline_run_stores_the_launch_file(tmp_path): def test_extra_values_can_be_merged_onto_a_launch_file(tmp_path): - """One parameter takes one form. A run launched from a file that also wants - a derived value adds it explicitly, rather than the SDK growing a second - config argument for a case that is not the common one.""" + """A run launched from a file that also wants structured values passes a + mapping carrying the source under ``CONFIG_SOURCE_KEY`` — what verifiers does.""" path = tmp_path / "eval.toml" path.write_text(EVAL_TOML) + config = { + "model": "deepseek/deepseek-v4-flash", + CONFIG_SOURCE_KEY: ConfigSource.from_file(path).to_dict(), + } - run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path), config=path) - run.update_config({"resolved_model": "deepseek/deepseek-v4-flash"}) + run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path), config=config) run.finish() state = json.loads((tmp_path / run.id / "run.json").read_text()) - assert state["config"]["resolved_model"] == "deepseek/deepseek-v4-flash" + assert state["config"]["model"] == "deepseek/deepseek-v4-flash" assert state["config"][CONFIG_SOURCE_KEY]["text"] == EVAL_TOML + assert run.config_source.filename == "eval.toml" def test_the_run_reports_its_own_source(tmp_path): @@ -272,14 +209,8 @@ def test_an_online_run_sends_the_source_in_create_metadata( handler = RecordingHandler(eval_routes) monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler)) - run = pr.init( - name="tb2", - environments=["gsm8k"], - api_key="test-key", - config=path, - traces=False, - samples=False, - ) + monkeypatch.setattr("prime_traces.TracesClient", lambda **_: object()) + run = pr.init(name="tb2", environments=["gsm8k"], api_key="test-key", config=path) run.finish() create = next(r for r in handler.requests if r.url.path == "/api/v1/evaluations/") diff --git a/packages/prime-runs/tests/test_evals_backend.py b/packages/prime-runs/tests/test_evals_backend.py index 09e1d735c..2c4839e92 100644 --- a/packages/prime-runs/tests/test_evals_backend.py +++ b/packages/prime-runs/tests/test_evals_backend.py @@ -8,10 +8,7 @@ from prime_runs.exceptions import ( ConfigurationError, EnvironmentResolutionError, - ForbiddenError, - PaymentRequiredError, RetryableAPIError, - UnauthorizedError, ) from prime_runs.models import EnvironmentRef, RunSpec, RunStatus @@ -75,9 +72,7 @@ def test_a_published_environment_slug_supplies_dataset_and_default_name( make_platform_client, eval_routes ): routes = dict(eval_routes) - routes["GET /api/v1/environmentshub/alice/gsm8k/@latest"] = { - "data": {"id": "env-published"} - } + routes["GET /api/v1/environmentshub/alice/gsm8k/@latest"] = {"data": {"id": "env-published"}} backend, handler = make_backend(make_platform_client, routes) backend.create(RunSpec(environments=[EnvironmentRef.coerce("alice/gsm8k")])) @@ -153,80 +148,22 @@ def test_an_ambiguous_finalize_failure_is_not_replayed(make_platform_client, eva assert handler.paths().count("POST /api/v1/evaluations/eval-abc/finalize") == 1 -def test_a_failed_run_falls_back_to_metadata_when_the_status_endpoint_is_missing( - make_platform_client, eval_routes, caplog -): - """The platform has no producer-facing way to fail an evaluation yet. - - Until it does, the run cannot leave RUNNING — but the failure must still be - recorded somewhere an operator and the dashboard can both read it, and the - SDK must say plainly that the run will keep showing as running. - """ +def test_a_failed_run_is_recorded_in_metadata(make_platform_client, eval_routes, caplog): + """The platform has no producer-facing way to fail an evaluation. The run + cannot leave RUNNING, but the failure is recorded where an operator and the + dashboard can read it, and the SDK says the run will keep showing as running.""" backend, handler = make_backend(make_platform_client, eval_routes) with caplog.at_level("WARNING"): backend.finalize("eval-abc", status=RunStatus.FAILED, error="boom") - assert "POST /api/v1/evaluations/eval-abc/status" in handler.paths() + assert "POST /api/v1/evaluations/eval-abc/finalize" not in handler.paths() terminal = handler.bodies_for("/api/v1/evaluations/eval-abc")[0]["metadata"]["prime_runs"] assert terminal["status"] == "failed" assert terminal["error"] == "boom" assert "keep showing as running" in caplog.text -def test_the_missing_status_endpoint_is_probed_once_per_backend(make_platform_client, eval_routes): - backend, handler = make_backend(make_platform_client, eval_routes) - - backend.finalize("eval-abc", status=RunStatus.FAILED, error="one") - backend.finalize("eval-abc", status=RunStatus.CRASHED, error="two") - - assert handler.paths().count("POST /api/v1/evaluations/eval-abc/status") == 1 - - -def test_a_status_endpoint_that_exists_is_used_instead_of_the_fallback( - make_platform_client, eval_routes -): - routes = dict(eval_routes) - routes["POST /api/v1/evaluations/eval-abc/status"] = {"evaluation_id": "eval-abc"} - backend, handler = make_backend(make_platform_client, routes) - - backend.finalize("eval-abc", status=RunStatus.FAILED, error="boom") - - assert handler.bodies_for("/api/v1/evaluations/eval-abc/status")[0] == { - "status": "FAILED", - "error": "boom", - } - assert "PUT /api/v1/evaluations/eval-abc" not in handler.paths() - - -def test_a_transient_status_failure_is_retried(make_platform_client, eval_routes): - routes = dict(eval_routes) - responses = [httpx.Response(503), httpx.Response(200, json={"evaluation_id": "eval-abc"})] - routes["POST /api/v1/evaluations/eval-abc/status"] = lambda request: responses.pop(0) - backend, handler = make_backend(make_platform_client, routes) - - backend.finalize("eval-abc", status=RunStatus.FAILED, error="boom") - - assert handler.paths().count("POST /api/v1/evaluations/eval-abc/status") == 2 - assert "PUT /api/v1/evaluations/eval-abc" not in handler.paths() - - -def test_exhausted_status_retries_fall_back_to_metadata(make_platform_client, eval_routes, caplog): - routes = dict(eval_routes) - routes["POST /api/v1/evaluations/eval-abc/status"] = lambda request: httpx.Response(503) - handler = RecordingHandler(routes) - client = make_platform_client(handler, max_attempts=2) - backend = EvalsBackend(client, frontend_url="https://app.example") - - with caplog.at_level("WARNING"): - backend.finalize("eval-abc", status=RunStatus.FAILED, error="boom") - - assert handler.paths().count("POST /api/v1/evaluations/eval-abc/status") == 2 - terminal = handler.bodies_for("/api/v1/evaluations/eval-abc")[0]["metadata"]["prime_runs"] - assert terminal["status"] == "failed" - assert "remained unavailable after retries" in caplog.text - - def test_update_sends_nothing_when_there_is_nothing_to_send(make_platform_client, eval_routes): backend, handler = make_backend(make_platform_client, eval_routes) @@ -235,41 +172,6 @@ def test_update_sends_nothing_when_there_is_nothing_to_send(make_platform_client assert handler.requests == [] -def test_attach_survives_a_read_failure(make_platform_client, eval_routes): - """Losing a run's name to a transient read is not worth failing a resume on.""" - routes = dict(eval_routes) - routes["GET /api/v1/evaluations/eval-abc"] = lambda request: httpx.Response( - 500, json={"detail": "nope"} - ) - backend, _ = make_backend(make_platform_client, routes) - - handle = backend.attach("eval-abc") - - assert handle.id == "eval-abc" - assert handle.url == "https://app.example/dashboard/evaluations/eval-abc" - - -@pytest.mark.parametrize( - ("status_code", "error_type"), - [ - (401, UnauthorizedError), - (402, PaymentRequiredError), - (403, ForbiddenError), - ], -) -def test_attach_propagates_permanent_access_failures( - make_platform_client, eval_routes, status_code, error_type -): - routes = dict(eval_routes) - routes["GET /api/v1/evaluations/eval-abc"] = lambda request: httpx.Response( - status_code, json={"detail": "denied"} - ) - backend, _ = make_backend(make_platform_client, routes) - - with pytest.raises(error_type): - backend.attach("eval-abc") - - def test_a_pinned_environment_version_reaches_the_api(make_platform_client, eval_routes): """The API's EnvironmentReference carries version_id. Dropping it attaches the run to whatever version the hub resolves today — the difference between diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index e94314727..15a8a9dc8 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -1,4 +1,4 @@ -"""``init()``: mode resolution, offline runs, online runs, rank handling.""" +"""``init()``: mode resolution, offline runs, online runs.""" import json import os @@ -8,24 +8,19 @@ from conftest import RecordingHandler import prime_runs as pr -from prime_runs.backends.offline import DEFAULT_DIR_ENV from prime_runs.exceptions import ConfigurationError from prime_runs.models import RunStatus -from prime_runs.run import MODE_ENV, RUN_ID_ENV, _exported_run_ids # ------------------------------------------------------------------- offline def test_an_offline_run_is_a_real_run(tmp_path): - """The reason producers can delete their ``--no-push`` branch: same ID, same - status, same calls — just a different destination.""" run = pr.init(name="local", environments=["gsm8k"], mode="offline", dir=str(tmp_path)) assert run.id.startswith("offline-") assert run.url == str((tmp_path / run.id).resolve()) assert run.mode == "offline" - run.log({"reward": 0.5}, step=1) run.log_traces([{"id": "t1"}]) run.finish(summary={"avg_reward": 0.5}) @@ -46,7 +41,7 @@ def test_offline_records_are_written_in_the_wire_format(tmp_path): records = [json.loads(line) for line in lines] assert [record["id"] for record in records] == ["t1", "t2"] - assert all(record["run"]["id"] == run.id for record in records) + assert all(record["run"] == {"id": run.id, "type": "eval"} for record in records) def test_episodes_are_written_to_their_own_file(tmp_path): @@ -57,17 +52,6 @@ def test_episodes_are_written_to_their_own_file(tmp_path): assert (tmp_path / run.id / "records" / "episode.jsonl").exists() -def test_offline_metrics_are_a_time_series(tmp_path): - run = pr.init(mode="offline", dir=str(tmp_path)) - run.log({"loss": 2.0}, step=1) - run.log({"loss": 1.0}, step=2) - run.flush() - run.finish() - - lines = (tmp_path / run.id / "metrics.jsonl").read_text().splitlines() - assert [json.loads(line)["loss"] for line in lines] == [2.0, 1.0] - - def test_a_record_that_already_names_a_run_is_left_alone(tmp_path): """Producers stamp the run themselves; two sources of truth for the run ID is how traces end up on the wrong run.""" @@ -109,7 +93,6 @@ def test_a_disabled_run_still_answers_every_call(tmp_path): """Same object shape, so producer code needs no branching.""" run = pr.init(mode="disabled", dir=str(tmp_path)) - run.log({"reward": 1.0}, step=1) run.log_traces([{"id": "t1"}]) run.finish(summary={"avg_reward": 1.0}) @@ -118,82 +101,9 @@ def test_a_disabled_run_still_answers_every_call(tmp_path): assert not list(tmp_path.iterdir()) -def test_training_runs_are_not_supported_yet(tmp_path): - with pytest.raises(ConfigurationError, match="training runs"): - pr.init(kind="train", api_key="test-key", environments=["gsm8k"]) - - -# --------------------------------------------------------------------- rank - - -def test_a_non_primary_rank_with_no_run_to_join_records_nothing(monkeypatch, tmp_path): - """Otherwise rank 3 creates a second run for the same job.""" - monkeypatch.setenv("RANK", "3") - - run = pr.init(name="local", api_key="test-key", environments=["gsm8k"]) - - assert run.mode == "disabled" - assert run.is_primary is False - run.finish() - - -def test_a_non_primary_offline_rank_without_a_run_to_join_records_nothing(monkeypatch, tmp_path): - """An offline rank must not create a run it is forbidden to finalize.""" - monkeypatch.setenv("RANK", "3") - - run = pr.init(mode="offline", dir=str(tmp_path)) - - assert run.mode == "disabled" - assert run.is_primary is False - run.finish() - assert not list(tmp_path.iterdir()) - - -def test_a_run_id_in_the_environment_is_joined_not_recreated(monkeypatch, tmp_path): - monkeypatch.setenv("DP_RANK", "2") - monkeypatch.setenv(RUN_ID_ENV, "offline-shared") - - run = pr.init(mode="offline", dir=str(tmp_path)) - - assert run.id == "offline-shared" - assert run.is_primary is False - run.finish() - - -def test_init_publishes_the_run_id_for_child_processes(tmp_path): - """Forked workers and subprocess launchers join the run their parent opened - instead of each opening their own.""" - run = pr.init(mode="offline", dir=str(tmp_path)) - - assert os.environ[RUN_ID_ENV] == run.id - run.finish() - - -def test_a_child_inherits_the_resolved_offline_mode_and_directory(monkeypatch, tmp_path): - """An API key must not make the child switch an explicit offline parent online.""" - monkeypatch.setenv("PRIME_API_KEY", "test-key") - monkeypatch.setattr( - "prime_runs.run.PlatformClient", - lambda **_: pytest.fail("the inherited child unexpectedly selected online mode"), - ) - parent = pr.init(mode="offline", dir=str(tmp_path)) - - assert os.environ[MODE_ENV] == "offline" - assert os.environ[DEFAULT_DIR_ENV] == str(tmp_path.resolve()) - - # Emulate the PID distinction a real child process inherits. - _exported_run_ids[parent.id] = os.getpid() - 1 - child = pr.init(handle_signals=False) - child.log_traces([{"id": "from-child"}]) - child.finish() - parent.finish() - - assert child.id == parent.id - assert child.mode == "offline" - assert (tmp_path / parent.id / "records" / "trace.jsonl").exists() - assert RUN_ID_ENV not in os.environ - assert MODE_ENV not in os.environ - assert DEFAULT_DIR_ENV not in os.environ +def test_online_without_an_api_key_is_a_configuration_error(): + with pytest.raises(ConfigurationError, match="needs an API key"): + pr.init(mode="online", environments=["gsm8k"]) # ------------------------------------------------------------------- online @@ -201,20 +111,20 @@ def test_a_child_inherits_the_resolved_offline_mode_and_directory(monkeypatch, t @pytest.fixture def online(monkeypatch, make_platform_client, eval_routes): - """``init(mode="online")`` wired to a MockTransport.""" + """``init(mode="online")`` wired to a MockTransport, traces sink off.""" def _init(routes=None, **kwargs): handler = RecordingHandler(routes or eval_routes) monkeypatch.setattr( "prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler) ) + monkeypatch.setattr("prime_runs.run.TracesSink", lambda **_: _NullSink()) run = pr.init( name="test-run", environments=["gsm8k"], model="Qwen3-8B", framework="verifiers", api_key="test-key", - traces=False, **kwargs, ) return run, handler @@ -222,55 +132,21 @@ def _init(routes=None, **kwargs): return _init -def test_samples_use_a_separate_client_from_run_finalization( - monkeypatch, make_platform_client, eval_routes -): - handler = RecordingHandler(eval_routes) - - class TrackingClient: - def __init__(self): - self._delegate = make_platform_client(handler) - self.closed = False - - def get(self, *args, **kwargs): - return self._delegate.get(*args, **kwargs) - - def post(self, *args, **kwargs): - return self._delegate.post(*args, **kwargs) - - def put(self, *args, **kwargs): - return self._delegate.put(*args, **kwargs) - - def close(self): - self.closed = True +class _NullSink: + name = "traces" + enabled = True - clients = [] + def start(self, run_id, context): + pass - def make_client(**kwargs): - client = TrackingClient() - clients.append(client) - return client + def write(self, records): + pass - monkeypatch.setattr("prime_runs.run.PlatformClient", make_client) - run = pr.init( - name="test-run", - environments=["gsm8k"], - api_key="test-key", - traces=False, - handle_signals=False, - ) - assert len(clients) == 2 + def flush(self): + pass - # Model UploadWorker.close() timing out: it intentionally leaves its sink - # open, while lifecycle finalization still closes the backend transport. - run._worker.close = lambda timeout=None: None - run.finish() - - backend_client, samples_client = clients - assert backend_client.closed is True - assert samples_client.closed is False - run._worker.sinks[0].close() - assert samples_client.closed is True + def close(self): + pass def test_an_online_run_returns_the_platforms_id_and_viewer_url(online): @@ -282,6 +158,20 @@ def test_an_online_run_returns_the_platforms_id_and_viewer_url(online): run.finish() +def test_an_online_run_has_both_transports_by_default( + monkeypatch, make_platform_client, eval_routes +): + """Traces is the system of record; the sample table is what the viewer reads.""" + handler = RecordingHandler(eval_routes) + monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler)) + monkeypatch.setattr("prime_traces.TracesClient", lambda **_: object()) + + run = pr.init(name="test-run", environments=["gsm8k"], api_key="test-key") + + assert [sink.name for sink in run._worker.sinks] == ["traces", "eval_samples"] + run.finish() + + def test_episodes_stream_to_the_sample_table_while_the_run_is_going(online): run, handler = online() @@ -306,7 +196,7 @@ def test_finishing_an_online_run_finalizes_it_with_its_metrics(online): def test_the_end_to_end_shape_a_producer_writes(online): - """The whole surface, in the order verifiers will call it.""" + """The whole surface, in the order verifiers calls it.""" from prime_runs import metrics episodes = [make_episode(f"ep-{n}", [make_trace(idx=n, reward=float(n))]) for n in range(3)] @@ -323,13 +213,7 @@ def test_the_end_to_end_shape_a_producer_writes(online): assert run.errors == [] -# ------------------------------------------------------- run id inheritance - - def test_a_second_init_in_one_process_opens_its_own_run(tmp_path): - """init() exports PRIME_RUN_ID for child processes. Reading our own export - back would silently attach the second eval to the first, and it would never - create or finalize a run of its own.""" first = pr.init(mode="offline", dir=str(tmp_path)) first.finish() @@ -340,109 +224,24 @@ def test_a_second_init_in_one_process_opens_its_own_run(tmp_path): assert (tmp_path / second.id / "run.json").exists() -def test_a_finished_run_stops_advertising_itself(tmp_path): - run = pr.init(mode="offline", dir=str(tmp_path)) - assert os.environ[RUN_ID_ENV] == run.id - - run.finish() - - assert RUN_ID_ENV not in os.environ - - -def test_an_id_inherited_from_a_parent_process_is_joined(monkeypatch, tmp_path): - """The env var without a matching PID belongs to an ancestor.""" - monkeypatch.setenv(RUN_ID_ENV, "offline-from-parent") - - run = pr.init(mode="offline", dir=str(tmp_path)) - - assert run.id == "offline-from-parent" - assert run.is_primary is False - run.finish() - - -def test_an_explicit_id_is_a_resume_and_still_finalizes(online): - """Resuming after a crash has to be able to close the run out; only an ID - picked up from the environment belongs to someone else.""" - run, handler = online(id="eval-abc") - - run.finish(summary={"avg_reward": 1.0}) - - assert "POST /api/v1/evaluations/eval-abc/finalize" in handler.paths() - - -def test_resuming_preserves_existing_config_and_summary(online, eval_routes): - routes = dict(eval_routes) - routes["GET /api/v1/evaluations/eval-abc"] = { - **routes["GET /api/v1/evaluations/eval-abc"], - "metadata": {"before_crash": True, "overridden": "old"}, - "metrics": {"old_reward": 0.5, "overridden": "old"}, - } - - run, handler = online(routes=routes, id="eval-abc", config={"overridden": "new"}) - run.update_config({"after_resume": True}) - run.log({"overridden": "new", "new_reward": 1.0}) - run.finish() - - update = handler.bodies_for("/api/v1/evaluations/eval-abc")[0] - assert update["metadata"] == { - "before_crash": True, - "overridden": "new", - "after_resume": True, - } - # Recovered metrics survive; anything logged after the resume wins. - assert update["metrics"] == { - "old_reward": 0.5, - "overridden": "new", - "new_reward": 1.0, - } - - -def test_init_forwards_the_finish_timeout(tmp_path): - run = pr.init(mode="offline", dir=str(tmp_path), finish_timeout=0.25) - - assert run._finish_timeout == 0.25 - run.finish() - - -def test_an_id_inherited_from_the_environment_does_not_finalize(monkeypatch, online): - monkeypatch.setenv(RUN_ID_ENV, "eval-abc") - - run, handler = online() - run.finish() - - assert "POST /api/v1/evaluations/eval-abc/finalize" not in handler.paths() +# --------------------------------------------------------------------- fork @pytest.mark.skipif(not hasattr(os, "fork"), reason="fork is POSIX-only") -# Forking a threaded process is exactly the situation under test — hosted evals -# do it, and the uploader thread is why the SDK needs a fork hook at all. @pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning") -def test_a_forked_child_joins_the_run_without_duplicating_the_parents_records(tmp_path): - """The end-to-end shape hosted evals actually hit. - - At fork time the parent has records in the upload queue and bytes in the - sink's write buffer. The child inherits copies of both; writing them would - put every one of those records in the file twice, and opening its own run - would split one job across two. - """ - import json - - run = pr.init(mode="offline", dir=str(tmp_path), handle_signals=False) +def test_a_forked_child_does_not_duplicate_the_parents_records_or_close_its_run(tmp_path): + """At fork time the parent has records in the upload queue. The child + inherits a copy; writing them would put every record in the file twice, + and the inherited atexit hook must not finalize the parent's run.""" + run = pr.init(mode="offline", dir=str(tmp_path)) run.log_traces([{"id": f"parent-{n}"} for n in range(5)]) pid = os.fork() if pid == 0: # pragma: no cover - asserted through the child's exit code code = 0 try: - child = pr.init(mode="offline", dir=str(tmp_path), handle_signals=False) - if child.id != run.id: - code = 1 - if child.is_primary: - code = 3 - child.log_traces([{"id": "child-1"}]) - child.finish() - # The original handle is inherited too. Its atexit callback must be - # harmless in the child: the parent still owns this lifecycle. + run.log_traces([{"id": "child-1"}]) + run.flush() run._on_process_exit() state = json.loads((tmp_path / run.id / "run.json").read_text()) if state["status"] != RunStatus.RUNNING.value: @@ -453,7 +252,7 @@ def test_a_forked_child_joins_the_run_without_duplicating_the_parents_records(tm os._exit(code) _, status = os.waitpid(pid, 0) - assert os.waitstatus_to_exitcode(status) == 0, "the child did not join the parent's run" + assert os.waitstatus_to_exitcode(status) == 0 run.finish() lines = (tmp_path / run.id / "records" / "trace.jsonl").read_text().splitlines() @@ -463,14 +262,7 @@ def test_a_forked_child_joins_the_run_without_duplicating_the_parents_records(tm def test_offline_records_are_on_disk_before_any_flush(tmp_path): - """Nothing may sit in a process-local write buffer. - - A buffered writer holds records in memory until it decides to flush, and a - fork copies that buffer — after which both processes eventually write it and - every buffered record lands in the file twice. Reading the file back through - a separate handle, with no flush and no close, is what proves the buffer is - not there to be copied. - """ + """Nothing may sit in a process-local write buffer a fork could copy.""" run = pr.init(mode="offline", dir=str(tmp_path)) run.log_traces([{"id": "t1"}]) run.flush() @@ -487,12 +279,7 @@ def test_offline_records_are_on_disk_before_any_flush(tmp_path): def test_offline_records_reject_nonfinite_json_instead_of_writing_invalid_jsonl(tmp_path): - run = pr.init( - mode="offline", - dir=str(tmp_path), - handle_signals=False, - on_error="raise", - ) + run = pr.init(mode="offline", dir=str(tmp_path), on_error="raise") run.log_traces([{"id": "bad", "reward": float("nan")}]) with pytest.raises(ValueError, match="Out of range float values"): diff --git a/packages/prime-runs/tests/test_projection.py b/packages/prime-runs/tests/test_projection.py index 3929a426a..34646fc1d 100644 --- a/packages/prime-runs/tests/test_projection.py +++ b/packages/prime-runs/tests/test_projection.py @@ -12,7 +12,6 @@ MAX_SAMPLES_PAYLOAD_BYTES, batch_samples, build_samples, - is_episode, summary_trace_index, trace_to_sample, ) @@ -136,6 +135,10 @@ def test_batch_samples_returns_nothing_for_no_samples(): assert batch_samples([]) == [] -def test_is_episode_distinguishes_episodes_from_traces(): +def test_is_episode_distinguishes_episodes_from_traces_in_either_form(): + from prime_runs.sinks import is_episode + assert is_episode(make_episode()) assert not is_episode(make_trace()) + assert is_episode({"id": "e", "traces": []}) + assert not is_episode({"id": "t"}) diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index b43d4b24a..f3805b961 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -1,6 +1,5 @@ -"""The run handle: lifecycle, containment, ranks, terminal status.""" +"""The run handle: lifecycle, containment, terminal status.""" -import signal import threading import time from concurrent.futures import ThreadPoolExecutor @@ -17,29 +16,20 @@ class FakeBackend: - def __init__(self, supports_step_metrics: bool = False, fail_on: Optional[str] = None) -> None: - self.kind = "eval" - self.supports_step_metrics = supports_step_metrics + def __init__(self, fail_on: Optional[str] = None) -> None: self.fail_on = fail_on self.updates: List[Dict[str, Any]] = [] - self.points: List[Any] = [] self.finalized: List[Dict[str, Any]] = [] self.closed = False def create(self, spec: RunSpec) -> RunHandle: return RunHandle(id="run-1", name=spec.name, url="https://app.example/run-1") - def attach(self, run_id: str) -> RunHandle: - return RunHandle(id=run_id) - def update(self, run_id, *, config=None, summary=None) -> None: if self.fail_on == "update": raise RuntimeError("update exploded") self.updates.append({"config": config, "summary": summary}) - def log_metrics(self, run_id, metrics, step=None) -> None: - self.points.append((metrics, step)) - def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> None: if self.fail_on == "finalize": raise RuntimeError("finalize exploded") @@ -51,9 +41,9 @@ def close(self) -> None: self.closed = True -def make_run(backend=None, sinks=None, **kwargs) -> Run: +def make_run(backend=None, sinks=None, config=None, **kwargs) -> Run: backend = backend or FakeBackend() - spec = RunSpec(name="test-run", kind="eval", framework="verifiers", model="Qwen3-8B") + spec = RunSpec(name="test-run", framework="verifiers", model="Qwen3-8B", config=config or {}) return Run( backend=backend, handle=backend.create(spec), @@ -69,7 +59,6 @@ def test_the_handle_exposes_what_a_producer_prints(): assert run.id == "run-1" assert run.url == "https://app.example/run-1" assert run.status is RunStatus.RUNNING - assert run.is_primary is True run.finish() @@ -91,11 +80,10 @@ def test_traces_reach_the_sinks_while_the_run_is_still_going(): sink = FakeSink() run = make_run(sinks=[sink]) - run.log_traces([{"id": "t1"}], step=2) + run.log_traces([{"id": "t1"}]) run.flush() - assert sink.batches[0][0] == [{"id": "t1"}] - assert sink.batches[0][2] == 2 + assert sink.batches == [[{"id": "t1"}]] assert not run.finished run.finish() @@ -111,84 +99,16 @@ def test_an_empty_batch_is_not_sent(): run.finish() -def test_metrics_land_in_the_summary_when_the_backend_has_no_time_series(): - backend = FakeBackend(supports_step_metrics=False) - run = make_run(backend, summary_flush_seconds=0.0) - - run.log({"reward": 0.5}, step=1) - run.log({"reward": 0.75}, step=2) - run.flush() - - assert run.summary["reward"] == 0.75 - assert backend.points == [] - assert backend.updates, "the summary was flushed" - run.finish() - - -def test_metrics_become_a_time_series_when_the_backend_has_one(): - backend = FakeBackend(supports_step_metrics=True) - run = make_run(backend) - - run.log({"loss": 2.0}, step=1) - run.flush() - - assert backend.points == [({"loss": 2.0}, 1)] - run.finish() - - -def test_commit_false_stages_without_writing(): - backend = FakeBackend(supports_step_metrics=True) - run = make_run(backend) - - run.log({"loss": 2.0}, step=1, commit=False) - run.flush() - - assert backend.points == [] - assert run.summary["loss"] == 2.0 - run.finish() - - -def test_commit_false_merges_staged_metrics_into_the_next_point(): - backend = FakeBackend(supports_step_metrics=True) - run = make_run(backend) - - run.log({"loss": 2.0}, step=7, commit=False) - run.log({"reward": 0.5}) - run.flush() - - assert backend.points == [({"loss": 2.0, "reward": 0.5}, 7)] - run.finish() - - -def test_periodic_summary_updates_run_on_the_uploader_thread(): - caller_thread = threading.get_ident() - update_threads = [] - - class ThreadRecordingBackend(FakeBackend): - def update(self, run_id, *, config=None, summary=None) -> None: - update_threads.append(threading.get_ident()) - super().update(run_id, config=config, summary=summary) - - backend = ThreadRecordingBackend(supports_step_metrics=False) - run = make_run(backend, summary_flush_seconds=0.0) - - run.log({"reward": 0.5}) - run.flush() - - assert update_threads - assert all(thread_id != caller_thread for thread_id in update_threads) - run.finish() - - -def test_non_finite_metrics_are_dropped_rather_than_failing_the_request(): +def test_non_finite_summary_values_are_dropped_rather_than_failing_the_request(): """A diverged loss serializes as bare ``NaN``, which strict JSON rejects — the whole request fails on a payload nobody can inspect.""" - run = make_run(summary_flush_seconds=0.0) + backend = FakeBackend() + run = make_run(backend) - run.log({"loss": float("nan"), "grad": float("inf"), "reward": 0.5}) + run.finish(summary={"loss": float("nan"), "grad": float("inf"), "reward": 0.5}) assert run.summary == {"reward": 0.5} - run.finish() + assert backend.finalized[0]["summary"] == {"reward": 0.5} def test_finish_flushes_records_before_reporting_the_terminal_status(): @@ -197,9 +117,9 @@ def test_finish_flushes_records_before_reporting_the_terminal_status(): order: List[str] = [] class OrderedSink(FakeSink): - def write(self, records, *, line_format=None, step=None) -> None: + def write(self, records) -> None: order.append("write") - super().write(records, line_format=line_format, step=step) + super().write(records) class OrderedBackend(FakeBackend): def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> None: @@ -235,9 +155,7 @@ def __init__(self) -> None: def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> None: self.finalize_started.set() assert self.release_finalize.wait(2.0) - super().finalize( - run_id, status=status, summary=summary, error=error, config=config - ) + super().finalize(run_id, status=status, summary=summary, error=error, config=config) backend = BlockingBackend() run = make_run(backend) @@ -283,8 +201,6 @@ def test_logging_after_finish_is_a_producer_bug(): run = make_run() run.finish() - with pytest.raises(RunFinishedError): - run.log({"reward": 1.0}) with pytest.raises(RunFinishedError): run.log_traces([{"id": "t1"}]) @@ -292,7 +208,7 @@ def test_logging_after_finish_is_a_producer_bug(): def test_the_context_manager_completes_a_clean_run(): backend = FakeBackend() with make_run(backend) as run: - run.log({"reward": 1.0}) + run.log_traces([{"id": "t1"}]) assert backend.finalized[0]["status"] is RunStatus.COMPLETED @@ -337,8 +253,7 @@ def interrupt_finish(*args, **kwargs): def test_an_interrupt_is_recorded_as_a_decision_not_a_fault(): - """Ctrl-C must not land in the same bucket as a broken eval — and it must - agree with the SIGINT handler, which normally gets there first.""" + """Ctrl-C must not land in the same bucket as a broken eval.""" backend = FakeBackend() with pytest.raises(KeyboardInterrupt): @@ -349,29 +264,11 @@ def test_an_interrupt_is_recorded_as_a_decision_not_a_fault(): assert backend.finalized[0]["error"] == "interrupted" -def test_a_termination_signal_reports_crashed_like_atexit_does(): - """The producer never said the run failed; it was stopped from outside.""" - backend = FakeBackend() - run = make_run(backend) - chained = [] - # Stand in for the handler the SDK displaced. Anything but SIG_DFL, which - # would re-raise the signal and take the test runner down with it. - run._previous_signal_handlers[signal.SIGTERM] = lambda *a: chained.append(a) - - run._handle_signal(signal.SIGTERM, None) - - assert backend.finalized[0]["status"] is RunStatus.CRASHED - assert "SIGTERM" in backend.finalized[0]["error"] - assert chained, "the displaced handler still runs" - signal.signal(signal.SIGTERM, signal.SIG_DFL) - - def test_finish_hands_the_full_config_to_finalize(): """The evaluations API replaces metadata wholesale, so a backend recording terminal state inside it needs the whole picture to merge into.""" backend = FakeBackend() - run = make_run(backend) - run.update_config({"num_rollouts": 4}) + run = make_run(backend, config={"num_rollouts": 4}) run.finish(status=RunStatus.FAILED, error="boom") @@ -423,6 +320,7 @@ def test_a_process_that_exits_without_finishing_reports_crashed(): def test_a_forked_handle_gets_a_fresh_lock_and_loses_lifecycle_ownership(): + """The child's inherited atexit hook must never finalize the parent's run.""" backend = FakeBackend() run = make_run(backend) inherited_lock = run._finish_lock @@ -430,9 +328,9 @@ def test_a_forked_handle_gets_a_fresh_lock_and_loses_lifecycle_ownership(): run.reset_after_fork() assert run._finish_lock is not inherited_lock - assert run.is_primary is False run.finish() assert backend.finalized == [] + assert backend.closed is True def test_a_backend_failure_does_not_escape_into_the_producer_by_default(): @@ -499,31 +397,6 @@ def test_a_sink_error_is_recorded_on_the_run(): run.finish() -def test_a_non_primary_rank_does_not_close_the_shared_run(): - """Eight ranks racing to finalize produce seven confusing failures.""" - backend = FakeBackend() - run = make_run(backend, is_primary=False) - - run.log({"reward": 1.0}) - run.finish() - - assert backend.finalized == [] - assert backend.updates == [] - assert run.is_primary is False - - -def test_a_non_primary_rank_still_uploads_its_own_records(): - """The point of eight ranks is that they contribute to one run.""" - sink = FakeSink() - run = make_run(sinks=[sink], is_primary=False) - - run.log_traces([{"id": "t1"}]) - run.flush() - - assert sink.batches - run.finish() - - def test_dropped_records_are_reported_on_the_handle(): run = make_run() run._worker.dropped = 3 @@ -569,9 +442,7 @@ def test_an_upload_failure_is_reported_once(): def test_a_gated_trace_upload_reaches_strict_callers_and_loss_accounting(): class GatedClient: def upload_records(self, records, **kwargs): - raise ForbiddenError( - "not in beta", status_code=403, code="service_not_enabled" - ) + raise ForbiddenError("not in beta", status_code=403, code="service_not_enabled") def close(self) -> None: pass @@ -583,88 +454,3 @@ def close(self) -> None: run.finish() assert run.failed_records == {"traces": 1} - - -def test_a_signal_interrupting_finish_is_chained_after_teardown(monkeypatch): - events = [] - - class SignallingBackend(FakeBackend): - def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> None: - events.append("finalize-start") - run._handle_signal(signal.SIGTERM, None) - events.append("finalize-end") - super().finalize( - run_id, status=status, summary=summary, error=error, config=config - ) - - backend = SignallingBackend() - run = make_run(backend) - run._previous_signal_handlers[signal.SIGTERM] = lambda *_: events.append("signal") - monkeypatch.setattr("prime_runs.run.signal.signal", lambda *_: None) - - run.finish() - - assert events == ["finalize-start", "finalize-end", "signal"] - assert backend.closed is True - - -def test_signal_handlers_are_restored_when_the_run_finishes(): - """`self._handle_signal` builds a new bound method on every access, so an - identity check against a fresh one never matches — leaving the handler - installed, pinning the finished run, and blocking the next run in the - process from installing its own.""" - original = signal.getsignal(signal.SIGTERM) - run = make_run() - run.install_signal_handlers() - assert signal.getsignal(signal.SIGTERM) is run._signal_handler - - run.finish() - - assert signal.getsignal(signal.SIGTERM) is original - - -def test_a_later_run_can_install_its_own_handlers(): - first = make_run() - first.install_signal_handlers() - first.finish() - - second = make_run() - second.install_signal_handlers() - - assert signal.getsignal(signal.SIGTERM) is second._signal_handler - second.finish() - assert signal.getsignal(signal.SIGTERM) is signal.SIG_DFL - - -def test_a_run_finished_in_an_executor_relinquishes_handlers_to_the_next_run(): - original = signal.getsignal(signal.SIGTERM) - first = make_run() - first.install_signal_handlers() - - with ThreadPoolExecutor(max_workers=1) as executor: - executor.submit(first.finish).result() - - # Python forbids signal.signal() off the main thread, so restoration is - # deferred rather than forgetting which handler was displaced. - assert signal.getsignal(signal.SIGTERM) is first._signal_handler - second = make_run() - second.install_signal_handlers() - assert signal.getsignal(signal.SIGTERM) is second._signal_handler - - second.finish() - assert signal.getsignal(signal.SIGTERM) is original - - -def test_a_child_run_can_replace_an_inherited_signal_handler(): - original = signal.getsignal(signal.SIGTERM) - inherited = make_run() - inherited.install_signal_handlers() - inherited.reset_after_fork() - - child = make_run() - child.install_signal_handlers() - - assert signal.getsignal(signal.SIGTERM) is child._signal_handler - child.finish() - inherited.finish() - assert signal.getsignal(signal.SIGTERM) is original diff --git a/packages/prime-runs/tests/test_samples_sink.py b/packages/prime-runs/tests/test_samples_sink.py index 379afde1a..6ea84556a 100644 --- a/packages/prime-runs/tests/test_samples_sink.py +++ b/packages/prime-runs/tests/test_samples_sink.py @@ -35,19 +35,6 @@ def test_rollout_numbering_is_continuous_across_streamed_batches(make_platform_c assert [body["samples"][0]["rollout_number"] for body in posted] == [1, 2] -def test_individual_traces_are_projected_with_continuous_rollout_numbers( - make_platform_client, eval_routes -): - sink, handler = make_sink(make_platform_client, eval_routes) - - sink.write([make_trace(trace_id="trace-1", idx=0)]) - sink.write([make_trace(trace_id="trace-2", idx=0)]) - - posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples") - assert [body["samples"][0]["sample_id"] for body in posted] == ["trace-1", "trace-2"] - assert [body["samples"][0]["rollout_number"] for body in posted] == [1, 2] - - def test_a_producer_that_already_speaks_v0_is_passed_through(make_platform_client, eval_routes): sink, handler = make_sink(make_platform_client, eval_routes) @@ -57,59 +44,13 @@ def test_a_producer_that_already_speaks_v0_is_passed_through(make_platform_clien assert body["samples"] == [{"sample_id": "s1", "reward": 1.0}] -def test_serialized_trace_records_are_projected(make_platform_client, eval_routes): - sink, handler = make_sink(make_platform_client, eval_routes) - - sink.write([make_trace(trace_id="serialized-trace", reward=0.75).to_record()]) - - body = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")[0] - assert body["samples"][0]["sample_id"] == "serialized-trace" - assert body["samples"][0]["reward"] == 0.75 - - -def test_serialized_episode_records_keep_the_native_wrapper(make_platform_client, eval_routes): - sink, handler = make_sink(make_platform_client, eval_routes) - record = make_episode("serialized-episode", [make_trace()]).to_record() - - sink.write([record]) - - sample = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")[0]["samples"][0] - assert sample["sample_id"] == "serialized-episode" - assert sample["info"]["native_wrapper"] == record - - -def test_serialized_message_graphs_recover_the_viewer_completion(make_platform_client, eval_routes): - sink, handler = make_sink(make_platform_client, eval_routes) - record = { - "id": "graph-trace", - "task": {"data": {"idx": 7, "answer": "42"}}, - "agent": {"name": "solver", "trainable": True}, - "nodes": [ - {"parent": None, "message": {"role": "user", "content": "6 * 7?"}}, - {"parent": 0, "message": {"role": "assistant", "content": "42"}}, - ], - "calls": [ - { - "node": 1, - "usage": {"prompt_tokens": 4, "completion_tokens": 1}, - } - ], - "rewards": {"correct": {"score": 1.0, "weight": 1.0}}, - } - - sink.write([record]) - - sample = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")[0]["samples"][0] - assert sample["example_id"] == 7 - assert sample["completion"][-1] == {"role": "assistant", "content": "42"} - assert sample["reward"] == 1.0 - - def test_records_this_sink_cannot_project_fail_explicitly(make_platform_client, eval_routes): sink, handler = make_sink(make_platform_client, eval_routes) - with pytest.raises(TypeError, match="non-empty 'id'"): + with pytest.raises(TypeError, match="cannot project"): sink.write([{"unrelated": True}]) + with pytest.raises(TypeError, match="cannot project"): + sink.write([make_trace()]) assert handler.requests == [] diff --git a/packages/prime-runs/tests/test_traces_sink.py b/packages/prime-runs/tests/test_traces_sink.py index f5948d7a8..ca0e4abef 100644 --- a/packages/prime-runs/tests/test_traces_sink.py +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -36,14 +36,13 @@ def test_records_go_out_with_provenance_but_not_the_join_key(): client = FakeTracesClient() sink = make_sink(client) - sink.write([{"id": "t1", "run": {"id": "run-1"}}], step=4) + sink.write([{"id": "t1", "run": {"id": "run-1"}}]) _, kwargs = client.calls[0] assert kwargs["context"] == { "source": "prime-runs", "run_kind": "eval", "framework": "verifiers", - "step": "4", } assert "run_id" not in kwargs["context"] @@ -54,18 +53,11 @@ def test_the_line_format_is_inferred_from_the_records(): sink.write([make_trace()]) sink.write([make_episode()]) + sink.write([{"id": "t", "traces": []}]) assert client.calls[0][1]["line_format"] is LineFormat.TRACE assert client.calls[1][1]["line_format"] is LineFormat.EPISODE - - -def test_an_explicit_line_format_wins(): - client = FakeTracesClient() - sink = make_sink(client) - - sink.write([{"id": "t1"}], line_format="episode") - - assert client.calls[0][1]["line_format"] is LineFormat.EPISODE + assert client.calls[2][1]["line_format"] is LineFormat.EPISODE def test_a_bare_mapping_gets_the_run_stamped_onto_a_copy(): @@ -149,9 +141,7 @@ def test_closing_the_sink_closes_the_client(): assert client.closed is True -def test_a_missing_traces_client_disables_the_sink_and_reports_the_failure( - monkeypatch, caplog -): +def test_a_missing_traces_client_disables_the_sink_and_reports_the_failure(monkeypatch, caplog): """The run applies warn/raise policy, so the sink must surface this failure.""" def explode(**kwargs): diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index a2a6f18c9..c9f00288c 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -5,7 +5,7 @@ from conftest import FakeSink -from prime_runs.worker import MetricItem, RunUpdateItem, UploadWorker, WriteItem +from prime_runs.worker import UploadWorker class BlockingSink(FakeSink): @@ -14,10 +14,10 @@ def __init__(self) -> None: self.entered = threading.Event() self.released = threading.Event() - def write(self, records, *, line_format=None, step=None) -> None: + def write(self, records) -> None: self.entered.set() self.released.wait(5.0) - super().write(records, line_format=line_format, step=step) + super().write(records) def drain(worker: UploadWorker) -> None: @@ -28,11 +28,11 @@ def test_records_reach_every_enabled_sink(): sinks = [FakeSink("a"), FakeSink("b")] worker = UploadWorker(sinks) - worker.submit(WriteItem(records=[{"id": 1}], line_format="trace", step=3)) + worker.submit([{"id": 1}]) drain(worker) for sink in sinks: - assert sink.batches == [([{"id": 1}], "trace", 3)] + assert sink.batches == [[{"id": 1}]] worker.close() @@ -41,7 +41,7 @@ def test_a_disabled_sink_is_skipped(): dead.enabled = False worker = UploadWorker([live, dead]) - worker.submit(WriteItem(records=[{"id": 1}])) + worker.submit([{"id": 1}]) drain(worker) assert live.batches and not dead.batches @@ -53,7 +53,7 @@ def test_one_sink_failing_does_not_stop_the_others(): reported = [] worker = UploadWorker([broken, healthy], on_error=lambda name, exc: reported.append(name)) - worker.submit(WriteItem(records=[{"id": 1}])) + worker.submit([{"id": 1}]) drain(worker) assert healthy.batches @@ -69,7 +69,7 @@ def test_a_failed_sink_is_not_called_again(): worker = UploadWorker([broken], on_error=lambda name, exc: reported.append(name)) for _ in range(3): - worker.submit(WriteItem(records=[{"id": 1}])) + worker.submit([{"id": 1}]) drain(worker) assert reported == ["broken"] @@ -82,11 +82,11 @@ def test_a_full_queue_drops_rather_than_blocking_the_producer(): worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.05) # First item is picked up and wedges the uploader inside sink.write(). - assert worker.submit(WriteItem(records=[{"id": 0}])) + assert worker.submit([{"id": 0}]) assert sink.entered.wait(5.0), "the uploader never reached the sink" # Second fills the one-slot queue; third has nowhere to go. - worker.submit(WriteItem(records=[{"id": 1}])) - accepted = worker.submit(WriteItem(records=[{"id": 2}, {"id": 3}])) + worker.submit([{"id": 1}]) + accepted = worker.submit([{"id": 2}, {"id": 3}]) assert accepted is False assert worker.dropped == 2 @@ -98,9 +98,9 @@ def test_a_full_queue_drops_rather_than_blocking_the_producer(): def test_flush_uses_the_drain_budget_to_get_behind_a_full_queue(): sink = BlockingSink() worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.01) - assert worker.submit(WriteItem(records=[{"id": 0}])) + assert worker.submit([{"id": 0}]) assert sink.entered.wait(1.0) - assert worker.submit(WriteItem(records=[{"id": 1}])) + assert worker.submit([{"id": 1}]) release = threading.Timer(0.05, sink.released.set) release.start() @@ -111,15 +111,15 @@ def test_flush_uses_the_drain_budget_to_get_behind_a_full_queue(): release.join() worker.close(timeout=1.0) - assert [batch[0][0]["id"] for batch in sink.batches] == [0, 1] + assert [batch[0]["id"] for batch in sink.batches] == [0, 1] def test_close_uses_its_budget_to_queue_the_stop_behind_pending_records(): sink = BlockingSink() worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.01) - assert worker.submit(WriteItem(records=[{"id": 0}])) + assert worker.submit([{"id": 0}]) assert sink.entered.wait(1.0) - assert worker.submit(WriteItem(records=[{"id": 1}])) + assert worker.submit([{"id": 1}]) release = threading.Timer(0.05, sink.released.set) release.start() @@ -129,54 +129,15 @@ def test_close_uses_its_budget_to_queue_the_stop_behind_pending_records(): sink.released.set() release.join() - assert [batch[0][0]["id"] for batch in sink.batches] == [0, 1] + assert [batch[0]["id"] for batch in sink.batches] == [0, 1] assert sink.closed is True -def test_metrics_ride_the_same_queue_when_the_backend_stores_a_time_series(): - points = [] - worker = UploadWorker([], metric_writer=lambda metrics, step: points.append((metrics, step))) - - worker.submit(MetricItem(metrics={"loss": 0.5}, step=7)) - drain(worker) - - assert points == [({"loss": 0.5}, 7)] - worker.close() - - -def test_run_updates_ride_the_uploader_queue(): - updates = [] - worker = UploadWorker( - [], update_writer=lambda config, summary: updates.append((config, summary)) - ) - - worker.submit(RunUpdateItem(config={"seed": 7}, summary={"reward": 0.5})) - drain(worker) - - assert updates == [({"seed": 7}, {"reward": 0.5})] - worker.close() - - -def test_a_metric_write_that_raises_does_not_kill_the_uploader(): - sink = FakeSink() - - def explode(metrics, step): - raise RuntimeError("nope") - - worker = UploadWorker([sink], metric_writer=explode) - worker.submit(MetricItem(metrics={"loss": 0.5}, step=1)) - worker.submit(WriteItem(records=[{"id": 1}])) - drain(worker) - - assert sink.batches, "the uploader survived the metric failure" - worker.close() - - def test_close_drains_then_closes_every_sink(): sink = FakeSink() worker = UploadWorker([sink]) - worker.submit(WriteItem(records=[{"id": 1}])) + worker.submit([{"id": 1}]) worker.close() assert sink.batches @@ -187,7 +148,7 @@ def test_submitting_after_close_is_refused(): worker = UploadWorker([FakeSink()]) worker.close() - assert worker.submit(WriteItem(records=[{"id": 1}])) is False + assert worker.submit([{"id": 1}]) is False def test_flush_without_a_running_thread_still_flushes_the_sinks(): @@ -206,7 +167,7 @@ def test_a_forked_child_starts_over_instead_of_re_uploading_the_parents_queue(): """ sink = FakeSink() worker = UploadWorker([sink], max_queue_size=4) - worker._queue.put(WriteItem(records=[{"id": "parents"}])) + worker._queue.put([{"id": "parents"}]) old_queue = worker._queue worker.reset_after_fork() @@ -227,12 +188,12 @@ def __init__(self) -> None: super().__init__("wedged") self.released = threading.Event() - def write(self, records, *, line_format=None, step=None) -> None: + def write(self, records) -> None: self.released.wait(10.0) sink = WedgedSink() worker = UploadWorker([sink]) - worker.submit(WriteItem(records=[{"id": 1}])) + worker.submit([{"id": 1}]) with caplog.at_level("WARNING"): worker.close(timeout=0.2) @@ -292,13 +253,13 @@ def test_a_transient_failure_drops_the_batch_but_keeps_the_sink(): from prime_runs.exceptions import RetryableAPIError class BlipSink(FakeSink): - def write(self, records, *, line_format=None, step=None) -> None: + def write(self, records) -> None: raise RetryableAPIError("bad gateway", status_code=502) sink = BlipSink("blippy") worker = UploadWorker([sink]) - worker.submit(WriteItem(records=[{"id": 1}, {"id": 2}])) + worker.submit([{"id": 1}, {"id": 2}]) drain(worker) assert sink.enabled is True @@ -315,14 +276,14 @@ def test_a_sustained_outage_eventually_retires_the_sink(): from prime_runs.worker import TRANSIENT_FAILURE_LIMIT class DeadSink(FakeSink): - def write(self, records, *, line_format=None, step=None) -> None: + def write(self, records) -> None: raise TransportError("connection refused") sink = DeadSink("dead") worker = UploadWorker([sink]) for _ in range(TRANSIENT_FAILURE_LIMIT): - worker.submit(WriteItem(records=[{"id": 1}])) + worker.submit([{"id": 1}]) drain(worker) assert sink.enabled is False @@ -342,18 +303,18 @@ def __init__(self) -> None: super().__init__("flaky") self.calls = 0 - def write(self, records, *, line_format=None, step=None) -> None: + def write(self, records) -> None: self.calls += 1 if self.calls % 2 == 1: raise RetryableAPIError("bad gateway", status_code=502) - super().write(records, line_format=line_format, step=step) + super().write(records) sink = FlakySink() worker = UploadWorker([sink]) # Far more failures than the limit, but never two in a row. for _ in range(TRANSIENT_FAILURE_LIMIT * 4): - worker.submit(WriteItem(records=[{"id": 1}])) + worker.submit([{"id": 1}]) drain(worker) assert sink.calls == TRANSIENT_FAILURE_LIMIT * 4 @@ -367,13 +328,13 @@ def test_a_permanent_failure_retires_the_sink_immediately(): from prime_runs.exceptions import UnauthorizedError class DeniedSink(FakeSink): - def write(self, records, *, line_format=None, step=None) -> None: + def write(self, records) -> None: raise UnauthorizedError("nope", status_code=401) sink = DeniedSink("denied") worker = UploadWorker([sink]) - worker.submit(WriteItem(records=[{"id": 1}])) + worker.submit([{"id": 1}]) drain(worker) assert sink.enabled is False @@ -387,13 +348,13 @@ def test_a_failed_batch_is_counted_once_per_sink_not_once_per_run(): from prime_runs.exceptions import RetryableAPIError class BlipSink(FakeSink): - def write(self, records, *, line_format=None, step=None) -> None: + def write(self, records) -> None: raise RetryableAPIError("bad gateway", status_code=502) broken, healthy = BlipSink("broken"), FakeSink("healthy") worker = UploadWorker([broken, healthy]) - worker.submit(WriteItem(records=[{"id": 1}, {"id": 2}])) + worker.submit([{"id": 1}, {"id": 2}]) drain(worker) assert worker.failed_records == {"broken": 2} From 5d087a03f556b2dea1867b920ac4d3b7a9617d7e Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Fri, 21 Aug 2026 09:31:58 -0700 Subject: [PATCH 21/39] refactor(runs): reuse prime_traces plumbing instead of copying it exceptions.py re-exports the prime_traces API error family and keeps only the SDK-local errors (PrimeRunsError, ConfigurationError, EnvironmentResolutionError, RunFinishedError). RunAPIError is gone: it was prime_traces.APIError under another name. is_transient and the traces sink's gating check lose their try-import dance since there is one family. config.Config subclasses prime_traces.core.Config and adds frontend_url. _http.PlatformClient uses raise_for_response, retry_delay and AMBIGUOUS_TRANSPORT_ERRORS from prime_traces.core.client; the local error mapping, Retry-After parsing and backoff ladder are deleted. What stays local is the per-call idempotent= replay policy, encode_json and the non-JSON-body guard. Backoff is now the shared jittered schedule. Everything used is in the released prime-traces 0.0.2; the pin is unchanged. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01W4S4DNVRX5BNtsLva7TZz1 --- packages/prime-runs/README.md | 3 +- .../prime-runs/src/prime_runs/__init__.py | 4 +- packages/prime-runs/src/prime_runs/_http.py | 161 +++++------------- .../src/prime_runs/backends/evals.py | 8 +- packages/prime-runs/src/prime_runs/config.py | 57 +------ .../prime-runs/src/prime_runs/exceptions.py | 121 +++++-------- .../prime-runs/src/prime_runs/sinks/traces.py | 5 +- packages/prime-runs/tests/test_http.py | 40 +++-- 8 files changed, 125 insertions(+), 274 deletions(-) diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index d07462720..bc4a547e5 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -84,7 +84,8 @@ degrades to offline with a warning. - **Contains its own errors.** With the default `on_error="warn"`, nothing the platform raises escapes into your loop. Use `on_error="raise"` in tests and CI, where a silent upload failure is the bug; failures surface from `flush()` - and `finish()`. + and `finish()`. Platform errors are the `prime_traces` exception family + (`pr.APIError` and friends), so one set of `except` clauses covers both SDKs. - **Applies backpressure.** The upload queue is bounded; a producer that durably outruns the uploader has records dropped and counted (`run.dropped_records`) rather than stalled. Per-sink losses are in diff --git a/packages/prime-runs/src/prime_runs/__init__.py b/packages/prime-runs/src/prime_runs/__init__.py index 5ce95149e..c92893a00 100644 --- a/packages/prime-runs/src/prime_runs/__init__.py +++ b/packages/prime-runs/src/prime_runs/__init__.py @@ -19,6 +19,7 @@ from . import metrics, projection from .exceptions import ( + APIError, ConfigurationError, EnvironmentResolutionError, ForbiddenError, @@ -26,7 +27,6 @@ PaymentRequiredError, PrimeRunsError, RetryableAPIError, - RunAPIError, RunFinishedError, TransportError, UnauthorizedError, @@ -52,7 +52,7 @@ "PrimeRunsError", "ConfigurationError", "EnvironmentResolutionError", - "RunAPIError", + "APIError", "RunFinishedError", "ForbiddenError", "NotFoundError", diff --git a/packages/prime-runs/src/prime_runs/_http.py b/packages/prime-runs/src/prime_runs/_http.py index 2dd31743f..3c6c9d61a 100644 --- a/packages/prime-runs/src/prime_runs/_http.py +++ b/packages/prime-runs/src/prime_runs/_http.py @@ -1,13 +1,12 @@ -"""Shared HTTP client for the platform run APIs. - -Maps status codes onto :mod:`prime_runs.exceptions` and retries 429/502/503/504 -and transport failures with backoff, honouring ``Retry-After``. - -Retry safety is decided per call. A failure is *ambiguous* when the request -may already have been processed (a 502/504, a read timeout). Replaying one is -fine for a GET or PUT and not for ``POST /evaluations/``, which would create a -second run. Callers declare intent with ``idempotent=``; unambiguous failures -(connect errors, 429) are replayed for every method. +"""HTTP client for the platform run APIs. + +Error mapping, backoff and the transport-failure classification come from +``prime_traces.core.client``; what is local is the retry *policy*, because it +is decided per call. A failure is *ambiguous* when the request may already have +been processed (a 502/504, a read timeout). Replaying one is fine for a GET or +PUT and not for ``POST /evaluations/``, which would create a second run. +Callers declare intent with ``idempotent=``; unambiguous failures (connect +errors, 429) are replayed for every method. """ import json @@ -16,27 +15,22 @@ from typing import Any, Dict, Mapping, Optional, Union import httpx +from prime_traces.core.client import ( + AMBIGUOUS_TRANSPORT_ERRORS, + raise_for_response, + retry_delay, +) from . import _fork -from .exceptions import ( - ForbiddenError, - NotFoundError, - PaymentRequiredError, - RetryableAPIError, - RunAPIError, - TransportError, - UnauthorizedError, -) +from .exceptions import APIError, APITimeoutError, RetryableAPIError, TransportError DEFAULT_TIMEOUT = httpx.Timeout(60.0, connect=10.0) # Sample batches are megabytes; uploads get a longer budget. UPLOAD_TIMEOUT = httpx.Timeout(300.0, connect=10.0) -RETRY_STATUS = frozenset({429, 502, 503, 504}) #: Refused before any work was done, so replaying cannot duplicate anything. #: 503 is excluded: it may come from an intermediary after forwarding. UNAMBIGUOUS_RETRY_STATUS = frozenset({429}) DEFAULT_MAX_ATTEMPTS = 5 -MAX_BACKOFF_SECONDS = 16.0 def _user_agent() -> str: @@ -46,24 +40,6 @@ def _user_agent() -> str: return f"prime-runs/{__version__} python/{py}" -def retry_delay(attempt: int, retry_after: Optional[float]) -> float: - """Seconds to wait before ``attempt`` (1-based). Server wins if it spoke.""" - if retry_after is not None and retry_after >= 0: - return min(retry_after, MAX_BACKOFF_SECONDS) - return min(2.0 ** (attempt - 1), MAX_BACKOFF_SECONDS) - - -def _parse_retry_after(response: httpx.Response) -> Optional[float]: - raw = response.headers.get("retry-after") - if not raw: - return None - try: - return float(raw) - except ValueError: - # HTTP-date form; fall back to the exponential schedule. - return None - - def normalize_base_url(url: str) -> str: """Strip a trailing ``/api/v1``; the client appends it itself.""" return url.rstrip("/").removesuffix("/api/v1") @@ -134,55 +110,42 @@ def request( """ url = f"{self.api_prefix}{path}" body = content if content is not None else (encode_json(json_body) if json_body else None) - headers = {"Content-Type": "application/json"} if body is not None else None + request_kwargs: Dict[str, Any] = { + "content": body, + "headers": {"Content-Type": "application/json"} if body is not None else None, + "params": dict(params) if params else None, + } + # ``None`` would disable httpx timeouts, not restore the default. + if timeout is not None: + request_kwargs["timeout"] = timeout attempts = max_attempts or self.max_attempts replayable = idempotent if idempotent is not None else method.upper() != "POST" - last_error: Optional[Exception] = None - for attempt in range(1, attempts + 1): - ambiguous = True + for attempt in range(attempts): + error: APIError try: - request_kwargs: Dict[str, Any] = { - "content": body, - "headers": headers, - "params": dict(params) if params else None, - } - # ``None`` would disable httpx timeouts, not restore the default. - if timeout is not None: - request_kwargs["timeout"] = timeout response = self._client.request(method, url, **request_kwargs) - except (httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout) as exc: - # No connection was ever established, so the server saw nothing. - ambiguous = False - last_error = TransportError(f"{method} {path} failed to connect: {exc}") except httpx.TimeoutException as exc: - last_error = TransportError(f"{method} {path} timed out: {exc}") + error = APITimeoutError(f"{method} {path} timed out: {exc}") + ambiguous = isinstance(exc, AMBIGUOUS_TRANSPORT_ERRORS) except httpx.RequestError as exc: - last_error = TransportError(f"{method} {path} failed: {type(exc).__name__}: {exc}") + error = TransportError(f"{method} {path} failed: {type(exc).__name__}: {exc}") + ambiguous = isinstance(exc, AMBIGUOUS_TRANSPORT_ERRORS) else: - if response.status_code in RETRY_STATUS: - ambiguous = response.status_code not in UNAMBIGUOUS_RETRY_STATUS - last_error = RetryableAPIError( - _error_message(response), - status_code=response.status_code, - code=_error_code(response), - retry_after=_parse_retry_after(response), - ) - elif response.is_error: - raise _map_error(response) + try: + raise_for_response(response) + except RetryableAPIError as exc: + error = exc + ambiguous = exc.status_code not in UNAMBIGUOUS_RETRY_STATUS else: return _decode(response) - if attempt == attempts: - break - if ambiguous and not replayable: + last = attempt == attempts - 1 + if last or (ambiguous and not replayable): # Possibly processed already; a duplicate cannot be undone. - break - after = getattr(last_error, "retry_after", None) - time.sleep(retry_delay(attempt, after)) - - assert last_error is not None - raise last_error + raise error + time.sleep(retry_delay(error, attempt)) + raise AssertionError("unreachable") # pragma: no cover def get(self, path: str, **kwargs: Any) -> Dict[str, Any]: return self.request("GET", path, **kwargs) @@ -210,52 +173,8 @@ def _decode(response: httpx.Response) -> Dict[str, Any]: try: payload = response.json() except ValueError as exc: - raise RunAPIError( + raise APIError( f"{response.request.method} {response.request.url.path} returned non-JSON " f"({response.status_code}): {response.text[:200]!r}" ) from exc return payload if isinstance(payload, dict) else {"data": payload} - - -def _error_body(response: httpx.Response) -> Dict[str, Any]: - try: - payload = response.json() - except ValueError: - return {} - return payload if isinstance(payload, dict) else {} - - -def _error_code(response: httpx.Response) -> Optional[str]: - body = _error_body(response) - code = body.get("code") or body.get("error_code") - return str(code) if code else None - - -def _error_message(response: httpx.Response) -> str: - body = _error_body(response) - detail = body.get("detail") or body.get("message") or body.get("error") - if detail is None: - detail = response.text[:200] or response.reason_phrase - return ( - f"HTTP {response.status_code} from " - f"{response.request.method} {response.request.url.path}: {detail}" - ) - - -def _map_error(response: httpx.Response) -> RunAPIError: - message = _error_message(response) - code = _error_code(response) - status = response.status_code - if status == 401: - return UnauthorizedError( - f"{message} — check PRIME_API_KEY or run `prime login`.", - status_code=status, - code=code, - ) - if status == 402: - return PaymentRequiredError(message, status_code=status, code=code) - if status == 403: - return ForbiddenError(message, status_code=status, code=code) - if status == 404: - return NotFoundError(message, status_code=status, code=code) - return RunAPIError(message, status_code=status, code=code) diff --git a/packages/prime-runs/src/prime_runs/backends/evals.py b/packages/prime-runs/src/prime_runs/backends/evals.py index 81398dc7c..ec2842a4c 100644 --- a/packages/prime-runs/src/prime_runs/backends/evals.py +++ b/packages/prime-runs/src/prime_runs/backends/evals.py @@ -13,7 +13,7 @@ from typing import Any, Dict, List, Optional from .._http import PlatformClient -from ..exceptions import ConfigurationError, EnvironmentResolutionError, RunAPIError +from ..exceptions import APIError, ConfigurationError, EnvironmentResolutionError from ..models import EnvironmentRef, RunHandle, RunSpec, RunStatus logger = logging.getLogger(__name__) @@ -62,7 +62,7 @@ def create(self, spec: RunSpec) -> RunHandle: response = self._client.post("/evaluations/", json_body=payload) run_id = response.get("evaluation_id") if not run_id: - raise RunAPIError( + raise APIError( f"POST /evaluations/ returned no evaluation_id (keys: {sorted(response)})" ) return RunHandle( @@ -152,7 +152,7 @@ def _lookup_environment(self, ref: EnvironmentRef) -> str: owner_slug, name = ref.slug.split("/", 1) try: response = self._client.get(f"/environmentshub/{owner_slug}/{name}/@latest") - except RunAPIError as exc: + except APIError as exc: raise EnvironmentResolutionError( f"Could not resolve environment {ref.slug!r}: {exc}" ) from exc @@ -169,7 +169,7 @@ def _lookup_environment(self, ref: EnvironmentRef) -> str: response = self._client.post( "/environmentshub/resolve", json_body=body, idempotent=True ) - except RunAPIError as exc: + except APIError as exc: raise EnvironmentResolutionError( f"Could not resolve environment {ref.name!r}: {exc}" ) from exc diff --git a/packages/prime-runs/src/prime_runs/config.py b/packages/prime-runs/src/prime_runs/config.py index 2326c8208..c25440a1f 100644 --- a/packages/prime-runs/src/prime_runs/config.py +++ b/packages/prime-runs/src/prime_runs/config.py @@ -1,63 +1,14 @@ -"""Configuration: ``~/.prime/config.json`` plus environment variables, env -taking precedence. Same shape as the other prime SDKs, plus ``frontend_url``.""" +"""Configuration: ``prime_traces.core.Config`` (``~/.prime/config.json`` plus +environment variables, env taking precedence) with the dashboard URL added.""" -import json import os -from pathlib import Path -from typing import Optional +from prime_traces.core import Config as _TracesConfig -class Config: - """Minimal configuration class for SDK packages. - Reads from ~/.prime/config.json and environment variables. - """ - - DEFAULT_BASE_URL: str = "https://api.primeintellect.ai" +class Config(_TracesConfig): DEFAULT_FRONTEND_URL: str = "https://app.primeintellect.ai" - def __init__(self) -> None: - self.config_dir = Path.home() / ".prime" - self.config_file = self.config_dir / "config.json" - self._load_config() - - def _load_config(self) -> None: - """Load configuration from file.""" - config_data: object = {} - if self.config_file.exists(): - try: - config_data = json.loads(self.config_file.read_text()) - except (json.JSONDecodeError, OSError, UnicodeDecodeError): - config_data = {} - # Valid JSON that is not an object (a list, a bare string) must degrade - # the same way invalid JSON does: every accessor below assumes a dict. - self.config = config_data if isinstance(config_data, dict) else {} - - @staticmethod - def _strip_api_v1(url: str) -> str: - return url.rstrip("/").removesuffix("/api/v1") - - @property - def api_key(self) -> str: - """API key with precedence: env > file > empty.""" - return os.getenv("PRIME_API_KEY") or self.config.get("api_key", "") - - @property - def team_id(self) -> Optional[str]: - """Team ID with precedence: env > file > None.""" - team_id = os.getenv("PRIME_TEAM_ID") - if team_id is not None: - return team_id - return self.config.get("team_id") or None - - @property - def base_url(self) -> str: - """Platform API base URL with precedence: env > file > default.""" - env_val = os.getenv("PRIME_API_BASE_URL") or os.getenv("PRIME_BASE_URL") - if env_val: - return self._strip_api_v1(env_val) - return self._strip_api_v1(self.config.get("base_url", self.DEFAULT_BASE_URL)) - @property def frontend_url(self) -> str: """Dashboard base URL; fallback when a create response omits ``viewer_url``.""" diff --git a/packages/prime-runs/src/prime_runs/exceptions.py b/packages/prime-runs/src/prime_runs/exceptions.py index c17ab94c2..adc5df8b9 100644 --- a/packages/prime-runs/src/prime_runs/exceptions.py +++ b/packages/prime-runs/src/prime_runs/exceptions.py @@ -1,83 +1,50 @@ -"""Exceptions for the Prime Runs SDK. Nothing here escapes into a producer -loop by default (``on_error="warn"``); callers opting into ``on_error="raise"`` -can branch on these types.""" - -from typing import Optional +"""Exceptions for the Prime Runs SDK. + +Two families. Errors raised here, before any request is made, derive from +:class:`PrimeRunsError`. Errors from the platform derive from +``prime_traces.APIError``: both SDKs talk to the same platform with the same +credential, and the uploader already handles the traces family, so one +vocabulary serves both. Nothing in either family escapes into a producer loop +under the default ``on_error="warn"``. To catch everything the SDK can raise +under ``on_error="raise"``:: + + except (pr.PrimeRunsError, pr.APIError): +""" + +from prime_traces.exceptions import ( + APIError, + APITimeoutError, + ForbiddenError, + NotFoundError, + PaymentRequiredError, + RetryableAPIError, + TransportError, + UnauthorizedError, +) + +__all__ = [ + "APIError", + "APITimeoutError", + "ConfigurationError", + "EnvironmentResolutionError", + "ForbiddenError", + "NotFoundError", + "PaymentRequiredError", + "PrimeRunsError", + "RetryableAPIError", + "RunFinishedError", + "TransportError", + "UnauthorizedError", + "is_transient", +] class PrimeRunsError(Exception): - """Base exception for the Prime Runs SDK.""" + """Base for errors raised by the SDK itself, before any request is made.""" class ConfigurationError(PrimeRunsError): - """Missing API key, unreadable config file, unknown mode. Raised before - any request is made.""" - - -class RunAPIError(PrimeRunsError): - """An HTTP error response from a run backend.""" - - def __init__( - self, - message: str, - *, - status_code: Optional[int] = None, - code: Optional[str] = None, - ): - self.status_code = status_code - self.code = code - super().__init__(message) - - -class UnauthorizedError(RunAPIError): - """401 — the credential was rejected. Stop rather than retry.""" - - -class PaymentRequiredError(RunAPIError): - """402 — payment required. Check billing status.""" - - -class ForbiddenError(RunAPIError): - """403 — authenticated, but not allowed: another owner's run, a team the - key cannot act for, or a feature gated to an allowlist. Named to match - ``prime_traces.ForbiddenError``.""" - - -class NotFoundError(RunAPIError): - """404 — the run, environment or evaluation does not exist for this owner.""" - - -class RetryableAPIError(RunAPIError): - """429/5xx — retry the same request after ``retry_after`` seconds.""" - - def __init__( - self, - message: str, - *, - status_code: Optional[int] = None, - code: Optional[str] = None, - retry_after: Optional[float] = None, - ): - super().__init__(message, status_code=status_code, code=code) - self.retry_after = retry_after - - -class TransportError(RunAPIError): - """The request failed below HTTP — connection refused, TLS failure, timeout.""" - - -def is_transient(exc: BaseException) -> bool: - """Whether a failure is about this moment (retry later) rather than this - run (stop). Decides whether a sink is retired. Covers the traces service's - exception family too, since both reach the uploader through one path.""" - if isinstance(exc, (RetryableAPIError, TransportError)): - return True - try: - from prime_traces.exceptions import RetryableAPIError as TracesRetryable - from prime_traces.exceptions import TransportError as TracesTransport - except ImportError: # pragma: no cover - dependency is declared - return False - return isinstance(exc, (TracesRetryable, TracesTransport)) + """Missing API key, unreadable config file, unknown mode.""" class EnvironmentResolutionError(PrimeRunsError): @@ -87,3 +54,9 @@ class EnvironmentResolutionError(PrimeRunsError): class RunFinishedError(PrimeRunsError): """A finished run was written to again: a producer bug.""" + + +def is_transient(exc: BaseException) -> bool: + """Whether a failure is about this moment (retry later) rather than this + run (stop). Decides whether a sink is retired.""" + return isinstance(exc, (RetryableAPIError, TransportError)) diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py index 61a6f5412..abc7f2067 100644 --- a/packages/prime-runs/src/prime_runs/sinks/traces.py +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -11,6 +11,7 @@ from typing import Any, Dict, Mapping, Optional, Sequence from .. import _fork +from ..exceptions import ForbiddenError from .base import Sink, is_episode, stamp_run logger = logging.getLogger(__name__) @@ -147,8 +148,4 @@ def _disable(self, reason: str) -> None: def _is_gated(exc: Exception) -> bool: """A 403 (``service_not_enabled``, or a write-only token) is not fixable at runtime, so the sink turns itself off instead of retrying.""" - try: - from prime_traces.exceptions import ForbiddenError - except ImportError: # pragma: no cover - dependency is declared - return False return isinstance(exc, ForbiddenError) diff --git a/packages/prime-runs/tests/test_http.py b/packages/prime-runs/tests/test_http.py index b4895b5bd..6d954f84d 100644 --- a/packages/prime-runs/tests/test_http.py +++ b/packages/prime-runs/tests/test_http.py @@ -3,13 +3,13 @@ import httpx import pytest -from prime_runs._http import PlatformClient, encode_json, retry_delay +from prime_runs._http import PlatformClient, encode_json from prime_runs.exceptions import ( + APIError, ForbiddenError, NotFoundError, PaymentRequiredError, RetryableAPIError, - RunAPIError, TransportError, UnauthorizedError, is_transient, @@ -32,18 +32,21 @@ def client_for(handler, *, base_url: str = "http://testserver", **kwargs) -> Pla (402, PaymentRequiredError), (403, ForbiddenError), (404, NotFoundError), - (400, RunAPIError), - (422, RunAPIError), + (400, APIError), + (422, APIError), ], ) def test_status_codes_map_to_types_callers_can_branch_on(status, expected): + """The classes are ``prime_traces``' own, so a producer that already + handles the traces client's errors handles these with the same clauses.""" client = client_for(lambda request: httpx.Response(status, json={"detail": "nope"})) with pytest.raises(expected) as caught: client.get("/evaluations/x") assert caught.value.status_code == status - assert "nope" in str(caught.value) + if status not in (401, 402): + assert "nope" in str(caught.value) def test_a_forbidden_response_is_permanent_so_a_sink_retires_on_it(): @@ -69,14 +72,16 @@ def test_retryable_statuses_are_retried_then_surface(no_sleep): def handler(request): attempts.append(request) - return httpx.Response(503, json={"code": "ingest_unavailable"}) + return httpx.Response(503, json={"detail": "overloaded"}) with pytest.raises(RetryableAPIError) as caught: client_for(handler, max_attempts=3).get("/evaluations/x") assert len(attempts) == 3 - assert caught.value.code == "ingest_unavailable" - assert no_sleep == [1.0, 2.0] + assert caught.value.status_code == 503 + # Two waits between three attempts, on prime_traces' jittered schedule. + assert len(no_sleep) == 2 + assert all(0.0 < delay <= 30.0 for delay in no_sleep) def test_a_retry_succeeds_without_bothering_the_caller(no_sleep): @@ -87,12 +92,17 @@ def test_a_retry_succeeds_without_bothering_the_caller(no_sleep): assert client.get("/evaluations/x") == {"ok": True} -def test_retry_after_beats_the_exponential_schedule(): - assert retry_delay(1, 7.5) == 7.5 - assert retry_delay(1, None) == 1.0 - assert retry_delay(4, None) == 8.0 - # Never wait longer than the ceiling, whatever the server asked for. - assert retry_delay(1, 900.0) == 16.0 +def test_retry_after_is_honoured(no_sleep): + """The schedule itself is ``prime_traces.core.client.retry_delay``; what is + ours is feeding it the server's header.""" + responses = [ + httpx.Response(429, headers={"Retry-After": "7.5"}), + httpx.Response(200, json={"ok": True}), + ] + + client_for(lambda request: responses.pop(0)).get("/evaluations/x") + + assert no_sleep == [7.5] def test_transport_failures_are_retried_and_typed(no_sleep): @@ -137,7 +147,7 @@ def test_an_empty_body_is_a_valid_response(): def test_a_non_json_body_names_the_request_that_produced_it(): client = client_for(lambda request: httpx.Response(200, text="gateway")) - with pytest.raises(RunAPIError, match="non-JSON"): + with pytest.raises(APIError, match="non-JSON"): client.get("/evaluations/x") From 8d46b27f63546a71d4c462f27115a618e7439c72 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Fri, 21 Aug 2026 09:45:39 -0700 Subject: [PATCH 22/39] refactor(runs): drop offline mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nothing consumed it: verifiers picks online or disabled. There was no sync path, and building one collides with the identity design — records carry the local run.id, which is the join key the traces service indexes, so a later upload would have to rewrite every record or the platform would have to accept client-issued IDs. If air-gapped evals are ever wanted, this comes back together with the sync command and that decision. Removes OfflineBackend, OfflineSink, dir=, PRIME_RUNS_DIR and mode="offline". A missing API key now resolves to disabled with a warning that the run will not be tracked, rather than silently writing a ./prime-runs/ directory into the cwd. Tests that used offline as a cheap real backend now use disabled or the online MockTransport fixture. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01W4S4DNVRX5BNtsLva7TZz1 --- packages/prime-runs/README.md | 6 +- .../src/prime_runs/backends/__init__.py | 9 +- .../src/prime_runs/backends/offline.py | 132 ---------------- packages/prime-runs/src/prime_runs/models.py | 6 +- packages/prime-runs/src/prime_runs/run.py | 41 +++-- .../src/prime_runs/sinks/__init__.py | 2 - .../src/prime_runs/sinks/offline.py | 82 ---------- packages/prime-runs/tests/conftest.py | 1 - .../prime-runs/tests/test_config_source.py | 75 ++++----- packages/prime-runs/tests/test_init.py | 147 +++++------------- 10 files changed, 100 insertions(+), 401 deletions(-) delete mode 100644 packages/prime-runs/src/prime_runs/backends/offline.py delete mode 100644 packages/prime-runs/src/prime_runs/sinks/offline.py diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index bc4a547e5..a3b2b50da 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -72,11 +72,10 @@ Nothing is redacted — keep credentials in the environment, not in the file. | mode | what happens | | --- | --- | | `online` | the run lives on the platform (default when an API key is present) | -| `offline` | the run lives in a local directory (`dir=` / `PRIME_RUNS_DIR`) | | `disabled` | every call is a no-op, with the same object shape | Set the mode explicitly, or through `$PRIME_RUNS_MODE`. A missing API key -degrades to offline with a warning. +disables the run with a warning — it never silently writes somewhere else. ## What the run handle does for you @@ -115,9 +114,8 @@ Resolved from environment variables first, then `~/.prime/config.json`: | platform API | `PRIME_API_BASE_URL` | `https://api.primeintellect.ai` | | dashboard | `PRIME_FRONTEND_URL` | `https://app.primeintellect.ai` | | traces service | `PRIME_TRACES_URL` | resolved by `prime-traces` | -| offline runs | `PRIME_RUNS_DIR` | `./prime-runs` | -Or pass `api_key=`, `base_url=`, `team_id=`, `dir=` to `init()`. +Or pass `api_key=`, `base_url=`, `team_id=` to `init()`. ## Transports diff --git a/packages/prime-runs/src/prime_runs/backends/__init__.py b/packages/prime-runs/src/prime_runs/backends/__init__.py index f94eb5343..30e6efc38 100644 --- a/packages/prime-runs/src/prime_runs/backends/__init__.py +++ b/packages/prime-runs/src/prime_runs/backends/__init__.py @@ -2,12 +2,5 @@ from .base import Backend from .evals import EvalsBackend -from .offline import OfflineBackend, default_dir, new_run_id -__all__ = [ - "Backend", - "EvalsBackend", - "OfflineBackend", - "default_dir", - "new_run_id", -] +__all__ = ["Backend", "EvalsBackend"] diff --git a/packages/prime-runs/src/prime_runs/backends/offline.py b/packages/prime-runs/src/prime_runs/backends/offline.py deleted file mode 100644 index 50703f5c7..000000000 --- a/packages/prime-runs/src/prime_runs/backends/offline.py +++ /dev/null @@ -1,132 +0,0 @@ -"""Offline runs: a local directory with the same lifecycle as a platform run. - -Layout, one directory per run:: - - //run.json spec + status + timestamps - //records/ whatever the offline sink wrote -""" - -import json -import logging -import os -import uuid -from dataclasses import asdict -from datetime import datetime, timezone -from pathlib import Path -from typing import Any, Dict, Optional, Union - -from ..models import RunHandle, RunSpec, RunStatus - -logger = logging.getLogger(__name__) - -DEFAULT_DIR_ENV = "PRIME_RUNS_DIR" -DEFAULT_DIR = "prime-runs" - - -def default_dir() -> Path: - return Path(os.getenv(DEFAULT_DIR_ENV) or DEFAULT_DIR) - - -def new_run_id() -> str: - """A locally issued run ID, visibly distinct from a platform one.""" - return f"offline-{uuid.uuid4().hex[:16]}" - - -class OfflineBackend: - """Run lifecycle recorded on the local filesystem.""" - - def __init__(self, directory: Union[str, Path, None] = None) -> None: - self.directory = Path(directory) if directory is not None else default_dir() - - def run_dir(self, run_id: str) -> Path: - return self.directory / run_id - - def create(self, spec: RunSpec) -> RunHandle: - run_id = new_run_id() - path = self.run_dir(run_id) - path.mkdir(parents=True, exist_ok=True) - run_name: str = spec.name or run_id - state: Dict[str, Any] = { - "id": run_id, - "name": run_name, - "status": RunStatus.RUNNING.value, - "created_at": _now(), - "spec": _spec_to_json(spec), - } - self._write_state(run_id, state) - return RunHandle(id=run_id, name=run_name, url=str(path.resolve())) - - def update( - self, - run_id: str, - *, - config: Optional[Dict[str, Any]] = None, - summary: Optional[Dict[str, Any]] = None, - ) -> None: - state = self._read_state(run_id) - if config: - state.setdefault("config", {}).update(config) - if summary: - state.setdefault("summary", {}).update(summary) - state["updated_at"] = _now() - self._write_state(run_id, state) - - def finalize( - self, - run_id: str, - *, - status: RunStatus, - summary: Optional[Dict[str, Any]] = None, - error: Optional[str] = None, - config: Optional[Dict[str, Any]] = None, - ) -> None: - state = self._read_state(run_id) - if config: - state.setdefault("config", {}).update(config) - state["status"] = status.value - state["finished_at"] = _now() - if error: - state["error"] = error - if summary: - state.setdefault("summary", {}).update(summary) - self._write_state(run_id, state) - - def close(self) -> None: - """Every write is already on disk.""" - - # ------------------------------------------------------------------ state - - def _state_path(self, run_id: str) -> Path: - return self.run_dir(run_id) / "run.json" - - def _read_state(self, run_id: str) -> Dict[str, Any]: - path = self._state_path(run_id) - if not path.exists(): - return {"id": run_id} - try: - state = json.loads(path.read_text(encoding="utf-8")) - except (json.JSONDecodeError, OSError, UnicodeDecodeError) as exc: - logger.warning("Could not read %s (%s); starting a fresh record", path, exc) - return {"id": run_id} - return state if isinstance(state, dict) else {"id": run_id} - - def _write_state(self, run_id: str, state: Dict[str, Any]) -> None: - path = self._state_path(run_id) - path.parent.mkdir(parents=True, exist_ok=True) - # Write-then-rename so a crash mid-write never truncates the record. - temp = path.with_suffix(".json.tmp") - temp.write_text(json.dumps(state, indent=2, ensure_ascii=False, default=str), "utf-8") - temp.replace(path) - - -def _now() -> str: - return datetime.now(timezone.utc).isoformat() - - -def _spec_to_json(spec: RunSpec) -> Dict[str, Any]: - data = asdict(spec) - data["environments"] = [ - {key: value for key, value in env.items() if value is not None} - for env in data.get("environments", []) - ] - return data diff --git a/packages/prime-runs/src/prime_runs/models.py b/packages/prime-runs/src/prime_runs/models.py index 15a5d7cb4..cf27c838c 100644 --- a/packages/prime-runs/src/prime_runs/models.py +++ b/packages/prime-runs/src/prime_runs/models.py @@ -12,9 +12,9 @@ from .exceptions import ConfigurationError -Mode = Literal["online", "offline", "disabled"] -"""``online`` talks to the platform, ``offline`` writes a local run directory, -``disabled`` makes every call a no-op with the same object shape.""" +Mode = Literal["online", "disabled"] +"""``online`` talks to the platform; ``disabled`` makes every call a no-op with +the same object shape.""" OnError = Literal["warn", "raise"] diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 2eb93db9e..11904622e 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -6,8 +6,7 @@ ``init()`` is called *before* rollouts start, and the ID it returns is *the* run ID everywhere — including inside every trace document the producer writes. -Nothing is re-stamped afterwards. Offline runs get a locally issued ID through -the same path. +Nothing is re-stamped afterwards. """ import atexit @@ -16,11 +15,12 @@ import os import threading import time +import uuid from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Union from . import _fork from ._http import DEFAULT_TIMEOUT, UPLOAD_TIMEOUT, PlatformClient -from .backends import Backend, EvalsBackend, OfflineBackend, new_run_id +from .backends import Backend, EvalsBackend from .config import Config from .exceptions import ConfigurationError, RunFinishedError from .models import ( @@ -34,7 +34,7 @@ RunSpec, RunStatus, ) -from .sinks import EvalSamplesSink, OfflineSink, Sink, TracesSink +from .sinks import EvalSamplesSink, Sink, TracesSink from .worker import UploadWorker logger = logging.getLogger(__name__) @@ -118,7 +118,7 @@ def name(self) -> Optional[str]: @property def url(self) -> Optional[str]: - """Where to open this run — a dashboard URL, or a local path offline.""" + """The dashboard URL; ``None`` when disabled.""" return self._handle.url @property @@ -365,7 +365,6 @@ def init( tags: Optional[Sequence[str]] = None, config: Optional[Any] = None, mode: Optional[Mode] = None, - dir: Optional[str] = None, team_id: Optional[str] = None, api_key: Optional[str] = None, base_url: Optional[str] = None, @@ -377,9 +376,9 @@ def init( in the run should carry, and the URL is what a producer prints. ``mode`` defaults to ``$PRIME_RUNS_MODE``, else online when there is an API - key and offline when there is not. ``config`` is what the run was configured - with: the path to the file it was launched from (stored byte for byte under - ``config_source``), or a mapping taken as given. + key and disabled (with a warning) when there is not. ``config`` is what the + run was configured with: the path to the file it was launched from (stored + byte for byte under ``config_source``), or a mapping taken as given. """ settings = Config() api_key = api_key if api_key is not None else settings.api_key @@ -402,18 +401,13 @@ def init( sinks: List[Sink] if resolved_mode == "disabled": backend = _DisabledBackend() - handle = RunHandle(id=new_run_id(), name=name) + handle = RunHandle(id=_disabled_run_id(), name=name) sinks = [] - elif resolved_mode == "offline": - offline = OfflineBackend(dir) - backend = offline - handle = offline.create(spec) - sinks = [OfflineSink(offline.directory)] else: if not api_key: raise ConfigurationError( 'mode="online" needs an API key. Set PRIME_API_KEY, run `prime login`, ' - 'or pass mode="offline".' + 'or pass mode="disabled".' ) client = PlatformClient(api_key=api_key, base_url=base_url, timeout=DEFAULT_TIMEOUT) backend = EvalsBackend(client, frontend_url=settings.frontend_url, team_id=team_id) @@ -435,11 +429,16 @@ def init( return run +def _disabled_run_id() -> str: + """A locally issued ID, visibly distinct from a platform one.""" + return f"disabled-{uuid.uuid4().hex[:16]}" + + class _DisabledBackend: """No-op lifecycle, so ``mode="disabled"`` needs no branching upstream.""" def create(self, spec: RunSpec) -> RunHandle: - return RunHandle(id=new_run_id()) + return RunHandle(id=_disabled_run_id()) def update(self, run_id: str, **kwargs: Any) -> None: return None @@ -456,17 +455,17 @@ def _resolve_mode(mode: Optional[Mode], *, api_key: str) -> Mode: env_mode = os.getenv(MODE_ENV) if env_mode: mode = env_mode.strip().lower() # type: ignore[assignment] - if mode not in (None, "online", "offline", "disabled"): - raise ConfigurationError(f"mode={mode!r} is not one of 'online', 'offline' or 'disabled'") + if mode not in (None, "online", "disabled"): + raise ConfigurationError(f"mode={mode!r} is not one of 'online' or 'disabled'") if mode is None: if api_key: mode = "online" else: logger.warning( "No API key found (set PRIME_API_KEY or run `prime login`); " - "recording this run offline instead." + "this run will not be tracked." ) - mode = "offline" + mode = "disabled" return mode # type: ignore[return-value] diff --git a/packages/prime-runs/src/prime_runs/sinks/__init__.py b/packages/prime-runs/src/prime_runs/sinks/__init__.py index d75fdb283..6f5a7ffdf 100644 --- a/packages/prime-runs/src/prime_runs/sinks/__init__.py +++ b/packages/prime-runs/src/prime_runs/sinks/__init__.py @@ -1,7 +1,6 @@ """Record transports. Independent of backends, and of each other.""" from .base import Sink, is_episode, stamp_run, to_mapping -from .offline import OfflineSink from .samples import EvalSamplesSink from .traces import TracesSink @@ -11,6 +10,5 @@ "stamp_run", "to_mapping", "EvalSamplesSink", - "OfflineSink", "TracesSink", ] diff --git a/packages/prime-runs/src/prime_runs/sinks/offline.py b/packages/prime-runs/src/prime_runs/sinks/offline.py deleted file mode 100644 index d63a72dc2..000000000 --- a/packages/prime-runs/src/prime_runs/sinks/offline.py +++ /dev/null @@ -1,82 +0,0 @@ -"""Local JSONL sink, written in the wire format Prime Traces accepts, so the -files can later be sent by ``prime_traces.TracesClient.upload_file`` untouched.""" - -import logging -from pathlib import Path -from typing import Any, BinaryIO, Mapping, Optional, Sequence, Union - -from .. import _fork -from .._http import encode_json -from .base import Sink, is_episode, stamp_run, to_mapping - -logger = logging.getLogger(__name__) - - -class OfflineSink(Sink): - """Appends records to ``//records/{trace,episode}.jsonl``.""" - - name = "offline" - - def __init__(self, directory: Union[str, Path]) -> None: - self.enabled = True - self.directory = Path(directory) - self._run_id: Optional[str] = None - self._handles: dict[str, BinaryIO] = {} - self.records_written = 0 - _fork.register(self) - - def reset_after_fork(self) -> None: - """Abandon inherited file handles. They are unbuffered (see ``_handle``), - so dropping them cannot flush a copy of the parent's buffer.""" - self._handles = {} - - def start(self, run_id: str, context: Mapping[str, str]) -> None: - self._run_id = run_id - self._records_dir.mkdir(parents=True, exist_ok=True) - - @property - def _records_dir(self) -> Path: - return self.directory / (self._run_id or "unknown") / "records" - - def write(self, records: Sequence[Any]) -> None: - if not self.enabled or not records: - return - handle = self._handle("episode" if is_episode(records[0]) else "trace") - for record in records: - mapping = to_mapping(record) - if self._run_id: - mapping = stamp_run(mapping, self._run_id) - # Same strict encoder as the online path: an archive that holds - # NaN cannot later be uploaded. - _write_all(handle, encode_json(mapping) + b"\n") - self.records_written += 1 - - def _handle(self, name: str) -> BinaryIO: - """An unbuffered append-mode handle per line format. Unbuffered so a - fork never copies pending records; ``O_APPEND`` keeps writers whole.""" - handle = self._handles.get(name) - if handle is None: - self._records_dir.mkdir(parents=True, exist_ok=True) - handle = open(self._records_dir / f"{name}.jsonl", "ab", buffering=0) - self._handles[name] = handle - return handle - - def flush(self) -> None: - """Every write already went to the file.""" - - def close(self) -> None: - for handle in self._handles.values(): - try: - handle.close() - except OSError as exc: # pragma: no cover - teardown must not raise - logger.debug("Error closing an offline record file: %s", exc) - self._handles.clear() - - -def _write_all(handle: BinaryIO, data: bytes) -> None: - """Write every byte. A raw handle may report a short write.""" - while data: - written = handle.write(data) - if not written: # pragma: no cover - only on a non-blocking handle - raise OSError("offline record write made no progress") - data = data[written:] diff --git a/packages/prime-runs/tests/conftest.py b/packages/prime-runs/tests/conftest.py index c266341c9..7ede20e6c 100644 --- a/packages/prime-runs/tests/conftest.py +++ b/packages/prime-runs/tests/conftest.py @@ -16,7 +16,6 @@ "PRIME_TRACES_URL", "PRIME_FRONTEND_URL", "PRIME_RUNS_MODE", - "PRIME_RUNS_DIR", ) diff --git a/packages/prime-runs/tests/test_config_source.py b/packages/prime-runs/tests/test_config_source.py index 8ef867fa7..a0120afe1 100644 --- a/packages/prime-runs/tests/test_config_source.py +++ b/packages/prime-runs/tests/test_config_source.py @@ -145,26 +145,33 @@ def test_a_mapping_that_looks_like_a_source_is_still_just_a_mapping(): # -------------------------------------------------------------- through init -def test_an_offline_run_stores_the_launch_file(tmp_path): +@pytest.fixture +def online(monkeypatch, make_platform_client, eval_routes): + handler = RecordingHandler(eval_routes) + monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler)) + monkeypatch.setattr("prime_traces.TracesClient", lambda **_: object()) + + def _init(**kwargs): + return pr.init(name="tb2", environments=["gsm8k"], api_key="test-key", **kwargs), handler + + return _init + + +def test_an_online_run_sends_the_source_in_create_metadata(tmp_path, online): path = tmp_path / "eval.toml" path.write_text(EVAL_TOML) - run = pr.init( - name="tb2", - environments=["gsm8k"], - mode="offline", - dir=str(tmp_path / "runs"), - config=path, - ) + run, handler = online(config=path) run.finish() - state = json.loads((tmp_path / "runs" / run.id / "run.json").read_text()) - stored = state["spec"]["config"] - assert stored[CONFIG_SOURCE_KEY]["text"] == EVAL_TOML - assert stored[CONFIG_SOURCE_KEY]["filename"] == "eval.toml" + create = next(r for r in handler.requests if r.url.path == "/api/v1/evaluations/") + metadata = json.loads(create.content)["metadata"] + assert metadata[CONFIG_SOURCE_KEY]["text"] == EVAL_TOML + assert metadata[CONFIG_SOURCE_KEY]["format"] == "toml" + assert metadata[CONFIG_SOURCE_KEY]["filename"] == "eval.toml" -def test_extra_values_can_be_merged_onto_a_launch_file(tmp_path): +def test_extra_values_can_be_merged_onto_a_launch_file(tmp_path, online): """A run launched from a file that also wants structured values passes a mapping carrying the source under ``CONFIG_SOURCE_KEY`` — what verifiers does.""" path = tmp_path / "eval.toml" @@ -174,12 +181,13 @@ def test_extra_values_can_be_merged_onto_a_launch_file(tmp_path): CONFIG_SOURCE_KEY: ConfigSource.from_file(path).to_dict(), } - run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path), config=config) + run, handler = online(config=config) run.finish() - state = json.loads((tmp_path / run.id / "run.json").read_text()) - assert state["config"]["model"] == "deepseek/deepseek-v4-flash" - assert state["config"][CONFIG_SOURCE_KEY]["text"] == EVAL_TOML + create = next(r for r in handler.requests if r.url.path == "/api/v1/evaluations/") + metadata = json.loads(create.content)["metadata"] + assert metadata["model"] == "deepseek/deepseek-v4-flash" + assert metadata[CONFIG_SOURCE_KEY]["text"] == EVAL_TOML assert run.config_source.filename == "eval.toml" @@ -187,49 +195,32 @@ def test_the_run_reports_its_own_source(tmp_path): path = tmp_path / "train.toml" path.write_text(EVAL_TOML) - run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path), config=path) + run = pr.init(environments=["gsm8k"], mode="disabled", config=path) assert run.config_source.filename == "train.toml" assert run.config_source.text == EVAL_TOML run.finish() -def test_a_run_without_a_source_reports_none(tmp_path): - run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path)) +def test_a_run_without_a_source_reports_none(): + run = pr.init(environments=["gsm8k"], mode="disabled") assert run.config_source is None run.finish() -def test_an_online_run_sends_the_source_in_create_metadata( - tmp_path, monkeypatch, make_platform_client, eval_routes -): - path = tmp_path / "eval.toml" - path.write_text(EVAL_TOML) - handler = RecordingHandler(eval_routes) - monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler)) - - monkeypatch.setattr("prime_traces.TracesClient", lambda **_: object()) - run = pr.init(name="tb2", environments=["gsm8k"], api_key="test-key", config=path) - run.finish() - - create = next(r for r in handler.requests if r.url.path == "/api/v1/evaluations/") - metadata = json.loads(create.content)["metadata"] - assert metadata[CONFIG_SOURCE_KEY]["text"] == EVAL_TOML - assert metadata[CONFIG_SOURCE_KEY]["format"] == "toml" - - -def test_the_source_survives_the_failure_fallback(tmp_path): +def test_the_source_survives_the_failure_fallback(tmp_path, online): """The fallback rewrites metadata to record a terminal state. It merges into the whole config, so the source must still be there afterwards.""" path = tmp_path / "eval.toml" path.write_text(EVAL_TOML) - run = pr.init(environments=["gsm8k"], mode="offline", dir=str(tmp_path), config=path) + run, handler = online(config=path) run.fail("something broke") - state = json.loads((tmp_path / run.id / "run.json").read_text()) - assert state["config"][CONFIG_SOURCE_KEY]["text"] == EVAL_TOML + update = handler.bodies_for("/api/v1/evaluations/eval-abc")[-1] + assert update["metadata"][CONFIG_SOURCE_KEY]["text"] == EVAL_TOML + assert update["metadata"]["prime_runs"]["status"] == "failed" def test_a_spec_defaults_to_no_source(): diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index 15a8a9dc8..c42165279 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -1,6 +1,5 @@ -"""``init()``: mode resolution, offline runs, online runs.""" +"""``init()``: mode resolution, disabled runs, online runs.""" -import json import os import pytest @@ -11,94 +10,46 @@ from prime_runs.exceptions import ConfigurationError from prime_runs.models import RunStatus -# ------------------------------------------------------------------- offline - - -def test_an_offline_run_is_a_real_run(tmp_path): - run = pr.init(name="local", environments=["gsm8k"], mode="offline", dir=str(tmp_path)) - - assert run.id.startswith("offline-") - assert run.url == str((tmp_path / run.id).resolve()) - assert run.mode == "offline" - - run.log_traces([{"id": "t1"}]) - run.finish(summary={"avg_reward": 0.5}) - - state = json.loads((tmp_path / run.id / "run.json").read_text()) - assert state["status"] == RunStatus.COMPLETED.value - assert state["summary"]["avg_reward"] == 0.5 - assert state["spec"]["environments"] == [{"name": "gsm8k"}] - - -def test_offline_records_are_written_in_the_wire_format(tmp_path): - """The archive is a deferred upload, not a debug dump: these bytes are what - ``TracesClient.upload_file`` sends, with the run already stamped.""" - run = pr.init(mode="offline", dir=str(tmp_path)) - run.log_traces([{"id": "t1"}, {"id": "t2"}]) - run.finish() - - lines = (tmp_path / run.id / "records" / "trace.jsonl").read_text().splitlines() - records = [json.loads(line) for line in lines] - - assert [record["id"] for record in records] == ["t1", "t2"] - assert all(record["run"] == {"id": run.id, "type": "eval"} for record in records) - - -def test_episodes_are_written_to_their_own_file(tmp_path): - run = pr.init(mode="offline", dir=str(tmp_path)) - run.log_traces([make_episode("ep-1", [make_trace()])]) - run.finish() - - assert (tmp_path / run.id / "records" / "episode.jsonl").exists() - - -def test_a_record_that_already_names_a_run_is_left_alone(tmp_path): - """Producers stamp the run themselves; two sources of truth for the run ID - is how traces end up on the wrong run.""" - run = pr.init(mode="offline", dir=str(tmp_path)) - run.log_traces([{"id": "t1", "run": {"id": "someone-elses-run"}}]) - run.finish() - - record = json.loads((tmp_path / run.id / "records" / "trace.jsonl").read_text()) - assert record["run"]["id"] == "someone-elses-run" - - # -------------------------------------------------------------------- modes -def test_no_api_key_degrades_to_offline_rather_than_skipping_the_run(tmp_path, caplog): +def test_no_api_key_disables_the_run_and_says_so(caplog): + """Loudly, not silently: a user who forgot ``prime login`` must not believe + the run was tracked.""" with caplog.at_level("WARNING"): - run = pr.init(name="local", dir=str(tmp_path)) + run = pr.init(name="local") - assert run.mode == "offline" - assert "offline" in caplog.text + assert run.mode == "disabled" + assert "not be tracked" in caplog.text run.finish() -def test_the_mode_can_be_set_from_the_environment(monkeypatch, tmp_path): +def test_the_mode_can_be_set_from_the_environment(monkeypatch): monkeypatch.setenv("PRIME_RUNS_MODE", "disabled") - run = pr.init(name="local", api_key="test-key", dir=str(tmp_path)) + run = pr.init(name="local", api_key="test-key") assert run.mode == "disabled" run.finish() -def test_an_unknown_mode_is_rejected(tmp_path): +def test_an_unknown_mode_is_rejected(): with pytest.raises(ConfigurationError, match="not one of"): - pr.init(mode="sideways", dir=str(tmp_path)) + pr.init(mode="sideways") def test_a_disabled_run_still_answers_every_call(tmp_path): - """Same object shape, so producer code needs no branching.""" - run = pr.init(mode="disabled", dir=str(tmp_path)) + """Same object shape, so producer code needs no branching — and nothing + touches the network or the filesystem.""" + run = pr.init(mode="disabled") run.log_traces([{"id": "t1"}]) run.finish(summary={"avg_reward": 1.0}) - assert run.id + assert run.id.startswith("disabled-") + assert run.url is None assert run.status is RunStatus.COMPLETED - assert not list(tmp_path.iterdir()) + assert not list(tmp_path.iterdir()) # tmp_path is $HOME here def test_online_without_an_api_key_is_a_configuration_error(): @@ -213,15 +164,14 @@ def test_the_end_to_end_shape_a_producer_writes(online): assert run.errors == [] -def test_a_second_init_in_one_process_opens_its_own_run(tmp_path): - first = pr.init(mode="offline", dir=str(tmp_path)) +def test_a_second_init_in_one_process_opens_its_own_run(): + first = pr.init(mode="disabled") first.finish() - second = pr.init(mode="offline", dir=str(tmp_path)) + second = pr.init(mode="disabled") second.finish() assert second.id != first.id - assert (tmp_path / second.id / "run.json").exists() # --------------------------------------------------------------------- fork @@ -229,22 +179,27 @@ def test_a_second_init_in_one_process_opens_its_own_run(tmp_path): @pytest.mark.skipif(not hasattr(os, "fork"), reason="fork is POSIX-only") @pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning") -def test_a_forked_child_does_not_duplicate_the_parents_records_or_close_its_run(tmp_path): - """At fork time the parent has records in the upload queue. The child - inherits a copy; writing them would put every record in the file twice, - and the inherited atexit hook must not finalize the parent's run.""" - run = pr.init(mode="offline", dir=str(tmp_path)) - run.log_traces([{"id": f"parent-{n}"} for n in range(5)]) +def test_a_forked_child_does_not_duplicate_the_parents_records_or_close_its_run(online): + """At fork time the parent may have records in the upload queue. The child + inherits a copy; writing them would upload every record twice, and the + inherited atexit hook must not finalize the parent's run.""" + run, handler = online() + run.log_traces([{"sample_id": f"parent-{n}"} for n in range(5)]) pid = os.fork() if pid == 0: # pragma: no cover - asserted through the child's exit code code = 0 try: - run.log_traces([{"id": "child-1"}]) + # ``handler`` is the child's copy: only what the child uploads lands here. + before = len(handler.bodies_for("/api/v1/evaluations/eval-abc/samples")) + run.log_traces([{"sample_id": "child-1"}]) run.flush() run._on_process_exit() - state = json.loads((tmp_path / run.id / "run.json").read_text()) - if state["status"] != RunStatus.RUNNING.value: + posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")[before:] + ids = [s["sample_id"] for body in posted for s in body["samples"]] + if ids != ["child-1"]: + code = 3 + if "POST /api/v1/evaluations/eval-abc/finalize" in handler.paths(): code = 4 except BaseException: code = 2 @@ -255,36 +210,16 @@ def test_a_forked_child_does_not_duplicate_the_parents_records_or_close_its_run( assert os.waitstatus_to_exitcode(status) == 0 run.finish() - lines = (tmp_path / run.id / "records" / "trace.jsonl").read_text().splitlines() - ids = [json.loads(line)["id"] for line in lines] - - assert sorted(ids) == sorted([f"parent-{n}" for n in range(5)] + ["child-1"]) - - -def test_offline_records_are_on_disk_before_any_flush(tmp_path): - """Nothing may sit in a process-local write buffer a fork could copy.""" - run = pr.init(mode="offline", dir=str(tmp_path)) - run.log_traces([{"id": "t1"}]) - run.flush() - - path = tmp_path / run.id / "records" / "trace.jsonl" - assert path.read_text().count('"t1"') == 1 - - run.log_traces([{"id": "t2"}]) - run.flush() - assert path.read_text().count('"t2"') == 1 - - run.finish() - assert len(path.read_text().splitlines()) == 2 + posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples") + ids = [s["sample_id"] for body in posted for s in body["samples"]] + assert sorted(ids) == [f"parent-{n}" for n in range(5)] + assert handler.paths().count("POST /api/v1/evaluations/eval-abc/finalize") == 1 -def test_offline_records_reject_nonfinite_json_instead_of_writing_invalid_jsonl(tmp_path): - run = pr.init(mode="offline", dir=str(tmp_path), on_error="raise") - run.log_traces([{"id": "bad", "reward": float("nan")}]) +def test_records_reject_nonfinite_json_instead_of_sending_an_opaque_400(online): + run, _ = online(on_error="raise") + run.log_traces([{"sample_id": "bad", "reward": float("nan")}]) with pytest.raises(ValueError, match="Out of range float values"): run.flush() - - path = tmp_path / run.id / "records" / "trace.jsonl" - assert path.read_text() == "" run.finish() From e31b00cb5bfc93ce9d874ce4e0ec71afb9a06ae7 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Fri, 21 Aug 2026 13:49:42 -0700 Subject: [PATCH 23/39] fix(runs): an account outside the traces beta is not a failed run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Prime Traces is gated to an owner allowlist; in production that is three internal teams. For everyone else the traces sink's first upload returns 403 service_not_enabled, and the sink re-raised it so the worker counted the batch as lost: three warnings per run, "N failed via traces" in the verifiers footer, and a ForbiddenError out of finish() under on_error="raise" — for a run that stored everything it was asked to. The sink now distinguishes the two 403s. service_not_enabled retires the sink at INFO and returns: nothing was lost, there was never anywhere for the records to go. forbidden (a token without the traces scope) is something the caller can fix, so it still disables the sink and raises for loss accounting. A 403 with no recognised code is treated as the latter, so the failure mode of being wrong is the old behaviour, not silent loss. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01W4S4DNVRX5BNtsLva7TZz1 --- packages/prime-runs/README.md | 6 ++-- .../prime-runs/src/prime_runs/sinks/traces.py | 35 +++++++++++++------ packages/prime-runs/tests/test_run.py | 34 +++++++++++++----- packages/prime-runs/tests/test_traces_sink.py | 29 ++++++++++++--- 4 files changed, 78 insertions(+), 26 deletions(-) diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index a3b2b50da..1bdfbb127 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -122,8 +122,10 @@ Or pass `api_key=`, `base_url=`, `team_id=` to `init()`. An online run writes every record to two sinks: Prime Traces (the system of record — streaming, episode-aware, content-addressed and therefore idempotent on retry) and the flat v0 sample table today's viewer reads. Both run because -Prime Traces is gated to an allowlist; when the viewer reads traces natively the -default sink list drops one entry and no producer changes. +Prime Traces is gated to an allowlist; an account outside it has the traces sink +turn itself off at the first upload — not counted as a failure, since nothing was +lost. When the viewer reads traces natively the default sink list drops one entry +and no producer changes. `prime_runs.projection` holds `trace_to_sample` / `build_samples`, the v0 projection moved here from verifiers. `prime_runs.metrics.from_episodes` is the diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py index abc7f2067..e75d772ff 100644 --- a/packages/prime-runs/src/prime_runs/sinks/traces.py +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -105,14 +105,20 @@ def write(self, records: Sequence[Any]) -> None: compress=self._compress, ) ) - except Exception as exc: - if self._is_gated(exc): - self._disable( + except ForbiddenError as exc: + # No runtime action fixes a 403, so the sink retires either way. + # What differs is whether the batch counts as lost. + if _is_not_enabled(exc): + # Outside the beta: there was never anywhere for these records + # to go, so nothing was lost. Not a warning, not a failure. + self._retire_quietly( f"Prime Traces is not enabled for this account ({exc}); " - "falling back to the remaining sinks" + "continuing with the remaining sinks" ) - # Re-raised either way so the worker's loss accounting and strict - # callers see the failed batch. + return + # A credential without the traces scope is something the caller + # can fix; raised so loss accounting and strict callers see it. + self._disable(f"this credential cannot write traces ({exc})") raise self.receipts_received += len(receipts) if self._receipt_history_size: @@ -144,8 +150,15 @@ def _disable(self, reason: str) -> None: logger.warning("Traces sink disabled: %s", reason) self.enabled = False - @staticmethod - def _is_gated(exc: Exception) -> bool: - """A 403 (``service_not_enabled``, or a write-only token) is not - fixable at runtime, so the sink turns itself off instead of retrying.""" - return isinstance(exc, ForbiddenError) + def _retire_quietly(self, reason: str) -> None: + if self.enabled: + logger.info("Traces sink off: %s", reason) + self.enabled = False + + +def _is_not_enabled(exc: ForbiddenError) -> bool: + """``service_not_enabled``: the account is outside the private beta. The + other 403, ``forbidden``, means the token lacks the ``traces`` scope.""" + from prime_traces import ErrorCode + + return exc.code == ErrorCode.SERVICE_NOT_ENABLED.value diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index f3805b961..0e3c3cc0d 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -439,18 +439,36 @@ def test_an_upload_failure_is_reported_once(): run.finish() -def test_a_gated_trace_upload_reaches_strict_callers_and_loss_accounting(): - class GatedClient: - def upload_records(self, records, **kwargs): - raise ForbiddenError("not in beta", status_code=403, code="service_not_enabled") +class _ForbiddenClient: + def __init__(self, code: str) -> None: + self.code = code - def close(self) -> None: - pass + def upload_records(self, records, **kwargs): + raise ForbiddenError("403", status_code=403, code=self.code) - run = make_run(sinks=[TracesSink(client=GatedClient())], on_error="raise") + def close(self) -> None: + pass + + +def test_an_account_outside_the_beta_is_not_a_failed_run(): + """Nothing was lost — the records went to every sink that applies to this + account — so the run finishes clean even under ``on_error="raise"``.""" + sink = TracesSink(client=_ForbiddenClient("service_not_enabled")) + run = make_run(sinks=[sink], on_error="raise") + run.log_traces([{"id": "t1"}]) + + run.finish() + + assert run.failed_records == {} + assert run.errors == [] + assert sink.enabled is False + + +def test_a_credential_without_the_traces_scope_reaches_strict_callers_and_loss_accounting(): + run = make_run(sinks=[TracesSink(client=_ForbiddenClient("forbidden"))], on_error="raise") run.log_traces([{"id": "t1"}]) - with pytest.raises(ForbiddenError, match="not in beta"): + with pytest.raises(ForbiddenError): run.finish() assert run.failed_records == {"traces": 1} diff --git a/packages/prime-runs/tests/test_traces_sink.py b/packages/prime-runs/tests/test_traces_sink.py index ca0e4abef..7e3952fe0 100644 --- a/packages/prime-runs/tests/test_traces_sink.py +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -85,20 +85,39 @@ def test_producer_objects_are_passed_through_untouched(): assert client.calls[0][0][0] is trace -def test_a_gated_account_disables_the_sink_and_reports_the_failed_batch(caplog): - """Prime Traces is in closed beta; no runtime action fixes a 403, so - retrying it for the rest of the run only produces noise.""" +def test_an_account_outside_the_beta_retires_the_sink_without_a_failure(caplog): + """Prime Traces is in closed beta. For everyone outside it there was never + anywhere for these records to go, so the sink turns itself off quietly: + no exception for the worker to count, nothing above INFO in the log.""" client = FakeTracesClient( raises=ForbiddenError("not in beta", status_code=403, code="service_not_enabled") ) sink = make_sink(client) + with caplog.at_level("INFO"): + sink.write([{"id": "t1"}]) + sink.write([{"id": "t2"}]) + + assert sink.enabled is False + assert len(client.calls) == 1 + assert "not enabled" in caplog.text + assert not [r for r in caplog.records if r.levelname == "WARNING"] + + +def test_a_credential_without_the_traces_scope_is_still_a_failure(caplog): + """The other 403 is fixable — mint a token with the scope — so it is raised + for loss accounting and strict callers, and the sink still retires.""" + client = FakeTracesClient( + raises=ForbiddenError("missing scope: traces", status_code=403, code="forbidden") + ) + sink = make_sink(client) + with caplog.at_level("WARNING"): - with pytest.raises(ForbiddenError, match="not in beta"): + with pytest.raises(ForbiddenError, match="missing scope"): sink.write([{"id": "t1"}]) assert sink.enabled is False - assert "not enabled" in caplog.text + assert "cannot write traces" in caplog.text def test_a_transient_failure_is_raised_so_the_worker_can_report_it(): From b53439fcca222639daf699ce84b4916ec0799df7 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Fri, 21 Aug 2026 14:52:40 -0700 Subject: [PATCH 24/39] refactor(runs): flatten backends/ into backend.py and drop leftover scaffolding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The backends/ package was shaped for two implementations; after the offline-mode cut it held one. Backend (protocol), EvalsBackend and the disabled no-op now share backend.py, next to the contract they satisfy. Also removed: @runtime_checkable on both protocols (no isinstance check anywhere), the normalize_base_url helper (inlined into PlatformClient — the prime_traces equivalent is private), and the top-level build_samples / trace_to_sample aliases (prime_runs.projection is the spelling verifiers already uses). The pyproject dependency comment no longer claims a hand-rolled retry loop. No behaviour change; 161 tests unchanged. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01W4S4DNVRX5BNtsLva7TZz1 --- packages/prime-runs/pyproject.toml | 8 +- .../prime-runs/src/prime_runs/__init__.py | 3 - packages/prime-runs/src/prime_runs/_http.py | 9 +- .../{backends/evals.py => backend.py} | 90 ++++++++++++++++--- .../src/prime_runs/backends/__init__.py | 6 -- .../src/prime_runs/backends/base.py | 48 ---------- packages/prime-runs/src/prime_runs/run.py | 28 +----- .../prime-runs/src/prime_runs/sinks/base.py | 3 +- .../prime-runs/tests/test_evals_backend.py | 2 +- 9 files changed, 92 insertions(+), 105 deletions(-) rename packages/prime-runs/src/prime_runs/{backends/evals.py => backend.py} (72%) delete mode 100644 packages/prime-runs/src/prime_runs/backends/__init__.py delete mode 100644 packages/prime-runs/src/prime_runs/backends/base.py diff --git a/packages/prime-runs/pyproject.toml b/packages/prime-runs/pyproject.toml index 165788c58..3f895ac21 100644 --- a/packages/prime-runs/pyproject.toml +++ b/packages/prime-runs/pyproject.toml @@ -14,10 +14,10 @@ authors = [ # consumer of this SDK. Nothing here may pull in prime, verifiers, typer, # rich or textual, directly or transitively. # -# Kept to what is actually imported. The other SDKs carry `pydantic` (they model -# response bodies) and `tenacity` (they retry through it); this package models -# nothing and hand-rolls its retry loop in `_http.py`, so neither belongs in the -# dependency tree of a package that lands inside verifiers and prime-rl. +# Kept to what is actually imported. This package models no response bodies +# (so no direct `pydantic`; it arrives transitively through prime-traces) and +# takes its error mapping and backoff from `prime_traces.core` rather than +# `tenacity`. Nothing is added to what verifiers and prime-rl already carry. dependencies = [ "httpx>=0.25.0", "prime-traces>=0.0.2", diff --git a/packages/prime-runs/src/prime_runs/__init__.py b/packages/prime-runs/src/prime_runs/__init__.py index c92893a00..fe66244e2 100644 --- a/packages/prime-runs/src/prime_runs/__init__.py +++ b/packages/prime-runs/src/prime_runs/__init__.py @@ -32,7 +32,6 @@ UnauthorizedError, ) from .models import CONFIG_SOURCE_KEY, ConfigSource, EnvironmentRef, RunStatus -from .projection import build_samples, trace_to_sample from .run import MODE_ENV, Run, init __version__ = "0.1.0" @@ -47,8 +46,6 @@ "MODE_ENV", "metrics", "projection", - "build_samples", - "trace_to_sample", "PrimeRunsError", "ConfigurationError", "EnvironmentResolutionError", diff --git a/packages/prime-runs/src/prime_runs/_http.py b/packages/prime-runs/src/prime_runs/_http.py index 3c6c9d61a..713ebed65 100644 --- a/packages/prime-runs/src/prime_runs/_http.py +++ b/packages/prime-runs/src/prime_runs/_http.py @@ -40,11 +40,6 @@ def _user_agent() -> str: return f"prime-runs/{__version__} python/{py}" -def normalize_base_url(url: str) -> str: - """Strip a trailing ``/api/v1``; the client appends it itself.""" - return url.rstrip("/").removesuffix("/api/v1") - - def encode_json(value: Any) -> bytes: """Compact UTF-8 JSON. ``allow_nan=False``: strict parsers server-side reject bare ``NaN`` and the failure is an opaque 400.""" @@ -65,7 +60,9 @@ def __init__( max_attempts: int = DEFAULT_MAX_ATTEMPTS, client: Optional[httpx.Client] = None, ) -> None: - self.base_url = normalize_base_url(base_url) + # Strip a trailing /api/v1 — the same normalization prime_traces applies — + # so an explicit base_url written with the suffix does not double it. + self.base_url = base_url.rstrip("/").removesuffix("/api/v1") self.api_prefix = f"{self.base_url}/api/v1" self.max_attempts = max(1, max_attempts) self._owns_client = client is None diff --git a/packages/prime-runs/src/prime_runs/backends/evals.py b/packages/prime-runs/src/prime_runs/backend.py similarity index 72% rename from packages/prime-runs/src/prime_runs/backends/evals.py rename to packages/prime-runs/src/prime_runs/backend.py index ec2842a4c..68e226c55 100644 --- a/packages/prime-runs/src/prime_runs/backends/evals.py +++ b/packages/prime-runs/src/prime_runs/backend.py @@ -1,24 +1,70 @@ -"""Eval runs over ``/api/v1/evaluations/*``, plus environment resolution -through the hub's get-or-create so a local run uploads without ``prime env push``. +"""Run backends: the contract, the evaluations backend, and the disabled no-op. -The eval API has no producer-facing way to mark a run failed: ``finalize`` -moves a run to COMPLETED and ``UpdateEvaluationRequest`` carries no status. A -failed or crashed run therefore keeps showing as running; the terminal state is -recorded in ``metadata.prime_runs`` so it is at least visible. +A backend owns the *lifecycle* of a run — creating it, updating what is known +about it, closing it out with a terminal status. It does not move records; +that is a sink's job (see :mod:`prime_runs.sinks`). + +:class:`EvalsBackend` works over ``/api/v1/evaluations/*``, resolving +environments through the hub's get-or-create so a local run uploads without +``prime env push``. The eval API has no producer-facing way to mark a run +failed: ``finalize`` moves a run to COMPLETED and ``UpdateEvaluationRequest`` +carries no status. A failed or crashed run therefore keeps showing as running; +the terminal state is recorded in ``metadata.prime_runs`` so it is at least +visible. """ import logging import uuid from datetime import datetime, timezone -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Protocol -from .._http import PlatformClient -from ..exceptions import APIError, ConfigurationError, EnvironmentResolutionError -from ..models import EnvironmentRef, RunHandle, RunSpec, RunStatus +from ._http import PlatformClient +from .exceptions import APIError, ConfigurationError, EnvironmentResolutionError +from .models import EnvironmentRef, RunHandle, RunSpec, RunStatus logger = logging.getLogger(__name__) +class Backend(Protocol): + def create(self, spec: RunSpec) -> RunHandle: + """Open a new run and return its identity.""" + ... + + def update( + self, + run_id: str, + *, + config: Optional[Dict[str, Any]] = None, + summary: Optional[Dict[str, Any]] = None, + ) -> None: + """Persist config (inputs) and/or summary (outputs). + + ``config`` is the run's *whole* config, not a patch: the evaluations API + replaces the stored metadata document. + """ + ... + + def finalize( + self, + run_id: str, + *, + status: RunStatus, + summary: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, + ) -> None: + """Close the run out. Called exactly once per run. ``config`` is passed + so a backend recording terminal state inside metadata can merge it.""" + ... + + def close(self) -> None: + """Release transport resources.""" + ... + + +# ------------------------------------------------------------------- evals + + class EvalsBackend: """Lifecycle for evaluation runs.""" @@ -197,3 +243,27 @@ def _default_name(spec: RunSpec) -> str: """The API requires a name; lead with the environment so runs sort together.""" stem = _first_environment_name(spec) or spec.framework or "eval" return f"{stem}-{uuid.uuid4().hex[:8]}" + + +# ---------------------------------------------------------------- disabled + + +def disabled_run_id() -> str: + """A locally issued ID, visibly distinct from a platform one.""" + return f"disabled-{uuid.uuid4().hex[:16]}" + + +class DisabledBackend: + """No-op lifecycle, so ``mode="disabled"`` needs no branching upstream.""" + + def create(self, spec: RunSpec) -> RunHandle: + return RunHandle(id=disabled_run_id()) + + def update(self, run_id: str, **kwargs: Any) -> None: + return None + + def finalize(self, run_id: str, **kwargs: Any) -> None: + return None + + def close(self) -> None: + return None diff --git a/packages/prime-runs/src/prime_runs/backends/__init__.py b/packages/prime-runs/src/prime_runs/backends/__init__.py deleted file mode 100644 index 30e6efc38..000000000 --- a/packages/prime-runs/src/prime_runs/backends/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Run lifecycle backends.""" - -from .base import Backend -from .evals import EvalsBackend - -__all__ = ["Backend", "EvalsBackend"] diff --git a/packages/prime-runs/src/prime_runs/backends/base.py b/packages/prime-runs/src/prime_runs/backends/base.py deleted file mode 100644 index 47f705faf..000000000 --- a/packages/prime-runs/src/prime_runs/backends/base.py +++ /dev/null @@ -1,48 +0,0 @@ -"""The contract a run backend implements. - -A backend owns the *lifecycle* of a run — creating it, updating what is known -about it, closing it out with a terminal status. It does not move records; -that is a sink's job (see :mod:`prime_runs.sinks`). -""" - -from typing import Any, Dict, Optional, Protocol, runtime_checkable - -from ..models import RunHandle, RunSpec, RunStatus - - -@runtime_checkable -class Backend(Protocol): - def create(self, spec: RunSpec) -> RunHandle: - """Open a new run and return its identity.""" - ... - - def update( - self, - run_id: str, - *, - config: Optional[Dict[str, Any]] = None, - summary: Optional[Dict[str, Any]] = None, - ) -> None: - """Persist config (inputs) and/or summary (outputs). - - ``config`` is the run's *whole* config, not a patch: the evaluations API - replaces the stored metadata document. - """ - ... - - def finalize( - self, - run_id: str, - *, - status: RunStatus, - summary: Optional[Dict[str, Any]] = None, - error: Optional[str] = None, - config: Optional[Dict[str, Any]] = None, - ) -> None: - """Close the run out. Called exactly once per run. ``config`` is passed - so a backend recording terminal state inside metadata can merge it.""" - ... - - def close(self) -> None: - """Release transport resources.""" - ... diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 11904622e..13bfc20f9 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -15,12 +15,11 @@ import os import threading import time -import uuid from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Union from . import _fork from ._http import DEFAULT_TIMEOUT, UPLOAD_TIMEOUT, PlatformClient -from .backends import Backend, EvalsBackend +from .backend import Backend, DisabledBackend, EvalsBackend, disabled_run_id from .config import Config from .exceptions import ConfigurationError, RunFinishedError from .models import ( @@ -400,8 +399,8 @@ def init( backend: Backend sinks: List[Sink] if resolved_mode == "disabled": - backend = _DisabledBackend() - handle = RunHandle(id=_disabled_run_id(), name=name) + backend = DisabledBackend() + handle = RunHandle(id=disabled_run_id(), name=name) sinks = [] else: if not api_key: @@ -429,27 +428,6 @@ def init( return run -def _disabled_run_id() -> str: - """A locally issued ID, visibly distinct from a platform one.""" - return f"disabled-{uuid.uuid4().hex[:16]}" - - -class _DisabledBackend: - """No-op lifecycle, so ``mode="disabled"`` needs no branching upstream.""" - - def create(self, spec: RunSpec) -> RunHandle: - return RunHandle(id=_disabled_run_id()) - - def update(self, run_id: str, **kwargs: Any) -> None: - return None - - def finalize(self, run_id: str, **kwargs: Any) -> None: - return None - - def close(self) -> None: - return None - - def _resolve_mode(mode: Optional[Mode], *, api_key: str) -> Mode: if mode is None: env_mode = os.getenv(MODE_ENV) diff --git a/packages/prime-runs/src/prime_runs/sinks/base.py b/packages/prime-runs/src/prime_runs/sinks/base.py index e8ec3edc1..03d7e39f5 100644 --- a/packages/prime-runs/src/prime_runs/sinks/base.py +++ b/packages/prime-runs/src/prime_runs/sinks/base.py @@ -9,12 +9,11 @@ and says why, once. """ -from typing import Any, Dict, Mapping, Protocol, Sequence, runtime_checkable +from typing import Any, Dict, Mapping, Protocol, Sequence from ..models import RUN_KIND -@runtime_checkable class Sink(Protocol): """A destination for run records.""" diff --git a/packages/prime-runs/tests/test_evals_backend.py b/packages/prime-runs/tests/test_evals_backend.py index 2c4839e92..4f9c53ff7 100644 --- a/packages/prime-runs/tests/test_evals_backend.py +++ b/packages/prime-runs/tests/test_evals_backend.py @@ -4,7 +4,7 @@ import pytest from conftest import RecordingHandler -from prime_runs.backends import EvalsBackend +from prime_runs.backend import EvalsBackend from prime_runs.exceptions import ( ConfigurationError, EnvironmentResolutionError, From 97d6482a21b31fcdb6bfdbcf7ccdb962e5659c56 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Fri, 21 Aug 2026 15:31:18 -0700 Subject: [PATCH 25/39] fix(runs): carry the run onto every member trace; count records lost to a retired sink MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both found by the first live e2e (verifiers cooper/prime-runs -> prod). 1. The traces service derives run_id from trace.run.id only and never reads the episode envelope's run (ingestion/extract.py); verifiers records the run on the Episode and its Trace has no run field. The sink passed producer episodes through untouched, so every row of an episode upload landed with an empty run_id — unqueryable by run. stamp_run now also stamps members that lack a run, and TracesSink runs producer objects through to_record() itself so the members are reachable (same bytes the transport would have produced). 2. After a sink was retired by an error, later batches skipped it with no accounting: five episodes lost to the traces sink reported as "1 failed via traces". The worker now counts records that skip a sink it retired. A sink that switches itself off without raising (service_not_enabled) is still not counted — nothing was lost. Verified: run lkcw5sfzgb6jlqli6qnpggns, 5/5 traces and 5/5 episodes carry run_id on prime-traces.pintel.dev. 166 tests. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01W4S4DNVRX5BNtsLva7TZz1 --- .../prime-runs/src/prime_runs/sinks/base.py | 28 ++++++++-- .../prime-runs/src/prime_runs/sinks/traces.py | 11 ++-- packages/prime-runs/src/prime_runs/worker.py | 17 +++++- packages/prime-runs/tests/test_traces_sink.py | 56 +++++++++++++++++-- packages/prime-runs/tests/test_worker.py | 34 +++++++++++ 5 files changed, 130 insertions(+), 16 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/sinks/base.py b/packages/prime-runs/src/prime_runs/sinks/base.py index 03d7e39f5..fee20ccbd 100644 --- a/packages/prime-runs/src/prime_runs/sinks/base.py +++ b/packages/prime-runs/src/prime_runs/sinks/base.py @@ -59,9 +59,25 @@ def is_episode(record: Any) -> bool: def stamp_run(mapping: Mapping[str, Any], run_id: str) -> Dict[str, Any]: - """A copy of ``mapping`` carrying ``run`` if it did not already. Producer - objects are never stamped — they carry their own ``run`` — but a bare dict - with no ``run.id`` is an orphaned, unqueryable upload.""" - if mapping.get("run"): - return dict(mapping) - return {**mapping, "run": {"id": run_id, "type": RUN_KIND}} + """A copy of ``mapping`` carrying ``run`` at the top level, and on every + member trace of an episode that lacks one. + + A record that already names a run keeps it; one without is stamped, since + an upload with no ``run.id`` is orphaned and unqueryable. Members matter + because the traces service derives ``run_id`` from ``trace.run.id`` only — + the episode envelope's ``run`` is never read — while producers (verifiers) + record the run on the episode and nowhere else. + """ + stamped = dict(mapping) + if not stamped.get("run"): + stamped["run"] = {"id": run_id, "type": RUN_KIND} + members = stamped.get("traces") + if isinstance(members, list): + run = stamped["run"] + stamped["traces"] = [ + {**member, "run": run} + if isinstance(member, Mapping) and not member.get("run") + else member + for member in members + ] + return stamped diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py index e75d772ff..faf66b4cb 100644 --- a/packages/prime-runs/src/prime_runs/sinks/traces.py +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -12,7 +12,7 @@ from .. import _fork from ..exceptions import ForbiddenError -from .base import Sink, is_episode, stamp_run +from .base import Sink, is_episode, stamp_run, to_mapping logger = logging.getLogger(__name__) @@ -126,10 +126,13 @@ def write(self, records: Sequence[Any]) -> None: del self.receipts[: -self._receipt_history_size] def _prepare(self, record: Any) -> Any: - """Stamp the run onto bare mappings; producer objects pass through.""" - if not isinstance(record, Mapping) or self._run_id is None: + """The wire mapping for a record, carrying the run on the envelope and + on every member trace. Producer objects go through their own + ``to_record()`` here rather than inside the transport — same bytes, + but the members are reachable for stamping.""" + if self._run_id is None: return record - return stamp_run(record, self._run_id) + return stamp_run(to_mapping(record), self._run_id) def flush(self) -> None: """Uploads are synchronous; nothing is held back here.""" diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index 2a352f836..e8137f4cf 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -35,6 +35,10 @@ class _Flush: event: threading.Event = field(default_factory=threading.Event) +def _sink_name(sink: Any) -> str: + return getattr(sink, "name", type(sink).__name__) + + def _deadline(timeout: Optional[float]) -> Optional[float]: if timeout is None: return None @@ -72,6 +76,10 @@ def __init__( #: from ``dropped``: another sink may well have stored them. self.failed_records: dict = {} self._transient_failures: dict = {} + #: Sinks this worker retired after a failure. Records that skip one of + #: these are lost to it and counted; a sink that switched itself off + #: without raising (nowhere for the records to go) is not in here. + self._retired: set = set() _fork.register(self) # ----------------------------------------------------------------- thread @@ -105,13 +113,16 @@ def _run(self) -> None: def _dispatch(self, records: Sequence[Any]) -> None: for sink in self.sinks: if not getattr(sink, "enabled", True): + name = _sink_name(sink) + if name in self._retired: + self.failed_records[name] = self.failed_records.get(name, 0) + len(records) continue try: sink.write(records) except Exception as exc: # noqa: BLE001 - one sink failing must not stop the others self._fail_sink(sink, exc, dropped=len(records)) else: - self._transient_failures.pop(getattr(sink, "name", id(sink)), None) + self._transient_failures.pop(_sink_name(sink), None) def _flush_sinks(self) -> None: for sink in self.sinks: @@ -126,7 +137,7 @@ def _fail_sink(self, sink: Any, exc: Exception, *, dropped: int = 0) -> None: """The batch is gone (the transports already retried). Decide whether the sink is too: permanent failures retire it at once, transient ones after ``TRANSIENT_FAILURE_LIMIT`` consecutive strikes.""" - name = getattr(sink, "name", type(sink).__name__) + name = _sink_name(sink) if dropped: self.failed_records[name] = self.failed_records.get(name, 0) + dropped @@ -147,6 +158,7 @@ def _fail_sink(self, sink: Any, exc: Exception, *, dropped: int = 0) -> None: self._notify(name, exc) return sink.enabled = False + self._retired.add(name) logger.warning( "Sink %s disabled after %d consecutive transient failures: %s: %s", name, @@ -156,6 +168,7 @@ def _fail_sink(self, sink: Any, exc: Exception, *, dropped: int = 0) -> None: ) else: sink.enabled = False + self._retired.add(name) logger.warning("Sink %s disabled after an error: %s: %s", name, type(exc).__name__, exc) self._notify(name, exc) diff --git a/packages/prime-runs/tests/test_traces_sink.py b/packages/prime-runs/tests/test_traces_sink.py index 7e3952fe0..e2cca6a81 100644 --- a/packages/prime-runs/tests/test_traces_sink.py +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -73,16 +73,64 @@ def test_a_bare_mapping_gets_the_run_stamped_onto_a_copy(): assert original == {"id": "t1"}, "the caller's dict was not mutated" -def test_producer_objects_are_passed_through_untouched(): - """Verifiers and prime-rl stamp the run at rollout time; rewriting their - objects here is how a second source of truth appears.""" +def test_a_record_that_names_its_run_keeps_it(): + """Producers stamp the run at rollout time; the sink never overrides a + ``run`` that is already there, which is how a second source of truth + would appear.""" + client = FakeTracesClient() + sink = make_sink(client) + + sink.write([{"id": "t1", "run": {"id": "theirs", "type": "eval"}}]) + + assert client.calls[0][0][0]["run"] == {"id": "theirs", "type": "eval"} + + +def test_an_episode_s_run_reaches_every_member_trace(): + """The traces service derives ``run_id`` from ``trace.run.id`` only and never + reads the envelope's ``run``, while verifiers records the run on the episode + and its ``Trace`` has no ``run`` field at all. Without this, every row of an + episode upload lands with an empty ``run_id`` (seen live, 2026-08-21).""" + client = FakeTracesClient() + sink = make_sink(client) + episode = make_episode("ep-1", [make_trace(trace_id="a"), make_trace(trace_id="b")]) + + sink.write([episode]) + + sent = client.calls[0][0][0] + assert sent["run"] == {"id": "run-1", "type": "eval"} + assert [member["run"] for member in sent["traces"]] == [sent["run"], sent["run"]] + assert [member["id"] for member in sent["traces"]] == ["a", "b"] + + +def test_a_member_that_names_its_own_run_is_left_alone(): + client = FakeTracesClient() + sink = make_sink(client) + theirs = {"id": "other-run", "type": "eval"} + episode = { + "id": "ep-1", + "run": {"id": "env-run", "type": "eval"}, + "traces": [{"id": "a"}, {"id": "b", "run": theirs}], + } + + sink.write([episode]) + + members = client.calls[0][0][0]["traces"] + assert members[0]["run"] == {"id": "env-run", "type": "eval"} + assert members[1]["run"] == theirs + assert episode["traces"][0] == {"id": "a"}, "the caller's members were not mutated" + + +def test_producer_objects_are_serialized_once_through_to_record(): + """The sink calls ``to_record()`` itself so members are reachable; the bytes + the transport sees are the same ones it would have produced.""" client = FakeTracesClient() sink = make_sink(client) trace = make_trace() sink.write([trace]) - assert client.calls[0][0][0] is trace + sent = client.calls[0][0][0] + assert sent == {**trace.to_record(), "run": {"id": "run-1", "type": "eval"}} def test_an_account_outside_the_beta_retires_the_sink_without_a_failure(caplog): diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index c9f00288c..d7c251708 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -76,6 +76,40 @@ def test_a_failed_sink_is_not_called_again(): worker.close() +def test_records_that_skip_a_sink_retired_by_error_are_counted_as_lost_to_it(): + """After the first failed batch the sink is off, but the producer keeps + logging; every later batch is just as lost to that sink as the first. Live + run 2026-08-21: five episodes, footer said "1 failed via traces".""" + broken = FakeSink("broken", fail_on_write=True) + worker = UploadWorker([broken]) + + for _ in range(5): + worker.submit([{"id": 1}]) + drain(worker) + + assert worker.failed_records == {"broken": 5} + worker.close() + + +def test_records_that_skip_a_sink_which_switched_itself_off_are_not_counted(): + """A sink that retires quietly (nowhere for the records to go, e.g. outside + the traces beta) lost nothing, and must not start a failure count.""" + + class QuietSink(FakeSink): + def write(self, records) -> None: + self.enabled = False # retires without raising, like service_not_enabled + + quiet = QuietSink("quiet") + worker = UploadWorker([quiet]) + + for _ in range(3): + worker.submit([{"id": 1}]) + drain(worker) + + assert worker.failed_records == {} + worker.close() + + def test_a_full_queue_drops_rather_than_blocking_the_producer(): """Stalling a training run to protect telemetry is the wrong trade.""" sink = BlockingSink() From 34faa4aaf4e45ed4b0fe387089585ddb698423fd Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Fri, 21 Aug 2026 18:20:12 -0700 Subject: [PATCH 26/39] =?UTF-8?q?refactor(runs):=20address=20review=20?= =?UTF-8?q?=E2=80=94=20Sink-typed=20worker,=20log=5Fepisodes,=20update=5Fs?= =?UTF-8?q?ummary?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review by kennethnym on #856: - UploadWorker takes Sequence[Sink] instead of List[Any]; the getattr fallbacks for `name` / `enabled` go with it (worker, run). - The traces sink imports ErrorCode / LineFormat / TracesClient at the top; prime-traces is a hard dependency and exceptions.py already imports it eagerly, so the lazy imports and the ImportError branch were dead. - Run.log_episodes() as the counterpart to log_traces(); both share one submit path, the sink still infers the line format. verifiers' call sites pass episodes, so they move to log_episodes. - Run.update_summary() merges outputs before finish() with the same non-finite filtering; metrics.from_episodes() returns a RunSummary TypedDict naming the three keys the dashboard reads. 169 tests; ty back at its 15-diagnostic baseline (FakeSink's start() now takes Mapping, matching the protocol). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01W4S4DNVRX5BNtsLva7TZz1 --- packages/prime-runs/README.md | 11 ++++--- packages/prime-runs/src/prime_runs/metrics.py | 16 +++++++--- packages/prime-runs/src/prime_runs/run.py | 32 ++++++++++++++++--- .../prime-runs/src/prime_runs/sinks/traces.py | 13 ++------ packages/prime-runs/src/prime_runs/worker.py | 27 +++++++--------- packages/prime-runs/tests/conftest.py | 6 ++-- packages/prime-runs/tests/test_run.py | 31 ++++++++++++++++++ packages/prime-runs/tests/test_traces_sink.py | 2 +- 8 files changed, 96 insertions(+), 42 deletions(-) diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index 1bdfbb127..5824f4117 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -21,7 +21,7 @@ run = pr.init( print(run.url) # https://app.primeintellect.ai/dashboard/evaluations/eval-... for episode in rollouts: # episodes carry run.id — see "Identity" below - run.log_traces([episode]) + run.log_episodes([episode]) # bare traces go through log_traces() run.finish(summary=pr.metrics.from_episodes(episodes)) ``` @@ -29,7 +29,10 @@ run.finish(summary=pr.metrics.from_episodes(episodes)) `init()` opens the run and returns a handle carrying its ID and dashboard URL. Records stream out on a background thread while the run proceeds, so the dashboard fills in as rollouts land. `finish()` closes the run out; a `with` -block does that for you, including when the body raises. +block does that for you, including when the body raises. Run-level outputs go +in `finish(summary=...)`, or incrementally through `update_summary()`; +`metrics.from_episodes()` returns them in the shape the dashboard reads +(`metrics.RunSummary`). ## Identity @@ -98,8 +101,8 @@ disables the run with a warning — it never silently writes somewhere else. ## From async code -`log_traces()` is a queue put, not a request, so it is safe to call from a -coroutine; it blocks only if the queue is full (up to 5s), which is the +`log_traces()` / `log_episodes()` are a queue put, not a request, so they are +safe to call from a coroutine; it blocks only if the queue is full (up to 5s), which is the backpressure. `init()` and `finish()` do network I/O — wrap them in `asyncio.to_thread` if a stall there would matter. diff --git a/packages/prime-runs/src/prime_runs/metrics.py b/packages/prime-runs/src/prime_runs/metrics.py index 7b31d9928..ee0f12670 100644 --- a/packages/prime-runs/src/prime_runs/metrics.py +++ b/packages/prime-runs/src/prime_runs/metrics.py @@ -2,12 +2,20 @@ so migrated runs keep identical dashboard numbers. Opt-in: pass the result to ``run.finish(summary=...)``. Duck-typed; no producer package is imported.""" -from typing import Any, Dict, Optional, Sequence +from typing import Any, Dict, Optional, Sequence, TypedDict -def from_episodes( - episodes: Sequence[Any], traces: Optional[Sequence[Any]] = None -) -> Dict[str, Any]: +class RunSummary(TypedDict): + """The run-level aggregates the eval dashboard reads. A summary may carry + more than this — ``finish(summary=...)`` stores whatever it is given — but + these three are what the dashboard renders.""" + + avg_reward: float + avg_metrics: Dict[str, float] + avg_error: float + + +def from_episodes(episodes: Sequence[Any], traces: Optional[Sequence[Any]] = None) -> RunSummary: """Run-level aggregates in the shape the eval dashboard reads. Rewards and metrics aggregate over the trainable traces only — fixed agents diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 13bfc20f9..758d0d3c1 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -94,7 +94,7 @@ def __init__( sink.start(handle.id, context) except Exception as exc: # noqa: BLE001 - a bad sink is not a bad run sink.enabled = False - self._note(f"starting sink {getattr(sink, 'name', sink)}", exc) + self._note(f"starting sink {sink.name}", exc) if self._on_error == "raise": # The backend may already have created a remote run; close # it out before the failure reaches the caller. finish() @@ -154,17 +154,39 @@ def __repr__(self) -> str: # ------------------------------------------------------------------- log - def log_traces(self, records: Iterable[Any]) -> None: - """Hand traces or episodes to the sinks. Returns immediately. + def log_traces(self, traces: Iterable[Any]) -> None: + """Hand bare traces to the sinks. Returns immediately. - Accepts verifiers ``Trace``/``Episode`` objects or plain JSON mappings. - Call this as rollouts complete; nothing is buffered until the end. + Accepts verifiers ``Trace`` objects or plain JSON mappings. Call this + as rollouts complete; nothing is buffered until the end. A rollout + that is an episode (a group of traces) goes through :meth:`log_episodes`. """ self._require_live("log_traces") + self._submit(traces) + + def log_episodes(self, episodes: Iterable[Any]) -> None: + """Hand episodes — grouped traces — to the sinks. Returns immediately. + + Accepts verifiers ``Episode`` objects or plain JSON mappings with a + ``traces`` list. The episode's ``run`` reaches every member trace. + """ + self._require_live("log_episodes") + self._submit(episodes) + + def _submit(self, records: Iterable[Any]) -> None: batch = list(records) if batch: self._worker.submit(batch) + def update_summary(self, values: Mapping[str, Any]) -> None: + """Merge run-level outputs into :attr:`summary` ahead of :meth:`finish`. + + Non-finite numbers are dropped here, as they are for + ``finish(summary=...)``; writing to ``summary`` directly skips that. + """ + self._require_live("update_summary") + self.summary.update(_clean_metrics(values)) + def flush(self, timeout: Optional[float] = 30.0) -> bool: """Block until queued records have been written. Under ``on_error="raise"`` this is the first place an upload failure surfaces.""" diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py index faf66b4cb..79b2831f9 100644 --- a/packages/prime-runs/src/prime_runs/sinks/traces.py +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -10,6 +10,8 @@ import logging from typing import Any, Dict, Mapping, Optional, Sequence +from prime_traces import ErrorCode, LineFormat, TracesClient + from .. import _fork from ..exceptions import ForbiddenError from .base import Sink, is_episode, stamp_run, to_mapping @@ -68,11 +70,6 @@ def _ensure_client(self) -> bool: exc = RuntimeError("an injected traces client cannot be reused after a fork") self._disable(str(exc)) raise exc - try: - from prime_traces import TracesClient - except ImportError as exc: # pragma: no cover - dependency is declared - self._disable(f"prime-traces is not installed ({exc})") - raise try: self._client = TracesClient(**self._client_kwargs) except Exception as exc: # noqa: BLE001 - the run applies its error policy @@ -90,8 +87,6 @@ def write(self, records: Sequence[Any]) -> None: if not self.enabled or not records or not self._ensure_client(): return - from prime_traces import LineFormat - # The same bytes under a different format are rejected as a conflict, # so infer from the first record, which is stable within a batch. line_format = LineFormat.EPISODE if is_episode(records[0]) else LineFormat.TRACE @@ -140,7 +135,7 @@ def flush(self) -> None: def close(self) -> None: client = self._client self._client = None - if client is not None and hasattr(client, "close"): + if client is not None: try: client.close() except Exception as exc: # noqa: BLE001 - teardown must not raise @@ -162,6 +157,4 @@ def _retire_quietly(self, reason: str) -> None: def _is_not_enabled(exc: ForbiddenError) -> bool: """``service_not_enabled``: the account is outside the private beta. The other 403, ``forbidden``, means the token lacks the ``traces`` scope.""" - from prime_traces import ErrorCode - return exc.code == ErrorCode.SERVICE_NOT_ENABLED.value diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index e8137f4cf..3391e66a9 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -15,10 +15,11 @@ import threading import time from dataclasses import dataclass, field -from typing import Any, Callable, List, Optional, Sequence +from typing import Any, Callable, Optional, Sequence from . import _fork from .exceptions import is_transient +from .sinks.base import Sink logger = logging.getLogger(__name__) @@ -35,10 +36,6 @@ class _Flush: event: threading.Event = field(default_factory=threading.Event) -def _sink_name(sink: Any) -> str: - return getattr(sink, "name", type(sink).__name__) - - def _deadline(timeout: Optional[float]) -> Optional[float]: if timeout is None: return None @@ -56,7 +53,7 @@ class UploadWorker: def __init__( self, - sinks: List[Any], + sinks: Sequence[Sink], *, max_queue_size: int = DEFAULT_QUEUE_SIZE, put_timeout: float = DEFAULT_PUT_TIMEOUT, @@ -112,32 +109,32 @@ def _run(self) -> None: def _dispatch(self, records: Sequence[Any]) -> None: for sink in self.sinks: - if not getattr(sink, "enabled", True): - name = _sink_name(sink) - if name in self._retired: - self.failed_records[name] = self.failed_records.get(name, 0) + len(records) + if not sink.enabled: + if sink.name in self._retired: + count = self.failed_records.get(sink.name, 0) + self.failed_records[sink.name] = count + len(records) continue try: sink.write(records) except Exception as exc: # noqa: BLE001 - one sink failing must not stop the others self._fail_sink(sink, exc, dropped=len(records)) else: - self._transient_failures.pop(_sink_name(sink), None) + self._transient_failures.pop(sink.name, None) def _flush_sinks(self) -> None: for sink in self.sinks: - if not getattr(sink, "enabled", True): + if not sink.enabled: continue try: sink.flush() except Exception as exc: # noqa: BLE001 self._fail_sink(sink, exc) - def _fail_sink(self, sink: Any, exc: Exception, *, dropped: int = 0) -> None: + def _fail_sink(self, sink: Sink, exc: Exception, *, dropped: int = 0) -> None: """The batch is gone (the transports already retried). Decide whether the sink is too: permanent failures retire it at once, transient ones after ``TRANSIENT_FAILURE_LIMIT`` consecutive strikes.""" - name = _sink_name(sink) + name = sink.name if dropped: self.failed_records[name] = self.failed_records.get(name, 0) + dropped @@ -246,7 +243,7 @@ def close(self, timeout: Optional[float] = 30.0) -> None: try: sink.close() except Exception as exc: # noqa: BLE001 - teardown must not raise - logger.debug("Error closing sink %s: %s", getattr(sink, "name", sink), exc) + logger.debug("Error closing sink %s: %s", sink.name, exc) # ------------------------------------------------------------------- fork diff --git a/packages/prime-runs/tests/conftest.py b/packages/prime-runs/tests/conftest.py index 7ede20e6c..9c7166c9c 100644 --- a/packages/prime-runs/tests/conftest.py +++ b/packages/prime-runs/tests/conftest.py @@ -1,7 +1,7 @@ """Shared fixtures. Every test is hermetic: no network, no real ~/.prime.""" from pathlib import Path -from typing import Any, Callable, Dict, List +from typing import Any, Callable, Dict, List, Mapping, Sequence import httpx import pytest @@ -123,10 +123,10 @@ def __init__(self, name: str = "fake", fail_on_write: bool = False) -> None: self.flushes = 0 self.closed = False - def start(self, run_id: str, context: Dict[str, str]) -> None: + def start(self, run_id: str, context: Mapping[str, str]) -> None: self.started.append((run_id, dict(context))) - def write(self, records) -> None: + def write(self, records: Sequence[Any]) -> None: if self.fail_on_write: raise RuntimeError("sink is broken") self.batches.append(list(records)) diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index 0e3c3cc0d..9eb6df6d3 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -99,6 +99,37 @@ def test_an_empty_batch_is_not_sent(): run.finish() +def test_episodes_take_the_same_path_as_traces(): + sink = FakeSink() + run = make_run(sinks=[sink]) + + run.log_episodes([{"id": "ep-1", "traces": [{"id": "t1"}]}]) + run.flush() + + assert sink.batches == [[{"id": "ep-1", "traces": [{"id": "t1"}]}]] + run.finish() + + +def test_summary_can_be_built_up_before_finish(): + backend = FakeBackend() + run = make_run(backend) + + run.update_summary({"avg_reward": 0.5, "loss": float("nan")}) + run.update_summary({"avg_error": 0.0}) + run.finish(summary={"avg_reward": 0.75}) + + assert run.summary == {"avg_reward": 0.75, "avg_error": 0.0} + assert backend.finalized[0]["summary"] == {"avg_reward": 0.75, "avg_error": 0.0} + + +def test_a_finished_run_refuses_more_summary(): + run = make_run() + run.finish() + + with pytest.raises(RunFinishedError): + run.update_summary({"late": 1.0}) + + def test_non_finite_summary_values_are_dropped_rather_than_failing_the_request(): """A diverged loss serializes as bare ``NaN``, which strict JSON rejects — the whole request fails on a payload nobody can inspect.""" diff --git a/packages/prime-runs/tests/test_traces_sink.py b/packages/prime-runs/tests/test_traces_sink.py index e2cca6a81..d3d15c7c1 100644 --- a/packages/prime-runs/tests/test_traces_sink.py +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -214,7 +214,7 @@ def test_a_missing_traces_client_disables_the_sink_and_reports_the_failure(monke def explode(**kwargs): raise RuntimeError("no credentials") - monkeypatch.setattr("prime_traces.TracesClient", explode) + monkeypatch.setattr("prime_runs.sinks.traces.TracesClient", explode) sink = TracesSink() with caplog.at_level("WARNING"): From defb9c67c4158483f10a2aeefe70cb8f46fae9cd Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Sat, 22 Aug 2026 08:37:42 -0700 Subject: [PATCH 27/39] fix(runs): a record the samples sink cannot project is skipped, not fatal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit log_episodes() documented plain JSON mappings, but EvalSamplesSink raised TypeError on any mapping without sample_id (and on bare trace objects), and the worker treats TypeError as permanent: the sink was retired for the rest of the run. Outside the traces beta that left the viewer with no samples at all. The v0 projection is attribute-based and the dict-based one was cut deliberately (#3), so such records genuinely have no row in the sample table. The sink now skips them — warned once per run, counted on `skipped` — and keeps storing what it can project. The log_traces / log_episodes docstrings say which records reach which table. Reported by Bugbot on #856. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01W4S4DNVRX5BNtsLva7TZz1 --- packages/prime-runs/src/prime_runs/run.py | 11 ++++++--- .../src/prime_runs/sinks/samples.py | 23 +++++++++++++++---- .../prime-runs/tests/test_samples_sink.py | 23 ++++++++++++++----- 3 files changed, 44 insertions(+), 13 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 758d0d3c1..201bed899 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -157,9 +157,11 @@ def __repr__(self) -> str: def log_traces(self, traces: Iterable[Any]) -> None: """Hand bare traces to the sinks. Returns immediately. - Accepts verifiers ``Trace`` objects or plain JSON mappings. Call this - as rollouts complete; nothing is buffered until the end. A rollout - that is an episode (a group of traces) goes through :meth:`log_episodes`. + Accepts verifiers ``Trace`` objects or plain JSON mappings. Both reach + Prime Traces; the v0 sample table is projected from *episodes*, so a + bare trace has no row there. A rollout that is an episode (a group of + traces) goes through :meth:`log_episodes`. Call this as rollouts + complete; nothing is buffered until the end. """ self._require_live("log_traces") self._submit(traces) @@ -169,6 +171,9 @@ def log_episodes(self, episodes: Iterable[Any]) -> None: Accepts verifiers ``Episode`` objects or plain JSON mappings with a ``traces`` list. The episode's ``run`` reaches every member trace. + Both reach Prime Traces; the v0 sample table (what today's viewer + reads) is projected from episode *objects* only — a JSON episode has + no row there, which the samples sink warns about once. """ self._require_live("log_episodes") self._submit(episodes) diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py index d99b0d719..9b75c12a2 100644 --- a/packages/prime-runs/src/prime_runs/sinks/samples.py +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -31,6 +31,9 @@ def __init__(self, client: PlatformClient) -> None: # same way a one-shot upload did. self._rollout_numbers: Dict[Any, int] = {} self.samples_written = 0 + #: Records this sink could not project and therefore did not store. + #: The traces sink takes them; the v0 table simply has no row for them. + self.skipped = 0 def start(self, run_id: str, context: Mapping[str, str]) -> None: self._run_id = run_id @@ -55,18 +58,30 @@ def write(self, records: Sequence[Any]) -> None: def _to_samples(self, records: Sequence[Any]) -> List[Dict[str, Any]]: """Episode objects are projected; v0 sample dicts (``sample_id``) pass - through. Anything else fails loudly rather than vanishing.""" + through. Anything else — a JSON episode, a bare trace — has no v0 + projection (it is attribute-based, see :mod:`prime_runs.projection`) + and is skipped: warned once, counted, and left to the traces sink. + Raising here would retire this sink for the rest of the run over one + record shape, which is worse than one missing row.""" samples: List[Dict[str, Any]] = [] + skipped = 0 for record in records: if isinstance(record, Mapping) and "sample_id" in record: samples.append(dict(record)) elif not isinstance(record, Mapping) and is_episode(record): samples.extend(build_samples([record], self._rollout_numbers)) else: - raise TypeError( - f"EvalSamplesSink cannot project {type(record).__name__}; expected an " - "episode object or a v0 sample dict" + skipped += 1 + if skipped: + if not self.skipped: + logger.warning( + "The v0 sample table is projected from episode objects; %d record(s) " + "in this batch (%s) have no projection and reach Prime Traces only. " + "Further skips are counted, not logged.", + skipped, + type(records[0]).__name__, ) + self.skipped += skipped return samples def flush(self) -> None: diff --git a/packages/prime-runs/tests/test_samples_sink.py b/packages/prime-runs/tests/test_samples_sink.py index 6ea84556a..cb23ad86d 100644 --- a/packages/prime-runs/tests/test_samples_sink.py +++ b/packages/prime-runs/tests/test_samples_sink.py @@ -1,6 +1,5 @@ """The legacy sample sink that keeps today's viewer working.""" -import pytest from _fakes import make_episode, make_trace from conftest import RecordingHandler @@ -44,15 +43,27 @@ def test_a_producer_that_already_speaks_v0_is_passed_through(make_platform_clien assert body["samples"] == [{"sample_id": "s1", "reward": 1.0}] -def test_records_this_sink_cannot_project_fail_explicitly(make_platform_client, eval_routes): +def test_records_this_sink_cannot_project_are_skipped_not_fatal( + make_platform_client, eval_routes, caplog +): + """A JSON episode or a bare trace has no v0 projection. Raising would have + the worker retire the sink for the run — outside the traces beta that is + an empty viewer — so the record is skipped, warned once, and counted, + while the episode object in the same batch is still stored.""" sink, handler = make_sink(make_platform_client, eval_routes) - with pytest.raises(TypeError, match="cannot project"): - sink.write([{"unrelated": True}]) - with pytest.raises(TypeError, match="cannot project"): + with caplog.at_level("WARNING"): + sink.write( + [{"id": "ep-json", "traces": [{"id": "t"}]}, make_episode("ep-1", [make_trace()])] + ) sink.write([make_trace()]) - assert handler.requests == [] + body = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")[0] + assert [row["sample_id"] for row in body["samples"]] == ["ep-1"] + assert len(handler.bodies_for("/api/v1/evaluations/eval-abc/samples")) == 1 + assert sink.skipped == 2 + assert sink.enabled is True + assert sum("no projection" in r.getMessage() for r in caplog.records) == 1 def test_an_empty_batch_makes_no_request(make_platform_client, eval_routes): From f44365e07c18a17ef8f4cd214ec23e1c7f80834d Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 27 Aug 2026 10:57:27 -0700 Subject: [PATCH 28/39] fix(runs): harden release source and run classification --- .github/workflows/release-runs.yml | 6 ++++++ packages/prime-runs/src/prime_runs/run.py | 2 +- packages/prime-runs/tests/test_run.py | 1 + packages/prime-runs/tests/test_traces_sink.py | 6 +++--- 4 files changed, 11 insertions(+), 4 deletions(-) diff --git a/.github/workflows/release-runs.yml b/.github/workflows/release-runs.yml index da263298a..05f989796 100644 --- a/.github/workflows/release-runs.yml +++ b/.github/workflows/release-runs.yml @@ -46,6 +46,12 @@ jobs: echo "Tag $TAG does not exist" fi + - name: Refuse untagged releases outside main + if: github.event_name == 'workflow_dispatch' && steps.check_tag.outputs.exists != 'true' && github.ref != 'refs/heads/main' + run: | + echo "::error title=Invalid release ref::New prime-runs releases must be dispatched from main, not ${{ github.ref }}." + exit 1 + # A forced release must build the tagged source, not whatever the # dispatch ref currently points at. - name: Check out existing tag diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 201bed899..4939d09ce 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -477,7 +477,7 @@ def _resolve_mode(mode: Optional[Mode], *, api_key: str) -> Mode: def _sink_context(spec: RunSpec) -> Dict[str, str]: """Upload-scoped provenance. Not the join key — that is ``run.id`` inside the trace document.""" - context = {"source": "prime-runs", "run_kind": RUN_KIND} + context = {"source": "prime-runs", "run_type": RUN_KIND} if spec.framework: context["framework"] = spec.framework if spec.model: diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index 9eb6df6d3..f41fb96af 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -71,6 +71,7 @@ def test_sinks_are_started_with_the_run_id_and_provenance(): assert run_id == "run-1" # Provenance only — the join key is run.id inside the trace document. assert context["source"] == "prime-runs" + assert context["run_type"] == "eval" assert context["framework"] == "verifiers" assert "evaluation_id" not in context run.finish() diff --git a/packages/prime-runs/tests/test_traces_sink.py b/packages/prime-runs/tests/test_traces_sink.py index d3d15c7c1..48c59fc3e 100644 --- a/packages/prime-runs/tests/test_traces_sink.py +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -26,7 +26,7 @@ def close(self) -> None: def make_sink(client=None, **kwargs) -> TracesSink: sink = TracesSink(client=client or FakeTracesClient(), **kwargs) - sink.start("run-1", {"source": "prime-runs", "run_kind": "eval", "framework": "verifiers"}) + sink.start("run-1", {"source": "prime-runs", "run_type": "eval", "framework": "verifiers"}) return sink @@ -41,7 +41,7 @@ def test_records_go_out_with_provenance_but_not_the_join_key(): _, kwargs = client.calls[0] assert kwargs["context"] == { "source": "prime-runs", - "run_kind": "eval", + "run_type": "eval", "framework": "verifiers", } assert "run_id" not in kwargs["context"] @@ -219,6 +219,6 @@ def explode(**kwargs): with caplog.at_level("WARNING"): with pytest.raises(RuntimeError, match="no credentials"): - sink.start("run-1", {"run_kind": "eval"}) + sink.start("run-1", {"run_type": "eval"}) assert sink.enabled is False From bb6e2d2bedac1d1a349782547a1466e573859f78 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 27 Aug 2026 10:58:58 -0700 Subject: [PATCH 29/39] fix(runs): report skipped sample record type --- packages/prime-runs/src/prime_runs/sinks/samples.py | 5 ++++- packages/prime-runs/tests/test_samples_sink.py | 13 +++++++++++-- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py index 9b75c12a2..b2e82828a 100644 --- a/packages/prime-runs/src/prime_runs/sinks/samples.py +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -65,6 +65,7 @@ def _to_samples(self, records: Sequence[Any]) -> List[Dict[str, Any]]: record shape, which is worse than one missing row.""" samples: List[Dict[str, Any]] = [] skipped = 0 + first_skipped_type: Optional[str] = None for record in records: if isinstance(record, Mapping) and "sample_id" in record: samples.append(dict(record)) @@ -72,6 +73,8 @@ def _to_samples(self, records: Sequence[Any]) -> List[Dict[str, Any]]: samples.extend(build_samples([record], self._rollout_numbers)) else: skipped += 1 + if first_skipped_type is None: + first_skipped_type = type(record).__name__ if skipped: if not self.skipped: logger.warning( @@ -79,7 +82,7 @@ def _to_samples(self, records: Sequence[Any]) -> List[Dict[str, Any]]: "in this batch (%s) have no projection and reach Prime Traces only. " "Further skips are counted, not logged.", skipped, - type(records[0]).__name__, + first_skipped_type, ) self.skipped += skipped return samples diff --git a/packages/prime-runs/tests/test_samples_sink.py b/packages/prime-runs/tests/test_samples_sink.py index cb23ad86d..e356e0d4c 100644 --- a/packages/prime-runs/tests/test_samples_sink.py +++ b/packages/prime-runs/tests/test_samples_sink.py @@ -54,7 +54,10 @@ def test_records_this_sink_cannot_project_are_skipped_not_fatal( with caplog.at_level("WARNING"): sink.write( - [{"id": "ep-json", "traces": [{"id": "t"}]}, make_episode("ep-1", [make_trace()])] + [ + make_episode("ep-1", [make_trace()]), + {"id": "ep-json", "traces": [{"id": "t"}]}, + ] ) sink.write([make_trace()]) @@ -63,7 +66,13 @@ def test_records_this_sink_cannot_project_are_skipped_not_fatal( assert len(handler.bodies_for("/api/v1/evaluations/eval-abc/samples")) == 1 assert sink.skipped == 2 assert sink.enabled is True - assert sum("no projection" in r.getMessage() for r in caplog.records) == 1 + warnings = [ + record.getMessage() + for record in caplog.records + if "no projection" in record.getMessage() + ] + assert len(warnings) == 1 + assert "(dict)" in warnings[0] def test_an_empty_batch_makes_no_request(make_platform_client, eval_routes): From 9d2a9dbde38b8afd5c463b0857618efb45d4f9eb Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 27 Aug 2026 11:05:48 -0700 Subject: [PATCH 30/39] fix(runs): close client when initialization fails --- packages/prime-runs/src/prime_runs/run.py | 14 ++++++++++++- packages/prime-runs/tests/test_init.py | 25 +++++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 4939d09ce..faf9ea0eb 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -437,7 +437,19 @@ def init( ) client = PlatformClient(api_key=api_key, base_url=base_url, timeout=DEFAULT_TIMEOUT) backend = EvalsBackend(client, frontend_url=settings.frontend_url, team_id=team_id) - handle = backend.create(spec) + try: + handle = backend.create(spec) + except BaseException: + # Ownership has not reached a Run yet, so nothing else can release + # the connection pool when environment resolution or creation fails. + try: + backend.close() + except Exception as close_error: # noqa: BLE001 - preserve the create failure + logger.debug( + "Error closing the platform client after run creation failed: %s", + close_error, + ) + raise # Both transports run during the transition: traces is the system of # record, the sample table is what today's viewer reads. sinks = [TracesSink(api_key=api_key, team_id=team_id), EvalSamplesSink(client)] diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index c42165279..54c9a087e 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -109,6 +109,31 @@ def test_an_online_run_returns_the_platforms_id_and_viewer_url(online): run.finish() +def test_a_failed_create_closes_the_platform_client(monkeypatch): + class FailingClient: + def __init__(self): + self.closed = False + + def post(self, *args, **kwargs): + raise pr.APIError("create failed") + + def close(self): + self.closed = True + + client = FailingClient() + monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: client) + + with pytest.raises(pr.APIError, match="create failed"): + pr.init( + name="test-run", + environments=[{"id": "env-123"}], + api_key="test-key", + mode="online", + ) + + assert client.closed is True + + def test_an_online_run_has_both_transports_by_default( monkeypatch, make_platform_client, eval_routes ): From 3ddaedd81326cca4c166224d446eda4411b17b5d Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 27 Aug 2026 11:48:27 -0700 Subject: [PATCH 31/39] fix(runs): keep sinks alive after record rejections --- .../prime-runs/src/prime_runs/exceptions.py | 29 +++++++++++++ packages/prime-runs/src/prime_runs/worker.py | 28 ++++++++++--- packages/prime-runs/tests/test_worker.py | 41 +++++++++++++++++++ 3 files changed, 92 insertions(+), 6 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/exceptions.py b/packages/prime-runs/src/prime_runs/exceptions.py index adc5df8b9..4649b673b 100644 --- a/packages/prime-runs/src/prime_runs/exceptions.py +++ b/packages/prime-runs/src/prime_runs/exceptions.py @@ -15,11 +15,14 @@ APIError, APITimeoutError, ForbiddenError, + LineFormatConflictError, NotFoundError, PaymentRequiredError, RetryableAPIError, + TraceTooLargeError, TransportError, UnauthorizedError, + ValidationRejectedError, ) __all__ = [ @@ -28,13 +31,17 @@ "ConfigurationError", "EnvironmentResolutionError", "ForbiddenError", + "LineFormatConflictError", "NotFoundError", "PaymentRequiredError", "PrimeRunsError", "RetryableAPIError", "RunFinishedError", + "TraceTooLargeError", "TransportError", "UnauthorizedError", + "ValidationRejectedError", + "is_record_rejection", "is_transient", ] @@ -56,6 +63,28 @@ class RunFinishedError(PrimeRunsError): """A finished run was written to again: a producer bug.""" +def is_record_rejection(exc: BaseException) -> bool: + """Whether a failed write is specific to the submitted record batch. + + A later batch can still succeed after serialization, size, or validation + failures, so these errors must not retire an otherwise healthy sink. + """ + if isinstance( + exc, + ( + TypeError, + ValueError, + TraceTooLargeError, + ValidationRejectedError, + LineFormatConflictError, + ), + ): + return True + # The legacy samples API returns framework-generated 413/422 responses, + # which the shared traces mapper intentionally leaves as plain APIError. + return isinstance(exc, APIError) and exc.status_code in (413, 422) + + def is_transient(exc: BaseException) -> bool: """Whether a failure is about this moment (retry later) rather than this run (stop). Decides whether a sink is retired.""" diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index 3391e66a9..fb14f2bd1 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -6,8 +6,9 @@ Fork safety: a forked child inherits the queue's memory but not the thread. The child starts over empty; the queued records belong to the parent. -Containment: a sink that raises is given a few consecutive transient strikes, -then disabled for the rest of the run. The thread never dies on one bad batch. +Containment: record-specific failures drop only their batch. A sink-wide +failure disables the sink, while transient failures get a few consecutive +strikes first. The thread never dies on one bad batch. """ import logging @@ -18,7 +19,7 @@ from typing import Any, Callable, Optional, Sequence from . import _fork -from .exceptions import is_transient +from .exceptions import is_record_rejection, is_transient from .sinks.base import Sink logger = logging.getLogger(__name__) @@ -131,13 +132,28 @@ def _flush_sinks(self) -> None: self._fail_sink(sink, exc) def _fail_sink(self, sink: Sink, exc: Exception, *, dropped: int = 0) -> None: - """The batch is gone (the transports already retried). Decide whether - the sink is too: permanent failures retire it at once, transient ones - after ``TRANSIENT_FAILURE_LIMIT`` consecutive strikes.""" + """Account for a failed operation and decide whether to retire its sink. + + Record-specific failures drop only the current batch. Sink-wide + permanent failures retire immediately; transient failures retire after + ``TRANSIENT_FAILURE_LIMIT`` consecutive strikes. + """ name = sink.name if dropped: self.failed_records[name] = self.failed_records.get(name, 0) + dropped + if dropped and is_record_rejection(exc): + logger.warning( + "Sink %s rejected a batch of %d record(s) (%s: %s); " + "still enabled for later records.", + name, + dropped, + type(exc).__name__, + exc, + ) + self._notify(name, exc) + return + if is_transient(exc): strikes = self._transient_failures.get(name, 0) + 1 self._transient_failures[name] = strikes diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index d7c251708..1671e0165 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -3,8 +3,10 @@ import queue import threading +import pytest from conftest import FakeSink +from prime_runs.exceptions import APIError, ValidationRejectedError from prime_runs.worker import UploadWorker @@ -304,6 +306,45 @@ def write(self, records) -> None: worker.close() +@pytest.mark.parametrize( + "error", + [ + ValueError("Out of range float values are not JSON compliant"), + ValidationRejectedError("invalid record", status_code=400), + APIError("unprocessable record", status_code=422), + ], + ids=["local-encoding", "traces-validation", "samples-validation"], +) +def test_a_record_rejection_drops_only_its_batch(error): + """Malformed content must not prevent later valid records from uploading.""" + + class RejectOnceSink(FakeSink): + def __init__(self) -> None: + super().__init__("reject-once") + self.calls = 0 + + def write(self, records) -> None: + self.calls += 1 + if self.calls == 1: + raise error + super().write(records) + + sink = RejectOnceSink() + reported = [] + worker = UploadWorker([sink], on_error=lambda name, exc: reported.append((name, exc))) + + worker.submit([{"id": "bad"}]) + worker.submit([{"id": "good"}]) + drain(worker) + + assert sink.enabled is True + assert sink.calls == 2 + assert sink.batches == [[{"id": "good"}]] + assert worker.failed_records == {"reject-once": 1} + assert reported == [("reject-once", error)] + worker.close() + + def test_a_sustained_outage_eventually_retires_the_sink(): """A blip is forgiven; hours of re-attempting every batch is not useful.""" from prime_runs.exceptions import TransportError From 57fe33f2dea4595029bcd6ab90a1ab7fb710aa49 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 27 Aug 2026 12:56:10 -0700 Subject: [PATCH 32/39] chore(runs): require prime-traces>=0.0.3 0.0.2 resolves the traces service URL to the platform API, which does not route /api/v1/traces, so every traces upload 404s and the traces sink retires with a warning on the first batch. 0.0.3 (#871, released in #872) defaults to https://prime-traces.pintel.dev. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01CyW4j57juAg5cheTX8PmUR --- packages/prime-runs/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/prime-runs/pyproject.toml b/packages/prime-runs/pyproject.toml index 3f895ac21..68682a320 100644 --- a/packages/prime-runs/pyproject.toml +++ b/packages/prime-runs/pyproject.toml @@ -20,7 +20,7 @@ authors = [ # `tenacity`. Nothing is added to what verifiers and prime-rl already carry. dependencies = [ "httpx>=0.25.0", - "prime-traces>=0.0.2", + "prime-traces>=0.0.3", ] keywords = ["evals", "evaluations", "training", "runs", "traces", "observability"] classifiers = [ From 9ca5a9093cceae6b4272a27a92b03d8b40903da1 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 27 Aug 2026 15:13:01 -0700 Subject: [PATCH 33/39] fix(runs): close upload lifecycle races --- packages/prime-runs/src/prime_runs/run.py | 24 ++++--- .../prime-runs/src/prime_runs/sinks/base.py | 16 +++++ .../src/prime_runs/sinks/samples.py | 71 ++++++++++++++----- packages/prime-runs/src/prime_runs/worker.py | 67 +++++++++-------- packages/prime-runs/tests/test_run.py | 50 +++++++++++++ .../prime-runs/tests/test_samples_sink.py | 61 ++++++++++++++++ packages/prime-runs/tests/test_worker.py | 6 +- 7 files changed, 238 insertions(+), 57 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index faf9ea0eb..3c2e73556 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -163,8 +163,7 @@ def log_traces(self, traces: Iterable[Any]) -> None: traces) goes through :meth:`log_episodes`. Call this as rollouts complete; nothing is buffered until the end. """ - self._require_live("log_traces") - self._submit(traces) + self._submit("log_traces", traces) def log_episodes(self, episodes: Iterable[Any]) -> None: """Hand episodes — grouped traces — to the sinks. Returns immediately. @@ -175,13 +174,17 @@ def log_episodes(self, episodes: Iterable[Any]) -> None: reads) is projected from episode *objects* only — a JSON episode has no row there, which the samples sink warns about once. """ - self._require_live("log_episodes") - self._submit(episodes) + self._submit("log_episodes", episodes) - def _submit(self, records: Iterable[Any]) -> None: - batch = list(records) - if batch: - self._worker.submit(batch) + def _submit(self, operation: str, records: Iterable[Any]) -> None: + # Finish holds the same lock through worker teardown. A log call either + # queues its batch first or observes that the run is already closing; + # it can never restart the worker between those two states. + with self._finish_lock: + self._require_live(operation) + batch = list(records) + if batch: + self._worker.submit(batch) def update_summary(self, values: Mapping[str, Any]) -> None: """Merge run-level outputs into :attr:`summary` ahead of :meth:`finish`. @@ -189,8 +192,9 @@ def update_summary(self, values: Mapping[str, Any]) -> None: Non-finite numbers are dropped here, as they are for ``finish(summary=...)``; writing to ``summary`` directly skips that. """ - self._require_live("update_summary") - self.summary.update(_clean_metrics(values)) + with self._finish_lock: + self._require_live("update_summary") + self.summary.update(_clean_metrics(values)) def flush(self, timeout: Optional[float] = 30.0) -> bool: """Block until queued records have been written. Under diff --git a/packages/prime-runs/src/prime_runs/sinks/base.py b/packages/prime-runs/src/prime_runs/sinks/base.py index fee20ccbd..3df8ba903 100644 --- a/packages/prime-runs/src/prime_runs/sinks/base.py +++ b/packages/prime-runs/src/prime_runs/sinks/base.py @@ -14,6 +14,22 @@ from ..models import RUN_KIND +class SinkWriteError(Exception): + """A sink stored only part of a submitted batch. + + ``cause`` drives the worker's retry/retirement policy; ``failed_records`` + lets it account for only the input records that were not fully stored. + """ + + def __init__(self, cause: Exception, *, failed_records: int) -> None: + self.cause = cause + self.failed_records = failed_records + super().__init__( + f"{failed_records} record(s) were not fully stored: " + f"{type(cause).__name__}: {cause}" + ) + + class Sink(Protocol): """A destination for run records.""" diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py index b2e82828a..a0520e5f4 100644 --- a/packages/prime-runs/src/prime_runs/sinks/samples.py +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -9,11 +9,12 @@ """ import logging -from typing import Any, Dict, List, Mapping, Optional, Sequence +from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple from .._http import UPLOAD_TIMEOUT, PlatformClient, encode_json +from ..exceptions import is_record_rejection, is_transient from ..projection import batch_samples, build_samples -from .base import Sink, is_episode +from .base import Sink, SinkWriteError, is_episode logger = logging.getLogger(__name__) @@ -44,33 +45,69 @@ def write(self, records: Sequence[Any]) -> None: if self._run_id is None: raise RuntimeError("EvalSamplesSink.write called before start()") - samples = self._to_samples(records) + samples, owners = self._to_samples(records) if not samples: return - for batch in batch_samples(samples): - self._client.post( - f"/evaluations/{self._run_id}/samples", - content=encode_json({"samples": batch}), - timeout=UPLOAD_TIMEOUT, - idempotent=False, # appends; a lost response must not duplicate rows - ) - self.samples_written += len(batch) - - def _to_samples(self, records: Sequence[Any]) -> List[Dict[str, Any]]: + batches = batch_samples(samples) + failed_owners: Set[int] = set() + reported_error: Optional[Exception] = None + offset = 0 + for batch_index, batch in enumerate(batches): + next_offset = offset + len(batch) + batch_owners = owners[offset:next_offset] + offset = next_offset + try: + self._client.post( + f"/evaluations/{self._run_id}/samples", + content=encode_json({"samples": batch}), + timeout=UPLOAD_TIMEOUT, + idempotent=False, # appends; a lost response must not duplicate rows + ) + except Exception as exc: + failed_owners.update(batch_owners) + recoverable = is_record_rejection(exc) or is_transient(exc) + if not recoverable: + # The same sink-wide failure will reject everything after + # this request, so do not send more traffic just to prove it. + for remaining_batch in batches[batch_index + 1 :]: + failed_owners.update(owners[offset : offset + len(remaining_batch)]) + offset += len(remaining_batch) + reported_error = exc + break + # Prefer a transient error for worker strike accounting when a + # write encountered both transient and record-local failures. + if reported_error is None or is_transient(exc): + reported_error = exc + else: + self.samples_written += len(batch) + + if reported_error is not None: + raise SinkWriteError( + reported_error, failed_records=len(failed_owners) + ) from reported_error + + def _to_samples( + self, records: Sequence[Any] + ) -> Tuple[List[Dict[str, Any]], List[int]]: """Episode objects are projected; v0 sample dicts (``sample_id``) pass through. Anything else — a JSON episode, a bare trace — has no v0 projection (it is attribute-based, see :mod:`prime_runs.projection`) and is skipped: warned once, counted, and left to the traces sink. Raising here would retire this sink for the rest of the run over one - record shape, which is worse than one missing row.""" + record shape, which is worse than one missing row. The parallel owner + list preserves input-record accounting when HTTP batching splits them.""" samples: List[Dict[str, Any]] = [] + owners: List[int] = [] skipped = 0 first_skipped_type: Optional[str] = None - for record in records: + for record_index, record in enumerate(records): if isinstance(record, Mapping) and "sample_id" in record: samples.append(dict(record)) + owners.append(record_index) elif not isinstance(record, Mapping) and is_episode(record): - samples.extend(build_samples([record], self._rollout_numbers)) + projected = build_samples([record], self._rollout_numbers) + samples.extend(projected) + owners.extend([record_index] * len(projected)) else: skipped += 1 if first_skipped_type is None: @@ -85,7 +122,7 @@ def _to_samples(self, records: Sequence[Any]) -> List[Dict[str, Any]]: first_skipped_type, ) self.skipped += skipped - return samples + return samples, owners def flush(self) -> None: """Writes are synchronous; the uploader thread owns the asynchrony.""" diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index fb14f2bd1..b107caf52 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -20,7 +20,7 @@ from . import _fork from .exceptions import is_record_rejection, is_transient -from .sinks.base import Sink +from .sinks.base import Sink, SinkWriteError logger = logging.getLogger(__name__) @@ -84,13 +84,18 @@ def __init__( def start(self) -> None: with self._lock: - if self._thread is not None and self._thread.is_alive(): - return - self._stopping.clear() - self._thread = threading.Thread( - target=self._run, name="prime-runs-uploader", daemon=True - ) - self._thread.start() + self._start_locked() + + def _start_locked(self) -> None: + """Start while ``_lock`` is held; a closed worker is one-shot.""" + if self._stopping.is_set(): + return + if self._thread is not None and self._thread.is_alive(): + return + self._thread = threading.Thread( + target=self._run, name="prime-runs-uploader", daemon=True + ) + self._thread.start() def _run(self) -> None: while True: @@ -117,6 +122,8 @@ def _dispatch(self, records: Sequence[Any]) -> None: continue try: sink.write(records) + except SinkWriteError as exc: + self._fail_sink(sink, exc.cause, dropped=exc.failed_records) except Exception as exc: # noqa: BLE001 - one sink failing must not stop the others self._fail_sink(sink, exc, dropped=len(records)) else: @@ -198,25 +205,26 @@ def _notify(self, name: str, exc: Exception) -> None: def submit(self, records: Sequence[Any]) -> bool: """Hand a batch to the uploader. ``False`` means it was dropped: the queue stayed full for the whole put timeout.""" - if self._stopping.is_set(): - return False - thread = self._thread - if thread is None or not thread.is_alive(): - self.start() - try: - self._queue.put(records, timeout=self.put_timeout) - return True - except queue.Full: - count = len(records) - self.dropped += count - logger.warning( - "Upload queue full after %.1fs; dropped %d item(s) (%d total). " - "The producer is outrunning the uploader.", - self.put_timeout, - count, - self.dropped, - ) - return False + # Serialize acceptance with close(): an accepted batch is queued before + # the stop sentinel, while a submission after close is refused. + with self._lock: + if self._stopping.is_set(): + return False + self._start_locked() + try: + self._queue.put(records, timeout=self.put_timeout) + return True + except queue.Full: + count = len(records) + self.dropped += count + logger.warning( + "Upload queue full after %.1fs; dropped %d item(s) (%d total). " + "The producer is outrunning the uploader.", + self.put_timeout, + count, + self.dropped, + ) + return False def flush(self, timeout: Optional[float] = None) -> bool: """Block until everything queued so far has been written.""" @@ -236,8 +244,9 @@ def flush(self, timeout: Optional[float] = None) -> bool: def close(self, timeout: Optional[float] = 30.0) -> None: """Drain, stop the thread, and close every sink.""" - self._stopping.set() - thread = self._thread + with self._lock: + self._stopping.set() + thread = self._thread if thread is not None and thread.is_alive(): deadline = _deadline(timeout) try: diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index f41fb96af..8b6c052e4 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -237,6 +237,56 @@ def test_logging_after_finish_is_a_producer_bug(): run.log_traces([{"id": "t1"}]) +def test_logging_that_started_first_is_queued_before_concurrent_finish(): + """Finish must not close the worker between the live check and enqueue.""" + + class SignalingBackend(FakeBackend): + def __init__(self) -> None: + super().__init__() + self.finalize_started = threading.Event() + + def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> None: + self.finalize_started.set() + super().finalize( + run_id, + status=status, + summary=summary, + error=error, + config=config, + ) + + backend = SignalingBackend() + sink = FakeSink() + run = make_run(backend, sinks=[sink]) + iteration_started = threading.Event() + release_record = threading.Event() + finish_called = threading.Event() + + def records(): + iteration_started.set() + assert release_record.wait(2.0) + yield {"id": "t1"} + + def finish_run() -> None: + finish_called.set() + run.finish() + + with ThreadPoolExecutor(max_workers=2) as executor: + logging = executor.submit(run.log_traces, records()) + assert iteration_started.wait(2.0) + finishing = executor.submit(finish_run) + assert finish_called.wait(2.0) + try: + assert not backend.finalize_started.wait(0.1) + finally: + release_record.set() + logging.result(timeout=2.0) + finishing.result(timeout=2.0) + + assert sink.batches == [[{"id": "t1"}]] + assert sink.closed is True + + def test_the_context_manager_completes_a_clean_run(): backend = FakeBackend() with make_run(backend) as run: diff --git a/packages/prime-runs/tests/test_samples_sink.py b/packages/prime-runs/tests/test_samples_sink.py index e356e0d4c..d32d632e8 100644 --- a/packages/prime-runs/tests/test_samples_sink.py +++ b/packages/prime-runs/tests/test_samples_sink.py @@ -1,9 +1,11 @@ """The legacy sample sink that keeps today's viewer working.""" +import httpx from _fakes import make_episode, make_trace from conftest import RecordingHandler from prime_runs.sinks import EvalSamplesSink +from prime_runs.worker import UploadWorker def make_sink(make_platform_client, eval_routes): @@ -81,3 +83,62 @@ def test_an_empty_batch_makes_no_request(make_platform_client, eval_routes): sink.write([]) assert handler.requests == [] + + +def test_a_rejected_sample_batch_does_not_abort_later_batches( + make_platform_client, eval_routes, monkeypatch +): + """One bad request must not hide valid rows later in the same worker batch.""" + statuses = iter([200, 422, 200, 200]) + + def respond(request): + status = next(statuses) + return httpx.Response(status, json={"detail": "invalid sample"} if status == 422 else {}) + + eval_routes["POST /api/v1/evaluations/eval-abc/samples"] = respond + monkeypatch.setattr( + "prime_runs.sinks.samples.batch_samples", + lambda samples: [[sample] for sample in samples], + ) + sink, handler = make_sink(make_platform_client, eval_routes) + worker = UploadWorker([sink]) + + worker.submit([{"sample_id": sample_id} for sample_id in ("a", "b", "c")]) + assert worker.flush(timeout=5.0) + worker.submit([{"sample_id": "d"}]) + assert worker.flush(timeout=5.0) + + posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples") + assert [body["samples"][0]["sample_id"] for body in posted] == ["a", "b", "c", "d"] + assert sink.samples_written == 3 + assert sink.enabled is True + assert worker.failed_records == {"eval_samples": 1} + worker.close() + + +def test_a_sink_wide_failure_counts_only_current_and_unattempted_records( + make_platform_client, eval_routes, monkeypatch +): + statuses = iter([200, 401]) + + def respond(request): + status = next(statuses) + return httpx.Response(status, json={"detail": "unauthorized"}) + + eval_routes["POST /api/v1/evaluations/eval-abc/samples"] = respond + monkeypatch.setattr( + "prime_runs.sinks.samples.batch_samples", + lambda samples: [[sample] for sample in samples], + ) + sink, handler = make_sink(make_platform_client, eval_routes) + worker = UploadWorker([sink]) + + worker.submit([{"sample_id": sample_id} for sample_id in ("a", "b", "c")]) + assert worker.flush(timeout=5.0) + + posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples") + assert [body["samples"][0]["sample_id"] for body in posted] == ["a", "b"] + assert sink.samples_written == 1 + assert sink.enabled is False + assert worker.failed_records == {"eval_samples": 2} + worker.close() diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index 1671e0165..d1cdb2e73 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -181,10 +181,14 @@ def test_close_drains_then_closes_every_sink(): def test_submitting_after_close_is_refused(): - worker = UploadWorker([FakeSink()]) + sink = FakeSink() + worker = UploadWorker([sink]) worker.close() + worker.start() assert worker.submit([{"id": 1}]) is False + assert worker._thread is None + assert sink.batches == [] def test_flush_without_a_running_thread_still_flushes_the_sinks(): From bce013024665f158dd6863bc34e838f1e64c6485 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Thu, 27 Aug 2026 15:27:40 -0700 Subject: [PATCH 34/39] fix(runs): reject writes promptly during finish --- packages/prime-runs/src/prime_runs/run.py | 27 ++++++---- packages/prime-runs/tests/test_run.py | 63 +++++++++++++++++++---- 2 files changed, 72 insertions(+), 18 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 3c2e73556..edbfa54ef 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -80,6 +80,10 @@ def __init__( self._deferred_error: Optional[BaseException] = None self._finish_timeout = DEFAULT_FINISH_TIMEOUT + # Guards only the transition from accepting writes to finishing. It is + # deliberately separate from _finish_lock, which is held through slow + # network teardown so concurrent finish callers wait for completion. + self._state_lock = threading.Lock() self._finish_lock = threading.Lock() self._finishing = False self._finished = False @@ -177,12 +181,14 @@ def log_episodes(self, episodes: Iterable[Any]) -> None: self._submit("log_episodes", episodes) def _submit(self, operation: str, records: Iterable[Any]) -> None: - # Finish holds the same lock through worker teardown. A log call either - # queues its batch first or observes that the run is already closing; - # it can never restart the worker between those two states. - with self._finish_lock: + # Refuse an already-finished run without consuming a lazy iterable. + self._require_live(operation) + batch = list(records) + # Finish holds this short-lived lock only while closing admission. A + # log call queues first or observes that teardown has begun, without + # waiting behind the slow flush/finalize path. + with self._state_lock: self._require_live(operation) - batch = list(records) if batch: self._worker.submit(batch) @@ -192,7 +198,7 @@ def update_summary(self, values: Mapping[str, Any]) -> None: Non-finite numbers are dropped here, as they are for ``finish(summary=...)``; writing to ``summary`` directly skips that. """ - with self._finish_lock: + with self._state_lock: self._require_live("update_summary") self.summary.update(_clean_metrics(values)) @@ -220,12 +226,14 @@ def finish( resolved = RunStatus(status) if not isinstance(status, RunStatus) else status if not resolved.is_terminal(): raise ValueError(f"finish() requires a terminal status, got {resolved.value!r}") - self._finishing = True + with self._state_lock: + self._finishing = True try: self._finish_once(summary, resolved, error) finally: - self._finishing = False - self._finished = True + with self._state_lock: + self._finishing = False + self._finished = True def _finish_once( self, @@ -329,6 +337,7 @@ def reset_after_fork(self) -> None: """Make an inherited handle safe in a forked child: fresh lock, and the parent keeps ownership of the lifecycle so the child's atexit hook cannot finalize a still-running run.""" + self._state_lock = threading.Lock() self._finish_lock = threading.Lock() self._finishing = False self._owns_lifecycle = False diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index 8b6c052e4..25673aa86 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -258,28 +258,31 @@ def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> backend = SignalingBackend() sink = FakeSink() run = make_run(backend, sinks=[sink]) - iteration_started = threading.Event() - release_record = threading.Event() + submit_started = threading.Event() + release_submit = threading.Event() finish_called = threading.Event() + original_submit = run._worker.submit - def records(): - iteration_started.set() - assert release_record.wait(2.0) - yield {"id": "t1"} + def blocking_submit(records): + submit_started.set() + assert release_submit.wait(2.0) + return original_submit(records) + + run._worker.submit = blocking_submit def finish_run() -> None: finish_called.set() run.finish() with ThreadPoolExecutor(max_workers=2) as executor: - logging = executor.submit(run.log_traces, records()) - assert iteration_started.wait(2.0) + logging = executor.submit(run.log_traces, [{"id": "t1"}]) + assert submit_started.wait(2.0) finishing = executor.submit(finish_run) assert finish_called.wait(2.0) try: assert not backend.finalize_started.wait(0.1) finally: - release_record.set() + release_submit.set() logging.result(timeout=2.0) finishing.result(timeout=2.0) @@ -287,6 +290,48 @@ def finish_run() -> None: assert sink.closed is True +@pytest.mark.parametrize( + "operation", + [ + lambda run: run.log_traces([{"id": "t1"}]), + lambda run: run.log_episodes([{"id": "ep1", "traces": []}]), + lambda run: run.update_summary({"reward": 1.0}), + ], + ids=["traces", "episodes", "summary"], +) +def test_writes_during_finish_fail_without_waiting_for_teardown(operation): + class BlockingBackend(FakeBackend): + def __init__(self) -> None: + super().__init__() + self.finalize_started = threading.Event() + self.release_finalize = threading.Event() + + def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> None: + self.finalize_started.set() + assert self.release_finalize.wait(2.0) + super().finalize( + run_id, + status=status, + summary=summary, + error=error, + config=config, + ) + + backend = BlockingBackend() + run = make_run(backend) + + with ThreadPoolExecutor(max_workers=2) as executor: + finishing = executor.submit(run.finish) + assert backend.finalize_started.wait(2.0) + writing = executor.submit(operation, run) + try: + with pytest.raises(RunFinishedError): + writing.result(timeout=0.2) + finally: + backend.release_finalize.set() + finishing.result(timeout=2.0) + + def test_the_context_manager_completes_a_clean_run(): backend = FakeBackend() with make_run(backend) as run: From 94c07918f9dd769713bd5d5950cabb7b6d73267a Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Fri, 28 Aug 2026 15:22:17 -0700 Subject: [PATCH 35/39] fix(runs): coalesce queued uploads, add CANCELLED and finish_timeout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The uploader dispatched one queue item per request. verifiers hands over one episode per rollout, so that was one POST /samples per episode against the platform's 300/min limit: any eval finishing more than five episodes a second got 429s, the retries backed the 256-item queue up, submit() blocked the rollout slot for 5s per episode, and records were then dropped from both sinks. The old one-shot upload never lost rows to rate limits. UploadWorker._run now drains everything already queued behind the item it took into one dispatch, so the request rate tracks upload latency rather than rollout throughput. Sinks already split by size (25 MB for samples, the traces client's own batching). Coalescing stops at a flush barrier (flush() promises everything before it has been written), at the stop sentinel, and at a batch of the other record kind (the traces sink infers the line format from the first record). A worker with no sinks — every disabled run — no longer copies batches into a queue or starts a thread for nothing. RunStatus gains CANCELLED: an interrupt is a decision, not a fault, and the platform's producer statuses are FAILED | CANCELLED. The context manager maps KeyboardInterrupt and asyncio.CancelledError to it (was CRASHED); the atexit hook still reports CRASHED. Added now so the enum does not change under consumers when the SDK adopts the status endpoint. init(finish_timeout=) and finish(timeout=) bound how long queued uploads get to drain before the run is closed out. An abort path can pass a few seconds so an interrupted producer is not pinned behind a slow platform for the full 300s. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01GBkJRDCFihXMjU1YpJTt14 --- packages/prime-runs/README.md | 22 +++-- packages/prime-runs/src/prime_runs/models.py | 7 +- packages/prime-runs/src/prime_runs/run.py | 37 ++++++-- packages/prime-runs/src/prime_runs/worker.py | 56 ++++++++++-- packages/prime-runs/tests/test_init.py | 13 ++- packages/prime-runs/tests/test_run.py | 50 +++++++++-- packages/prime-runs/tests/test_worker.py | 90 +++++++++++++++++++- 7 files changed, 243 insertions(+), 32 deletions(-) diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index 5824f4117..f18762751 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -83,6 +83,8 @@ disables the run with a warning — it never silently writes somewhere else. ## What the run handle does for you - **Streams instead of buffering.** Records go out as they are produced. + Whatever queues up while one upload is in flight goes out as the next one, + so the request rate tracks upload latency, not rollout throughput. - **Contains its own errors.** With the default `on_error="warn"`, nothing the platform raises escapes into your loop. Use `on_error="raise"` in tests and CI, where a silent upload failure is the bug; failures surface from `flush()` @@ -93,11 +95,15 @@ disables the run with a warning — it never silently writes somewhere else. (`run.dropped_records`) rather than stalled. Per-sink losses are in `run.failed_records`. - **Waits for its own uploads.** `finish()` gives queued records the same - budget a single upload gets (300s) and warns if they do not drain. + budget a single upload gets (300s) and warns if they do not drain. Set + `init(finish_timeout=...)` for the run, or `finish(timeout=...)` per call — an + abort path can pass a few seconds so an interrupted producer is not pinned + behind a slow platform. - **Reports a terminal status.** The context manager and an `atexit` hook both route to the same idempotent `finish()`. A run you said failed is `failed`; - one that stopped without saying — Ctrl-C inside a `with` block, an exit that - never reached `finish()` — is `crashed`. + one somebody stopped on purpose — Ctrl-C or a cancelled task inside a `with` + block — is `cancelled`; one that stopped without saying — an exit that never + reached `finish()` — is `crashed`. ## From async code @@ -141,6 +147,10 @@ package is imported — this is a leaf package, because the `prime` CLI depends Eval runs only. Training runs will arrive with a backend over `/api/v1/rft/external-runs`, designed against prime-rl's actual needs. -There is currently no producer-facing way to mark an evaluation **failed**. The -SDK records the terminal state under `metadata.prime_runs` and warns that the -run will keep showing as running on the dashboard. +The evaluations API this version targets has no producer-facing way to mark an +evaluation **failed** or **cancelled**: the SDK records the terminal state under +`metadata.prime_runs` and warns that the run will keep showing as running on +the dashboard. The platform is adding `PUT /evaluations/{id}` with a +`FAILED` / `CANCELLED` status; a follow-up release adopts it (keeping the +metadata fallback for older backends), which is why `RunStatus.CANCELLED` +already exists. diff --git a/packages/prime-runs/src/prime_runs/models.py b/packages/prime-runs/src/prime_runs/models.py index cf27c838c..c1c11537a 100644 --- a/packages/prime-runs/src/prime_runs/models.py +++ b/packages/prime-runs/src/prime_runs/models.py @@ -25,13 +25,16 @@ class RunStatus(str, Enum): """Terminal state a producer can report. - ``failed`` means the producer said the run failed; ``crashed`` means the - process exited without saying (only the SDK's atexit hook reports it). + ``failed`` means the producer said the run failed; ``cancelled`` means + somebody stopped it on purpose (Ctrl-C, a cancelled task) — a decision, + not a fault; ``crashed`` means the process exited without saying (only the + SDK's atexit hook reports it). """ RUNNING = "running" COMPLETED = "completed" FAILED = "failed" + CANCELLED = "cancelled" CRASHED = "crashed" def is_terminal(self) -> bool: diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index edbfa54ef..f0ceb251c 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -9,6 +9,7 @@ Nothing is re-stamped afterwards. """ +import asyncio import atexit import logging import math @@ -62,6 +63,7 @@ def __init__( sinks: Optional[List[Sink]] = None, mode: Mode = "online", on_error: OnError = "warn", + finish_timeout: Optional[float] = None, ) -> None: self._backend = backend self._handle = handle @@ -79,7 +81,9 @@ def __init__( # and re-raised from flush() or finish(), where the caller is looking. self._deferred_error: Optional[BaseException] = None - self._finish_timeout = DEFAULT_FINISH_TIMEOUT + self._finish_timeout = ( + DEFAULT_FINISH_TIMEOUT if finish_timeout is None else max(0.0, float(finish_timeout)) + ) # Guards only the transition from accepting writes to finishing. It is # deliberately separate from _finish_lock, which is held through slow # network teardown so concurrent finish callers wait for completion. @@ -217,9 +221,17 @@ def finish( *, status: Union[RunStatus, str] = RunStatus.COMPLETED, error: Optional[str] = None, + timeout: Optional[float] = None, ) -> None: """Flush everything and close the run out. Idempotent: the first caller - reports the status, concurrent callers wait for that teardown.""" + reports the status, concurrent callers wait for that teardown. + + ``timeout`` overrides the run's ``finish_timeout`` for this call: how + long queued uploads get to drain before the run is closed out anyway. + An abort path can pass a short budget so an interrupted producer is + not pinned behind a slow platform. The terminal status is still + reported (with its own retries) after the budget runs out. + """ with self._finish_lock: if self._finished: return @@ -229,7 +241,7 @@ def finish( with self._state_lock: self._finishing = True try: - self._finish_once(summary, resolved, error) + self._finish_once(summary, resolved, error, timeout) finally: with self._state_lock: self._finishing = False @@ -240,12 +252,14 @@ def _finish_once( summary: Optional[Mapping[str, Any]], resolved: RunStatus, error: Optional[str], + timeout: Optional[float] = None, ) -> None: if summary: self.summary.update(_clean_metrics(summary)) self._status = resolved - deadline = time.monotonic() + max(0.0, self._finish_timeout) + budget = self._finish_timeout if timeout is None else max(0.0, float(timeout)) + deadline = time.monotonic() + budget def remaining() -> float: return max(0.0, deadline - time.monotonic()) @@ -257,7 +271,7 @@ def remaining() -> float: "Run %s: uploads did not drain within %ss; finalizing anyway. " "Some records may be missing from this run.", self.id, - self._finish_timeout, + budget, ) self._worker.close(timeout=remaining()) @@ -295,9 +309,9 @@ def remaining() -> float: # Last, so a run whose teardown failed is still closed out first. self._raise_deferred() - def fail(self, error: Union[str, BaseException]) -> None: + def fail(self, error: Union[str, BaseException], *, timeout: Optional[float] = None) -> None: """Close the run out as failed.""" - self.finish(status=RunStatus.FAILED, error=_describe(error)) + self.finish(status=RunStatus.FAILED, error=_describe(error), timeout=timeout) # -------------------------------------------------------- context manager @@ -309,9 +323,9 @@ def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: self.finish() return False - if isinstance(exc, KeyboardInterrupt): + if isinstance(exc, (KeyboardInterrupt, asyncio.CancelledError)): # An interrupt is a decision, not a fault. - status, error = RunStatus.CRASHED, "interrupted" + status, error = RunStatus.CANCELLED, "interrupted" else: status, error = RunStatus.FAILED, _describe(exc) @@ -408,6 +422,7 @@ def init( api_key: Optional[str] = None, base_url: Optional[str] = None, on_error: OnError = "warn", + finish_timeout: Optional[float] = None, ) -> Run: """Start a run and return a handle to it. @@ -418,6 +433,9 @@ def init( key and disabled (with a warning) when there is not. ``config`` is what the run was configured with: the path to the file it was launched from (stored byte for byte under ``config_source``), or a mapping taken as given. + ``finish_timeout`` is how long :meth:`Run.finish` lets queued uploads drain + before closing the run out anyway (default: the upload timeout, 300 s); + ``finish(timeout=...)`` overrides it per call. """ settings = Config() api_key = api_key if api_key is not None else settings.api_key @@ -474,6 +492,7 @@ def init( sinks=sinks, mode=resolved_mode, on_error=on_error, + finish_timeout=finish_timeout, ) if run.url: logger.info("Run %s: %s", run.id, run.url) diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index b107caf52..560fea593 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -20,7 +20,7 @@ from . import _fork from .exceptions import is_record_rejection, is_transient -from .sinks.base import Sink, SinkWriteError +from .sinks.base import Sink, SinkWriteError, is_episode logger = logging.getLogger(__name__) @@ -37,6 +37,10 @@ class _Flush: event: threading.Event = field(default_factory=threading.Event) +#: "Nothing pulled ahead" — distinct from ``None``, which is the stop sentinel. +_NOTHING: Any = object() + + def _deadline(timeout: Optional[float]) -> Optional[float]: if timeout is None: return None @@ -92,14 +96,18 @@ def _start_locked(self) -> None: return if self._thread is not None and self._thread.is_alive(): return - self._thread = threading.Thread( - target=self._run, name="prime-runs-uploader", daemon=True - ) + self._thread = threading.Thread(target=self._run, name="prime-runs-uploader", daemon=True) self._thread.start() def _run(self) -> None: + # An item pulled off the queue while coalescing that could not join the + # batch (a control item — the stop sentinel is None, hence the marker — + # or records of the other kind); it is the next thing to process, ahead + # of anything queued after it. + pending: Any = _NOTHING while True: - item = self._queue.get() + item = pending if pending is not _NOTHING else self._queue.get() + pending = _NOTHING try: if item is None: return @@ -107,12 +115,43 @@ def _run(self) -> None: self._flush_sinks() item.event.set() else: - self._dispatch(item) + batch, pending = self._coalesce(item) + self._dispatch(batch) except Exception as exc: # noqa: BLE001 - the thread must outlive one bad batch logger.debug("Uploader iteration failed: %s", exc) finally: self._queue.task_done() + def _coalesce(self, first: Sequence[Any]) -> "tuple[list, Any]": + """Merge everything already queued behind ``first`` into one batch. + + A producer hands over one episode at a time as rollouts finish, and + each batch costs every sink a request: one ``POST /samples`` per + episode runs into the platform's per-minute limit on any fast eval, + and the retries then back the queue up until records are dropped. + Draining what has accumulated while the previous request was in + flight makes the request rate track upload latency instead of rollout + throughput. Sinks split a large batch by size themselves. + + Stops at a control item (a flush barrier must not be reordered past + the records queued before it) and at a batch of the other record kind + (the traces sink infers the line format from the first record). + Returns the merged batch and the item that stopped it, if any. + """ + batch = list(first) + kind = is_episode(batch[0]) if batch else None + while True: + try: + item = self._queue.get_nowait() + except queue.Empty: + return batch, _NOTHING + if item is None or isinstance(item, _Flush): + return batch, item + if item and (kind is None or is_episode(item[0]) != kind): + return batch, item + batch.extend(item) + self._queue.task_done() + def _dispatch(self, records: Sequence[Any]) -> None: for sink in self.sinks: if not sink.enabled: @@ -205,6 +244,11 @@ def _notify(self, name: str, exc: Exception) -> None: def submit(self, records: Sequence[Any]) -> bool: """Hand a batch to the uploader. ``False`` means it was dropped: the queue stayed full for the whole put timeout.""" + # No sinks (a disabled run): there is nowhere for the records to go, so + # do not copy them into a queue or start a thread to find that out. + # Retired sinks are not the same case — their losses are counted. + if not self.sinks: + return True # Serialize acceptance with close(): an accepted batch is queued before # the stop sentinel, while a submission after close is refused. with self._lock: diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index 54c9a087e..b7a8f3273 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -44,11 +44,14 @@ def test_a_disabled_run_still_answers_every_call(tmp_path): run = pr.init(mode="disabled") run.log_traces([{"id": "t1"}]) + # No sinks, so nothing was queued and no uploader thread was started for it. + assert run._worker._thread is None run.finish(summary={"avg_reward": 1.0}) assert run.id.startswith("disabled-") assert run.url is None assert run.status is RunStatus.COMPLETED + assert run.dropped_records == 0 assert not list(tmp_path.iterdir()) # tmp_path is $HOME here @@ -182,7 +185,15 @@ def test_the_end_to_end_shape_a_producer_writes(online): run.log_traces([episode]) run.finish(summary=metrics.from_episodes(episodes)) - assert handler.paths().count("POST /api/v1/evaluations/eval-abc/samples") == 3 + # Episodes queued while a request is in flight go out together, so the + # request count depends on timing; every row landing exactly once does not. + posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples") + assert 1 <= len(posted) <= 3 + assert sorted(s["sample_id"] for body in posted for s in body["samples"]) == [ + "ep-0", + "ep-1", + "ep-2", + ] finalize = handler.bodies_for("/api/v1/evaluations/eval-abc/finalize")[0] assert finalize["metrics"]["avg_reward"] == 1.0 assert run.status is RunStatus.COMPLETED diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index 25673aa86..16b67ab23 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -1,5 +1,6 @@ """The run handle: lifecycle, containment, terminal status.""" +import asyncio import threading import time from concurrent.futures import ThreadPoolExecutor @@ -379,18 +380,57 @@ def interrupt_finish(*args, **kwargs): run.finish() -def test_an_interrupt_is_recorded_as_a_decision_not_a_fault(): - """Ctrl-C must not land in the same bucket as a broken eval.""" +@pytest.mark.parametrize("interrupt", [KeyboardInterrupt, asyncio.CancelledError]) +def test_an_interrupt_is_recorded_as_a_decision_not_a_fault(interrupt): + """Ctrl-C must not land in the same bucket as a broken eval, nor in the + one for a process that vanished: the platform has a status for it.""" backend = FakeBackend() - with pytest.raises(KeyboardInterrupt): + with pytest.raises(interrupt): with make_run(backend): - raise KeyboardInterrupt + raise interrupt - assert backend.finalized[0]["status"] is RunStatus.CRASHED + assert backend.finalized[0]["status"] is RunStatus.CANCELLED assert backend.finalized[0]["error"] == "interrupted" +def test_cancelled_is_terminal_and_finalizes_like_failed(): + backend = FakeBackend() + run = make_run(backend) + + assert RunStatus.CANCELLED.is_terminal() + run.finish(status="cancelled", error="operator stopped it") + + assert run.status is RunStatus.CANCELLED + assert backend.finalized[0]["status"] is RunStatus.CANCELLED + assert backend.finalized[0]["error"] == "operator stopped it" + + +def test_finish_timeout_is_set_per_run_and_overridden_per_call(): + """An abort path passes a short budget so an interrupted producer is not + pinned behind a slow platform for the full drain.""" + run = make_run(finish_timeout=1.5) + observed = {} + + def record_flush(timeout=None): + observed["flush"] = timeout + return True + + run._worker.flush = record_flush + run.finish(timeout=0.25) + + assert run._finish_timeout == 1.5 + assert 0.0 < observed["flush"] <= 0.25 + + +def test_finish_timeout_default_is_the_upload_budget(): + from prime_runs.run import DEFAULT_FINISH_TIMEOUT + + run = make_run() + assert run._finish_timeout == DEFAULT_FINISH_TIMEOUT == 300.0 + run.finish() + + def test_finish_hands_the_full_config_to_finalize(): """The evaluations API replaces metadata wholesale, so a backend recording terminal state inside it needs the whole picture to merge into.""" diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index d1cdb2e73..870ba2543 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -337,7 +337,10 @@ def write(self, records) -> None: reported = [] worker = UploadWorker([sink], on_error=lambda name, exc: reported.append((name, exc))) + # Drained between the two: queued batches coalesce into one write, and a + # rejection is per write. worker.submit([{"id": "bad"}]) + drain(worker) worker.submit([{"id": "good"}]) drain(worker) @@ -361,9 +364,10 @@ def write(self, records) -> None: sink = DeadSink("dead") worker = UploadWorker([sink]) + # One write per strike: drained between submits so they do not coalesce. for _ in range(TRANSIENT_FAILURE_LIMIT): worker.submit([{"id": 1}]) - drain(worker) + drain(worker) assert sink.enabled is False worker.close() @@ -391,10 +395,11 @@ def write(self, records) -> None: sink = FlakySink() worker = UploadWorker([sink]) - # Far more failures than the limit, but never two in a row. + # Far more failures than the limit, but never two in a row. Drained one at + # a time: queued batches coalesce, and this is about consecutive *writes*. for _ in range(TRANSIENT_FAILURE_LIMIT * 4): worker.submit([{"id": 1}]) - drain(worker) + drain(worker) assert sink.calls == TRANSIENT_FAILURE_LIMIT * 4 assert sink.enabled is True @@ -440,3 +445,82 @@ def write(self, records) -> None: assert worker.dropped == 0 assert healthy.batches, "the healthy sink stored them" worker.close() + + +def test_records_queued_during_an_upload_go_out_as_one_batch(): + """One request per episode runs into the platform's per-minute limit on a + fast eval; whatever accumulates while a request is in flight must go out + together, so the request rate tracks upload latency instead.""" + sink = BlockingSink() + worker = UploadWorker([sink]) + + assert worker.submit([{"id": 0}]) + assert sink.entered.wait(1.0), "the uploader never reached the sink" + for n in (1, 2, 3): + assert worker.submit([{"id": n}]) + sink.released.set() + drain(worker) + + assert [[r["id"] for r in batch] for batch in sink.batches] == [[0], [1, 2, 3]] + worker.close() + + +def test_coalescing_keeps_traces_and_episodes_apart(): + """The traces sink infers the line format from a batch's first record, so + mixing episodes into a batch of bare traces would misfile every record.""" + sink = BlockingSink() + worker = UploadWorker([sink]) + + assert worker.submit([{"id": "t0"}]) + assert sink.entered.wait(1.0) + worker.submit([{"id": "t1"}]) + worker.submit([{"id": "e1", "traces": []}]) + worker.submit([{"id": "e2", "traces": []}]) + worker.submit([{"id": "t2"}]) + sink.released.set() + drain(worker) + + assert [[r["id"] for r in batch] for batch in sink.batches] == [ + ["t0"], + ["t1"], + ["e1", "e2"], + ["t2"], + ] + worker.close() + + +def test_a_flush_barrier_is_not_reordered_past_the_records_before_it(): + """Coalescing must stop at the barrier: flush() promises that everything + submitted before it has been written when it returns.""" + sink = BlockingSink() + worker = UploadWorker([sink]) + + assert worker.submit([{"id": 0}]) + assert sink.entered.wait(1.0) + worker.submit([{"id": 1}]) + release = threading.Timer(0.05, sink.released.set) + release.start() + try: + assert worker.flush(timeout=2.0) is True + finally: + release.join() + worker.submit([{"id": 2}]) + + assert [[r["id"] for r in batch] for batch in sink.batches] == [[0], [1]] + drain(worker) + assert [[r["id"] for r in batch] for batch in sink.batches] == [[0], [1], [2]] + worker.close() + + +def test_a_worker_with_no_sinks_never_starts_a_thread(): + """A disabled run submits every batch; copying them into a queue and + starting a thread to find there is nowhere to put them is pure cost.""" + worker = UploadWorker([]) + + assert worker.submit([{"id": 1}]) is True + assert worker.submit([{"id": 2}]) is True + + assert worker._thread is None + assert worker.dropped == 0 + assert worker.flush(timeout=1.0) is True + worker.close() From d07c245af845f136172b3d2a44040997b904b275 Mon Sep 17 00:00:00 2001 From: Cooper Miller <44559144+kcoopermiller@users.noreply.github.com> Date: Tue, 1 Sep 2026 16:46:17 -0700 Subject: [PATCH 36/39] prime-runs: RftBackend for training runs (#873) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(runs): RftBackend — training runs over /api/v1/rft/external-runs 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=, environments=[hub ids], training=TrainingSpec(max_steps, batch_size, rollouts_per_example, seq_len, wandb_*), 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 Claude-Session: https://claude.ai/code/session_01GBkJRDCFihXMjU1YpJTt14 * fix(runs): make the training samples sink fork-safe 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 Claude-Session: https://claude.ai/code/session_01GBkJRDCFihXMjU1YpJTt14 * fix(runs): harden training finalization and metrics * fix(runs): number training samples per upload and let training sinks 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 Claude-Session: https://claude.ai/code/session_01Fk5fBAAaTiPfJ5UpNyqxXR * fix(runs): keep a forked child's sample_id ranges apart from its parent'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 Claude-Session: https://claude.ai/code/session_01Fk5fBAAaTiPfJ5UpNyqxXR --------- Co-authored-by: Claude Fable 5 --- packages/prime-runs/README.md | 69 ++- packages/prime-runs/pyproject.toml | 5 + .../prime-runs/src/prime_runs/__init__.py | 13 +- packages/prime-runs/src/prime_runs/backend.py | 185 +++++++- packages/prime-runs/src/prime_runs/models.py | 32 +- .../prime-runs/src/prime_runs/projection.py | 121 +++++ packages/prime-runs/src/prime_runs/run.py | 197 ++++++-- .../src/prime_runs/sinks/__init__.py | 5 + .../prime-runs/src/prime_runs/sinks/base.py | 7 +- .../src/prime_runs/sinks/metrics.py | 79 +++ .../src/prime_runs/sinks/samples.py | 4 +- .../prime-runs/src/prime_runs/sinks/traces.py | 3 +- .../src/prime_runs/sinks/train_samples.py | 302 ++++++++++++ packages/prime-runs/src/prime_runs/worker.py | 50 +- packages/prime-runs/tests/_fakes.py | 55 +++ packages/prime-runs/tests/conftest.py | 52 ++ packages/prime-runs/tests/test_init.py | 212 ++++++++- .../prime-runs/tests/test_metrics_sink.py | 102 ++++ packages/prime-runs/tests/test_rft_backend.py | 240 ++++++++++ packages/prime-runs/tests/test_run.py | 90 ++++ .../prime-runs/tests/test_samples_sink.py | 4 +- .../tests/test_train_samples_sink.py | 448 ++++++++++++++++++ packages/prime-runs/tests/test_worker.py | 61 +++ uv.lock | 8 +- 24 files changed, 2282 insertions(+), 62 deletions(-) create mode 100644 packages/prime-runs/src/prime_runs/sinks/metrics.py create mode 100644 packages/prime-runs/src/prime_runs/sinks/train_samples.py create mode 100644 packages/prime-runs/tests/test_metrics_sink.py create mode 100644 packages/prime-runs/tests/test_rft_backend.py create mode 100644 packages/prime-runs/tests/test_train_samples_sink.py diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index f18762751..bc44079c7 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -1,6 +1,6 @@ # Prime Runs SDK -Track eval runs on the Prime Intellect platform. +Track eval and training runs on the Prime Intellect platform. ```bash pip install prime-runs @@ -105,6 +105,69 @@ disables the run with a warning — it never silently writes somewhere else. block — is `cancelled`; one that stopped without saying — an exit that never reached `finish()` — is `crashed`. +## Training runs + +```python +import prime_runs as pr + +run = pr.init( + kind="train", + name="qwen3-8b-gsm8k-rl", + model="Qwen/Qwen3-8B", # the base model + environments=["primeintellect/gsm8k"], # hub ids, passed through as given + training=pr.TrainingSpec(max_steps=1000, batch_size=64, rollouts_per_example=8), + config=train_config.model_dump(), # or the path to the launched TOML + team_id="team_...", # external runs are created for a team +) + +for step, (episodes, metrics) in enumerate(training_loop): + run.log_episodes(episodes) # carry run.work.step (verifiers TrainRunInfo) + run.log_metrics(metrics, step=step) + +run.finish() +``` + +A training run lives on `/api/v1/rft/external-runs` — what prime-rl's monitor +did by hand. Where it differs from an eval run: + +- **A team is required.** The platform enables external runs per team. A team + outside the allowlist gets a `ForbiddenError` from `init()` with the + platform's reason; the producer decides whether that means "run locally". +- **Metrics.** `log_metrics(values, step=...)` streams one row per call on its + own uploader, so a slow sample upload never holds a step's metrics back. + `_timestamp` is stamped on every row; non-finite values are dropped. The + endpoint allows 60 rows a minute per token. A row whose response was lost + is re-sent: the platform keeps one row per step, so nothing doubles. +- **Samples.** Every record reaches Prime Traces. The training viewer's sample + table is a Parquet object per upload, keyed by the training step an episode + was *dispatched* at — `run.work.step` on verifiers' `TrainRunInfo`, which + prime-rl stamps on every dispatched episode — and uploaded every 10th step + (prime-rl's cadence) for training-work episodes. Objects are additive: a + step logged in several `log_episodes` calls (an off-policy episode landing + in a later batch) gets one object per call, each numbering its `sample_id`s + in its own range (a forked child draws its ranges apart from its parent's), + and the viewer shows their union. So "step N" in the + sample viewer means dispatched at N, while the metrics logged at N describe + the batch trained at N. Encoding needs pyarrow: `pip install + 'prime-runs[train]'`. Without it the table is skipped with a warning and + traces still flow. +- **Attach.** `init(kind="train", id=os.environ["RUN_ID"])` joins a run a + launcher created: nothing is registered, the platform keeps the run's + failure marking, and a clean `finish()` still completes it. The id must be + an external run's (one created through `POST /rft/external-runs`): the + platform's monitoring endpoints answer 400 for a hosted run. +- **Status.** The RFT vocabulary is `completed | failed`. `cancelled` and + `crashed` are reported as `failed` with the reason in `error_message`; a + clean `finish()` is an idempotent finalize, safe to replay. +- **Outages.** A sink that strikes out on transient failures (three in a row, + each already retried) is paused for five minutes and then tried again, not + retired for the run: a platform deploy costs a multi-day run a window of + rows, never the rest of its curves. Records that arrive while a sink is + paused are counted in `run.failed_records` for that sink. +- **No config update.** `config=` is registered once as `run_config` and there + is no summary document, so `finish(summary=...)` is not sent for training + runs — log final numbers with `log_metrics`. + ## From async code `log_traces()` / `log_episodes()` are a queue put, not a request, so they are @@ -144,8 +207,8 @@ package is imported — this is a leaf package, because the `prime` CLI depends ## Status -Eval runs only. Training runs will arrive with a backend over -`/api/v1/rft/external-runs`, designed against prime-rl's actual needs. +Eval runs over `/api/v1/evaluations`, training runs over +`/api/v1/rft/external-runs` (see "Training runs"). The evaluations API this version targets has no producer-facing way to mark an evaluation **failed** or **cancelled**: the SDK records the terminal state under diff --git a/packages/prime-runs/pyproject.toml b/packages/prime-runs/pyproject.toml index 68682a320..0c6aa5fca 100644 --- a/packages/prime-runs/pyproject.toml +++ b/packages/prime-runs/pyproject.toml @@ -43,6 +43,11 @@ Documentation = "https://github.com/PrimeIntellect-ai/prime/tree/main/packages/p Repository = "https://github.com/PrimeIntellect-ai/prime.git" [project.optional-dependencies] +# The training sample table is Parquet. Optional: an eval producer never needs +# it, and prime-rl already carries pyarrow. +train = [ + "pyarrow>=15.0.0", +] dev = [ "pytest>=7.0.0", "pytest-asyncio>=0.23.0", diff --git a/packages/prime-runs/src/prime_runs/__init__.py b/packages/prime-runs/src/prime_runs/__init__.py index fe66244e2..6e8984062 100644 --- a/packages/prime-runs/src/prime_runs/__init__.py +++ b/packages/prime-runs/src/prime_runs/__init__.py @@ -1,6 +1,6 @@ """Prime Intellect Runs SDK. -Track eval runs on the Prime platform:: +Track eval and training runs on the Prime platform:: import prime_runs as pr @@ -31,7 +31,14 @@ TransportError, UnauthorizedError, ) -from .models import CONFIG_SOURCE_KEY, ConfigSource, EnvironmentRef, RunStatus +from .models import ( + CONFIG_SOURCE_KEY, + ConfigSource, + EnvironmentRef, + RunKind, + RunStatus, + TrainingSpec, +) from .run import MODE_ENV, Run, init __version__ = "0.1.0" @@ -40,6 +47,8 @@ "init", "Run", "RunStatus", + "RunKind", + "TrainingSpec", "ConfigSource", "CONFIG_SOURCE_KEY", "EnvironmentRef", diff --git a/packages/prime-runs/src/prime_runs/backend.py b/packages/prime-runs/src/prime_runs/backend.py index 68e226c55..76e08210f 100644 --- a/packages/prime-runs/src/prime_runs/backend.py +++ b/packages/prime-runs/src/prime_runs/backend.py @@ -20,7 +20,7 @@ from ._http import PlatformClient from .exceptions import APIError, ConfigurationError, EnvironmentResolutionError -from .models import EnvironmentRef, RunHandle, RunSpec, RunStatus +from .models import EnvironmentRef, RunHandle, RunSpec, RunStatus, TrainingSpec logger = logging.getLogger(__name__) @@ -267,3 +267,186 @@ def finalize(self, run_id: str, **kwargs: Any) -> None: def close(self) -> None: return None + + +# ------------------------------------------------------------------- rft + +#: The RFT API's status vocabulary is ``completed | failed``; the SDK's finer +#: distinctions travel in ``error_message``. +RFT_TERMINAL_STATUS = { + RunStatus.COMPLETED: "completed", + RunStatus.FAILED: "failed", + RunStatus.CANCELLED: "failed", + RunStatus.CRASHED: "failed", +} + + +class RftBackend: + """Lifecycle for external training runs over ``/api/v1/rft/*``. + + What prime-rl's ``TrainRun`` did, in the SDK: register a run (or attach to + one a managed launch created), close it out. External runs are created + *for a team* and the platform enables them per team, so ``team_id`` is + required and a team outside the allowlist gets a 403 from :meth:`create`; + the producer decides whether that means "run locally" (the way verifiers + treats a missing key) or "stop". + + The API has no config/summary update — ``run_config`` travels with + :meth:`create` and :meth:`update` is a no-op — and its status vocabulary + is ``completed | failed``: ``cancelled`` and ``crashed`` are reported as + ``failed`` with the reason in ``error_message``. + """ + + def __init__( + self, + client: PlatformClient, + *, + frontend_url: str, + team_id: Optional[str] = None, + ) -> None: + self._client = client + self._frontend_url = frontend_url.rstrip("/") + self._team_id = team_id + + # ------------------------------------------------------------------ create + + def create(self, spec: RunSpec) -> RunHandle: + if not spec.model: + raise ConfigurationError("A training run needs model= — the base model being trained.") + team_id = spec.team_id or self._team_id + if not team_id: + raise ConfigurationError( + "A training run needs a team: pass team_id= or set PRIME_TEAM_ID. " + "External training runs are created for a team, and the platform " + "enables them per team." + ) + training = spec.training or TrainingSpec() + payload: Dict[str, Any] = { + "base_model": spec.model, + "max_steps": int(training.max_steps), + # Environment ids are passed through: training environments are + # named by hub id, and the RFT API does its own resolution. + "environments": [_training_environment(ref) for ref in spec.environments], + "team_id": team_id, + } + _set_if(payload, "name", spec.name) + _set_if(payload, "batch_size", training.batch_size) + _set_if(payload, "rollouts_per_example", training.rollouts_per_example) + _set_if(payload, "seq_len", training.seq_len) + _set_if(payload, "run_config", spec.config or None) + _set_if(payload, "wandb_project", training.wandb_project) + _set_if(payload, "wandb_entity", training.wandb_entity) + _set_if(payload, "wandb_run_name", training.wandb_run_name) + + # Not replayable: a retry after a lost response would register a + # second run and only the second would be tracked. + response = self._client.post("/rft/external-runs", json_body=payload) + run = response.get("run") or {} + run_id = run.get("id") + if not run_id: + raise APIError(f"POST /rft/external-runs returned no run id (keys: {sorted(response)})") + return RunHandle( + id=str(run_id), + name=str(run.get("name") or spec.name or "") or None, + url=self.url_for(str(run_id)), + ) + + def attach(self, run_id: str) -> RunHandle: + """A handle for a run the platform already created — a launcher that + injected its id. No request: the first write proves access, and it + must be an *external* run (``POST /rft/external-runs``): the v1 + monitoring endpoints answer 400 for a hosted run's id.""" + return RunHandle(id=run_id, url=self.url_for(run_id)) + + def url_for(self, run_id: str) -> str: + return f"{self._frontend_url}/dashboard/training/{run_id}" + + # ------------------------------------------------------------------ update + + def update( + self, + run_id: str, + *, + config: Optional[Dict[str, Any]] = None, + summary: Optional[Dict[str, Any]] = None, + ) -> None: + """No-op: an RFT run takes its config at creation and has no summary + document. Run-level outputs go through ``log_metrics``.""" + if config or summary: + logger.debug("Run %s: the RFT API has no config/summary update; nothing sent", run_id) + + # ---------------------------------------------------------------- finalize + + def finalize( + self, + run_id: str, + *, + status: RunStatus, + summary: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, + ) -> None: + if status is RunStatus.COMPLETED: + # Compare-and-set on the server, and "already finalized" is a + # success, so a lost response is safe to replay. + try: + self._client.post( + "/rft/finalize", + json_body={"run_id": run_id, "exit_code": 0}, + idempotent=True, + ) + except APIError as exc: + # Preserve TrainRun's fallback: finalization may be unavailable + # even though the status endpoint can still close the run out. + logger.warning( + "Run %s could not be finalized (%s); falling back to a status update", + run_id, + exc, + ) + self._set_terminal_status(run_id, status=status) + return + + message = error or status.value + if status is not RunStatus.FAILED: + message = f"{status.value}: {message}" + self._set_terminal_status(run_id, status=status, error_message=message) + + def _set_terminal_status( + self, + run_id: str, + *, + status: RunStatus, + error_message: Optional[str] = None, + ) -> None: + payload: Dict[str, Any] = {"status": RFT_TERMINAL_STATUS[status]} + _set_if(payload, "error_message", error_message) + try: + self._client.put( + f"/rft/external-runs/{run_id}/status", + json_body=payload, + ) + except APIError as exc: + if exc.status_code == 409: + # Already closed out (a managed launcher, or an earlier attempt + # whose response was lost). The outcome is recorded either way. + logger.info( + "Run %s is already closed on the platform; not marking it %s", + run_id, + status.value, + ) + return + raise + + def log_metrics(self, run_id: str, metrics: Dict[str, Any]) -> None: + """One step's metrics, synchronously. The run handle batches these + through its uploader; this is the direct call.""" + self._client.post("/rft/metrics", json_body={"run_id": run_id, "metrics": metrics}) + + def close(self) -> None: + self._client.close() + + +def _training_environment(ref: EnvironmentRef) -> Dict[str, Any]: + entry: Dict[str, Any] = {"id": ref.id or ref.slug or ref.name} + _set_if(entry, "version_id", ref.version_id) + return entry diff --git a/packages/prime-runs/src/prime_runs/models.py b/packages/prime-runs/src/prime_runs/models.py index c1c11537a..ed7757b44 100644 --- a/packages/prime-runs/src/prime_runs/models.py +++ b/packages/prime-runs/src/prime_runs/models.py @@ -18,8 +18,14 @@ OnError = Literal["warn", "raise"] -RUN_KIND = "eval" -"""Stamped as ``run.type`` on records and sent as upload provenance.""" +RunKind = Literal["eval", "train"] +"""What a run is: a standalone evaluation (``/api/v1/evaluations``) or an +external training run (``/api/v1/rft/external-runs``). Stamped as ``run.type`` +on records and sent as upload provenance — the traces service's vocabulary, +which is also what verifiers' ``EvalRunInfo`` / ``TrainRunInfo`` carry.""" + +RUN_KIND: RunKind = "eval" +"""The default kind.""" class RunStatus(str, Enum): @@ -191,6 +197,26 @@ def __post_init__(self) -> None: ) +@dataclass +class TrainingSpec: + """What a training run is registered with, beyond what every run has + (``model`` is the base model, ``environments`` the training environments, + ``config`` the run config). All of it is for display on the dashboard. + ``max_steps`` is required by the platform; ``0`` means unknown.""" + + max_steps: int = 0 + batch_size: Optional[int] = None + rollouts_per_example: Optional[int] = None + seq_len: Optional[int] = None + wandb_project: Optional[str] = None + wandb_entity: Optional[str] = None + wandb_run_name: Optional[str] = None + + def __post_init__(self) -> None: + if self.max_steps < 0: + raise ValueError("max_steps must be non-negative") + + @dataclass class RunSpec: """Everything a backend needs to open a run: ``init()``'s arguments after @@ -205,6 +231,8 @@ class RunSpec: tags: List[str] = field(default_factory=list) team_id: Optional[str] = None config: Dict[str, Any] = field(default_factory=dict) + kind: RunKind = RUN_KIND + training: Optional[TrainingSpec] = None @dataclass diff --git a/packages/prime-runs/src/prime_runs/projection.py b/packages/prime-runs/src/prime_runs/projection.py index 1474bb3ae..1d3108ba3 100644 --- a/packages/prime-runs/src/prime_runs/projection.py +++ b/packages/prime-runs/src/prime_runs/projection.py @@ -179,3 +179,124 @@ def batch_samples(samples: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any] if batch: batches.append(batch) return batches + + +# ------------------------------------------------------------ training table + + +def parquet_available() -> bool: + """Whether the training sample table can be encoded (``prime-runs[train]``).""" + try: + import pyarrow # noqa: F401 + except ImportError: + return False + return True + + +def train_sample_schema() -> Any: + """The training sample table's Parquet schema (prime-rl's ``SAMPLE_SCHEMA``).""" + import pyarrow as pa + + return 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 episodes_to_parquet_bytes( + episodes: Sequence[Any], run_id: str, step: int, *, sample_id_offset: int = 0 +) -> Optional[bytes]: + """One training step's episodes as the viewer's Parquet table, one row per + episode. Moved here from prime-rl's ``monitors/prime.py``. + + Sample construction is shared with the eval projection (:func:`build_samples`: + the complete native episode in ``info.native_wrapper``, a 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. ``None`` when no episode has a trajectory. + + ``sample_id`` numbers the rows from ``sample_id_offset``: the viewer looks + samples up by (step, sample_id), and a step uploaded as more than one + object (see :mod:`prime_runs.sinks.train_samples`) needs each object to + number its rows in its own range. + """ + import io + import json + from datetime import datetime, timezone + + import pyarrow as pa + import pyarrow.parquet as pq + + advantages: Dict[Any, Optional[float]] = {} + env_names: Dict[Any, str] = {} + for episode in episodes: + if not episode.traces: + continue + summary_trace = episode.traces[summary_trace_index(episode)] + advantages[episode.id] = (summary_trace.info or {}).get("advantage") + env = getattr(episode, "env", None) + env_names[episode.id] = str(getattr(env, "id", None) or "") + + now = datetime.now(timezone.utc) + rows: List[Dict[str, Any]] = [] + for sample_id, sample in enumerate(build_samples(episodes), start=sample_id_offset): + trajectory = sample["trajectory"] + if not trajectory: # no branches: 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=train_sample_schema()) + buffer = io.BytesIO() + pq.write_table(table, buffer, compression="snappy", use_dictionary=True, write_statistics=True) + return buffer.getvalue() diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index f0ceb251c..1e4ecf75b 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -20,7 +20,7 @@ from . import _fork from ._http import DEFAULT_TIMEOUT, UPLOAD_TIMEOUT, PlatformClient -from .backend import Backend, DisabledBackend, EvalsBackend, disabled_run_id +from .backend import Backend, DisabledBackend, EvalsBackend, RftBackend, disabled_run_id from .config import Config from .exceptions import ConfigurationError, RunFinishedError from .models import ( @@ -31,10 +31,12 @@ Mode, OnError, RunHandle, + RunKind, RunSpec, RunStatus, + TrainingSpec, ) -from .sinks import EvalSamplesSink, Sink, TracesSink +from .sinks import EvalSamplesSink, RftMetricsSink, RftSamplesSink, Sink, TracesSink from .worker import UploadWorker logger = logging.getLogger(__name__) @@ -44,6 +46,11 @@ #: Derived from the upload timeout: a single in-flight sample POST may take #: this long, and a shorter budget would abandon an upload about to succeed. DEFAULT_FINISH_TIMEOUT = float(UPLOAD_TIMEOUT.read or 300.0) +#: A training run lasts days and a platform deploy lasts minutes: a training +#: sink that struck out on transient failures is tried again after this many +#: seconds instead of being retired for the rest of the run. Eval runs are +#: minutes long and keep the permanent retirement. +TRAIN_RETIRE_COOLDOWN = 300.0 class Run: @@ -64,6 +71,9 @@ def __init__( mode: Mode = "online", on_error: OnError = "warn", finish_timeout: Optional[float] = None, + metrics_sinks: Optional[List[Sink]] = None, + attached: bool = False, + retire_cooldown: Optional[float] = None, ) -> None: self._backend = backend self._handle = handle @@ -73,6 +83,10 @@ def __init__( self._status = RunStatus.RUNNING # A forked child inherits this handle but must not close the parent's run. self._owns_lifecycle = True + # A run the platform created and handed to this process (a managed + # launch): the platform owns its failure marking, this process only + # completes it. + self._attached = attached self.config: Dict[str, Any] = dict(spec.config) self.summary: Dict[str, Any] = {} @@ -95,9 +109,18 @@ def __init__( _fork.register(self) sinks = sinks or [] - self._worker = UploadWorker(sinks, on_error=self._record_sink_error) + metrics_sinks = metrics_sinks or [] + # Records and metrics drain on separate uploaders: a metrics dict is + # not a record (no sink would take both), and a slow sample upload + # must not hold a step's metrics behind it. + self._worker = UploadWorker( + sinks, on_error=self._record_sink_error, retire_cooldown=retire_cooldown + ) + self._metrics_worker = UploadWorker( + metrics_sinks, on_error=self._record_sink_error, retire_cooldown=retire_cooldown + ) context = _sink_context(spec) - for sink in sinks: + for sink in [*sinks, *metrics_sinks]: try: sink.start(handle.id, context) except Exception as exc: # noqa: BLE001 - a bad sink is not a bad run @@ -146,16 +169,26 @@ def status(self) -> RunStatus: def finished(self) -> bool: return self._finished + @property + def kind(self) -> RunKind: + """``eval`` or ``train``.""" + return self._spec.kind + + @property + def attached(self) -> bool: + """Whether this process joined a run the platform had already created.""" + return self._attached + @property def dropped_records(self) -> int: - """Records that reached no sink because the queue was full.""" - return self._worker.dropped + """Records (and metrics) that reached no sink because a queue was full.""" + return self._worker.dropped + self._metrics_worker.dropped @property def failed_records(self) -> Dict[str, int]: """Records each sink could not store, by sink name. Per sink because another sink may still hold them.""" - return dict(self._worker.failed_records) + return {**self._worker.failed_records, **self._metrics_worker.failed_records} def __repr__(self) -> str: return f"" @@ -184,6 +217,24 @@ def log_episodes(self, episodes: Iterable[Any]) -> None: """ self._submit("log_episodes", episodes) + def log_metrics(self, values: Mapping[str, Any], *, step: Optional[int] = None) -> None: + """Hand one step's metrics to the run. Returns immediately. + + Training runs only (an eval run has no metrics sink; the call is a + no-op there). ``step`` is recorded as ``values["step"]`` unless the + mapping already carries one; every row is stamped with ``_timestamp`` + so step-less rows (inference metrics) keep a time anchor. Non-finite + numbers are dropped, as they are for :meth:`update_summary`. + """ + self._require_live("log_metrics") + record = _clean_metrics(values) + if step is not None: + record.setdefault("step", step) + record.setdefault("_timestamp", time.time()) + with self._state_lock: + self._require_live("log_metrics") + self._metrics_worker.submit([record]) + def _submit(self, operation: str, records: Iterable[Any]) -> None: # Refuse an already-finished run without consuming a lazy iterable. self._require_live(operation) @@ -209,7 +260,13 @@ def update_summary(self, values: Mapping[str, Any]) -> None: def flush(self, timeout: Optional[float] = 30.0) -> bool: """Block until queued records have been written. Under ``on_error="raise"`` this is the first place an upload failure surfaces.""" - flushed = self._worker.flush(timeout=timeout) + deadline = None if timeout is None else time.monotonic() + max(0.0, timeout) + + def remaining() -> Optional[float]: + return None if deadline is None else max(0.0, deadline - time.monotonic()) + + flushed = self._worker.flush(timeout=remaining()) + flushed = self._metrics_worker.flush(timeout=remaining()) and flushed self._raise_deferred() return flushed @@ -266,7 +323,9 @@ def remaining() -> float: # Records first, so a dashboard reacting to the terminal status never # sees a finished run with samples still landing. - if not self._worker.flush(timeout=remaining()): + drained = self._worker.flush(timeout=remaining()) + drained = self._metrics_worker.flush(timeout=remaining()) and drained + if not drained: logger.warning( "Run %s: uploads did not drain within %ss; finalizing anyway. " "Some records may be missing from this run.", @@ -274,6 +333,7 @@ def remaining() -> float: budget, ) self._worker.close(timeout=remaining()) + self._metrics_worker.close(timeout=remaining()) if self._owns_lifecycle: self._teardown_step( @@ -282,27 +342,37 @@ def remaining() -> float: self.id, config=self.config or None, summary=self.summary or None ), ) - self._teardown_step( - "finalizing the run", - lambda: self._backend.finalize( + if self._attached and resolved is not RunStatus.COMPLETED: + # A managed launch marks its own run failed when this process + # goes away; reporting it from here would race that. + logger.info( + "Run %s: %s, but the platform owns the status of an attached run; " + "leaving it to the launcher", self.id, - status=resolved, - summary=self.summary or None, - error=error or (self.errors[0] if self.errors else None), - config=self.config or None, - ), - ) + resolved.value, + ) + else: + self._teardown_step( + "finalizing the run", + lambda: self._backend.finalize( + self.id, + status=resolved, + summary=self.summary or None, + error=error or (self.errors[0] if self.errors else None), + config=self.config or None, + ), + ) self._teardown_step("closing the backend", self._backend.close) atexit.unregister(self._atexit_hook) - if self._worker.dropped: + if self.dropped_records: logger.warning( "Run %s finished with %d record(s) that reached no sink; the producer " "outran the uploader.", self.id, - self._worker.dropped, + self.dropped_records, ) - for sink_name, count in self._worker.failed_records.items(): + for sink_name, count in self.failed_records.items(): logger.warning( "Run %s: the %s sink could not store %d record(s)", self.id, sink_name, count ) @@ -423,6 +493,9 @@ def init( base_url: Optional[str] = None, on_error: OnError = "warn", finish_timeout: Optional[float] = None, + kind: RunKind = "eval", + id: Optional[str] = None, + training: Optional[TrainingSpec] = None, ) -> Run: """Start a run and return a handle to it. @@ -436,12 +509,33 @@ def init( ``finish_timeout`` is how long :meth:`Run.finish` lets queued uploads drain before closing the run out anyway (default: the upload timeout, 300 s); ``finish(timeout=...)`` overrides it per call. + + ``kind="train"`` opens an external training run instead of an evaluation: + ``model`` is the base model, ``environments`` the training environments + (hub ids, passed through), ``training`` the rest of what the dashboard + shows, and a team is required. ``id`` attaches to a training run the + platform already created (a launcher that injected ``$RUN_ID``): nothing + is registered, the platform keeps ownership of the run's failure marking, + and a clean :meth:`Run.finish` still completes it. The run must itself be + an external run — one created through ``POST /rft/external-runs`` — since + the platform's monitoring endpoints answer 400 for a hosted (managed) run. + Training sinks that strike out on transient failures are tried again after + ``TRAIN_RETIRE_COOLDOWN`` seconds rather than retired for the run. """ settings = Config() api_key = api_key if api_key is not None else settings.api_key base_url = base_url or settings.base_url team_id = team_id if team_id is not None else settings.team_id + if kind not in ("eval", "train"): + raise ConfigurationError(f"kind={kind!r} is not one of 'eval' or 'train'") + if id is not None and kind != "train": + raise ConfigurationError( + "id= attaches to an existing training run; an eval run is always created here." + ) + if training is not None and kind != "train": + raise ConfigurationError("training= only applies to kind='train'") + spec = RunSpec( name=name, environments=[EnvironmentRef.coerce(entry) for entry in (environments or [])], @@ -451,14 +545,18 @@ def init( tags=list(tags or []), team_id=team_id, config=_normalize_config(config), + kind=kind, + training=training, ) resolved_mode = _resolve_mode(mode, api_key=api_key) backend: Backend sinks: List[Sink] + metrics_sinks: List[Sink] = [] + attached = False if resolved_mode == "disabled": backend = DisabledBackend() - handle = RunHandle(id=disabled_run_id(), name=name) + handle = RunHandle(id=id or disabled_run_id(), name=name) sinks = [] else: if not api_key: @@ -467,9 +565,17 @@ def init( 'or pass mode="disabled".' ) client = PlatformClient(api_key=api_key, base_url=base_url, timeout=DEFAULT_TIMEOUT) - backend = EvalsBackend(client, frontend_url=settings.frontend_url, team_id=team_id) + if kind == "train": + rft = RftBackend(client, frontend_url=settings.frontend_url, team_id=team_id) + backend = rft + else: + backend = EvalsBackend(client, frontend_url=settings.frontend_url, team_id=team_id) try: - handle = backend.create(spec) + if kind == "train" and id is not None: + handle = rft.attach(id) + attached = True + else: + handle = backend.create(spec) except BaseException: # Ownership has not reached a Run yet, so nothing else can release # the connection pool when environment resolution or creation fails. @@ -483,7 +589,12 @@ def init( raise # Both transports run during the transition: traces is the system of # record, the sample table is what today's viewer reads. - sinks = [TracesSink(api_key=api_key, team_id=team_id), EvalSamplesSink(client)] + traces = TracesSink(api_key=api_key, team_id=team_id) + if kind == "train": + sinks = [traces, RftSamplesSink(client)] + metrics_sinks = [RftMetricsSink(client)] + else: + sinks = [traces, EvalSamplesSink(client)] run = Run( backend=backend, @@ -493,6 +604,9 @@ def init( mode=resolved_mode, on_error=on_error, finish_timeout=finish_timeout, + metrics_sinks=metrics_sinks, + attached=attached, + retire_cooldown=TRAIN_RETIRE_COOLDOWN if kind == "train" else None, ) if run.url: logger.info("Run %s: %s", run.id, run.url) @@ -521,7 +635,7 @@ def _resolve_mode(mode: Optional[Mode], *, api_key: str) -> Mode: def _sink_context(spec: RunSpec) -> Dict[str, str]: """Upload-scoped provenance. Not the join key — that is ``run.id`` inside the trace document.""" - context = {"source": "prime-runs", "run_type": RUN_KIND} + context = {"source": "prime-runs", "run_type": spec.kind or RUN_KIND} if spec.framework: context["framework"] = spec.framework if spec.model: @@ -555,21 +669,42 @@ def _describe(error: Union[str, BaseException]) -> str: return str(error) -def _clean_metrics(metrics: Mapping[str, Any]) -> Dict[str, Any]: - """Drop NaN/infinity: strict JSON rejects them, and the failure would - surface as an opaque 400 on the whole request.""" +def _clean_metrics(metrics: Mapping[str, Any], *, _path: str = "") -> Dict[str, Any]: + """Recursively drop NaN/infinity: strict JSON rejects them, and the + failure would surface as an opaque 400 on the whole request.""" cleaned: Dict[str, Any] = {} for key, value in metrics.items(): + path = f"{_path}.{key}" if _path else str(key) if isinstance(value, float) and not math.isfinite(value): - logger.debug("Dropping non-finite metric %s=%r", key, value) + logger.debug("Dropping non-finite metric %s=%r", path, value) continue if isinstance(value, Mapping): - nested = _clean_metrics(value) + nested = _clean_metrics(value, _path=path) if nested: cleaned[key] = nested continue + if isinstance(value, (list, tuple)): + cleaned[key] = _clean_metric_sequence(value, path) + continue cleaned[key] = value return cleaned +def _clean_metric_sequence(values: Sequence[Any], path: str) -> List[Any]: + """Clean JSON-array-shaped metric values while preserving their order.""" + cleaned: List[Any] = [] + for index, value in enumerate(values): + item_path = f"{path}[{index}]" + if isinstance(value, float) and not math.isfinite(value): + logger.debug("Dropping non-finite metric %s=%r", item_path, value) + continue + if isinstance(value, Mapping): + cleaned.append(_clean_metrics(value, _path=item_path)) + elif isinstance(value, (list, tuple)): + cleaned.append(_clean_metric_sequence(value, item_path)) + else: + cleaned.append(value) + return cleaned + + __all__ = ["Run", "init", "MODE_ENV"] diff --git a/packages/prime-runs/src/prime_runs/sinks/__init__.py b/packages/prime-runs/src/prime_runs/sinks/__init__.py index 6f5a7ffdf..080793f71 100644 --- a/packages/prime-runs/src/prime_runs/sinks/__init__.py +++ b/packages/prime-runs/src/prime_runs/sinks/__init__.py @@ -1,8 +1,10 @@ """Record transports. Independent of backends, and of each other.""" from .base import Sink, is_episode, stamp_run, to_mapping +from .metrics import RftMetricsSink from .samples import EvalSamplesSink from .traces import TracesSink +from .train_samples import RftSamplesSink, training_step __all__ = [ "Sink", @@ -10,5 +12,8 @@ "stamp_run", "to_mapping", "EvalSamplesSink", + "RftMetricsSink", + "RftSamplesSink", "TracesSink", + "training_step", ] diff --git a/packages/prime-runs/src/prime_runs/sinks/base.py b/packages/prime-runs/src/prime_runs/sinks/base.py index 3df8ba903..a9c74c62e 100644 --- a/packages/prime-runs/src/prime_runs/sinks/base.py +++ b/packages/prime-runs/src/prime_runs/sinks/base.py @@ -25,8 +25,7 @@ def __init__(self, cause: Exception, *, failed_records: int) -> None: self.cause = cause self.failed_records = failed_records super().__init__( - f"{failed_records} record(s) were not fully stored: " - f"{type(cause).__name__}: {cause}" + f"{failed_records} record(s) were not fully stored: {type(cause).__name__}: {cause}" ) @@ -74,7 +73,7 @@ def is_episode(record: Any) -> bool: return hasattr(record, "traces") -def stamp_run(mapping: Mapping[str, Any], run_id: str) -> Dict[str, Any]: +def stamp_run(mapping: Mapping[str, Any], run_id: str, run_type: str = RUN_KIND) -> Dict[str, Any]: """A copy of ``mapping`` carrying ``run`` at the top level, and on every member trace of an episode that lacks one. @@ -86,7 +85,7 @@ def stamp_run(mapping: Mapping[str, Any], run_id: str) -> Dict[str, Any]: """ stamped = dict(mapping) if not stamped.get("run"): - stamped["run"] = {"id": run_id, "type": RUN_KIND} + stamped["run"] = {"id": run_id, "type": run_type} members = stamped.get("traces") if isinstance(members, list): run = stamped["run"] diff --git a/packages/prime-runs/src/prime_runs/sinks/metrics.py b/packages/prime-runs/src/prime_runs/sinks/metrics.py new file mode 100644 index 000000000..3a6f38036 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/metrics.py @@ -0,0 +1,79 @@ +"""Training metrics sink: one ``POST /rft/metrics`` per logged step. + +The RFT metrics endpoint takes a single step's metrics per request, so a +coalesced batch costs one request per record. The platform allows 60 a minute +per token; a 429 is retried with the server's ``Retry-After`` before it counts +as a strike. + +Requests are replayed after an ambiguous failure (a 502/504, a read timeout): +the platform keeps one row per step, so a duplicate collapses on its side, +while a lost row is a hole in the training curves for good. +""" + +import logging +from typing import Any, Mapping, Optional, Sequence + +from .._http import PlatformClient +from ..exceptions import is_record_rejection, is_transient +from .base import Sink, SinkWriteError + +logger = logging.getLogger(__name__) + + +class RftMetricsSink(Sink): + """Posts per-step metrics dicts to the RFT metrics endpoint.""" + + name = "rft_metrics" + + def __init__(self, client: PlatformClient) -> None: + self.enabled = True + self._client = client + self._run_id: Optional[str] = None + self.steps_written = 0 + + def start(self, run_id: str, context: Mapping[str, str]) -> None: + self._run_id = run_id + + def write(self, records: Sequence[Any]) -> None: + if not self.enabled or not records: + return + if self._run_id is None: + raise RuntimeError("RftMetricsSink.write called before start()") + + failed = 0 + reported: Optional[Exception] = None + for index, record in enumerate(records): + if not isinstance(record, Mapping): + failed += 1 + if reported is None: + reported = TypeError(f"metrics must be a mapping, got {type(record).__name__}") + continue + try: + # Replayable: the platform merges rows by step, so a retry after + # a lost response cannot duplicate anything visible. + self._client.post( + "/rft/metrics", + json_body={"run_id": self._run_id, "metrics": dict(record)}, + idempotent=True, + ) + except Exception as exc: # noqa: BLE001 - classified below + failed += 1 + if not (is_record_rejection(exc) or is_transient(exc)): + # Sink-wide: the rest would fail the same way. + failed += len(records) - index - 1 + reported = exc + break + # Prefer a transient error for the worker's strike accounting. + if reported is None or is_transient(exc): + reported = exc + else: + self.steps_written += 1 + + if reported is not None: + raise SinkWriteError(reported, failed_records=failed) from reported + + def flush(self) -> None: + """Writes are synchronous; the uploader thread owns the asynchrony.""" + + def close(self) -> None: + """The client is shared with the backend, which closes it.""" diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py index a0520e5f4..cf6a6465b 100644 --- a/packages/prime-runs/src/prime_runs/sinks/samples.py +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -86,9 +86,7 @@ def write(self, records: Sequence[Any]) -> None: reported_error, failed_records=len(failed_owners) ) from reported_error - def _to_samples( - self, records: Sequence[Any] - ) -> Tuple[List[Dict[str, Any]], List[int]]: + def _to_samples(self, records: Sequence[Any]) -> Tuple[List[Dict[str, Any]], List[int]]: """Episode objects are projected; v0 sample dicts (``sample_id``) pass through. Anything else — a JSON episode, a bare trace — has no v0 projection (it is attribute-based, see :mod:`prime_runs.projection`) diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py index 79b2831f9..8da688eba 100644 --- a/packages/prime-runs/src/prime_runs/sinks/traces.py +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -14,6 +14,7 @@ from .. import _fork from ..exceptions import ForbiddenError +from ..models import RUN_KIND from .base import Sink, is_episode, stamp_run, to_mapping logger = logging.getLogger(__name__) @@ -127,7 +128,7 @@ def _prepare(self, record: Any) -> Any: but the members are reachable for stamping.""" if self._run_id is None: return record - return stamp_run(to_mapping(record), self._run_id) + return stamp_run(to_mapping(record), self._run_id, self._context.get("run_type", RUN_KIND)) def flush(self) -> None: """Uploads are synchronous; nothing is held back here.""" diff --git a/packages/prime-runs/src/prime_runs/sinks/train_samples.py b/packages/prime-runs/src/prime_runs/sinks/train_samples.py new file mode 100644 index 000000000..5f3472b6c --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/train_samples.py @@ -0,0 +1,302 @@ +"""Training samples sink: the step-keyed Parquet table the training viewer reads. + +What prime-rl's ``TrainRun.upload_samples`` did, in the SDK: for each training +step, encode the step's episodes to one Parquet object and send it through the +platform's presign -> PUT to object storage -> confirm flow. Like the eval +samples sink this serves *today's* viewer; Prime Traces is the system of record +and this sink leaves the default list once the viewer reads traces natively. + +Which episodes to upload is read off the records: a verifiers episode carries +``run.work`` (``TrainWorkInfo(step=...)``) once its producer has stamped it, and +prime-rl stamps every dispatched episode. Episodes from eval work, bare traces +and JSON episodes have no row here and go to Prime Traces only. + +Objects are additive per step: every upload mints its own +``step_{step}_{uuid}.parquet`` key and the viewer unions every object under a +step, so a step's episodes may arrive across several ``log_episodes`` calls — +prime-rl hands over a step's batch at ship time, and an off-policy episode +dispatched at step N can land in a later batch. Each call becomes one more +object; ``sample_id`` is numbered from a per-upload offset so (step, +sample_id), which the viewer looks samples up by, stays unique across them. +Log a step in one call to keep it to one object. +""" + +import logging +import secrets +import time +from typing import Any, Dict, List, Mapping, Optional, Protocol, Sequence, Tuple + +import httpx +from prime_traces.core.client import retry_delay + +from .. import _fork +from .._http import UPLOAD_TIMEOUT, PlatformClient +from ..exceptions import ( + APIError, + RetryableAPIError, + TransportError, + is_record_rejection, + is_transient, +) +from .base import Sink, SinkWriteError, is_episode + +logger = logging.getLogger(__name__) + +#: Upload every Nth training step (prime-rl's cadence, chosen so the sample +#: table does not receive every rollout of every step). +DEFAULT_STEP_INTERVAL = 10 +#: Attempts for the object-storage PUT, which is idempotent. +UPLOAD_ATTEMPTS = 3 + +#: ``sample_id = range * SAMPLE_ID_STRIDE + row``. In the process that opened +#: the run the range is the step's upload index: the first object numbers its +#: rows from 0 (what a single-object step always did), a straggler object from +#: 2**32, and so on; no object holds 2**32 rows. A forked child inherits a copy +#: of those counters and would hand out the same ranges as its parent, so it +#: draws its ranges at random from the upper half of the range space instead, +#: which the parent's counter never reaches. Every id stays below 2**53: the +#: viewer parses ``sample_id`` as a JSON number. +SAMPLE_ID_STRIDE = 1 << 32 +SAMPLE_ID_RANGES = 1 << 21 + + +class Encoder(Protocol): + """``(episodes, run_id, step, sample_id_offset=...) -> parquet bytes``, or + ``None`` when there is nothing to upload for the step. Rows are numbered + from ``sample_id_offset``.""" + + def __call__( + self, episodes: Sequence[Any], run_id: str, step: int, *, sample_id_offset: int = 0 + ) -> Optional[bytes]: ... + + +def _field(obj: Any, name: str) -> Any: + return obj.get(name) if isinstance(obj, Mapping) else getattr(obj, name, None) + + +def training_step(record: Any) -> Optional[int]: + """The training step an episode was dispatched at, from its ``run.work`` + (verifiers ``TrainRunInfo`` / ``TrainWorkInfo``). ``None`` for anything + else: eval work, a record without provenance, a bare trace.""" + run = _field(record, "run") + work = _field(run, "work") if run is not None else None + if _field(run, "type") != "train" or _field(work, "type") != "train": + return None + step = _field(work, "step") + return step if isinstance(step, int) and not isinstance(step, bool) else None + + +class RftSamplesSink(Sink): + """Encodes each training step's episodes to Parquet and uploads it.""" + + name = "rft_samples" + + def __init__( + self, + client: PlatformClient, + *, + encoder: Optional[Encoder] = None, + step_interval: int = DEFAULT_STEP_INTERVAL, + upload_client: Optional[httpx.Client] = None, + ) -> None: + if step_interval < 1: + raise ValueError("step_interval must be at least 1") + self.enabled = True + self._client = client + self._encoder = encoder + self._step_interval = step_interval + self._upload_client = upload_client + self._owns_upload_client = upload_client is None + # An inherited socket belongs to the parent; see reset_after_fork. + self._forked_with_injected_client = False + self._run_id: Optional[str] = None + self.steps_written = 0 + #: Records with no row in this table (eval work, bare traces, JSON). + self.skipped = 0 + #: Training episodes left out by the step cadence. Not a loss. + self.sampled_out = 0 + #: Uploads attempted per step; each numbers its rows in its own range. + self._uploads_per_step: Dict[int, int] = {} + #: A forked child cannot share the parent's counters; see SAMPLE_ID_RANGES. + self._forked = False + _fork.register(self) + + # ------------------------------------------------------------------ setup + + def reset_after_fork(self) -> None: + """Drop (not close) the inherited storage client: its socket is the + parent's. An owned client is rebuilt on the next upload; an injected + one cannot be, so the child's uploads fail rather than share it. The + inherited upload counters are the parent's too: from here on ranges + are drawn at random, apart from anything the parent hands out.""" + if self._owns_upload_client: + self._upload_client = None + else: + self._forked_with_injected_client = True + self._forked = True + + def start(self, run_id: str, context: Mapping[str, str]) -> None: + self._run_id = run_id + if self._encoder is None: + from ..projection import episodes_to_parquet_bytes, parquet_available + + if not parquet_available(): + # Not a failure: nothing was lost, the episodes still reach + # Prime Traces. Say once how to turn the table on. + logger.warning( + "Training samples sink off: pyarrow is not installed " + "(pip install 'prime-runs[train]'); episodes reach Prime Traces only" + ) + self.enabled = False + return + self._encoder = episodes_to_parquet_bytes + + # ------------------------------------------------------------------ write + + def write(self, records: Sequence[Any]) -> None: + if not self.enabled or not records: + return + if self._run_id is None: + raise RuntimeError("RftSamplesSink.write called before start()") + assert self._encoder is not None # start() set it or disabled the sink + + steps = self._group(records) + failed = 0 + reported: Optional[Exception] = None + for position, (step, episodes) in enumerate(steps): + try: + uploaded = self._upload_step(step, episodes) + except Exception as exc: # noqa: BLE001 - classified below + failed += len(episodes) + if not (is_record_rejection(exc) or is_transient(exc)): + # Sink-wide: the remaining steps would fail the same way. + failed += sum(len(rest) for _, rest in steps[position + 1 :]) + reported = exc + break + if reported is None or is_transient(exc): + reported = exc + else: + if uploaded: + self.steps_written += 1 + + if reported is not None: + raise SinkWriteError(reported, failed_records=failed) from reported + + def _group(self, records: Sequence[Any]) -> List[Tuple[int, List[Any]]]: + """Training episodes by step, on the upload cadence, in step order.""" + by_step: Dict[int, List[Any]] = {} + skipped = 0 + for record in records: + step = training_step(record) + # The encoder projects attributes (``build_samples``), so a JSON + # episode has no row here either. + if step is None or isinstance(record, Mapping) or not is_episode(record): + skipped += 1 + continue + if step % self._step_interval: + self.sampled_out += 1 + continue + by_step.setdefault(step, []).append(record) + if skipped: + if not self.skipped: + logger.info( + "The training sample table holds training-work episodes only; %d " + "record(s) in this batch have no row there and reach Prime Traces only.", + skipped, + ) + self.skipped += skipped + return sorted(by_step.items()) + + def _upload_step(self, step: int, episodes: List[Any]) -> bool: + """Encode and upload one step. ``False`` when the step had nothing to + show (no trajectories) and no request was made.""" + assert self._encoder is not None and self._run_id is not None + # Every upload of a step numbers its rows in its own range, counted by + # attempt: an upload whose confirm was lost may still have landed. + payload = self._encoder( + episodes, self._run_id, step, sample_id_offset=self._next_range(step) * SAMPLE_ID_STRIDE + ) + if payload is None: + return False + + # Replayable: a presign only mints a URL. + presign = self._client.post( + "/rft/samples/presign", + json_body={"run_id": self._run_id, "step": step}, + idempotent=True, + ) + data = presign.get("data") or presign + url = data.get("presignedUrl") or data.get("presigned_url") + key = data.get("s3Key") or data.get("s3_key") + if not url or not key: + raise APIError( + f"POST /rft/samples/presign returned no upload URL/key (keys: {sorted(data)})" + ) + + self._put_object(url, payload) + + # Records the object under the step. Replayable: confirm checks the + # key and the object, then refreshes the run's progress; a second + # confirm of the same key changes nothing. + self._client.post( + "/rft/samples/confirm", + json_body={"run_id": self._run_id, "step": step, "s3_key": key}, + idempotent=True, + ) + return True + + def _next_range(self, step: int) -> int: + """The ``sample_id`` range for the next upload of ``step``.""" + if self._forked: + half = SAMPLE_ID_RANGES // 2 + return half + secrets.randbelow(half) + upload_index = self._uploads_per_step.get(step, 0) + self._uploads_per_step[step] = upload_index + 1 + return upload_index + + def _put_object(self, url: str, payload: bytes) -> None: + """PUT the object to storage. Bare client: the presigned URL carries its + own credentials and rejects the platform's auth headers. Idempotent, so + transport failures and 5xx are retried.""" + client = self._upload() + for attempt in range(UPLOAD_ATTEMPTS): + error: APIError + try: + response = client.put( + url, + content=payload, + headers={"Content-Type": "application/parquet"}, + timeout=UPLOAD_TIMEOUT, + ) + except httpx.RequestError as exc: + error = TransportError(f"uploading samples failed: {type(exc).__name__}: {exc}") + else: + if response.is_success: + return + message = f"uploading samples failed: HTTP {response.status_code}" + if response.status_code >= 500 or response.status_code == 429: + error = RetryableAPIError(message, status_code=response.status_code) + else: + raise APIError(message, status_code=response.status_code) + if attempt == UPLOAD_ATTEMPTS - 1: + raise error + time.sleep(retry_delay(error, attempt)) + + def _upload(self) -> httpx.Client: + if self._forked_with_injected_client: + raise RuntimeError("an injected upload client cannot be reused after a fork") + if self._upload_client is None: + self._upload_client = httpx.Client(follow_redirects=False) + return self._upload_client + + def flush(self) -> None: + """Uploads are synchronous; nothing is held back here.""" + + def close(self) -> None: + client = self._upload_client + self._upload_client = None + if client is not None and self._owns_upload_client: + try: + client.close() + except Exception as exc: # noqa: BLE001 - teardown must not raise + logger.debug("Error closing the upload client: %s", exc) diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index 560fea593..3a912d0e3 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -8,7 +8,9 @@ Containment: record-specific failures drop only their batch. A sink-wide failure disables the sink, while transient failures get a few consecutive -strikes first. The thread never dies on one bad batch. +strikes first — and, with a ``retire_cooldown``, the sink is tried again once +the cooldown has passed, so an outage costs a long run a window of records +rather than the rest of the run. The thread never dies on one bad batch. """ import logging @@ -63,11 +65,15 @@ def __init__( max_queue_size: int = DEFAULT_QUEUE_SIZE, put_timeout: float = DEFAULT_PUT_TIMEOUT, on_error: Optional[Callable[[str, Exception], None]] = None, + retire_cooldown: Optional[float] = None, ) -> None: self.sinks = sinks self.max_queue_size = max_queue_size self.put_timeout = put_timeout self._on_error = on_error + #: Seconds after which a sink retired on transient failures is tried + #: again; ``None`` retires it for the rest of the run. + self._retire_cooldown = retire_cooldown self._queue: "queue.Queue[Any]" = queue.Queue(maxsize=max_queue_size) self._thread: Optional[threading.Thread] = None self._stopping = threading.Event() @@ -82,6 +88,8 @@ def __init__( #: these are lost to it and counted; a sink that switched itself off #: without raising (nowhere for the records to go) is not in here. self._retired: set = set() + #: Sink name -> monotonic time at which a paused sink is tried again. + self._cooldowns: dict = {} _fork.register(self) # ----------------------------------------------------------------- thread @@ -154,6 +162,7 @@ def _coalesce(self, first: Sequence[Any]) -> "tuple[list, Any]": def _dispatch(self, records: Sequence[Any]) -> None: for sink in self.sinks: + self._maybe_revive(sink) if not sink.enabled: if sink.name in self._retired: count = self.failed_records.get(sink.name, 0) @@ -168,6 +177,17 @@ def _dispatch(self, records: Sequence[Any]) -> None: else: self._transient_failures.pop(sink.name, None) + def _maybe_revive(self, sink: Sink) -> None: + """Re-enable a sink whose cooldown has passed. Its strikes start over.""" + until = self._cooldowns.get(sink.name) + if until is None or time.monotonic() < until: + return + del self._cooldowns[sink.name] + self._retired.discard(sink.name) + self._transient_failures.pop(sink.name, None) + sink.enabled = True + logger.info("Sink %s re-enabled after its cooldown", sink.name) + def _flush_sinks(self) -> None: for sink in self.sinks: if not sink.enabled: @@ -182,7 +202,8 @@ def _fail_sink(self, sink: Sink, exc: Exception, *, dropped: int = 0) -> None: Record-specific failures drop only the current batch. Sink-wide permanent failures retire immediately; transient failures retire after - ``TRANSIENT_FAILURE_LIMIT`` consecutive strikes. + ``TRANSIENT_FAILURE_LIMIT`` consecutive strikes — for good, or until + the worker's ``retire_cooldown`` has passed. """ name = sink.name if dropped: @@ -218,13 +239,24 @@ def _fail_sink(self, sink: Sink, exc: Exception, *, dropped: int = 0) -> None: return sink.enabled = False self._retired.add(name) - logger.warning( - "Sink %s disabled after %d consecutive transient failures: %s: %s", - name, - strikes, - type(exc).__name__, - exc, - ) + if self._retire_cooldown is None: + logger.warning( + "Sink %s disabled after %d consecutive transient failures: %s: %s", + name, + strikes, + type(exc).__name__, + exc, + ) + else: + self._cooldowns[name] = time.monotonic() + self._retire_cooldown + logger.warning( + "Sink %s paused for %.0fs after %d consecutive transient failures: %s: %s", + name, + self._retire_cooldown, + strikes, + type(exc).__name__, + exc, + ) else: sink.enabled = False self._retired.add(name) diff --git a/packages/prime-runs/tests/_fakes.py b/packages/prime-runs/tests/_fakes.py index 4cef94549..00522ae49 100644 --- a/packages/prime-runs/tests/_fakes.py +++ b/packages/prime-runs/tests/_fakes.py @@ -74,11 +74,41 @@ def to_record(self) -> Dict[str, Any]: return {"id": self.id, "reward": self.reward} +@dataclass +class EnvInfo: + id: str = "gsm8k" + name: Optional[str] = None + + +@dataclass +class TrainWorkInfo: + step: int + type: str = "train" + + +@dataclass +class EvalWorkInfo: + step: int + type: str = "eval" + + +@dataclass +class TrainRunInfo: + """verifiers stamps this on every episode a training run dispatches.""" + + id: str + work: Any + type: str = "train" + name: Optional[str] = None + + @dataclass class Episode: id: str = "episode-1" traces: List[Trace] = field(default_factory=list) ok: bool = True + env: EnvInfo = field(default_factory=EnvInfo) + run: Optional[Any] = None def to_record(self) -> Dict[str, Any]: return {"id": self.id, "traces": [trace.to_record() for trace in self.traces]} @@ -113,3 +143,28 @@ def make_episode( episode_id: str = "episode-1", traces: Optional[List[Trace]] = None, ok: bool = True ) -> Episode: return Episode(id=episode_id, traces=traces if traces is not None else [make_trace()], ok=ok) + + +def make_train_episode( + episode_id: str = "episode-1", + *, + step: int = 10, + work: str = "train", + idx: int = 0, + reward: float = 1.0, + advantage: Optional[float] = None, + env_id: str = "gsm8k", + run_id: str = "run-1", +) -> Episode: + """An episode the way prime-rl hands it over: stamped with the run and the + training step it was dispatched at, ``advantage`` in the trace's info.""" + trace = make_trace(trace_id=f"{episode_id}-t", idx=idx, reward=reward) + if advantage is not None: + trace.info["advantage"] = advantage + work_info = TrainWorkInfo(step=step) if work == "train" else EvalWorkInfo(step=step) + return Episode( + id=episode_id, + traces=[trace], + env=EnvInfo(id=env_id), + run=TrainRunInfo(id=run_id, work=work_info), + ) diff --git a/packages/prime-runs/tests/conftest.py b/packages/prime-runs/tests/conftest.py index 9c7166c9c..a24a62587 100644 --- a/packages/prime-runs/tests/conftest.py +++ b/packages/prime-runs/tests/conftest.py @@ -141,3 +141,55 @@ def close(self) -> None: @pytest.fixture def fake_sink() -> Callable[..., FakeSink]: return FakeSink + + +@pytest.fixture +def rft_routes() -> Dict[str, Any]: + """The happy path for a training run: register, metrics, samples, finalize.""" + + def presign(request: httpx.Request) -> httpx.Response: + import json + + step = json.loads(request.content)["step"] + return httpx.Response( + 200, + json={ + "data": { + "presignedUrl": f"http://storage.test/run-abc/step-{step}.parquet", + "s3Key": f"runs/run-abc/step-{step}.parquet", + "expiresIn": 900, + } + }, + ) + + return { + "POST /api/v1/rft/external-runs": lambda request: httpx.Response( + 201, + json={"run": {"id": "run-abc", "name": "test-run", "status": "RUNNING"}}, + ), + "POST /api/v1/rft/metrics": {"data": {"status": "success"}}, + "POST /api/v1/rft/samples/presign": presign, + "POST /api/v1/rft/samples/confirm": {"data": {"status": "success"}}, + "POST /api/v1/rft/finalize": { + "data": {"status": "success", "message": "Run finalized successfully"} + }, + "PUT /api/v1/rft/external-runs/run-abc/status": { + "run": {"id": "run-abc", "status": "FAILED"} + }, + } + + +class StorageHandler: + """Object storage for presigned PUTs: records every upload, answers 200.""" + + def __init__(self, status_codes: Sequence[int] = ()) -> None: + self.uploads: List[httpx.Request] = [] + self._status_codes = list(status_codes) + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.uploads.append(request) + code = self._status_codes.pop(0) if self._status_codes else 200 + return httpx.Response(code) + + def client(self) -> httpx.Client: + return httpx.Client(transport=httpx.MockTransport(self)) diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index b7a8f3273..ca0debdc0 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -8,7 +8,7 @@ import prime_runs as pr from prime_runs.exceptions import ConfigurationError -from prime_runs.models import RunStatus +from prime_runs.models import RunSpec, RunStatus # -------------------------------------------------------------------- modes @@ -112,6 +112,13 @@ def test_an_online_run_returns_the_platforms_id_and_viewer_url(online): run.finish() +def test_eval_uploaders_keep_the_permanent_retirement(online): + run, _ = online() + + assert run._worker._retire_cooldown is None + run.finish() + + def test_a_failed_create_closes_the_platform_client(monkeypatch): class FailingClient: def __init__(self): @@ -259,3 +266,206 @@ def test_records_reject_nonfinite_json_instead_of_sending_an_opaque_400(online): with pytest.raises(ValueError, match="Out of range float values"): run.flush() run.finish() + + +# ----------------------------------------------------------------- training + + +@pytest.fixture +def online_train(monkeypatch, make_platform_client, rft_routes): + """``init(kind="train", mode="online")`` wired to a MockTransport: traces + sink off, sample uploads to an in-memory store, a recording encoder.""" + from conftest import StorageHandler + + from prime_runs.sinks import RftSamplesSink + + storage = StorageHandler() + encoder_calls = [] + + def encoder(episodes, run_id, step, sample_id_offset=0): + encoder_calls.append((len(episodes), run_id, step)) + return b"parquet" + + def samples_sink(client, **kwargs): + return RftSamplesSink(client, encoder=encoder, upload_client=storage.client(), **kwargs) + + def _init(routes=None, **kwargs): + handler = RecordingHandler(routes or rft_routes) + monkeypatch.setattr( + "prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler) + ) + monkeypatch.setattr("prime_runs.run.TracesSink", lambda **_: _NullSink()) + monkeypatch.setattr("prime_runs.run.RftSamplesSink", samples_sink) + # The RFT create response carries no viewer URL; the SDK builds one. + monkeypatch.setenv("PRIME_FRONTEND_URL", "https://app.example") + kwargs.setdefault("name", "test-run") + kwargs.setdefault("model", "Qwen/Qwen3-8B") + kwargs.setdefault("environments", ["primeintellect/vf-math"]) + kwargs.setdefault("team_id", "team-1") + run = pr.init(kind="train", api_key="test-key", **kwargs) + return run, handler, storage, encoder_calls + + return _init + + +def test_a_training_run_registers_with_the_rft_api(online_train): + run, handler, _, _ = online_train( + training=pr.TrainingSpec(max_steps=50, batch_size=32), + config={"trainer": {"lr": 1e-6}}, + ) + + assert run.kind == "train" + assert run.id == "run-abc" + assert run.url == "https://app.example/dashboard/training/run-abc" + assert run.attached is False + body = handler.bodies_for("/api/v1/rft/external-runs")[0] + assert body["base_model"] == "Qwen/Qwen3-8B" + assert body["max_steps"] == 50 + assert body["batch_size"] == 32 + assert body["environments"] == [{"id": "primeintellect/vf-math"}] + assert body["run_config"] == {"trainer": {"lr": 1e-6}} + assert body["team_id"] == "team-1" + assert [sink.name for sink in run._worker.sinks] == ["traces", "rft_samples"] + assert [sink.name for sink in run._metrics_worker.sinks] == ["rft_metrics"] + run.finish() + + +def test_a_training_run_streams_metrics_and_step_samples(online_train): + from _fakes import make_train_episode + + run, handler, storage, encoder_calls = online_train() + + run.log_metrics({"loss": 0.5}, step=10) + run.log_episodes([make_train_episode("e1", step=10), make_train_episode("e2", step=10)]) + run.finish() + + metrics = handler.bodies_for("/api/v1/rft/metrics") + assert len(metrics) == 1 + assert metrics[0]["run_id"] == "run-abc" + assert metrics[0]["metrics"]["loss"] == 0.5 + assert metrics[0]["metrics"]["step"] == 10 + assert "_timestamp" in metrics[0]["metrics"] + assert encoder_calls == [(2, "run-abc", 10)] + assert len(storage.uploads) == 1 + assert handler.bodies_for("/api/v1/rft/samples/confirm")[0]["step"] == 10 + # Closed out through the idempotent finalize, not the status PUT. + assert handler.bodies_for("/api/v1/rft/finalize") == [{"run_id": "run-abc", "exit_code": 0}] + assert "PUT /api/v1/rft/external-runs/run-abc/status" not in handler.paths() + assert run.status is RunStatus.COMPLETED + assert run.errors == [] + + +def test_a_failed_training_run_is_marked_failed_with_the_reason(online_train): + run, handler, _, _ = online_train() + + with pytest.raises(ValueError): + with run: + raise ValueError("loss is NaN") + + body = handler.bodies_for("/api/v1/rft/external-runs/run-abc/status")[0] + assert body == {"status": "failed", "error_message": "ValueError: loss is NaN"} + assert "POST /api/v1/rft/finalize" not in handler.paths() + + +def test_training_uploaders_pause_rather_than_retire_after_an_outage(online_train): + """A training run lasts days; a sink that struck out on a platform blip + is tried again after a cooldown instead of losing the rest of the run.""" + from prime_runs.run import TRAIN_RETIRE_COOLDOWN + + run, _, _, _ = online_train() + + assert run._worker._retire_cooldown == TRAIN_RETIRE_COOLDOWN == 300.0 + assert run._metrics_worker._retire_cooldown == TRAIN_RETIRE_COOLDOWN + run.finish() + + +def test_attaching_to_a_managed_run_registers_nothing(online_train): + run, handler, _, _ = online_train(id="run-managed") + + assert run.id == "run-managed" + assert run.attached is True + assert run.url == "https://app.example/dashboard/training/run-managed" + assert "POST /api/v1/rft/external-runs" not in handler.paths() + run.finish() + # A clean exit still completes it. + assert handler.bodies_for("/api/v1/rft/finalize") == [{"run_id": "run-managed", "exit_code": 0}] + + +def test_an_attached_run_leaves_failure_marking_to_the_launcher(online_train, caplog): + """A managed launch marks its own run failed when the process goes away; + reporting it from here would race that (prime-rl's contract).""" + run, handler, _, _ = online_train(id="run-managed") + + with caplog.at_level("INFO"): + run.fail("boom") + + assert not any("status" in path for path in handler.paths()) + assert "POST /api/v1/rft/finalize" not in handler.paths() + assert "leaving it to the launcher" in caplog.text + assert run.status is RunStatus.FAILED + + +def test_a_training_run_needs_a_team(monkeypatch, make_platform_client, rft_routes): + handler = RecordingHandler(rft_routes) + monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler)) + + with pytest.raises(ConfigurationError, match="team"): + pr.init(kind="train", model="Qwen/Qwen3-8B", api_key="test-key", mode="online") + assert handler.paths() == [] + + +def test_a_team_outside_the_allowlist_surfaces_the_platforms_reason( + monkeypatch, make_platform_client, rft_routes +): + import httpx + + routes = dict(rft_routes) + routes["POST /api/v1/rft/external-runs"] = lambda request: httpx.Response( + 403, json={"detail": "External training runs are not enabled for this team"} + ) + handler = RecordingHandler(routes) + monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler)) + + with pytest.raises(pr.ForbiddenError, match="not enabled for this team"): + pr.init(kind="train", model="m", team_id="team-1", api_key="test-key", mode="online") + + +def test_a_disabled_training_run_keeps_an_attached_id(tmp_path): + run = pr.init(kind="train", mode="disabled", id="run-managed", model="m") + + assert run.id == "run-managed" + assert run.kind == "train" + run.log_metrics({"loss": 0.1}, step=1) + assert run._metrics_worker._thread is None + run.finish() + assert run.status is RunStatus.COMPLETED + + +@pytest.mark.parametrize( + "kwargs", + [ + {"kind": "eval", "id": "eval-1"}, + {"kind": "eval", "training": pr.TrainingSpec()}, + {"kind": "serve"}, + ], + ids=["eval-with-id", "eval-with-training", "unknown-kind"], +) +def test_kind_arguments_are_checked_before_anything_else(kwargs): + with pytest.raises(ConfigurationError): + pr.init(mode="disabled", **kwargs) + + +def test_training_records_are_stamped_with_the_train_run_type(): + """Bare dicts get ``run.type`` from the run's kind, matching what verifiers' + ``TrainRunInfo`` carries, so the traces service files them as training.""" + from prime_runs.run import _sink_context + from prime_runs.sinks.base import stamp_run + + spec = RunSpec(kind="train", model="m") + context = _sink_context(spec) + + assert context["run_type"] == "train" + assert stamp_run({"id": "t1"}, "run-abc", context["run_type"])["run"] == { + "id": "run-abc", + "type": "train", + } diff --git a/packages/prime-runs/tests/test_metrics_sink.py b/packages/prime-runs/tests/test_metrics_sink.py new file mode 100644 index 000000000..624d7dbe0 --- /dev/null +++ b/packages/prime-runs/tests/test_metrics_sink.py @@ -0,0 +1,102 @@ +"""The training metrics sink: one request per logged step.""" + +import httpx +import pytest +from conftest import RecordingHandler + +from prime_runs.exceptions import RetryableAPIError, UnauthorizedError +from prime_runs.sinks import RftMetricsSink +from prime_runs.sinks.base import SinkWriteError + + +def make_sink(make_platform_client, routes): + handler = RecordingHandler(routes) + sink = RftMetricsSink(make_platform_client(handler)) + sink.start("run-abc", {}) + return sink, handler + + +def test_each_metrics_dict_is_one_request(make_platform_client, rft_routes): + sink, handler = make_sink(make_platform_client, rft_routes) + + sink.write([{"step": 1, "loss": 0.9}, {"step": 2, "loss": 0.8}]) + + assert handler.bodies_for("/api/v1/rft/metrics") == [ + {"run_id": "run-abc", "metrics": {"step": 1, "loss": 0.9}}, + {"run_id": "run-abc", "metrics": {"step": 2, "loss": 0.8}}, + ] + assert sink.steps_written == 2 + + +def test_an_ambiguous_failure_is_replayed(make_platform_client, rft_routes, no_sleep): + """The platform keeps one row per step, so a retry after a lost response + cannot double anything, while a lost row is a hole in the curves for good.""" + calls = [] + + def flaky(request: httpx.Request) -> httpx.Response: + calls.append(request) + if len(calls) == 1: + return httpx.Response(502) + return httpx.Response(200, json={"data": {"status": "success"}}) + + routes = dict(rft_routes) + routes["POST /api/v1/rft/metrics"] = flaky + sink, handler = make_sink(make_platform_client, routes) + + sink.write([{"step": 1}, {"step": 2}]) + + assert len(calls) == 3 + assert sink.steps_written == 2 + + +def test_a_persistent_transient_failure_loses_only_its_step( + make_platform_client, rft_routes, no_sleep +): + import json + + def down_for_step_one(request: httpx.Request) -> httpx.Response: + if json.loads(request.content)["metrics"]["step"] == 1: + return httpx.Response(502) + return httpx.Response(200, json={"data": {"status": "success"}}) + + routes = dict(rft_routes) + routes["POST /api/v1/rft/metrics"] = down_for_step_one + sink, handler = make_sink(make_platform_client, routes) + + with pytest.raises(SinkWriteError) as info: + sink.write([{"step": 1}, {"step": 2}]) + + assert info.value.failed_records == 1 + assert isinstance(info.value.cause, RetryableAPIError) + assert sink.steps_written == 1 + + +def test_a_sink_wide_failure_stops_the_batch_and_counts_the_rest(make_platform_client, rft_routes): + routes = dict(rft_routes) + routes["POST /api/v1/rft/metrics"] = lambda request: httpx.Response(401) + sink, handler = make_sink(make_platform_client, routes) + + with pytest.raises(SinkWriteError) as info: + sink.write([{"step": 1}, {"step": 2}, {"step": 3}]) + + assert info.value.failed_records == 3 + assert isinstance(info.value.cause, UnauthorizedError) + assert len(handler.paths()) == 1 + + +def test_a_record_that_is_not_a_mapping_is_counted_not_sent(make_platform_client, rft_routes): + sink, handler = make_sink(make_platform_client, rft_routes) + + with pytest.raises(SinkWriteError) as info: + sink.write(["loss=0.5", {"step": 1}]) + + assert info.value.failed_records == 1 + assert isinstance(info.value.cause, TypeError) + assert len(handler.bodies_for("/api/v1/rft/metrics")) == 1 + + +def test_write_before_start_is_a_producer_bug(make_platform_client, rft_routes): + sink = RftMetricsSink(make_platform_client(RecordingHandler(rft_routes))) + + with pytest.raises(RuntimeError, match="before start"): + sink.write([{"step": 1}]) diff --git a/packages/prime-runs/tests/test_rft_backend.py b/packages/prime-runs/tests/test_rft_backend.py new file mode 100644 index 000000000..22c2e9ef8 --- /dev/null +++ b/packages/prime-runs/tests/test_rft_backend.py @@ -0,0 +1,240 @@ +"""Training run lifecycle against ``/api/v1/rft/*``.""" + +import httpx +import pytest +from conftest import RecordingHandler + +from prime_runs.backend import RftBackend +from prime_runs.exceptions import ConfigurationError, ForbiddenError, RetryableAPIError +from prime_runs.models import EnvironmentRef, RunSpec, RunStatus, TrainingSpec + + +def make_backend(make_platform_client, routes, **kwargs): + handler = RecordingHandler(routes) + client = make_platform_client(handler) + kwargs.setdefault("team_id", "team-1") + backend = RftBackend(client, frontend_url="https://app.example", **kwargs) + return backend, handler + + +def train_spec(**overrides) -> RunSpec: + values = dict( + name="test-run", + kind="train", + model="Qwen/Qwen3-8B", + environments=[EnvironmentRef.coerce("primeintellect/vf-math")], + training=TrainingSpec(max_steps=100, batch_size=64, rollouts_per_example=8, seq_len=4096), + config={"trainer": {"lr": 1e-6}}, + ) + values.update(overrides) + return RunSpec(**values) + + +def test_create_registers_an_external_run_for_the_team(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes) + + handle = backend.create(train_spec()) + + assert handler.paths() == ["POST /api/v1/rft/external-runs"] + body = handler.bodies_for("/api/v1/rft/external-runs")[0] + assert body["base_model"] == "Qwen/Qwen3-8B" + assert body["max_steps"] == 100 + assert body["batch_size"] == 64 + assert body["rollouts_per_example"] == 8 + assert body["seq_len"] == 4096 + assert body["run_config"] == {"trainer": {"lr": 1e-6}} + assert body["team_id"] == "team-1" + assert body["name"] == "test-run" + assert handle.id == "run-abc" + assert handle.url == "https://app.example/dashboard/training/run-abc" + + +@pytest.mark.parametrize( + ("environment", "expected"), + [ + ("primeintellect/vf-math", {"id": "primeintellect/vf-math"}), + ("agentic-judge+gsm8k", {"id": "agentic-judge+gsm8k"}), + ({"id": "env-123", "version_id": "v7"}, {"id": "env-123", "version_id": "v7"}), + ], + ids=["slug", "name", "id"], +) +def test_training_environments_are_passed_through_not_resolved( + make_platform_client, rft_routes, environment, expected +): + """The RFT API names environments by hub id and resolves them itself; a + hub round-trip here would be a second opinion.""" + backend, handler = make_backend(make_platform_client, rft_routes) + + backend.create(train_spec(environments=[EnvironmentRef.coerce(environment)])) + + assert handler.bodies_for("/api/v1/rft/external-runs")[0]["environments"] == [expected] + assert "POST /api/v1/environmentshub/resolve" not in handler.paths() + + +def test_the_spec_team_wins_over_the_backend_default(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes, team_id="team-default") + + backend.create(train_spec(team_id="team-explicit")) + + assert handler.bodies_for("/api/v1/rft/external-runs")[0]["team_id"] == "team-explicit" + + +def test_create_needs_a_base_model(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes) + + with pytest.raises(ConfigurationError, match="model="): + backend.create(train_spec(model=None)) + assert handler.paths() == [] + + +def test_create_needs_a_team(make_platform_client, rft_routes): + """External runs are created for a team; the API refuses without one, so + say so before the request rather than surfacing its 403.""" + backend, handler = make_backend(make_platform_client, rft_routes, team_id=None) + + with pytest.raises(ConfigurationError, match="team"): + backend.create(train_spec()) + assert handler.paths() == [] + + +def test_a_team_outside_the_allowlist_is_a_forbidden_error(make_platform_client, rft_routes): + """The platform enables external runs per team. The producer decides what + a refusal means; the SDK passes the platform's reason through.""" + routes = dict(rft_routes) + routes["POST /api/v1/rft/external-runs"] = lambda request: httpx.Response( + 403, json={"detail": "External training runs are not enabled for this team"} + ) + backend, handler = make_backend(make_platform_client, routes) + + with pytest.raises(ForbiddenError, match="not enabled for this team"): + backend.create(train_spec()) + assert len(handler.paths()) == 1 + + +def test_create_is_not_replayed_after_an_ambiguous_failure(make_platform_client, rft_routes): + """A retry after a lost response would register a second run.""" + routes = dict(rft_routes) + routes["POST /api/v1/rft/external-runs"] = lambda request: httpx.Response(502) + backend, handler = make_backend(make_platform_client, routes) + + with pytest.raises(RetryableAPIError): + backend.create(train_spec()) + assert handler.paths() == ["POST /api/v1/rft/external-runs"] + + +def test_attach_makes_no_request(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes) + + handle = backend.attach("run-managed") + + assert handle.id == "run-managed" + assert handle.url == "https://app.example/dashboard/training/run-managed" + assert handler.paths() == [] + + +def test_update_is_a_no_op(make_platform_client, rft_routes): + """The RFT API takes the config at creation and has no summary document.""" + backend, handler = make_backend(make_platform_client, rft_routes) + + backend.update("run-abc", config={"a": 1}, summary={"loss": 0.1}) + + assert handler.paths() == [] + + +def test_a_completed_run_is_finalized_with_exit_code_zero(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes) + + backend.finalize("run-abc", status=RunStatus.COMPLETED, summary={"loss": 0.1}) + + assert handler.paths() == ["POST /api/v1/rft/finalize"] + assert handler.bodies_for("/api/v1/rft/finalize")[0] == {"run_id": "run-abc", "exit_code": 0} + + +def test_finalize_is_replayed_after_an_ambiguous_failure(make_platform_client, rft_routes): + """Compare-and-set on the server, and "already finalized" is a success: + a lost response costs nothing to replay, unlike the eval finalize.""" + attempts = [] + + def flaky(request: httpx.Request) -> httpx.Response: + attempts.append(request) + if len(attempts) == 1: + return httpx.Response(504) + return httpx.Response(200, json={"data": {"status": "success"}}) + + routes = dict(rft_routes) + routes["POST /api/v1/rft/finalize"] = flaky + backend, handler = make_backend(make_platform_client, routes) + + backend.finalize("run-abc", status=RunStatus.COMPLETED) + + assert len(attempts) == 2 + + +def test_finalize_falls_back_to_a_completed_status_update(make_platform_client, rft_routes): + """A finalize-specific failure must not strand an otherwise completed run.""" + routes = dict(rft_routes) + routes["POST /api/v1/rft/finalize"] = lambda request: httpx.Response( + 400, json={"detail": "finalize unavailable"} + ) + backend, handler = make_backend(make_platform_client, routes) + + backend.finalize("run-abc", status=RunStatus.COMPLETED) + + assert handler.paths() == [ + "POST /api/v1/rft/finalize", + "PUT /api/v1/rft/external-runs/run-abc/status", + ] + assert handler.bodies_for("/api/v1/rft/external-runs/run-abc/status") == [ + {"status": "completed"} + ] + + +@pytest.mark.parametrize( + ("status", "error", "expected_message"), + [ + (RunStatus.FAILED, "ValueError: boom", "ValueError: boom"), + (RunStatus.CANCELLED, "interrupted", "cancelled: interrupted"), + (RunStatus.CRASHED, "process exited without finishing the run", "crashed: process exited"), + (RunStatus.FAILED, None, "failed"), + ], + ids=["failed", "cancelled", "crashed", "failed-no-reason"], +) +def test_other_terminal_states_are_reported_as_failed_with_the_reason( + make_platform_client, rft_routes, status, error, expected_message +): + """The RFT vocabulary is completed | failed; the finer distinction rides + in error_message so an operator can still tell a Ctrl-C from a crash.""" + backend, handler = make_backend(make_platform_client, rft_routes) + + backend.finalize("run-abc", status=status, error=error) + + assert handler.paths() == ["PUT /api/v1/rft/external-runs/run-abc/status"] + body = handler.bodies_for("/api/v1/rft/external-runs/run-abc/status")[0] + assert body["status"] == "failed" + assert body["error_message"].startswith(expected_message) + assert "POST /api/v1/rft/finalize" not in handler.paths() + + +def test_marking_an_already_closed_run_failed_is_not_an_error(make_platform_client, rft_routes): + """409: a managed launcher or an earlier attempt got there first. The + outcome is recorded either way, so teardown must not raise over it.""" + routes = dict(rft_routes) + routes["PUT /api/v1/rft/external-runs/run-abc/status"] = lambda request: httpx.Response( + 409, json={"detail": "Cannot update status of run in 'COMPLETED' state."} + ) + backend, handler = make_backend(make_platform_client, routes) + + backend.finalize("run-abc", status=RunStatus.FAILED, error="boom") + + assert handler.paths() == ["PUT /api/v1/rft/external-runs/run-abc/status"] + + +def test_log_metrics_posts_one_step(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes) + + backend.log_metrics("run-abc", {"step": 3, "loss": 0.5}) + + assert handler.bodies_for("/api/v1/rft/metrics")[0] == { + "run_id": "run-abc", + "metrics": {"step": 3, "loss": 0.5}, + } diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index 16b67ab23..2f9f06b46 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -639,3 +639,93 @@ def test_a_credential_without_the_traces_scope_reaches_strict_callers_and_loss_a run.finish() assert run.failed_records == {"traces": 1} + + +# ------------------------------------------------------------------ metrics + + +def test_log_metrics_goes_to_the_metrics_sinks_only(): + records, metrics = FakeSink("records"), FakeSink("metrics") + run = make_run(sinks=[records], metrics_sinks=[metrics]) + + run.log_metrics({"loss": 0.5, "nan": float("nan")}, step=3) + run.log_traces([{"id": "t1"}]) + run.flush() + + assert records.batches == [[{"id": "t1"}]] + (batch,) = metrics.batches + (row,) = batch + assert row["loss"] == 0.5 + assert row["step"] == 3 + assert "nan" not in row + assert isinstance(row["_timestamp"], float) + run.finish() + + +def test_log_metrics_drops_non_finite_values_inside_sequences(): + metrics = FakeSink("metrics") + run = make_run(metrics_sinks=[metrics]) + + run.log_metrics( + { + "quantiles": [0.1, float("nan"), 0.9], + "nested": [{"loss": float("inf"), "reward": 1.0}, [float("-inf"), 2.0]], + } + ) + run.flush() + + row = metrics.batches[0][0] + assert row["quantiles"] == [0.1, 0.9] + assert row["nested"] == [{"reward": 1.0}, [2.0]] + run.finish() + + +def test_log_metrics_keeps_a_step_the_producer_already_set(): + metrics = FakeSink("metrics") + run = make_run(metrics_sinks=[metrics]) + + run.log_metrics({"step": 7, "loss": 0.5}, step=8) + run.flush() + + assert metrics.batches[0][0]["step"] == 7 + run.finish() + + +def test_log_metrics_without_a_metrics_sink_is_a_no_op(): + run = make_run() + + run.log_metrics({"loss": 0.5}, step=1) + + assert run._metrics_worker._thread is None + assert run.dropped_records == 0 + run.finish() + + +def test_log_metrics_on_a_finished_run_is_a_producer_bug(): + run = make_run(metrics_sinks=[FakeSink("metrics")]) + run.finish() + + with pytest.raises(RunFinishedError): + run.log_metrics({"loss": 0.5}) + + +def test_finish_drains_and_closes_the_metrics_uploader_too(): + metrics = FakeSink("metrics") + run = make_run(metrics_sinks=[metrics]) + + run.log_metrics({"loss": 0.5}, step=1) + run.finish() + + assert len(metrics.batches) == 1 + assert metrics.closed is True + + +def test_losses_are_reported_across_both_uploaders(): + broken = FakeSink("metrics", fail_on_write=True) + run = make_run(metrics_sinks=[broken]) + + run.log_metrics({"loss": 0.5}, step=1) + run.flush() + + assert run.failed_records == {"metrics": 1} + run.finish() diff --git a/packages/prime-runs/tests/test_samples_sink.py b/packages/prime-runs/tests/test_samples_sink.py index d32d632e8..2a46ff2e9 100644 --- a/packages/prime-runs/tests/test_samples_sink.py +++ b/packages/prime-runs/tests/test_samples_sink.py @@ -69,9 +69,7 @@ def test_records_this_sink_cannot_project_are_skipped_not_fatal( assert sink.skipped == 2 assert sink.enabled is True warnings = [ - record.getMessage() - for record in caplog.records - if "no projection" in record.getMessage() + record.getMessage() for record in caplog.records if "no projection" in record.getMessage() ] assert len(warnings) == 1 assert "(dict)" in warnings[0] diff --git a/packages/prime-runs/tests/test_train_samples_sink.py b/packages/prime-runs/tests/test_train_samples_sink.py new file mode 100644 index 000000000..6c5aa8946 --- /dev/null +++ b/packages/prime-runs/tests/test_train_samples_sink.py @@ -0,0 +1,448 @@ +"""The training samples sink: step-keyed Parquet through presign -> PUT -> confirm.""" + +import httpx +import pytest +from _fakes import make_episode, make_trace, make_train_episode +from conftest import RecordingHandler, StorageHandler + +from prime_runs.exceptions import APIError, RetryableAPIError, TransportError +from prime_runs.sinks import RftSamplesSink, training_step +from prime_runs.sinks.base import SinkWriteError + + +class Encoder: + """Records what it was asked to encode; returns a recognizable payload.""" + + def __init__(self, payload=b"parquet-bytes"): + self.calls = [] + self.offsets = [] + self.payload = payload + + def __call__(self, episodes, run_id, step, sample_id_offset=0): + self.calls.append((list(episodes), run_id, step)) + self.offsets.append(sample_id_offset) + return self.payload + + +def make_sink(make_platform_client, routes, *, encoder=None, storage=None, **kwargs): + handler = RecordingHandler(routes) + storage = storage or StorageHandler() + encoder = encoder or Encoder() + sink = RftSamplesSink( + make_platform_client(handler), encoder=encoder, upload_client=storage.client(), **kwargs + ) + sink.start("run-abc", {}) + return sink, handler, storage, encoder + + +def test_each_step_on_the_cadence_is_one_upload(make_platform_client, rft_routes): + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes) + episodes = [ + make_train_episode("e1", step=10), + make_train_episode("e2", step=20), + make_train_episode("e3", step=10), + make_train_episode("e4", step=11), # off-cadence: not a loss + ] + + sink.write(episodes) + + assert [(step, [e.id for e in eps]) for eps, _, step in encoder.calls] == [ + (10, ["e1", "e3"]), + (20, ["e2"]), + ] + assert [body["step"] for body in handler.bodies_for("/api/v1/rft/samples/presign")] == [10, 20] + assert [str(u.url) for u in storage.uploads] == [ + "http://storage.test/run-abc/step-10.parquet", + "http://storage.test/run-abc/step-20.parquet", + ] + assert storage.uploads[0].content == b"parquet-bytes" + assert storage.uploads[0].headers["content-type"] == "application/parquet" + confirms = handler.bodies_for("/api/v1/rft/samples/confirm") + assert confirms[0] == { + "run_id": "run-abc", + "step": 10, + "s3_key": "runs/run-abc/step-10.parquet", + } + assert sink.steps_written == 2 + assert sink.sampled_out == 1 + assert sink.skipped == 0 + + +def test_the_cadence_is_configurable(make_platform_client, rft_routes): + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes, step_interval=1) + + sink.write([make_train_episode("e1", step=7)]) + + assert [step for _, _, step in encoder.calls] == [7] + with pytest.raises(ValueError): + RftSamplesSink(sink._client, step_interval=0) + + +def test_records_without_training_work_have_no_row_here(make_platform_client, rft_routes): + """Eval-work episodes, bare traces and JSON episodes reach Prime Traces + only; here they are counted, not uploaded, and not a loss.""" + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes) + + sink.write( + [ + make_train_episode("eval-work", step=10, work="eval"), + make_episode("unstamped", [make_trace()]), + make_trace(), + { + "id": "json-episode", + "traces": [], + "run": {"type": "train", "work": {"type": "train", "step": 10}}, + }, + ] + ) + + assert encoder.calls == [] + assert handler.paths() == [] + assert sink.skipped == 4 + + +def test_nothing_to_encode_makes_no_request(make_platform_client, rft_routes): + sink, handler, storage, encoder = make_sink( + make_platform_client, rft_routes, encoder=Encoder(payload=None) + ) + + sink.write([make_train_episode("e1", step=10)]) + + assert encoder.calls and handler.paths() == [] + assert sink.steps_written == 0 + + +def test_a_failed_presign_loses_that_step_only(make_platform_client, rft_routes): + """Presign is replayable (it only mints a URL), so a blip is retried; a + step whose presign keeps failing is lost on its own, not the batch.""" + import json + + def flaky(request: httpx.Request) -> httpx.Response: + if json.loads(request.content)["step"] == 10: + return httpx.Response(502) + return rft_routes["POST /api/v1/rft/samples/presign"](request) + + routes = dict(rft_routes) + routes["POST /api/v1/rft/samples/presign"] = flaky + sink, handler, storage, encoder = make_sink(make_platform_client, routes) + + with pytest.raises(SinkWriteError) as info: + sink.write( + [ + make_train_episode("e1", step=10), + make_train_episode("e2", step=10), + make_train_episode("e3", step=20), + ] + ) + + # Step 10 (two episodes) was lost; step 20 went through. + assert info.value.failed_records == 2 + assert isinstance(info.value.cause, RetryableAPIError) + assert sink.steps_written == 1 + assert [str(u.url) for u in storage.uploads] == ["http://storage.test/run-abc/step-20.parquet"] + + +def test_the_storage_put_is_retried_then_reported_as_transient( + make_platform_client, rft_routes, no_sleep +): + storage = StorageHandler(status_codes=[503, 503, 503]) + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes, storage=storage) + + with pytest.raises(SinkWriteError) as info: + sink.write([make_train_episode("e1", step=10)]) + + assert len(storage.uploads) == 3 + assert isinstance(info.value.cause, RetryableAPIError) + assert handler.paths() == ["POST /api/v1/rft/samples/presign"] # never confirmed + + +def test_a_storage_put_that_recovers_is_confirmed(make_platform_client, rft_routes, no_sleep): + storage = StorageHandler(status_codes=[500, 200]) + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes, storage=storage) + + sink.write([make_train_episode("e1", step=10)]) + + assert len(storage.uploads) == 2 + assert "POST /api/v1/rft/samples/confirm" in handler.paths() + + +def test_a_storage_rejection_is_not_retried(make_platform_client, rft_routes): + storage = StorageHandler(status_codes=[403]) + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes, storage=storage) + + with pytest.raises(SinkWriteError) as info: + sink.write([make_train_episode("e1", step=10)]) + + assert len(storage.uploads) == 1 + assert isinstance(info.value.cause, APIError) + assert not isinstance(info.value.cause, RetryableAPIError) + + +def test_a_dropped_connection_to_storage_is_transient(make_platform_client, rft_routes, no_sleep): + def refuse(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("connection refused") + + storage = StorageHandler() + storage.__call__ = refuse # type: ignore[method-assign] + upload = httpx.Client(transport=httpx.MockTransport(refuse)) + handler = RecordingHandler(rft_routes) + sink = RftSamplesSink(make_platform_client(handler), encoder=Encoder(), upload_client=upload) + sink.start("run-abc", {}) + + with pytest.raises(SinkWriteError) as info: + sink.write([make_train_episode("e1", step=10)]) + + assert isinstance(info.value.cause, TransportError) + + +def test_an_encoder_failure_is_a_record_rejection(make_platform_client, rft_routes): + def broken(episodes, run_id, step, sample_id_offset=0): + raise ValueError("cannot serialize") + + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes, encoder=broken) + + with pytest.raises(SinkWriteError) as info: + sink.write([make_train_episode("e1", step=10)]) + + assert isinstance(info.value.cause, ValueError) + assert handler.paths() == [] + + +def test_a_presign_without_a_url_is_an_api_error(make_platform_client, rft_routes): + routes = dict(rft_routes) + routes["POST /api/v1/rft/samples/presign"] = {"data": {"expiresIn": 900}} + sink, handler, storage, encoder = make_sink(make_platform_client, routes) + + with pytest.raises(SinkWriteError) as info: + sink.write([make_train_episode("e1", step=10)]) + + assert isinstance(info.value.cause, APIError) + assert "presign" in str(info.value.cause) + + +def test_a_step_logged_again_gets_its_own_sample_id_range(make_platform_client, rft_routes): + """Objects are additive per step on the platform (each presign mints a new + key and the viewer unions them), and the viewer looks samples up by + (step, sample_id): a straggler batch must not reuse the first object's ids.""" + from prime_runs.sinks.train_samples import SAMPLE_ID_STRIDE + + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes) + + sink.write([make_train_episode("e1", step=10), make_train_episode("e2", step=20)]) + sink.write([make_train_episode("e3", step=10)]) # dispatched at 10, shipped later + + assert [step for _, _, step in encoder.calls] == [10, 20, 10] + assert encoder.offsets == [0, 0, SAMPLE_ID_STRIDE] + presigned = [body["step"] for body in handler.bodies_for("/api/v1/rft/samples/presign")] + assert presigned == [10, 20, 10] + assert sink.steps_written == 3 + + +def test_an_upload_that_failed_still_advances_the_sample_id_range(make_platform_client, rft_routes): + """An attempt that failed late may still have landed an object; the next + one for the step must not overlap it.""" + from prime_runs.sinks.train_samples import SAMPLE_ID_STRIDE + + storage = StorageHandler(status_codes=[403]) + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes, storage=storage) + + with pytest.raises(SinkWriteError): + sink.write([make_train_episode("e1", step=10)]) + sink.write([make_train_episode("e2", step=10)]) + + assert encoder.offsets == [0, SAMPLE_ID_STRIDE] + + +def test_a_forked_child_numbers_its_rows_apart_from_the_parent(make_platform_client, rft_routes): + """The child inherits a copy of the parent's upload counters; continuing + them would hand out the parent's next range. The child's ranges come from + a band the parent's counter never reaches, and every id stays a JSON-safe + integer for the viewer.""" + from prime_runs.sinks.train_samples import SAMPLE_ID_RANGES, SAMPLE_ID_STRIDE + + assert SAMPLE_ID_RANGES * SAMPLE_ID_STRIDE <= 2**53 + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes) + sink.write([make_train_episode("e1", step=10)]) + + sink.reset_after_fork() # the child's copy of the sink... + # ...with its own storage client, as an owned one is rebuilt on the next upload. + sink._upload_client, sink._forked_with_injected_client = storage.client(), False + sink.write([make_train_episode("e2", step=10)]) + sink.write([make_train_episode("e3", step=10)]) + + parent, first, second = encoder.offsets + assert parent == 0 + band = range(SAMPLE_ID_RANGES // 2 * SAMPLE_ID_STRIDE, SAMPLE_ID_RANGES * SAMPLE_ID_STRIDE) + assert first in band and second in band and first != second + assert first % SAMPLE_ID_STRIDE == 0 + assert sink._uploads_per_step == {10: 1}, "the inherited counters are left to the parent" + + +def test_a_confirm_whose_response_was_lost_is_replayed(make_platform_client, rft_routes, no_sleep): + """Confirm validates the key and refreshes progress; a second confirm of + the same key changes nothing, so an ambiguous failure is retried rather + than costing the step.""" + attempts = [] + + def flaky(request: httpx.Request) -> httpx.Response: + attempts.append(request) + if len(attempts) == 1: + return httpx.Response(504) + return httpx.Response(200, json={"data": {"status": "success"}}) + + routes = dict(rft_routes) + routes["POST /api/v1/rft/samples/confirm"] = flaky + sink, handler, storage, encoder = make_sink(make_platform_client, routes) + + sink.write([make_train_episode("e1", step=10)]) + + assert len(attempts) == 2 + assert len(storage.uploads) == 1 + assert sink.steps_written == 1 + + +def test_without_pyarrow_the_sink_turns_itself_off_quietly( + make_platform_client, rft_routes, monkeypatch, caplog +): + """No loss: the episodes still reach Prime Traces. Say how to turn it on.""" + monkeypatch.setattr("prime_runs.projection.parquet_available", lambda: False) + handler = RecordingHandler(rft_routes) + sink = RftSamplesSink(make_platform_client(handler)) + + with caplog.at_level("WARNING"): + sink.start("run-abc", {}) + + assert sink.enabled is False + assert "prime-runs[train]" in caplog.text + sink.write([make_train_episode("e1", step=10)]) + assert handler.paths() == [] + + +def test_close_releases_only_an_owned_upload_client(make_platform_client, rft_routes): + handler = RecordingHandler(rft_routes) + injected = StorageHandler().client() + sink = RftSamplesSink(make_platform_client(handler), encoder=Encoder(), upload_client=injected) + sink.close() + assert not injected.is_closed + + owned = RftSamplesSink(make_platform_client(handler), encoder=Encoder()) + owned.start("run-abc", {}) + client = owned._upload() + owned.close() + assert client.is_closed + + +@pytest.mark.parametrize( + ("record", "expected"), + [ + (make_train_episode("e", step=12), 12), + (make_train_episode("e", step=12, work="eval"), None), + (make_episode("e"), None), + ({"run": {"type": "train", "work": {"type": "train", "step": 3}}}, 3), + ({"run": {"type": "eval"}}, None), + ({"run": {"type": "train", "work": {"type": "train", "step": True}}}, None), + ("not a record", None), + ], + ids=["train", "eval-work", "unstamped", "json", "json-eval", "bool-step", "string"], +) +def test_training_step_reads_the_episode_provenance(record, expected): + assert training_step(record) == expected + + +def test_the_default_encoder_writes_the_viewer_table(): + """The real Parquet encoder, when pyarrow is installed: one row per episode + with the RFT-only columns layered on the shared sample projection.""" + pa = pytest.importorskip("pyarrow") + pq = pytest.importorskip("pyarrow.parquet") + import io + import json + + from prime_runs.projection import episodes_to_parquet_bytes, train_sample_schema + + episodes = [ + make_train_episode("e1", step=10, idx=3, reward=1.0, advantage=0.25, env_id="gsm8k"), + make_train_episode("e2", step=10, idx=4, reward=0.0, advantage=-0.25, env_id="gsm8k"), + ] + + payload = episodes_to_parquet_bytes(episodes, "run-abc", 10) + + assert payload is not None + table = pq.read_table(io.BytesIO(payload)) + assert table.schema.equals(train_sample_schema()) + rows = table.to_pylist() + assert [row["run_id"] for row in rows] == ["run-abc", "run-abc"] + assert [row["step"] for row in rows] == [10, 10] + assert [row["problem_id"] for row in rows] == [3, 4] + assert [row["sample_id"] for row in rows] == [0, 1] + assert [row["reward"] for row in rows] == [1.0, 0.0] + assert [row["advantage"] for row in rows] == [0.25, -0.25] + assert [row["env_name"] for row in rows] == ["gsm8k", "gsm8k"] + trajectory = json.loads(rows[0]["trajectory"]) + assert trajectory[-1]["advantage"] == 0.25 + assert rows[0]["num_output_tokens"] == 5 + info = json.loads(rows[0]["info"]) + assert info["native_wrapper"]["id"] == "e1" + assert isinstance(pa.Table, type) + + +def test_the_default_encoder_numbers_rows_from_the_offset(): + pytest.importorskip("pyarrow") + import io + + import pyarrow.parquet as pq + + from prime_runs.projection import episodes_to_parquet_bytes + + episodes = [make_train_episode("e1", step=10, idx=3), make_train_episode("e2", step=10, idx=4)] + + payload = episodes_to_parquet_bytes(episodes, "run-abc", 10, sample_id_offset=7) + + rows = pq.read_table(io.BytesIO(payload)).to_pylist() + assert [row["sample_id"] for row in rows] == [7, 8] + assert [row["problem_id"] for row in rows] == [3, 4] + + +def test_the_default_encoder_skips_episodes_without_a_trajectory(): + pytest.importorskip("pyarrow") + from _fakes import Episode, Trace + + from prime_runs.projection import episodes_to_parquet_bytes + + bare = Episode(id="no-branches", traces=[Trace(id="t", branches=[])]) + assert episodes_to_parquet_bytes([bare], "run-abc", 10) is None + assert episodes_to_parquet_bytes([], "run-abc", 10) is None + + +def test_a_forked_child_drops_the_inherited_storage_client(make_platform_client, rft_routes): + """The socket belongs to the parent: the child must neither reuse it nor + close it (a close would send close_notify on the parent's connection).""" + handler = RecordingHandler(rft_routes) + owned = RftSamplesSink(make_platform_client(handler), encoder=Encoder()) + owned.start("run-abc", {}) + inherited = owned._upload() + + owned.reset_after_fork() + + assert owned._upload_client is None + assert not inherited.is_closed # dropped, not closed + assert owned._upload() is not inherited + owned.close() + + injected = RftSamplesSink( + make_platform_client(handler), encoder=Encoder(), upload_client=StorageHandler().client() + ) + injected.start("run-abc", {}) + injected.reset_after_fork() + with pytest.raises(SinkWriteError) as info: + injected.write([make_train_episode("e1", step=10)]) + assert isinstance(info.value.cause, RuntimeError) + assert "after a fork" in str(info.value.cause) + + +def test_the_sink_is_registered_for_fork_resets(): + from prime_runs import _fork + from prime_runs.sinks import RftSamplesSink + + sink = RftSamplesSink(object(), encoder=Encoder()) # type: ignore[arg-type] + + assert sink in _fork._registry diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index 870ba2543..0b17e8e8e 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -373,6 +373,67 @@ def write(self, records) -> None: worker.close() +def test_a_cooldown_gives_a_retired_sink_another_chance(): + """A training run outlives any platform deploy: after the cooldown the + sink is tried again, and records that arrived meanwhile are counted as + lost to it, not silently forgotten.""" + from prime_runs.exceptions import TransportError + from prime_runs.worker import TRANSIENT_FAILURE_LIMIT + + class OutageSink(FakeSink): + def __init__(self) -> None: + super().__init__("outage") + self.down = True + + def write(self, records) -> None: + if self.down: + raise TransportError("connection refused") + super().write(records) + + sink = OutageSink() + worker = UploadWorker([sink], retire_cooldown=0.05) + + for _ in range(TRANSIENT_FAILURE_LIMIT): + worker.submit([{"id": "lost"}]) + drain(worker) + assert sink.enabled is False + worker.submit([{"id": "during"}]) + drain(worker) + assert sink.batches == [] + + # A real wait: the autouse no_sleep fixture stubs out time.sleep. + threading.Event().wait(0.06) + sink.down = False + worker.submit([{"id": "after"}]) + drain(worker) + + assert sink.enabled is True + assert sink.batches == [[{"id": "after"}]] + assert worker.failed_records == {"outage": TRANSIENT_FAILURE_LIMIT + 1} + worker.close() + + +def test_a_permanent_failure_is_not_revived_by_the_cooldown(): + from prime_runs.exceptions import UnauthorizedError + + class DeniedSink(FakeSink): + def write(self, records) -> None: + raise UnauthorizedError("nope", status_code=401) + + sink = DeniedSink("denied") + worker = UploadWorker([sink], retire_cooldown=0.01) + + worker.submit([{"id": 1}]) + drain(worker) + threading.Event().wait(0.02) + worker.submit([{"id": 2}]) + drain(worker) + + assert sink.enabled is False + assert worker.failed_records == {"denied": 2} + worker.close() + + def test_a_success_forgives_earlier_blips(): """Strikes are consecutive: an intermittent gateway must never accumulate its way to a retirement across an otherwise healthy run.""" diff --git a/uv.lock b/uv.lock index cd3ce2368..991446b92 100644 --- a/uv.lock +++ b/uv.lock @@ -8,7 +8,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-12T18:00:50.512194Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P7D" [options.exclude-newer-package] @@ -1461,16 +1461,20 @@ dev = [ { name = "pytest-asyncio" }, { name = "ruff" }, ] +train = [ + { name = "pyarrow" }, +] [package.metadata] requires-dist = [ { name = "httpx", specifier = ">=0.25.0" }, { name = "prime-traces", editable = "packages/prime-traces" }, + { name = "pyarrow", marker = "extra == 'train'", specifier = ">=15.0.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.23.0" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.13.1" }, ] -provides-extras = ["dev"] +provides-extras = ["dev", "train"] [[package]] name = "prime-sandboxes" From e4ca6569326ee734d13cf1110a28fe4dbbbbca9b Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Tue, 1 Sep 2026 17:10:03 -0700 Subject: [PATCH 37/39] refactor(runs): trim the SDK and write non-finite numbers as null - docstrings and comments cut to the contract or the non-obvious reason; section banners removed - one send_each write loop in sinks/base.py replaces three copies; _finite replaces the two metric-cleaning functions; the worker's deadline helpers are shared with run.py - unused surface removed: RftBackend.log_metrics, the PlatformClient context manager, the mapping form of ConfigSource.coerce - tests: docstrings stripped, 28 redundant tests removed, the after-finish checks parametrized - README rewritten in the prime-traces shape - projection: NaN and infinity become null in projected samples and in the Parquet JSON columns, so a non-finite reward or metric no longer loses the eval batch or the training step it sits in Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Fk5fBAAaTiPfJ5UpNyqxXR --- packages/prime-runs/README.md | 262 +++++----------- .../prime-runs/src/prime_runs/__init__.py | 19 +- packages/prime-runs/src/prime_runs/_fork.py | 23 +- packages/prime-runs/src/prime_runs/_http.py | 70 ++--- packages/prime-runs/src/prime_runs/backend.py | 173 +++------- packages/prime-runs/src/prime_runs/config.py | 5 +- .../prime-runs/src/prime_runs/exceptions.py | 29 +- packages/prime-runs/src/prime_runs/metrics.py | 23 +- packages/prime-runs/src/prime_runs/models.py | 83 ++--- .../prime-runs/src/prime_runs/projection.py | 145 ++++----- packages/prime-runs/src/prime_runs/run.py | 295 +++++------------- .../src/prime_runs/sinks/__init__.py | 2 +- .../prime-runs/src/prime_runs/sinks/base.py | 84 ++--- .../src/prime_runs/sinks/metrics.py | 65 ++-- .../src/prime_runs/sinks/samples.py | 96 ++---- .../prime-runs/src/prime_runs/sinks/traces.py | 60 ++-- .../src/prime_runs/sinks/train_samples.py | 144 +++------ packages/prime-runs/src/prime_runs/worker.py | 112 ++----- packages/prime-runs/tests/_fakes.py | 9 +- .../prime-runs/tests/test_config_source.py | 36 +-- .../prime-runs/tests/test_evals_backend.py | 21 -- packages/prime-runs/tests/test_http.py | 38 +-- packages/prime-runs/tests/test_init.py | 34 -- packages/prime-runs/tests/test_metrics.py | 16 +- .../prime-runs/tests/test_metrics_sink.py | 9 - packages/prime-runs/tests/test_projection.py | 39 +-- packages/prime-runs/tests/test_rft_backend.py | 26 -- packages/prime-runs/tests/test_run.py | 89 +----- .../prime-runs/tests/test_samples_sink.py | 6 - packages/prime-runs/tests/test_traces_sink.py | 51 --- .../tests/test_train_samples_sink.py | 71 ++--- packages/prime-runs/tests/test_worker.py | 98 ------ 32 files changed, 569 insertions(+), 1664 deletions(-) diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index bc44079c7..1e7a3bb9f 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -1,219 +1,123 @@ # Prime Runs SDK -Track eval and training runs on the Prime Intellect platform. +Track evaluation and training runs on the Prime Intellect platform: `init()` +opens the run, records stream out while it proceeds, and `finish()` closes it +out with a terminal status. + +## Install ```bash -pip install prime-runs +uv add prime-runs # or: pip install prime-runs +uv add 'prime-runs[train]' # training runs: adds pyarrow for the sample table ``` -## Quick start +## Eval runs ```python import prime_runs as pr +from verifiers.v1 import EvalRunInfo run = pr.init( name="gsm8k-qwen3-8b", - environments=["gsm8k"], + environments=["gsm8k"], # hub names (get-or-create) or owner/name slugs model="Qwen/Qwen3-8B", framework="verifiers", - config="eval.toml", # the file the run was launched from + config="eval.toml", # the launched file, stored byte for byte ) -print(run.url) # https://app.primeintellect.ai/dashboard/evaluations/eval-... +print(run.url) # https://app.primeintellect.ai/dashboard/evaluations/... -for episode in rollouts: # episodes carry run.id — see "Identity" below - run.log_episodes([episode]) # bare traces go through log_traces() +for episode in rollouts: + episode.record_run(EvalRunInfo(id=run.id)) # every trace carries the run id + run.log_episodes([episode]) # a queue put; bare traces: log_traces() run.finish(summary=pr.metrics.from_episodes(episodes)) ``` -`init()` opens the run and returns a handle carrying its ID and dashboard URL. -Records stream out on a background thread while the run proceeds, so the -dashboard fills in as rollouts land. `finish()` closes the run out; a `with` -block does that for you, including when the body raises. Run-level outputs go -in `finish(summary=...)`, or incrementally through `update_summary()`; -`metrics.from_episodes()` returns them in the shape the dashboard reads -(`metrics.RunSummary`). - -## Identity - -`init()` is called **before** the first rollout, and the ID it returns is *the* -run ID everywhere — including inside every trace document you write: - -```python -run = pr.init(...) -trace.record_run(EvalRunInfo(id=run.id)) # verifiers -``` +`init()` is called before the first rollout, and the id it returns is the run +id everywhere, including inside every trace document. A `with run:` block +finishes for you: an exception marks the run `failed`, Ctrl-C `cancelled`, and +a process that exits without finishing is reported `crashed` by an atexit hook. -The ingestion service indexes `run.id` from the trace body, so "every trace for -this run" is a fast query. Bare dicts without a `run` key are stamped for you; -producer objects are passed through untouched. - -## Config - -`config=` takes the path to the file the run was launched from, or a mapping. - -| you pass | what is stored | -| --- | --- | -| a path | the file, byte for byte, under `config_source` | -| a mapping | exactly as given | - -The file is the run's real configuration — comments, key order and section -grouping included — so it is stored verbatim, not parsed. A `str` or `Path` is -always a path; use `pr.ConfigSource(text=...)` for contents already in memory. -To send structured values *and* the file, put the file under -`pr.CONFIG_SOURCE_KEY` in the mapping: - -```python -config = {**cfg.model_dump(exclude_unset=True), - pr.CONFIG_SOURCE_KEY: pr.ConfigSource.from_file("eval.toml").to_dict()} -``` - -Nothing is redacted — keep credentials in the environment, not in the file. - -## Modes - -| mode | what happens | -| --- | --- | -| `online` | the run lives on the platform (default when an API key is present) | -| `disabled` | every call is a no-op, with the same object shape | - -Set the mode explicitly, or through `$PRIME_RUNS_MODE`. A missing API key -disables the run with a warning — it never silently writes somewhere else. - -## What the run handle does for you - -- **Streams instead of buffering.** Records go out as they are produced. - Whatever queues up while one upload is in flight goes out as the next one, - so the request rate tracks upload latency, not rollout throughput. -- **Contains its own errors.** With the default `on_error="warn"`, nothing the - platform raises escapes into your loop. Use `on_error="raise"` in tests and - CI, where a silent upload failure is the bug; failures surface from `flush()` - and `finish()`. Platform errors are the `prime_traces` exception family - (`pr.APIError` and friends), so one set of `except` clauses covers both SDKs. -- **Applies backpressure.** The upload queue is bounded; a producer that - durably outruns the uploader has records dropped and counted - (`run.dropped_records`) rather than stalled. Per-sink losses are in - `run.failed_records`. -- **Waits for its own uploads.** `finish()` gives queued records the same - budget a single upload gets (300s) and warns if they do not drain. Set - `init(finish_timeout=...)` for the run, or `finish(timeout=...)` per call — an - abort path can pass a few seconds so an interrupted producer is not pinned - behind a slow platform. -- **Reports a terminal status.** The context manager and an `atexit` hook both - route to the same idempotent `finish()`. A run you said failed is `failed`; - one somebody stopped on purpose — Ctrl-C or a cancelled task inside a `with` - block — is `cancelled`; one that stopped without saying — an exit that never - reached `finish()` — is `crashed`. +`config=` takes the path to the launched file (kept verbatim under +`config_source`, comments and all) or a mapping stored as given; put a file +under `pr.CONFIG_SOURCE_KEY` in the mapping to send both. Nothing is redacted. ## Training runs ```python -import prime_runs as pr - run = pr.init( kind="train", name="qwen3-8b-gsm8k-rl", - model="Qwen/Qwen3-8B", # the base model - environments=["primeintellect/gsm8k"], # hub ids, passed through as given + model="Qwen/Qwen3-8B", # the base model + environments=["primeintellect/gsm8k"], # hub ids, passed through training=pr.TrainingSpec(max_steps=1000, batch_size=64, rollouts_per_example=8), - config=train_config.model_dump(), # or the path to the launched TOML - team_id="team_...", # external runs are created for a team + config=train_config.model_dump(), + team_id="team_...", # external runs belong to a team ) for step, (episodes, metrics) in enumerate(training_loop): - run.log_episodes(episodes) # carry run.work.step (verifiers TrainRunInfo) + run.log_episodes(episodes) # episodes carry run.work.step (TrainRunInfo) run.log_metrics(metrics, step=step) run.finish() ``` -A training run lives on `/api/v1/rft/external-runs` — what prime-rl's monitor -did by hand. Where it differs from an eval run: - -- **A team is required.** The platform enables external runs per team. A team - outside the allowlist gets a `ForbiddenError` from `init()` with the - platform's reason; the producer decides whether that means "run locally". -- **Metrics.** `log_metrics(values, step=...)` streams one row per call on its - own uploader, so a slow sample upload never holds a step's metrics back. - `_timestamp` is stamped on every row; non-finite values are dropped. The - endpoint allows 60 rows a minute per token. A row whose response was lost - is re-sent: the platform keeps one row per step, so nothing doubles. -- **Samples.** Every record reaches Prime Traces. The training viewer's sample - table is a Parquet object per upload, keyed by the training step an episode - was *dispatched* at — `run.work.step` on verifiers' `TrainRunInfo`, which - prime-rl stamps on every dispatched episode — and uploaded every 10th step - (prime-rl's cadence) for training-work episodes. Objects are additive: a - step logged in several `log_episodes` calls (an off-policy episode landing - in a later batch) gets one object per call, each numbering its `sample_id`s - in its own range (a forked child draws its ranges apart from its parent's), - and the viewer shows their union. So "step N" in the - sample viewer means dispatched at N, while the metrics logged at N describe - the batch trained at N. Encoding needs pyarrow: `pip install - 'prime-runs[train]'`. Without it the table is skipped with a warning and - traces still flow. -- **Attach.** `init(kind="train", id=os.environ["RUN_ID"])` joins a run a - launcher created: nothing is registered, the platform keeps the run's - failure marking, and a clean `finish()` still completes it. The id must be - an external run's (one created through `POST /rft/external-runs`): the - platform's monitoring endpoints answer 400 for a hosted run. -- **Status.** The RFT vocabulary is `completed | failed`. `cancelled` and - `crashed` are reported as `failed` with the reason in `error_message`; a - clean `finish()` is an idempotent finalize, safe to replay. -- **Outages.** A sink that strikes out on transient failures (three in a row, - each already retried) is paused for five minutes and then tried again, not - retired for the run: a platform deploy costs a multi-day run a window of - rows, never the rest of its curves. Records that arrive while a sink is - paused are counted in `run.failed_records` for that sink. -- **No config update.** `config=` is registered once as `run_config` and there - is no summary document, so `finish(summary=...)` is not sent for training - runs — log final numbers with `log_metrics`. - -## From async code - -`log_traces()` / `log_episodes()` are a queue put, not a request, so they are -safe to call from a coroutine; it blocks only if the queue is full (up to 5s), which is the -backpressure. `init()` and `finish()` do network I/O — wrap them in -`asyncio.to_thread` if a stall there would matter. +- The platform enables external runs per team; a team outside the allowlist + gets a `ForbiddenError` from `init()`. +- `init(kind="train", id=os.environ["RUN_ID"])` attaches to an external run a + launcher already created: nothing is registered, the platform keeps the run's + failure marking, and a clean `finish()` still completes it. +- Metrics are one row per `log_metrics` call, on their own uploader. The sample + table gets one Parquet object per upload, every 10th step, keyed by the step + an episode was dispatched at; a step logged in several calls gets several + objects, and the viewer shows their union. +- The status vocabulary is `completed | failed`; `cancelled` and `crashed` + arrive as `failed` with the reason in `error_message`. + +## How it behaves + +- **Streams.** Records go out on a background thread as they are logged; + whatever queues up during one request goes out as the next. +- **Contains its errors.** With the default `on_error="warn"` nothing the + platform raises escapes into your loop; `on_error="raise"` surfaces the first + failure from `flush()` or `finish()`, for tests and CI. Platform errors are + the `prime_traces` exception family. +- **Degrades.** A transient failure costs its batch, three in a row retire the + sink (a training run pauses it for five minutes instead), and a full queue + drops records rather than stalling the run. Losses are counted in + `run.dropped_records` and `run.failed_records`. +- **Drains on exit.** `finish()` gives queued uploads up to `finish_timeout` + (300 s) before closing the run out; an abort path can pass + `finish(timeout=...)`. + +An online run writes to Prime Traces (the system of record, gated to an +allowlist; outside it that sink turns itself off quietly) and to the sample +table today's viewer reads. `log_*()` are queue puts, safe inside a coroutine; +`init()` and `finish()` do network I/O. ## Configuration -Resolved from environment variables first, then `~/.prime/config.json`: - -| setting | env var | default | -| --- | --- | --- | -| API key | `PRIME_API_KEY` | — | -| team | `PRIME_TEAM_ID` | — | -| platform API | `PRIME_API_BASE_URL` | `https://api.primeintellect.ai` | -| dashboard | `PRIME_FRONTEND_URL` | `https://app.primeintellect.ai` | -| traces service | `PRIME_TRACES_URL` | resolved by `prime-traces` | - -Or pass `api_key=`, `base_url=`, `team_id=` to `init()`. - -## Transports - -An online run writes every record to two sinks: Prime Traces (the system of -record — streaming, episode-aware, content-addressed and therefore idempotent -on retry) and the flat v0 sample table today's viewer reads. Both run because -Prime Traces is gated to an allowlist; an account outside it has the traces sink -turn itself off at the first upload — not counted as a failure, since nothing was -lost. When the viewer reads traces natively the default sink list drops one entry -and no producer changes. - -`prime_runs.projection` holds `trace_to_sample` / `build_samples`, the v0 -projection moved here from verifiers. `prime_runs.metrics.from_episodes` is the -run-level aggregation the eval dashboard reads. Both are duck-typed; no producer -package is imported — this is a leaf package, because the `prime` CLI depends on -`verifiers` and verifiers depends on this. - -## Status - -Eval runs over `/api/v1/evaluations`, training runs over -`/api/v1/rft/external-runs` (see "Training runs"). - -The evaluations API this version targets has no producer-facing way to mark an -evaluation **failed** or **cancelled**: the SDK records the terminal state under -`metadata.prime_runs` and warns that the run will keep showing as running on -the dashboard. The platform is adding `PUT /evaluations/{id}` with a -`FAILED` / `CANCELLED` status; a follow-up release adopts it (keeping the -metadata fallback for older backends), which is why `RunStatus.CANCELLED` -already exists. +| Source | Meaning | +| ---------------------- | ---------------------------------------------------------------------- | +| `PRIME_API_KEY` | Platform API token | +| `PRIME_TEAM_ID` | Team context; required for training runs | +| `PRIME_API_BASE_URL` | Platform API; defaults to `https://api.primeintellect.ai` | +| `PRIME_FRONTEND_URL` | Dashboard; defaults to `https://app.primeintellect.ai` | +| `PRIME_TRACES_URL` | Prime Traces service, resolved by `prime-traces` | +| `PRIME_RUNS_MODE` | `online` or `disabled`; unset means online when there is an API key | +| `~/.prime/config.json` | Shared prime CLI config (`api_key`, `team_id`, `base_url`) | + +Precedence is `init()` argument → environment variable → config file. A missing +API key disables the run with a warning. + +## Not yet available + +- **Failed or cancelled evaluations on the dashboard.** The evaluations API has + no producer-facing status endpoint yet, so the terminal state is recorded + under `metadata.prime_runs` and the run keeps showing as running. + +## Related packages + +- [prime-traces](https://github.com/PrimeIntellect-ai/prime/tree/main/packages/prime-traces) — Prime Traces SDK +- [prime](https://github.com/PrimeIntellect-ai/prime/tree/main/packages/prime) — Prime CLI +- [prime-evals](https://github.com/PrimeIntellect-ai/prime/tree/main/packages/prime-evals) — Evals SDK diff --git a/packages/prime-runs/src/prime_runs/__init__.py b/packages/prime-runs/src/prime_runs/__init__.py index 6e8984062..ac86cfd66 100644 --- a/packages/prime-runs/src/prime_runs/__init__.py +++ b/packages/prime-runs/src/prime_runs/__init__.py @@ -1,20 +1,7 @@ -"""Prime Intellect Runs SDK. +"""Prime Intellect Runs SDK: track eval and training runs on the Prime platform. -Track eval and training runs on the Prime platform:: - - import prime_runs as pr - - run = pr.init(name="gsm8k-qwen3-8b", environments=["gsm8k"], model=model) - print(run.url) - - for episode in rollouts: # stamp run.id onto the traces - run.log_traces([episode]) - - run.finish(summary=pr.metrics.from_episodes(episodes)) - -This is a leaf package: the ``prime`` CLI depends on ``verifiers``, so verifiers -can never depend on ``prime``. Nothing here imports a producer package; records -are duck-typed through ``to_record()``. +A leaf package: the ``prime`` CLI depends on ``verifiers`` and verifiers depends +on this, so nothing here imports a producer package; records are duck-typed. """ from . import metrics, projection diff --git a/packages/prime-runs/src/prime_runs/_fork.py b/packages/prime-runs/src/prime_runs/_fork.py index ff2eefb2f..f44d29c79 100644 --- a/packages/prime-runs/src/prime_runs/_fork.py +++ b/packages/prime-runs/src/prime_runs/_fork.py @@ -1,16 +1,8 @@ -"""One process-wide ``os.register_at_fork`` hook, shared by everything stateful. - -A forked child inherits copies of every open socket and buffered file handle. -Anything holding a connection or a file registers here and gets told to start -over in the child. Two rules for ``reset_after_fork``: - -- **Drop, do not close.** Closing an inherited transport can send bytes down a - socket the parent is still using. -- **Do not flush.** An inherited buffer holds records the parent will write. - -Registration is weak and the hook is installed once: ``register_at_fork`` -cannot be undone, so per-instance registration would pin every run forever. -""" +"""One process-wide ``os.register_at_fork`` hook for everything holding a +connection. A child must *drop* inherited transports (closing one can write to +the parent's socket) and never flush inherited buffers (those records are the +parent's). Registration is weak and the hook is installed once, since +``register_at_fork`` cannot be undone.""" import logging import os @@ -37,11 +29,8 @@ def register(obj: Any) -> None: def _reset_all() -> None: - # A fresh lock: the child inherits the parent's, which may have been held by - # a thread that does not exist here. Nothing else runs at this point, so the - # swap is safe. global _lock - _lock = threading.Lock() + _lock = threading.Lock() # the inherited one may be held by a thread that is not here for obj in list(_registry): try: obj.reset_after_fork() diff --git a/packages/prime-runs/src/prime_runs/_http.py b/packages/prime-runs/src/prime_runs/_http.py index 713ebed65..9e0fd725b 100644 --- a/packages/prime-runs/src/prime_runs/_http.py +++ b/packages/prime-runs/src/prime_runs/_http.py @@ -1,12 +1,10 @@ -"""HTTP client for the platform run APIs. - -Error mapping, backoff and the transport-failure classification come from -``prime_traces.core.client``; what is local is the retry *policy*, because it -is decided per call. A failure is *ambiguous* when the request may already have -been processed (a 502/504, a read timeout). Replaying one is fine for a GET or -PUT and not for ``POST /evaluations/``, which would create a second run. -Callers declare intent with ``idempotent=``; unambiguous failures (connect -errors, 429) are replayed for every method. +"""Authenticated client for ``{base_url}/api/v1`` with a per-call retry policy. + +A failure is *ambiguous* when the server may already have processed the request +(a 502/504, a read timeout). Replaying one is safe for a GET or PUT and not for +``POST /evaluations/``, which would create a second run; callers declare intent +with ``idempotent=``. Connect errors and 429 are replayed for every method. +Error mapping and backoff come from ``prime_traces.core.client``. """ import json @@ -15,20 +13,16 @@ from typing import Any, Dict, Mapping, Optional, Union import httpx -from prime_traces.core.client import ( - AMBIGUOUS_TRANSPORT_ERRORS, - raise_for_response, - retry_delay, -) +from prime_traces.core.client import AMBIGUOUS_TRANSPORT_ERRORS, raise_for_response, retry_delay from . import _fork from .exceptions import APIError, APITimeoutError, RetryableAPIError, TransportError DEFAULT_TIMEOUT = httpx.Timeout(60.0, connect=10.0) -# Sample batches are megabytes; uploads get a longer budget. +#: Sample batches are megabytes. UPLOAD_TIMEOUT = httpx.Timeout(300.0, connect=10.0) -#: Refused before any work was done, so replaying cannot duplicate anything. -#: 503 is excluded: it may come from an intermediary after forwarding. +#: Refused before any work was done. 503 is excluded: it may come from an +#: intermediary after forwarding. UNAMBIGUOUS_RETRY_STATUS = frozenset({429}) DEFAULT_MAX_ATTEMPTS = 5 @@ -41,16 +35,13 @@ def _user_agent() -> str: def encode_json(value: Any) -> bytes: - """Compact UTF-8 JSON. ``allow_nan=False``: strict parsers server-side - reject bare ``NaN`` and the failure is an opaque 400.""" + """Compact UTF-8 JSON; bare ``NaN`` is refused here rather than as a 400.""" return json.dumps(value, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode( "utf-8" ) class PlatformClient: - """Minimal authenticated client for ``{base_url}/api/v1``.""" - def __init__( self, *, @@ -60,32 +51,23 @@ def __init__( max_attempts: int = DEFAULT_MAX_ATTEMPTS, client: Optional[httpx.Client] = None, ) -> None: - # Strip a trailing /api/v1 — the same normalization prime_traces applies — - # so an explicit base_url written with the suffix does not double it. self.base_url = base_url.rstrip("/").removesuffix("/api/v1") self.api_prefix = f"{self.base_url}/api/v1" self.max_attempts = max(1, max_attempts) self._owns_client = client is None - self._headers = { - "Authorization": f"Bearer {api_key}", - "User-Agent": _user_agent(), - } + self._headers = {"Authorization": f"Bearer {api_key}", "User-Agent": _user_agent()} self._timeout = timeout self._client = client or self._new_client() - if self._owns_client: - # Only a pool we opened ourselves is ours to rebuild after a fork. + if self._owns_client: # only a pool we opened is ours to rebuild after a fork _fork.register(self) def _new_client(self) -> httpx.Client: return httpx.Client( - headers=dict(self._headers), - follow_redirects=True, - timeout=self._timeout, + headers=dict(self._headers), follow_redirects=True, timeout=self._timeout ) def reset_after_fork(self) -> None: - """Rebuild the pool in a forked child. The old one is dropped, not - closed: closing could send ``close_notify`` on the parent's socket.""" + """Rebuild the pool; the old one is dropped, not closed.""" self._client = self._new_client() def request( @@ -100,11 +82,8 @@ def request( max_attempts: Optional[int] = None, idempotent: Optional[bool] = None, ) -> Dict[str, Any]: - """Send one request, retrying transient failures. Returns the JSON body. - - ``idempotent`` defaults to ``method != "POST"``; a POST that is safe to - replay (get-or-create) passes ``idempotent=True`` explicitly. - """ + """Send one request, retrying transient failures; returns the JSON body. + ``idempotent`` defaults to ``method != "POST"``.""" url = f"{self.api_prefix}{path}" body = content if content is not None else (encode_json(json_body) if json_body else None) request_kwargs: Dict[str, Any] = { @@ -112,8 +91,7 @@ def request( "headers": {"Content-Type": "application/json"} if body is not None else None, "params": dict(params) if params else None, } - # ``None`` would disable httpx timeouts, not restore the default. - if timeout is not None: + if timeout is not None: # None would disable httpx timeouts, not restore the default request_kwargs["timeout"] = timeout attempts = max_attempts or self.max_attempts replayable = idempotent if idempotent is not None else method.upper() != "POST" @@ -137,9 +115,7 @@ def request( else: return _decode(response) - last = attempt == attempts - 1 - if last or (ambiguous and not replayable): - # Possibly processed already; a duplicate cannot be undone. + if attempt == attempts - 1 or (ambiguous and not replayable): raise error time.sleep(retry_delay(error, attempt)) raise AssertionError("unreachable") # pragma: no cover @@ -157,12 +133,6 @@ def close(self) -> None: if self._owns_client: self._client.close() - def __enter__(self) -> "PlatformClient": - return self - - def __exit__(self, *exc_info: Any) -> None: - self.close() - def _decode(response: httpx.Response) -> Dict[str, Any]: if not response.content: diff --git a/packages/prime-runs/src/prime_runs/backend.py b/packages/prime-runs/src/prime_runs/backend.py index 76e08210f..f10acaed4 100644 --- a/packages/prime-runs/src/prime_runs/backend.py +++ b/packages/prime-runs/src/prime_runs/backend.py @@ -1,16 +1,8 @@ -"""Run backends: the contract, the evaluations backend, and the disabled no-op. - -A backend owns the *lifecycle* of a run — creating it, updating what is known -about it, closing it out with a terminal status. It does not move records; -that is a sink's job (see :mod:`prime_runs.sinks`). - -:class:`EvalsBackend` works over ``/api/v1/evaluations/*``, resolving -environments through the hub's get-or-create so a local run uploads without -``prime env push``. The eval API has no producer-facing way to mark a run -failed: ``finalize`` moves a run to COMPLETED and ``UpdateEvaluationRequest`` -carries no status. A failed or crashed run therefore keeps showing as running; -the terminal state is recorded in ``metadata.prime_runs`` so it is at least -visible. +"""Run lifecycles: create, update, finalize. Sinks, not backends, move records. + +``EvalsBackend`` speaks ``/api/v1/evaluations``; ``RftBackend`` speaks +``/api/v1/rft/external-runs``; ``DisabledBackend`` is the no-op behind +``mode="disabled"``. """ import logging @@ -26,9 +18,7 @@ class Backend(Protocol): - def create(self, spec: RunSpec) -> RunHandle: - """Open a new run and return its identity.""" - ... + def create(self, spec: RunSpec) -> RunHandle: ... def update( self, @@ -37,11 +27,7 @@ def update( config: Optional[Dict[str, Any]] = None, summary: Optional[Dict[str, Any]] = None, ) -> None: - """Persist config (inputs) and/or summary (outputs). - - ``config`` is the run's *whole* config, not a patch: the evaluations API - replaces the stored metadata document. - """ + """Persist the run's whole config (not a patch) and/or summary.""" ... def finalize( @@ -53,20 +39,17 @@ def finalize( error: Optional[str] = None, config: Optional[Dict[str, Any]] = None, ) -> None: - """Close the run out. Called exactly once per run. ``config`` is passed - so a backend recording terminal state inside metadata can merge it.""" + """Close the run out, once per run.""" ... - def close(self) -> None: - """Release transport resources.""" - ... - - -# ------------------------------------------------------------------- evals + def close(self) -> None: ... class EvalsBackend: - """Lifecycle for evaluation runs.""" + """Environments are resolved through the hub's get-or-create, so a local run + uploads without ``prime env push``. The API has no producer-facing way to mark + a run failed, so a non-completed terminal state is recorded in + ``metadata.prime_runs`` and the run keeps showing as running.""" def __init__( self, @@ -79,8 +62,6 @@ def __init__( self._frontend_url = frontend_url.rstrip("/") self._team_id = team_id - # ------------------------------------------------------------------ create - def create(self, spec: RunSpec) -> RunHandle: environments = self._resolve_environments(spec.environments) if not environments: @@ -88,7 +69,6 @@ def create(self, spec: RunSpec) -> RunHandle: "An eval run needs at least one environment. Pass " 'environments=["my-env"] to init().' ) - run_name: str = spec.name or _default_name(spec) payload: Dict[str, Any] = { "name": run_name, @@ -96,15 +76,13 @@ def create(self, spec: RunSpec) -> RunHandle: "tags": list(spec.tags), } _set_if(payload, "model_name", spec.model) - # The API's `dataset` column is always the environment under another name. - _set_if(payload, "dataset", _first_environment_name(spec)) + _set_if(payload, "dataset", _first_environment_name(spec)) # the env under another name _set_if(payload, "framework", spec.framework) _set_if(payload, "description", spec.description) _set_if(payload, "metadata", spec.config or None) _set_if(payload, "team_id", spec.team_id or self._team_id) - # Not replayable (the POST default): a retry after a lost response would - # create a second run and only the second would be tracked. + # Not replayable: a retry after a lost response would create a second run. response = self._client.post("/evaluations/", json_body=payload) run_id = response.get("evaluation_id") if not run_id: @@ -120,8 +98,6 @@ def create(self, spec: RunSpec) -> RunHandle: def url_for(self, run_id: str) -> str: return f"{self._frontend_url}/dashboard/evaluations/{run_id}" - # ------------------------------------------------------------------ update - def update( self, run_id: str, @@ -129,16 +105,11 @@ def update( config: Optional[Dict[str, Any]] = None, summary: Optional[Dict[str, Any]] = None, ) -> None: - """``config`` must be the whole config: the service writes metadata with - a document-level ``$set``.""" payload: Dict[str, Any] = {} _set_if(payload, "metadata", config or None) _set_if(payload, "metrics", summary or None) - if not payload: - return - self._client.put(f"/evaluations/{run_id}", json_body=payload) - - # ---------------------------------------------------------------- finalize + if payload: + self._client.put(f"/evaluations/{run_id}", json_body=payload) def finalize( self, @@ -152,20 +123,12 @@ def finalize( if status is RunStatus.COMPLETED: body: Dict[str, Any] = {} _set_if(body, "metrics", summary or None) - # Finalization enqueues the platform's statistics task; replaying an - # ambiguous failure could enqueue it twice, so this stays non-idempotent. - self._client.post( - f"/evaluations/{run_id}/finalize", - json_body=body or {"metrics": {}}, - ) + # Not replayable: finalization enqueues the platform's statistics task. + self._client.post(f"/evaluations/{run_id}/finalize", json_body=body or {"metrics": {}}) return - # No status endpoint exists yet. Record the terminal state in metadata, - # merged into the full config because this PUT replaces the document. - terminal = { - "status": status.value, - "finished_at": datetime.now(timezone.utc).isoformat(), - } + # Merged into the whole config: the PUT replaces the metadata document. + terminal = {"status": status.value, "finished_at": datetime.now(timezone.utc).isoformat()} if error: terminal["error"] = error self.update(run_id, config={**(config or {}), "prime_runs": terminal}, summary=summary) @@ -179,13 +142,10 @@ def finalize( def close(self) -> None: self._client.close() - # ----------------------------------------------------------- environments - def _resolve_environments(self, refs: List[EnvironmentRef]) -> List[Dict[str, Any]]: - """Hub references as the API's ``EnvironmentReference`` objects, carrying - a pinned ``version_id`` through. A reference that cannot be resolved - raises rather than being skipped: a run silently attached to the wrong - environments looks like a successful upload.""" + """Hub references as the API's ``EnvironmentReference`` objects. One that + cannot be resolved raises rather than being skipped: a run silently + attached to the wrong environments looks like a successful upload.""" resolved: List[Dict[str, Any]] = [] for ref in refs: entry: Dict[str, Any] = {"id": ref.id or self._lookup_environment(ref)} @@ -202,8 +162,7 @@ def _lookup_environment(self, ref: EnvironmentRef) -> str: raise EnvironmentResolutionError( f"Could not resolve environment {ref.slug!r}: {exc}" ) from exc - details = response.get("data") or response - environment_id = details.get("id") + environment_id = (response.get("data") or response).get("id") if not environment_id: raise EnvironmentResolutionError(f"Hub returned no id for environment {ref.slug!r}") return str(environment_id) @@ -240,22 +199,15 @@ def _first_environment_name(spec: RunSpec) -> Optional[str]: def _default_name(spec: RunSpec) -> str: - """The API requires a name; lead with the environment so runs sort together.""" stem = _first_environment_name(spec) or spec.framework or "eval" return f"{stem}-{uuid.uuid4().hex[:8]}" -# ---------------------------------------------------------------- disabled - - def disabled_run_id() -> str: - """A locally issued ID, visibly distinct from a platform one.""" return f"disabled-{uuid.uuid4().hex[:16]}" class DisabledBackend: - """No-op lifecycle, so ``mode="disabled"`` needs no branching upstream.""" - def create(self, spec: RunSpec) -> RunHandle: return RunHandle(id=disabled_run_id()) @@ -269,10 +221,8 @@ def close(self) -> None: return None -# ------------------------------------------------------------------- rft - -#: The RFT API's status vocabulary is ``completed | failed``; the SDK's finer -#: distinctions travel in ``error_message``. +#: The RFT API's vocabulary is ``completed | failed``; the finer distinctions +#: travel in ``error_message``. RFT_TERMINAL_STATUS = { RunStatus.COMPLETED: "completed", RunStatus.FAILED: "failed", @@ -282,20 +232,10 @@ def close(self) -> None: class RftBackend: - """Lifecycle for external training runs over ``/api/v1/rft/*``. - - What prime-rl's ``TrainRun`` did, in the SDK: register a run (or attach to - one a managed launch created), close it out. External runs are created - *for a team* and the platform enables them per team, so ``team_id`` is - required and a team outside the allowlist gets a 403 from :meth:`create`; - the producer decides whether that means "run locally" (the way verifiers - treats a missing key) or "stop". - - The API has no config/summary update — ``run_config`` travels with - :meth:`create` and :meth:`update` is a no-op — and its status vocabulary - is ``completed | failed``: ``cancelled`` and ``crashed`` are reported as - ``failed`` with the reason in ``error_message``. - """ + """External training runs are created for a team, and the platform enables + them per team: ``team_id`` is required and a team outside the allowlist gets + a 403 from :meth:`create`. The API takes the config at creation and has no + summary document, so :meth:`update` is a no-op.""" def __init__( self, @@ -308,8 +248,6 @@ def __init__( self._frontend_url = frontend_url.rstrip("/") self._team_id = team_id - # ------------------------------------------------------------------ create - def create(self, spec: RunSpec) -> RunHandle: if not spec.model: raise ConfigurationError("A training run needs model= — the base model being trained.") @@ -324,8 +262,7 @@ def create(self, spec: RunSpec) -> RunHandle: payload: Dict[str, Any] = { "base_model": spec.model, "max_steps": int(training.max_steps), - # Environment ids are passed through: training environments are - # named by hub id, and the RFT API does its own resolution. + # Passed through: the RFT API resolves hub ids itself. "environments": [_training_environment(ref) for ref in spec.environments], "team_id": team_id, } @@ -338,8 +275,7 @@ def create(self, spec: RunSpec) -> RunHandle: _set_if(payload, "wandb_entity", training.wandb_entity) _set_if(payload, "wandb_run_name", training.wandb_run_name) - # Not replayable: a retry after a lost response would register a - # second run and only the second would be tracked. + # Not replayable: a retry after a lost response would register a second run. response = self._client.post("/rft/external-runs", json_body=payload) run = response.get("run") or {} run_id = run.get("id") @@ -352,17 +288,14 @@ def create(self, spec: RunSpec) -> RunHandle: ) def attach(self, run_id: str) -> RunHandle: - """A handle for a run the platform already created — a launcher that - injected its id. No request: the first write proves access, and it - must be an *external* run (``POST /rft/external-runs``): the v1 - monitoring endpoints answer 400 for a hosted run's id.""" + """A handle for a run a launcher already created. No request: the first + write proves access. The run must be an external one; the monitoring + endpoints answer 400 for a hosted run's id.""" return RunHandle(id=run_id, url=self.url_for(run_id)) def url_for(self, run_id: str) -> str: return f"{self._frontend_url}/dashboard/training/{run_id}" - # ------------------------------------------------------------------ update - def update( self, run_id: str, @@ -370,13 +303,9 @@ def update( config: Optional[Dict[str, Any]] = None, summary: Optional[Dict[str, Any]] = None, ) -> None: - """No-op: an RFT run takes its config at creation and has no summary - document. Run-level outputs go through ``log_metrics``.""" if config or summary: logger.debug("Run %s: the RFT API has no config/summary update; nothing sent", run_id) - # ---------------------------------------------------------------- finalize - def finalize( self, run_id: str, @@ -387,17 +316,13 @@ def finalize( config: Optional[Dict[str, Any]] = None, ) -> None: if status is RunStatus.COMPLETED: - # Compare-and-set on the server, and "already finalized" is a - # success, so a lost response is safe to replay. + # Compare-and-set server-side, and "already finalized" is a success. try: self._client.post( - "/rft/finalize", - json_body={"run_id": run_id, "exit_code": 0}, - idempotent=True, + "/rft/finalize", json_body={"run_id": run_id, "exit_code": 0}, idempotent=True ) except APIError as exc: - # Preserve TrainRun's fallback: finalization may be unavailable - # even though the status endpoint can still close the run out. + # The status endpoint may still close the run out. logger.warning( "Run %s could not be finalized (%s); falling back to a status update", run_id, @@ -412,23 +337,14 @@ def finalize( self._set_terminal_status(run_id, status=status, error_message=message) def _set_terminal_status( - self, - run_id: str, - *, - status: RunStatus, - error_message: Optional[str] = None, + self, run_id: str, *, status: RunStatus, error_message: Optional[str] = None ) -> None: payload: Dict[str, Any] = {"status": RFT_TERMINAL_STATUS[status]} _set_if(payload, "error_message", error_message) try: - self._client.put( - f"/rft/external-runs/{run_id}/status", - json_body=payload, - ) + self._client.put(f"/rft/external-runs/{run_id}/status", json_body=payload) except APIError as exc: - if exc.status_code == 409: - # Already closed out (a managed launcher, or an earlier attempt - # whose response was lost). The outcome is recorded either way. + if exc.status_code == 409: # already closed: the outcome is recorded either way logger.info( "Run %s is already closed on the platform; not marking it %s", run_id, @@ -437,11 +353,6 @@ def _set_terminal_status( return raise - def log_metrics(self, run_id: str, metrics: Dict[str, Any]) -> None: - """One step's metrics, synchronously. The run handle batches these - through its uploader; this is the direct call.""" - self._client.post("/rft/metrics", json_body={"run_id": run_id, "metrics": metrics}) - def close(self) -> None: self._client.close() diff --git a/packages/prime-runs/src/prime_runs/config.py b/packages/prime-runs/src/prime_runs/config.py index c25440a1f..2a4caba54 100644 --- a/packages/prime-runs/src/prime_runs/config.py +++ b/packages/prime-runs/src/prime_runs/config.py @@ -1,5 +1,5 @@ -"""Configuration: ``prime_traces.core.Config`` (``~/.prime/config.json`` plus -environment variables, env taking precedence) with the dashboard URL added.""" +"""``prime_traces.core.Config`` (env vars, then ``~/.prime/config.json``) plus the +dashboard URL.""" import os @@ -11,7 +11,6 @@ class Config(_TracesConfig): @property def frontend_url(self) -> str: - """Dashboard base URL; fallback when a create response omits ``viewer_url``.""" env_val = os.getenv("PRIME_FRONTEND_URL") if env_val: return env_val.rstrip("/") diff --git a/packages/prime-runs/src/prime_runs/exceptions.py b/packages/prime-runs/src/prime_runs/exceptions.py index 4649b673b..37658eeaf 100644 --- a/packages/prime-runs/src/prime_runs/exceptions.py +++ b/packages/prime-runs/src/prime_runs/exceptions.py @@ -1,15 +1,6 @@ -"""Exceptions for the Prime Runs SDK. - -Two families. Errors raised here, before any request is made, derive from -:class:`PrimeRunsError`. Errors from the platform derive from -``prime_traces.APIError``: both SDKs talk to the same platform with the same -credential, and the uploader already handles the traces family, so one -vocabulary serves both. Nothing in either family escapes into a producer loop -under the default ``on_error="warn"``. To catch everything the SDK can raise -under ``on_error="raise"``:: - - except (pr.PrimeRunsError, pr.APIError): -""" +"""Errors raised before any request derive from ``PrimeRunsError``; errors from +the platform are ``prime_traces``' own, so one set of ``except`` clauses covers +both SDKs. Under ``on_error="warn"`` none of them escape into a producer loop.""" from prime_traces.exceptions import ( APIError, @@ -55,8 +46,7 @@ class ConfigurationError(PrimeRunsError): class EnvironmentResolutionError(PrimeRunsError): - """An environment named in ``init()`` could not be resolved to a hub ID — - usually a typo or a permissions problem, not an outage.""" + """An environment named in ``init()`` could not be resolved to a hub id.""" class RunFinishedError(PrimeRunsError): @@ -64,11 +54,7 @@ class RunFinishedError(PrimeRunsError): def is_record_rejection(exc: BaseException) -> bool: - """Whether a failed write is specific to the submitted record batch. - - A later batch can still succeed after serialization, size, or validation - failures, so these errors must not retire an otherwise healthy sink. - """ + """A failure specific to the submitted batch: a later batch can still succeed.""" if isinstance( exc, ( @@ -81,11 +67,10 @@ def is_record_rejection(exc: BaseException) -> bool: ): return True # The legacy samples API returns framework-generated 413/422 responses, - # which the shared traces mapper intentionally leaves as plain APIError. + # which the shared traces mapper leaves as plain APIError. return isinstance(exc, APIError) and exc.status_code in (413, 422) def is_transient(exc: BaseException) -> bool: - """Whether a failure is about this moment (retry later) rather than this - run (stop). Decides whether a sink is retired.""" + """A failure about this moment (retry later) rather than this run (stop).""" return isinstance(exc, (RetryableAPIError, TransportError)) diff --git a/packages/prime-runs/src/prime_runs/metrics.py b/packages/prime-runs/src/prime_runs/metrics.py index ee0f12670..de5f31f0c 100644 --- a/packages/prime-runs/src/prime_runs/metrics.py +++ b/packages/prime-runs/src/prime_runs/metrics.py @@ -1,30 +1,21 @@ -"""Run-level aggregates over native episodes — the aggregation verifiers used, -so migrated runs keep identical dashboard numbers. Opt-in: pass the result to -``run.finish(summary=...)``. Duck-typed; no producer package is imported.""" +"""Run-level aggregates over native episodes, in the shape the eval dashboard +reads; the aggregation verifiers used, so migrated runs keep their numbers. +Pass the result to ``run.finish(summary=...)``.""" from typing import Any, Dict, Optional, Sequence, TypedDict class RunSummary(TypedDict): - """The run-level aggregates the eval dashboard reads. A summary may carry - more than this — ``finish(summary=...)`` stores whatever it is given — but - these three are what the dashboard renders.""" - avg_reward: float avg_metrics: Dict[str, float] avg_error: float def from_episodes(episodes: Sequence[Any], traces: Optional[Sequence[Any]] = None) -> RunSummary: - """Run-level aggregates in the shape the eval dashboard reads. - - Rewards and metrics aggregate over the trainable traces only — fixed agents - (a judge, a modeled user) often carry no rewards and would dilute every mean - with structural zeros — falling back to all traces when none are trainable, - the same rule the dashboard applies. ``avg_error`` is the share of EPISODES - that aren't ok: a hook failure counts even when its traces are clean or it - left none behind. - """ + """Rewards and metrics aggregate over the trainable traces only (a judge or + a modeled user would dilute every mean with structural zeros), falling back + to all traces when none are trainable, as the dashboard does. ``avg_error`` + is the share of episodes that are not ok.""" if traces is None: traces = [trace for episode in episodes for trace in episode.traces] scored = [trace for trace in traces if trace.agent.trainable] or list(traces) diff --git a/packages/prime-runs/src/prime_runs/models.py b/packages/prime-runs/src/prime_runs/models.py index ed7757b44..c2419d3f9 100644 --- a/packages/prime-runs/src/prime_runs/models.py +++ b/packages/prime-runs/src/prime_runs/models.py @@ -1,8 +1,5 @@ -"""Types shared across backends, sinks and the ``Run`` handle. - -Response bodies are deliberately not modeled: backends pull the two or three -fields they need and hand back a ``RunHandle``. -""" +"""Types shared across backends, sinks and the ``Run`` handle. Response bodies +are not modeled: backends pull the fields they need into a ``RunHandle``.""" import os from dataclasses import dataclass, field @@ -13,29 +10,19 @@ from .exceptions import ConfigurationError Mode = Literal["online", "disabled"] -"""``online`` talks to the platform; ``disabled`` makes every call a no-op with -the same object shape.""" - OnError = Literal["warn", "raise"] RunKind = Literal["eval", "train"] -"""What a run is: a standalone evaluation (``/api/v1/evaluations``) or an -external training run (``/api/v1/rft/external-runs``). Stamped as ``run.type`` -on records and sent as upload provenance — the traces service's vocabulary, -which is also what verifiers' ``EvalRunInfo`` / ``TrainRunInfo`` carry.""" +"""Stamped as ``run.type`` on records and sent as upload provenance; the traces +service's vocabulary, matching verifiers' ``EvalRunInfo`` / ``TrainRunInfo``.""" RUN_KIND: RunKind = "eval" -"""The default kind.""" class RunStatus(str, Enum): - """Terminal state a producer can report. - - ``failed`` means the producer said the run failed; ``cancelled`` means - somebody stopped it on purpose (Ctrl-C, a cancelled task) — a decision, - not a fault; ``crashed`` means the process exited without saying (only the - SDK's atexit hook reports it). - """ + """``failed``: the producer said so. ``cancelled``: stopped on purpose (an + interrupt). ``crashed``: the process exited without saying (only the + atexit hook reports it).""" RUNNING = "running" COMPLETED = "completed" @@ -49,11 +36,8 @@ def is_terminal(self) -> bool: @dataclass class EnvironmentRef: - """An environment as a producer names it, before hub resolution. - - ``id`` short-circuits resolution; ``name`` goes through the hub's - get-or-create; ``slug`` looks up a published ``owner/name`` environment. - """ + """``id`` short-circuits resolution; ``name`` goes through the hub's + get-or-create; ``slug`` looks up a published ``owner/name`` environment.""" name: Optional[str] = None id: Optional[str] = None @@ -88,29 +72,17 @@ def __post_init__(self) -> None: CONFIG_SOURCE_KEY = "config_source" -"""Where a config file lands inside a run's config, as a :class:`ConfigSource` -dict. Present means "render this verbatim"; absent means "render the structure".""" +"""Where a config file lands inside a run's config, as a ``ConfigSource`` dict.""" MAX_CONFIG_SOURCE_BYTES = 256 * 1024 -"""A hand-written run config is kilobytes; anything past this is refused at -``init()`` rather than truncated.""" -_CONFIG_SOURCE_FORMATS = { - ".toml": "toml", - ".json": "json", - ".yaml": "yaml", - ".yml": "yaml", -} +_CONFIG_SOURCE_FORMATS = {".toml": "toml", ".json": "json", ".yaml": "yaml", ".yml": "yaml"} @dataclass class ConfigSource: - """The config file a run was started from, kept byte-for-byte. - - That file *is* the run's configuration — comments, key order and section - grouping included — where a resolved model dump is a different artifact. - Nothing here is redacted: keep credentials in the environment, not the file. - """ + """The config file a run was started from, kept byte for byte. Nothing is + redacted: keep credentials in the environment, not the file.""" text: str format: str = "toml" @@ -124,7 +96,7 @@ def to_dict(self) -> Dict[str, Any]: @classmethod def from_mapping(cls, value: Mapping[str, Any]) -> Optional["ConfigSource"]: - """Rebuild from stored metadata. ``None`` if the mapping is not one of ours.""" + """Rebuild from stored metadata; ``None`` if the mapping is not one of ours.""" text = value.get("text") if not isinstance(text, str): return None @@ -138,7 +110,6 @@ def from_mapping(cls, value: Mapping[str, Any]) -> Optional["ConfigSource"]: @classmethod def from_file(cls, path: Union[str, "os.PathLike[str]"]) -> "ConfigSource": - """Read a config file, inferring its format from the suffix.""" resolved = Path(path) try: raw = resolved.read_bytes() @@ -168,23 +139,13 @@ def from_file(cls, path: Union[str, "os.PathLike[str]"]) -> "ConfigSource": @classmethod def coerce(cls, value: Any) -> Optional["ConfigSource"]: - """Normalize the config-file form of ``init(config=...)``. - - A ``str`` or ``PathLike`` is a *path*, never inline text; inline text - goes through ``ConfigSource(text=...)`` explicitly. - """ + """A ``str`` or ``PathLike`` is a *path*, never inline text.""" if value is None or isinstance(value, cls): return value - if isinstance(value, Mapping): - source = cls.from_mapping(value) - if source is None: - raise ValueError("a config-source mapping must contain a 'text' string") - return source if isinstance(value, (str, os.PathLike)): return cls.from_file(value) raise TypeError( - "a config source must be a path, a ConfigSource or a mapping, " - f"got {type(value).__name__}" + f"a config source must be a path or a ConfigSource, got {type(value).__name__}" ) def __post_init__(self) -> None: @@ -199,10 +160,8 @@ def __post_init__(self) -> None: @dataclass class TrainingSpec: - """What a training run is registered with, beyond what every run has - (``model`` is the base model, ``environments`` the training environments, - ``config`` the run config). All of it is for display on the dashboard. - ``max_steps`` is required by the platform; ``0`` means unknown.""" + """Display fields a training run is registered with. ``max_steps`` is + required by the platform; ``0`` means unknown.""" max_steps: int = 0 batch_size: Optional[int] = None @@ -219,9 +178,7 @@ def __post_init__(self) -> None: @dataclass class RunSpec: - """Everything a backend needs to open a run: ``init()``'s arguments after - normalization. ``config`` is the run's inputs; outputs accumulate on the - handle as ``summary``.""" + """``init()``'s arguments after normalization.""" name: Optional[str] = None environments: List[EnvironmentRef] = field(default_factory=list) @@ -237,8 +194,6 @@ class RunSpec: @dataclass class RunHandle: - """What a backend returns once the run exists on the other side.""" - id: str name: Optional[str] = None url: Optional[str] = None diff --git a/packages/prime-runs/src/prime_runs/projection.py b/packages/prime-runs/src/prime_runs/projection.py index 1d3108ba3..958b751f5 100644 --- a/packages/prime-runs/src/prime_runs/projection.py +++ b/packages/prime-runs/src/prime_runs/projection.py @@ -1,28 +1,23 @@ -"""Projection of native traces onto the platform's v0 eval-sample format. - -Moved here from ``verifiers.v1.utils.platform``: this is knowledge about a -platform wire format, so it lives in the client for that wire. Duck-typed — -verifiers ``Trace``/``Episode`` satisfy it structurally and are not imported. - -The projection serves the *current* viewer, which reads the flat sample table; -it leaves the default path once the Viewer API reads traces natively. -""" +"""Projections of native episodes onto the platform's sample tables: the v0 +eval-sample format (moved here from ``verifiers.v1.utils.platform``) and the +training viewer's Parquet table (moved here from prime-rl). Duck-typed: verifiers +``Trace``/``Episode`` satisfy it structurally and are not imported.""" +import json import logging +import math from typing import Any, Dict, Iterable, List, Optional, Sequence from ._http import encode_json logger = logging.getLogger(__name__) -# Repeated /samples posts append; this is the platform's per-request ceiling. +#: The platform's per-request ceiling on ``POST /samples``. MAX_SAMPLES_PAYLOAD_BYTES = 25 * 1024 * 1024 -# The bytes an empty {"samples":[]} envelope costs, counted against every batch. ENVELOPE_BYTES = len(b'{"samples":[]}') def json_bytes(value: Any) -> int: - """Serialized size of ``value`` under the encoding actually sent.""" return len(encode_json(value)) @@ -30,13 +25,25 @@ def _dump(items: Iterable[Any]) -> List[Dict[str, Any]]: return [item.model_dump(mode="json", exclude_none=True) for item in items] -def summary_trace_index(episode: Any) -> int: - """Index of the trace whose flat projection represents the episode. +def _nullify_non_finite(value: Any) -> Any: + """NaN and infinity are not JSON. The tables show them as ``null`` + rather than losing the row (or the batch) they sit in.""" + if isinstance(value, float) and not math.isfinite(value): + return None + if isinstance(value, dict): + return {key: _nullify_non_finite(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_nullify_non_finite(item) for item in value] + return value - The first trainable trace, else the first trace. Shared by the projection - and the rollout counter so the two can never disagree about which trace - (and therefore which ``task.data.idx``) an episode is numbered under. - """ + +def _json_column(value: Any) -> str: + return json.dumps(_nullify_non_finite(value), allow_nan=False) + + +def summary_trace_index(episode: Any) -> int: + """The first trainable trace, else the first trace: the one whose flat + projection (and ``task.data.idx``) represents the episode.""" return next( (index for index, trace in enumerate(episode.traces) if trace.agent.trainable), 0, @@ -46,15 +53,8 @@ def summary_trace_index(episode: Any) -> int: def trace_to_sample( trace: Any, rollout_number: int = 1, episode_id: Optional[str] = None ) -> Dict[str, Any]: - """One trace -> the platform's sample dict (the v0 eval-sample format). - - The hub table stays flat — one row per trace; its episode is denormalized - onto the row (``episode_id`` from the envelope, plus the trace's own - ``agent``/``trainable``), so a multi-trace rollout's grouping travels with - each row without a nested schema. No prompt/completion split (meaningless - mid-branch): ``completion`` is the final branch's messages, ``trajectory`` - one message list per branch. - """ + """One trace as a v0 sample row. ``completion`` is the final branch's + messages, ``trajectory`` one message list per branch.""" task = trace.task.data.model_dump(mode="json", exclude_none=True) branches = trace.branches sample: Dict[str, Any] = { @@ -68,7 +68,6 @@ def trace_to_sample( "prompt": [], "completion": _dump(branches[-1].messages) if branches else [], "answer": task.get("answer"), - # Keyed `tool_defs` because the v0 sample format already carries it there. "tool_defs": _dump(trace.tools) if trace.tools else None, "reward": trace.reward, "timing": trace.timing.model_dump(mode="json", exclude_none=True), @@ -92,41 +91,33 @@ def trace_to_sample( else None, "info": dict(trace.info) or None, } - # Flatten sub-rewards to top-level keys the way v0 does (raw scores, as v0's - # per-function outputs were); env metrics stay nested. + # Sub-rewards are flattened to top-level keys the way v0 did; env metrics stay nested. for name, reward in trace.rewards.items(): if reward is not None: sample.setdefault(name, reward.score) - return sample + return _nullify_non_finite(sample) def episode_to_samples(episode: Any, rollout_number: int) -> List[Dict[str, Any]]: - """One episode -> the sample rows the platform should store for it. - - Normally a single row: the native episode in ``info.native_wrapper`` is - authoritative and carries every trace, while one trainable trace (or the - first) supplies the flat summary older consumers read, identified by - ``native_trace_index``. An episode too large for one request falls back to - one projected row per trace, which loses the native wrapper but keeps the - run visible rather than dropping it. - """ + """Normally one row: the summary trace's projection, with the whole native + episode in ``info.native_wrapper``. An episode too large for one request + falls back to one projected row per trace.""" if not episode.traces: return [] summary_index = summary_trace_index(episode) - summary_trace = episode.traces[summary_index] - sample = trace_to_sample(summary_trace, rollout_number, episode.id) + sample = trace_to_sample(episode.traces[summary_index], rollout_number, episode.id) sample["sample_id"] = episode.id - sample["info"] = { - **(sample["info"] or {}), - "native_wrapper": episode.to_record(), - "native_trace_index": summary_index, - } + sample["info"] = _nullify_non_finite( + { + **(sample["info"] or {}), + "native_wrapper": episode.to_record(), + "native_trace_index": summary_index, + } + ) if ENVELOPE_BYTES + json_bytes(sample) <= MAX_SAMPLES_PAYLOAD_BYTES: return [sample] - logger.warning( - "Episode %s exceeds the platform sample limit; uploading projected traces", - episode.id, + "Episode %s exceeds the platform sample limit; uploading projected traces", episode.id ) return [trace_to_sample(trace, rollout_number, episode.id) for trace in episode.traces] @@ -134,13 +125,8 @@ def episode_to_samples(episode: Any, rollout_number: int) -> List[Dict[str, Any] def build_samples( episodes: Sequence[Any], rollout_numbers: Optional[Dict[Any, int]] = None ) -> List[Dict[str, Any]]: - """Project episodes to platform samples. - - ``rollout_numbers`` carries the per-example counter across calls. Streaming - producers pass the same dict every time so rollout numbering stays - consistent when episodes arrive in batches instead of one final list; a - one-shot caller can ignore it. - """ + """Project episodes to v0 samples. A streaming producer passes the same + ``rollout_numbers`` dict every call so numbering stays consistent.""" counts = rollout_numbers if rollout_numbers is not None else {} samples: List[Dict[str, Any]] = [] for episode in episodes: @@ -153,11 +139,8 @@ def build_samples( def batch_samples(samples: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any]]]: - """Split samples into request-sized batches. - - Raises ``ValueError`` on a sample too large to send alone: silently dropping - it would report a successful run that is missing rows. - """ + """Split samples into request-sized batches; a sample too large to send + alone raises rather than being silently dropped.""" batches: List[List[Dict[str, Any]]] = [] batch: List[Dict[str, Any]] = [] payload_bytes = ENVELOPE_BYTES @@ -168,8 +151,7 @@ def batch_samples(samples: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any] f"sample {index} is too large to upload " f"({ENVELOPE_BYTES + sample_bytes} > {MAX_SAMPLES_PAYLOAD_BYTES} bytes)" ) - # The +1 is the comma that joins this sample to the previous one. - next_bytes = payload_bytes + (1 if batch else 0) + sample_bytes + next_bytes = payload_bytes + (1 if batch else 0) + sample_bytes # +1: the joining comma if batch and next_bytes > MAX_SAMPLES_PAYLOAD_BYTES: batches.append(batch) batch = [] @@ -181,11 +163,7 @@ def batch_samples(samples: Sequence[Dict[str, Any]]) -> List[List[Dict[str, Any] return batches -# ------------------------------------------------------------ training table - - def parquet_available() -> bool: - """Whether the training sample table can be encoded (``prime-runs[train]``).""" try: import pyarrow # noqa: F401 except ImportError: @@ -194,7 +172,6 @@ def parquet_available() -> bool: def train_sample_schema() -> Any: - """The training sample table's Parquet schema (prime-rl's ``SAMPLE_SCHEMA``).""" import pyarrow as pa return pa.schema( @@ -226,21 +203,9 @@ def episodes_to_parquet_bytes( episodes: Sequence[Any], run_id: str, step: int, *, sample_id_offset: int = 0 ) -> Optional[bytes]: """One training step's episodes as the viewer's Parquet table, one row per - episode. Moved here from prime-rl's ``monitors/prime.py``. - - Sample construction is shared with the eval projection (:func:`build_samples`: - the complete native episode in ``info.native_wrapper``, a 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. ``None`` when no episode has a trajectory. - - ``sample_id`` numbers the rows from ``sample_id_offset``: the viewer looks - samples up by (step, sample_id), and a step uploaded as more than one - object (see :mod:`prime_runs.sinks.train_samples`) needs each object to - number its rows in its own range. - """ + episode (the v0 projection plus the RFT-only columns), numbered from + ``sample_id_offset``. ``None`` when no episode has a trajectory.""" import io - import json from datetime import datetime, timezone import pyarrow as pa @@ -260,7 +225,7 @@ def episodes_to_parquet_bytes( rows: List[Dict[str, Any]] = [] for sample_id, sample in enumerate(build_samples(episodes), start=sample_id_offset): trajectory = sample["trajectory"] - if not trajectory: # no branches: an episode that errored before any message + if not trajectory: continue advantage = advantages.get(sample["episode_id"]) trajectory = [{**branch, "advantage": advantage} for branch in trajectory] @@ -278,16 +243,16 @@ def episodes_to_parquet_bytes( "problem_id": problem_id, "sample_id": sample_id, "prompt": "", - "completion": json.dumps(sample["completion"]), - "trajectory": json.dumps(trajectory), + "completion": _json_column(sample["completion"]), + "trajectory": _json_column(trajectory), "answer": "", "env_name": env_names.get(sample["episode_id"], ""), - "task": json.dumps(sample["task"]), - "info": json.dumps(sample["info"]), + "task": _json_column(sample["task"]), + "info": _json_column(sample["info"]), "reward": sample["reward"], "advantage": advantage, - "metrics": json.dumps(sample["metrics"]), - "timing": json.dumps(sample["timing"]), + "metrics": _json_column(sample["metrics"]), + "timing": _json_column(sample["timing"]), "num_input_tokens": trajectory[-1]["num_input_tokens"], "num_output_tokens": trajectory[-1]["num_output_tokens"], "created_at": now, diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index 1e4ecf75b..aaf425f49 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -1,12 +1,7 @@ -"""The run handle, and ``init()`` that produces one. +"""The run handle, and ``init()`` that opens one. -A run is a long-lived thing with a status, so it is an object rather than three -stateless calls: records stream instead of buffering, errors are contained, and -a process that dies still reports a terminal status. - -``init()`` is called *before* rollouts start, and the ID it returns is *the* run -ID everywhere — including inside every trace document the producer writes. -Nothing is re-stamped afterwards. +``init()`` is called before the first rollout; the id it returns is the run id +everywhere, including inside every trace document the producer writes. """ import asyncio @@ -25,7 +20,6 @@ from .exceptions import ConfigurationError, RunFinishedError from .models import ( CONFIG_SOURCE_KEY, - RUN_KIND, ConfigSource, EnvironmentRef, Mode, @@ -37,28 +31,25 @@ TrainingSpec, ) from .sinks import EvalSamplesSink, RftMetricsSink, RftSamplesSink, Sink, TracesSink -from .worker import UploadWorker +from .worker import UploadWorker, deadline_after, time_left logger = logging.getLogger(__name__) MODE_ENV = "PRIME_RUNS_MODE" -#: How long ``finish()`` gives queued uploads to drain before finalizing anyway. -#: Derived from the upload timeout: a single in-flight sample POST may take -#: this long, and a shorter budget would abandon an upload about to succeed. +#: How long ``finish()`` lets queued uploads drain: one in-flight sample POST +#: may take this long, and a shorter budget would abandon it about to succeed. DEFAULT_FINISH_TIMEOUT = float(UPLOAD_TIMEOUT.read or 300.0) -#: A training run lasts days and a platform deploy lasts minutes: a training -#: sink that struck out on transient failures is tried again after this many -#: seconds instead of being retired for the rest of the run. Eval runs are -#: minutes long and keep the permanent retirement. +#: Training runs last days: a training sink that struck out on transient +#: failures is tried again after this long instead of retired for the run. TRAIN_RETIRE_COOLDOWN = 300.0 class Run: - """A live run: an ID, a URL, somewhere to put traces, a summary. + """A live run: an id, a URL, somewhere to put records, a summary. - With the default ``on_error="warn"`` nothing raised by the platform escapes - into a producer's loop; ``on_error="raise"`` surfaces the first failure - from :meth:`flush` or :meth:`finish`, for tests and CI. + With ``on_error="warn"`` (the default) nothing the platform raises escapes + into the producer's loop; ``on_error="raise"`` surfaces the first failure + from :meth:`flush` or :meth:`finish`. """ def __init__( @@ -81,26 +72,24 @@ def __init__( self._mode: Mode = mode self._on_error: OnError = on_error self._status = RunStatus.RUNNING - # A forked child inherits this handle but must not close the parent's run. + # A forked child inherits the handle but must not close the parent's run. self._owns_lifecycle = True - # A run the platform created and handed to this process (a managed - # launch): the platform owns its failure marking, this process only - # completes it. + # A run a launcher created and handed over: the platform owns its + # failure marking, this process only completes it. self._attached = attached self.config: Dict[str, Any] = dict(spec.config) self.summary: Dict[str, Any] = {} self.errors: List[str] = [] - # Under on_error="raise", a failure on the uploader thread is held here - # and re-raised from flush() or finish(), where the caller is looking. + # on_error="raise": the first uploader-thread failure, re-raised from + # flush() or finish() where the caller is looking. self._deferred_error: Optional[BaseException] = None self._finish_timeout = ( DEFAULT_FINISH_TIMEOUT if finish_timeout is None else max(0.0, float(finish_timeout)) ) - # Guards only the transition from accepting writes to finishing. It is - # deliberately separate from _finish_lock, which is held through slow - # network teardown so concurrent finish callers wait for completion. + # _state_lock guards admission only; _finish_lock is held through the + # slow teardown so concurrent finish() callers wait for it. self._state_lock = threading.Lock() self._finish_lock = threading.Lock() self._finishing = False @@ -110,9 +99,7 @@ def __init__( sinks = sinks or [] metrics_sinks = metrics_sinks or [] - # Records and metrics drain on separate uploaders: a metrics dict is - # not a record (no sink would take both), and a slow sample upload - # must not hold a step's metrics behind it. + # Separate uploaders: a slow sample upload must not hold metrics back. self._worker = UploadWorker( sinks, on_error=self._record_sink_error, retire_cooldown=retire_cooldown ) @@ -127,19 +114,14 @@ def __init__( sink.enabled = False self._note(f"starting sink {sink.name}", exc) if self._on_error == "raise": - # The backend may already have created a remote run; close - # it out before the failure reaches the caller. finish() - # re-raises the error noted above. + # Close the remote run before the failure reaches the caller. self.finish(status=RunStatus.FAILED, error=_describe(exc)) raise exc atexit.register(self._atexit_hook) - # -------------------------------------------------------------- identity - @property def id(self) -> str: - """The run ID. Stamp this onto every trace the run produces.""" return self._handle.id @property @@ -148,12 +130,10 @@ def name(self) -> Optional[str]: @property def url(self) -> Optional[str]: - """The dashboard URL; ``None`` when disabled.""" return self._handle.url @property def config_source(self) -> Optional[ConfigSource]: - """The config file this run was launched from, if one was given.""" raw = self.config.get(CONFIG_SOURCE_KEY) return ConfigSource.from_mapping(raw) if isinstance(raw, Mapping) else None @@ -171,61 +151,39 @@ def finished(self) -> bool: @property def kind(self) -> RunKind: - """``eval`` or ``train``.""" return self._spec.kind @property def attached(self) -> bool: - """Whether this process joined a run the platform had already created.""" return self._attached @property def dropped_records(self) -> int: - """Records (and metrics) that reached no sink because a queue was full.""" + """Records that reached no sink because a queue was full.""" return self._worker.dropped + self._metrics_worker.dropped @property def failed_records(self) -> Dict[str, int]: - """Records each sink could not store, by sink name. Per sink because - another sink may still hold them.""" + """Records each sink could not store, by sink name.""" return {**self._worker.failed_records, **self._metrics_worker.failed_records} def __repr__(self) -> str: return f"" - # ------------------------------------------------------------------- log - def log_traces(self, traces: Iterable[Any]) -> None: - """Hand bare traces to the sinks. Returns immediately. - - Accepts verifiers ``Trace`` objects or plain JSON mappings. Both reach - Prime Traces; the v0 sample table is projected from *episodes*, so a - bare trace has no row there. A rollout that is an episode (a group of - traces) goes through :meth:`log_episodes`. Call this as rollouts - complete; nothing is buffered until the end. - """ + """Queue bare traces: verifiers ``Trace`` objects or JSON mappings.""" self._submit("log_traces", traces) def log_episodes(self, episodes: Iterable[Any]) -> None: - """Hand episodes — grouped traces — to the sinks. Returns immediately. - - Accepts verifiers ``Episode`` objects or plain JSON mappings with a - ``traces`` list. The episode's ``run`` reaches every member trace. - Both reach Prime Traces; the v0 sample table (what today's viewer - reads) is projected from episode *objects* only — a JSON episode has - no row there, which the samples sink warns about once. - """ + """Queue episodes: verifiers ``Episode`` objects or JSON mappings with a + ``traces`` list. The sample tables are projected from episode objects + only; a JSON episode reaches Prime Traces alone.""" self._submit("log_episodes", episodes) def log_metrics(self, values: Mapping[str, Any], *, step: Optional[int] = None) -> None: - """Hand one step's metrics to the run. Returns immediately. - - Training runs only (an eval run has no metrics sink; the call is a - no-op there). ``step`` is recorded as ``values["step"]`` unless the - mapping already carries one; every row is stamped with ``_timestamp`` - so step-less rows (inference metrics) keep a time anchor. Non-finite - numbers are dropped, as they are for :meth:`update_summary`. - """ + """Queue one step's metrics (training runs; a no-op for evals). ``step`` + becomes ``values["step"]`` unless already set, every row gets a + ``_timestamp``, and non-finite numbers are dropped.""" self._require_live("log_metrics") record = _clean_metrics(values) if step is not None: @@ -236,23 +194,15 @@ def log_metrics(self, values: Mapping[str, Any], *, step: Optional[int] = None) self._metrics_worker.submit([record]) def _submit(self, operation: str, records: Iterable[Any]) -> None: - # Refuse an already-finished run without consuming a lazy iterable. - self._require_live(operation) + self._require_live(operation) # before consuming a lazy iterable batch = list(records) - # Finish holds this short-lived lock only while closing admission. A - # log call queues first or observes that teardown has begun, without - # waiting behind the slow flush/finalize path. with self._state_lock: self._require_live(operation) if batch: self._worker.submit(batch) def update_summary(self, values: Mapping[str, Any]) -> None: - """Merge run-level outputs into :attr:`summary` ahead of :meth:`finish`. - - Non-finite numbers are dropped here, as they are for - ``finish(summary=...)``; writing to ``summary`` directly skips that. - """ + """Merge run-level outputs into :attr:`summary` ahead of :meth:`finish`.""" with self._state_lock: self._require_live("update_summary") self.summary.update(_clean_metrics(values)) @@ -260,18 +210,12 @@ def update_summary(self, values: Mapping[str, Any]) -> None: def flush(self, timeout: Optional[float] = 30.0) -> bool: """Block until queued records have been written. Under ``on_error="raise"`` this is the first place an upload failure surfaces.""" - deadline = None if timeout is None else time.monotonic() + max(0.0, timeout) - - def remaining() -> Optional[float]: - return None if deadline is None else max(0.0, deadline - time.monotonic()) - - flushed = self._worker.flush(timeout=remaining()) - flushed = self._metrics_worker.flush(timeout=remaining()) and flushed + deadline = deadline_after(timeout) + flushed = self._worker.flush(timeout=time_left(deadline)) + flushed = self._metrics_worker.flush(timeout=time_left(deadline)) and flushed self._raise_deferred() return flushed - # ---------------------------------------------------------------- finish - def finish( self, summary: Optional[Mapping[str, Any]] = None, @@ -280,15 +224,9 @@ def finish( error: Optional[str] = None, timeout: Optional[float] = None, ) -> None: - """Flush everything and close the run out. Idempotent: the first caller + """Drain the uploads and close the run out. Idempotent: the first caller reports the status, concurrent callers wait for that teardown. - - ``timeout`` overrides the run's ``finish_timeout`` for this call: how - long queued uploads get to drain before the run is closed out anyway. - An abort path can pass a short budget so an interrupted producer is - not pinned behind a slow platform. The terminal status is still - reported (with its own retries) after the budget runs out. - """ + ``timeout`` overrides the run's ``finish_timeout`` for this call.""" with self._finish_lock: if self._finished: return @@ -309,22 +247,17 @@ def _finish_once( summary: Optional[Mapping[str, Any]], resolved: RunStatus, error: Optional[str], - timeout: Optional[float] = None, + timeout: Optional[float], ) -> None: if summary: self.summary.update(_clean_metrics(summary)) self._status = resolved budget = self._finish_timeout if timeout is None else max(0.0, float(timeout)) - deadline = time.monotonic() + budget - - def remaining() -> float: - return max(0.0, deadline - time.monotonic()) - - # Records first, so a dashboard reacting to the terminal status never - # sees a finished run with samples still landing. - drained = self._worker.flush(timeout=remaining()) - drained = self._metrics_worker.flush(timeout=remaining()) and drained + deadline = deadline_after(budget) + # Records first, so the dashboard never sees a finished run with samples landing. + drained = self._worker.flush(timeout=time_left(deadline)) + drained = self._metrics_worker.flush(timeout=time_left(deadline)) and drained if not drained: logger.warning( "Run %s: uploads did not drain within %ss; finalizing anyway. " @@ -332,8 +265,8 @@ def remaining() -> float: self.id, budget, ) - self._worker.close(timeout=remaining()) - self._metrics_worker.close(timeout=remaining()) + self._worker.close(timeout=time_left(deadline)) + self._metrics_worker.close(timeout=time_left(deadline)) if self._owns_lifecycle: self._teardown_step( @@ -343,8 +276,7 @@ def remaining() -> float: ), ) if self._attached and resolved is not RunStatus.COMPLETED: - # A managed launch marks its own run failed when this process - # goes away; reporting it from here would race that. + # The launcher marks its own run failed; reporting it here would race that. logger.info( "Run %s: %s, but the platform owns the status of an attached run; " "leaving it to the launcher", @@ -376,15 +308,11 @@ def remaining() -> float: logger.warning( "Run %s: the %s sink could not store %d record(s)", self.id, sink_name, count ) - # Last, so a run whose teardown failed is still closed out first. - self._raise_deferred() + self._raise_deferred() # last, so the run is closed out even when this raises def fail(self, error: Union[str, BaseException], *, timeout: Optional[float] = None) -> None: - """Close the run out as failed.""" self.finish(status=RunStatus.FAILED, error=_describe(error), timeout=timeout) - # -------------------------------------------------------- context manager - def __enter__(self) -> "Run": return self @@ -392,19 +320,15 @@ def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: if exc_type is None: self.finish() return False - if isinstance(exc, (KeyboardInterrupt, asyncio.CancelledError)): - # An interrupt is a decision, not a fault. - status, error = RunStatus.CANCELLED, "interrupted" + status, error = RunStatus.CANCELLED, "interrupted" # a decision, not a fault else: status, error = RunStatus.FAILED, _describe(exc) - try: self.finish(status=status, error=error) except Exception as finish_error: - # The producer's exception is why this block is unwinding; a - # telemetry teardown error must not replace it. Control-flow - # exceptions (KeyboardInterrupt, SystemExit) deliberately pass. + # The producer's exception is why this block is unwinding; a teardown + # error must not replace it. Control-flow exceptions pass through. logger.warning( "Run %s: finishing after %s also failed: %s: %s", self.id, @@ -415,12 +339,9 @@ def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: ) return False - # ------------------------------------------------------------- internals - def reset_after_fork(self) -> None: - """Make an inherited handle safe in a forked child: fresh lock, and the - parent keeps ownership of the lifecycle so the child's atexit hook - cannot finalize a still-running run.""" + """Fresh locks, and the parent keeps the lifecycle: the child's atexit + hook must not finalize a still-running run.""" self._state_lock = threading.Lock() self._finish_lock = threading.Lock() self._finishing = False @@ -428,8 +349,7 @@ def reset_after_fork(self) -> None: self._deferred_error = None def _on_process_exit(self) -> None: - """The process is exiting and nobody called ``finish()``: report CRASHED - rather than FAILED, since the producer never said the run failed.""" + """atexit: nobody called ``finish()``, so report CRASHED rather than FAILED.""" if self._finished: return logger.warning("Run %s was never finished; reporting it as crashed", self.id) @@ -446,11 +366,9 @@ def _require_live(self, operation: str) -> None: ) def _record_sink_error(self, sink_name: str, exc: Exception) -> None: - """Called on the uploader thread when a sink gives up on a batch.""" self._note(f"writing to the {sink_name} sink", exc) def _teardown_step(self, what: str, call: Any) -> None: - """Run one teardown step without letting it skip the steps after it.""" try: call() except Exception as exc: # noqa: BLE001 - teardown must continue @@ -467,7 +385,6 @@ def _note(self, what: str, exc: BaseException) -> None: logger.warning("Run %s: %s", self.id, message) def _raise_deferred(self) -> None: - """Re-raise the first held failure, once.""" exc = self._deferred_error if exc is None: return @@ -475,9 +392,6 @@ def _raise_deferred(self) -> None: raise exc -# --------------------------------------------------------------------- init - - def init( *, name: Optional[str] = None, @@ -497,30 +411,19 @@ def init( id: Optional[str] = None, training: Optional[TrainingSpec] = None, ) -> Run: - """Start a run and return a handle to it. - - Call this *before* the first rollout: the ID it returns is what every trace - in the run should carry, and the URL is what a producer prints. + """Open a run and return its handle. Call it before the first rollout. ``mode`` defaults to ``$PRIME_RUNS_MODE``, else online when there is an API - key and disabled (with a warning) when there is not. ``config`` is what the - run was configured with: the path to the file it was launched from (stored - byte for byte under ``config_source``), or a mapping taken as given. - ``finish_timeout`` is how long :meth:`Run.finish` lets queued uploads drain - before closing the run out anyway (default: the upload timeout, 300 s); - ``finish(timeout=...)`` overrides it per call. - - ``kind="train"`` opens an external training run instead of an evaluation: - ``model`` is the base model, ``environments`` the training environments - (hub ids, passed through), ``training`` the rest of what the dashboard - shows, and a team is required. ``id`` attaches to a training run the - platform already created (a launcher that injected ``$RUN_ID``): nothing - is registered, the platform keeps ownership of the run's failure marking, - and a clean :meth:`Run.finish` still completes it. The run must itself be - an external run — one created through ``POST /rft/external-runs`` — since - the platform's monitoring endpoints answer 400 for a hosted (managed) run. - Training sinks that strike out on transient failures are tried again after - ``TRAIN_RETIRE_COOLDOWN`` seconds rather than retired for the run. + key and disabled (with a warning) when there is not. ``config`` is the path + to the file the run was launched from (stored byte for byte under + ``config_source``) or a mapping taken as given. ``finish_timeout`` bounds + the drain in :meth:`Run.finish`. + + ``kind="train"`` opens an external training run: ``model`` is the base + model, ``environments`` the hub ids, ``training`` the display fields, and a + team is required. ``id`` attaches to an external run a launcher already + created (``$RUN_ID``): nothing is registered, the platform keeps the run's + failure marking, and a clean finish still completes it. """ settings = Config() api_key = api_key if api_key is not None else settings.api_key @@ -577,15 +480,11 @@ def init( else: handle = backend.create(spec) except BaseException: - # Ownership has not reached a Run yet, so nothing else can release - # the connection pool when environment resolution or creation fails. + # Nothing else can release the connection pool yet. try: backend.close() except Exception as close_error: # noqa: BLE001 - preserve the create failure - logger.debug( - "Error closing the platform client after run creation failed: %s", - close_error, - ) + logger.debug("Error closing the platform client: %s", close_error) raise # Both transports run during the transition: traces is the system of # record, the sample table is what today's viewer reads. @@ -633,9 +532,8 @@ def _resolve_mode(mode: Optional[Mode], *, api_key: str) -> Mode: def _sink_context(spec: RunSpec) -> Dict[str, str]: - """Upload-scoped provenance. Not the join key — that is ``run.id`` inside - the trace document.""" - context = {"source": "prime-runs", "run_type": spec.kind or RUN_KIND} + """Upload-scoped provenance; the join key is ``run.id`` inside the document.""" + context = {"source": "prime-runs", "run_type": spec.kind} if spec.framework: context["framework"] = spec.framework if spec.model: @@ -644,12 +542,6 @@ def _sink_context(spec: RunSpec) -> Dict[str, str]: def _normalize_config(value: Any) -> Dict[str, Any]: - """A producer's config as a plain dict. - - A path is the file the run was launched from, kept byte for byte under - ``config_source`` (stored, not parsed — the platform can read TOML). A - mapping is taken exactly as given. - """ if value is None: return {} if isinstance(value, Mapping): @@ -669,42 +561,21 @@ def _describe(error: Union[str, BaseException]) -> str: return str(error) -def _clean_metrics(metrics: Mapping[str, Any], *, _path: str = "") -> Dict[str, Any]: - """Recursively drop NaN/infinity: strict JSON rejects them, and the - failure would surface as an opaque 400 on the whole request.""" - cleaned: Dict[str, Any] = {} - for key, value in metrics.items(): - path = f"{_path}.{key}" if _path else str(key) - if isinstance(value, float) and not math.isfinite(value): - logger.debug("Dropping non-finite metric %s=%r", path, value) - continue - if isinstance(value, Mapping): - nested = _clean_metrics(value, _path=path) - if nested: - cleaned[key] = nested - continue - if isinstance(value, (list, tuple)): - cleaned[key] = _clean_metric_sequence(value, path) - continue - cleaned[key] = value - return cleaned - - -def _clean_metric_sequence(values: Sequence[Any], path: str) -> List[Any]: - """Clean JSON-array-shaped metric values while preserving their order.""" - cleaned: List[Any] = [] - for index, value in enumerate(values): - item_path = f"{path}[{index}]" - if isinstance(value, float) and not math.isfinite(value): - logger.debug("Dropping non-finite metric %s=%r", item_path, value) - continue - if isinstance(value, Mapping): - cleaned.append(_clean_metrics(value, _path=item_path)) - elif isinstance(value, (list, tuple)): - cleaned.append(_clean_metric_sequence(value, item_path)) - else: - cleaned.append(value) - return cleaned +_DROP: Any = object() + + +def _finite(value: Any) -> Any: + """``value`` without NaN/infinity anywhere inside it (strict JSON rejects + them as an opaque 400), or ``_DROP``. Empty nested mappings are dropped.""" + if isinstance(value, float) and not math.isfinite(value): + return _DROP + if isinstance(value, Mapping): + items = ((key, _finite(item)) for key, item in value.items()) + return {key: item for key, item in items if item is not _DROP and item != {}} + if isinstance(value, (list, tuple)): + return [item for item in (_finite(item) for item in value) if item is not _DROP] + return value -__all__ = ["Run", "init", "MODE_ENV"] +def _clean_metrics(metrics: Mapping[str, Any]) -> Dict[str, Any]: + return _finite(dict(metrics)) diff --git a/packages/prime-runs/src/prime_runs/sinks/__init__.py b/packages/prime-runs/src/prime_runs/sinks/__init__.py index 080793f71..ea393e379 100644 --- a/packages/prime-runs/src/prime_runs/sinks/__init__.py +++ b/packages/prime-runs/src/prime_runs/sinks/__init__.py @@ -1,4 +1,4 @@ -"""Record transports. Independent of backends, and of each other.""" +"""Record transports, independent of backends and of each other.""" from .base import Sink, is_episode, stamp_run, to_mapping from .metrics import RftMetricsSink diff --git a/packages/prime-runs/src/prime_runs/sinks/base.py b/packages/prime-runs/src/prime_runs/sinks/base.py index a9c74c62e..5418e16d0 100644 --- a/packages/prime-runs/src/prime_runs/sinks/base.py +++ b/packages/prime-runs/src/prime_runs/sinks/base.py @@ -1,25 +1,14 @@ -"""The contract a record sink implements, plus the record helpers sinks share. +"""The sink contract, the shared write loop, and the record helpers.""" -A sink moves records — traces, episodes — to wherever they are stored. It knows -nothing about run lifecycle. Sinks are independent of each other: during the -transition the traces sink and the legacy eval-samples sink both run, and -retiring one is a change to the default sink list, not to any producer. - -Every sink must be degradable: one that cannot write sets ``enabled = False`` -and says why, once. -""" - -from typing import Any, Dict, Mapping, Protocol, Sequence +from typing import Any, Callable, Dict, Mapping, Optional, Protocol, Sequence, Tuple +from ..exceptions import is_record_rejection, is_transient from ..models import RUN_KIND class SinkWriteError(Exception): - """A sink stored only part of a submitted batch. - - ``cause`` drives the worker's retry/retirement policy; ``failed_records`` - lets it account for only the input records that were not fully stored. - """ + """A sink stored only part of a batch. ``cause`` drives the worker's retry + and retirement policy, ``failed_records`` its loss accounting.""" def __init__(self, cause: Exception, *, failed_records: int) -> None: self.cause = cause @@ -30,31 +19,50 @@ def __init__(self, cause: Exception, *, failed_records: int) -> None: class Sink(Protocol): - """A destination for run records.""" + """A destination for run records. ``write`` runs on the uploader thread. A + sink that cannot write sets ``enabled = False`` and says why, once.""" name: str enabled: bool - def start(self, run_id: str, context: Mapping[str, str]) -> None: - """Bind the sink to a run before the first write.""" - ... + def start(self, run_id: str, context: Mapping[str, str]) -> None: ... + + def write(self, records: Sequence[Any]) -> None: ... + + def flush(self) -> None: ... + + def close(self) -> None: ... + - def write(self, records: Sequence[Any]) -> None: - """Send one batch. Called from the uploader thread, never inline.""" - ... +Unit = Tuple[int, Callable[[], Any]] +"""``(record count, send)``: one request's worth of a batch.""" - def flush(self) -> None: - """Block until everything handed over so far has been written.""" - ... - def close(self) -> None: - """Release transport resources.""" - ... +def send_each(units: Sequence[Unit]) -> None: + """Send a batch one unit at a time, then raise ``SinkWriteError`` if any + failed. A record rejection or a transient failure costs its own unit; a + sink-wide failure stops the batch and costs the rest too. A transient + error is what gets reported, so the worker's strike accounting is not + fooled by a record-local one in the same batch.""" + failed = 0 + reported: Optional[Exception] = None + for position, (count, send) in enumerate(units): + try: + send() + except Exception as exc: # noqa: BLE001 - classified below + failed += count + if not (is_record_rejection(exc) or is_transient(exc)): + failed += sum(rest for rest, _ in units[position + 1 :]) + reported = exc + break + if reported is None or is_transient(exc): + reported = exc + if reported is not None: + raise SinkWriteError(reported, failed_records=failed) from reported def to_mapping(record: Any) -> Mapping[str, Any]: - """The JSON mapping for a record: a dict passes through, anything else must - implement ``to_record()`` (verifiers ``Trace``/``Episode`` do).""" + """A dict passes through; anything else must implement ``to_record()``.""" if isinstance(record, Mapping): return record to_record = getattr(record, "to_record", None) @@ -67,22 +75,16 @@ def to_mapping(record: Any) -> Mapping[str, Any]: def is_episode(record: Any) -> bool: - """Whether a record (object or mapping) is an episode rather than a trace.""" if isinstance(record, Mapping): return "traces" in record return hasattr(record, "traces") def stamp_run(mapping: Mapping[str, Any], run_id: str, run_type: str = RUN_KIND) -> Dict[str, Any]: - """A copy of ``mapping`` carrying ``run`` at the top level, and on every - member trace of an episode that lacks one. - - A record that already names a run keeps it; one without is stamped, since - an upload with no ``run.id`` is orphaned and unqueryable. Members matter - because the traces service derives ``run_id`` from ``trace.run.id`` only — - the episode envelope's ``run`` is never read — while producers (verifiers) - record the run on the episode and nowhere else. - """ + """A copy of ``mapping`` with ``run`` on the envelope and on every member + trace that lacks one. A record that already names a run keeps it. Members + matter because the traces service reads ``run_id`` from ``trace.run.id`` + only, while verifiers records the run on the episode.""" stamped = dict(mapping) if not stamped.get("run"): stamped["run"] = {"id": run_id, "type": run_type} diff --git a/packages/prime-runs/src/prime_runs/sinks/metrics.py b/packages/prime-runs/src/prime_runs/sinks/metrics.py index 3a6f38036..f71c675e9 100644 --- a/packages/prime-runs/src/prime_runs/sinks/metrics.py +++ b/packages/prime-runs/src/prime_runs/sinks/metrics.py @@ -1,28 +1,17 @@ """Training metrics sink: one ``POST /rft/metrics`` per logged step. -The RFT metrics endpoint takes a single step's metrics per request, so a -coalesced batch costs one request per record. The platform allows 60 a minute -per token; a 429 is retried with the server's ``Retry-After`` before it counts -as a strike. - -Requests are replayed after an ambiguous failure (a 502/504, a read timeout): -the platform keeps one row per step, so a duplicate collapses on its side, -while a lost row is a hole in the training curves for good. +Requests are replayed after an ambiguous failure: the platform keeps one row per +step, so a duplicate collapses server-side, while a lost row is a hole in the +curves for good. """ -import logging -from typing import Any, Mapping, Optional, Sequence +from typing import Any, Callable, Mapping, Optional, Sequence from .._http import PlatformClient -from ..exceptions import is_record_rejection, is_transient -from .base import Sink, SinkWriteError - -logger = logging.getLogger(__name__) +from .base import Sink, send_each class RftMetricsSink(Sink): - """Posts per-step metrics dicts to the RFT metrics endpoint.""" - name = "rft_metrics" def __init__(self, client: PlatformClient) -> None: @@ -39,41 +28,23 @@ def write(self, records: Sequence[Any]) -> None: return if self._run_id is None: raise RuntimeError("RftMetricsSink.write called before start()") + send_each([(1, self._sender(record)) for record in records]) - failed = 0 - reported: Optional[Exception] = None - for index, record in enumerate(records): + def _sender(self, record: Any) -> Callable[[], None]: + def send() -> None: if not isinstance(record, Mapping): - failed += 1 - if reported is None: - reported = TypeError(f"metrics must be a mapping, got {type(record).__name__}") - continue - try: - # Replayable: the platform merges rows by step, so a retry after - # a lost response cannot duplicate anything visible. - self._client.post( - "/rft/metrics", - json_body={"run_id": self._run_id, "metrics": dict(record)}, - idempotent=True, - ) - except Exception as exc: # noqa: BLE001 - classified below - failed += 1 - if not (is_record_rejection(exc) or is_transient(exc)): - # Sink-wide: the rest would fail the same way. - failed += len(records) - index - 1 - reported = exc - break - # Prefer a transient error for the worker's strike accounting. - if reported is None or is_transient(exc): - reported = exc - else: - self.steps_written += 1 + raise TypeError(f"metrics must be a mapping, got {type(record).__name__}") + self._client.post( + "/rft/metrics", + json_body={"run_id": self._run_id, "metrics": dict(record)}, + idempotent=True, + ) + self.steps_written += 1 - if reported is not None: - raise SinkWriteError(reported, failed_records=failed) from reported + return send def flush(self) -> None: - """Writes are synchronous; the uploader thread owns the asynchrony.""" + """Writes are synchronous.""" def close(self) -> None: - """The client is shared with the backend, which closes it.""" + """The client belongs to the backend.""" diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py index cf6a6465b..b9f384481 100644 --- a/packages/prime-runs/src/prime_runs/sinks/samples.py +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -1,26 +1,21 @@ -"""Legacy sink: the flat eval-sample table behind today's viewer. +"""Legacy sink: the flat eval-sample table today's viewer reads. -Prime Traces is gated to an allowlist and the viewer reads the v0 sample table, -so both sinks run until the viewer reads traces natively; retiring this one is -a change to the default sink list. - -``POST /samples`` appends, so a request whose response was lost is not -replayed: a lost batch is recoverable, duplicated rows skew every average. +``POST /samples`` appends, so a request whose response was lost is not replayed: +a lost batch is recoverable, duplicated rows skew every average. """ import logging -from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple +from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple from .._http import UPLOAD_TIMEOUT, PlatformClient, encode_json -from ..exceptions import is_record_rejection, is_transient from ..projection import batch_samples, build_samples -from .base import Sink, SinkWriteError, is_episode +from .base import Sink, Unit, is_episode, send_each logger = logging.getLogger(__name__) class EvalSamplesSink(Sink): - """Projects episodes to v0 samples and pushes them to the evaluations API.""" + """Projects episodes to v0 samples and posts them to the evaluations API.""" name = "eval_samples" @@ -28,12 +23,12 @@ def __init__(self, client: PlatformClient) -> None: self.enabled = True self._client = client self._run_id: Optional[str] = None - # Carried across calls so a streaming producer numbers rollouts the - # same way a one-shot upload did. + # Per-example rollout counter, carried across calls so streamed batches + # number rollouts the way a one-shot upload did. self._rollout_numbers: Dict[Any, int] = {} self.samples_written = 0 - #: Records this sink could not project and therefore did not store. - #: The traces sink takes them; the v0 table simply has no row for them. + #: Records with no v0 projection (JSON episodes, bare traces). They + #: reach Prime Traces only and are not a loss. self.skipped = 0 def start(self, run_id: str, context: Mapping[str, str]) -> None: @@ -44,56 +39,31 @@ def write(self, records: Sequence[Any]) -> None: return if self._run_id is None: raise RuntimeError("EvalSamplesSink.write called before start()") - samples, owners = self._to_samples(records) - if not samples: - return - batches = batch_samples(samples) - failed_owners: Set[int] = set() - reported_error: Optional[Exception] = None + units: List[Unit] = [] offset = 0 - for batch_index, batch in enumerate(batches): - next_offset = offset + len(batch) - batch_owners = owners[offset:next_offset] - offset = next_offset - try: - self._client.post( - f"/evaluations/{self._run_id}/samples", - content=encode_json({"samples": batch}), - timeout=UPLOAD_TIMEOUT, - idempotent=False, # appends; a lost response must not duplicate rows - ) - except Exception as exc: - failed_owners.update(batch_owners) - recoverable = is_record_rejection(exc) or is_transient(exc) - if not recoverable: - # The same sink-wide failure will reject everything after - # this request, so do not send more traffic just to prove it. - for remaining_batch in batches[batch_index + 1 :]: - failed_owners.update(owners[offset : offset + len(remaining_batch)]) - offset += len(remaining_batch) - reported_error = exc - break - # Prefer a transient error for worker strike accounting when a - # write encountered both transient and record-local failures. - if reported_error is None or is_transient(exc): - reported_error = exc - else: - self.samples_written += len(batch) - - if reported_error is not None: - raise SinkWriteError( - reported_error, failed_records=len(failed_owners) - ) from reported_error + for batch in batch_samples(samples): + batch_owners = set(owners[offset : offset + len(batch)]) + offset += len(batch) + units.append((len(batch_owners), self._sender(batch))) + send_each(units) + + def _sender(self, batch: List[Dict[str, Any]]) -> Callable[[], None]: + def send() -> None: + self._client.post( + f"/evaluations/{self._run_id}/samples", + content=encode_json({"samples": batch}), + timeout=UPLOAD_TIMEOUT, + idempotent=False, + ) + self.samples_written += len(batch) + + return send def _to_samples(self, records: Sequence[Any]) -> Tuple[List[Dict[str, Any]], List[int]]: - """Episode objects are projected; v0 sample dicts (``sample_id``) pass - through. Anything else — a JSON episode, a bare trace — has no v0 - projection (it is attribute-based, see :mod:`prime_runs.projection`) - and is skipped: warned once, counted, and left to the traces sink. - Raising here would retire this sink for the rest of the run over one - record shape, which is worse than one missing row. The parallel owner - list preserves input-record accounting when HTTP batching splits them.""" + """Project episode objects; pass v0 sample dicts through; skip the rest, + warned once. ``owners`` maps each sample back to its input record so + loss accounting survives HTTP batching.""" samples: List[Dict[str, Any]] = [] owners: List[int] = [] skipped = 0 @@ -123,7 +93,7 @@ def _to_samples(self, records: Sequence[Any]) -> Tuple[List[Dict[str, Any]], Lis return samples, owners def flush(self) -> None: - """Writes are synchronous; the uploader thread owns the asynchrony.""" + """Writes are synchronous.""" def close(self) -> None: - """The client is shared with the backend, which closes it.""" + """The client belongs to the backend.""" diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py index 8da688eba..90858f9ac 100644 --- a/packages/prime-runs/src/prime_runs/sinks/traces.py +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -1,10 +1,8 @@ -"""Streaming sink over Prime Traces — the primary record transport. +"""Streaming sink over Prime Traces, the system of record. -Uploads are content-addressed, so a retry after a lost response replays rather -than duplicates, and episode-aware, so a multi-trace rollout keeps its grouping. - -The join key is ``run.id`` inside the trace document, which the ingestion -service indexes; ``context`` is upload-scoped provenance only. +Uploads are content-addressed (a retry after a lost response replays rather +than duplicates) and episode-aware. The join key is ``run.id`` inside the trace +document; ``context`` is upload-scoped provenance only. """ import logging @@ -23,8 +21,6 @@ class TracesSink(Sink): - """Uploads records through the Prime Traces service.""" - name = "traces" def __init__( @@ -41,8 +37,7 @@ def __init__( self.enabled = True self._client = client self._injected_client = client is not None - # The traces service has its own URL (PRIME_TRACES_URL / `traces_url`), - # which prime-traces resolves itself. + # The traces service URL (PRIME_TRACES_URL) is resolved by prime-traces. self._client_kwargs: Dict[str, Any] = {} if api_key is not None: self._client_kwargs["api_key"] = api_key @@ -56,17 +51,15 @@ def __init__( self._receipt_history_size = receipt_history_size _fork.register(self) - # ------------------------------------------------------------------ setup - def start(self, run_id: str, context: Mapping[str, str]) -> None: self._run_id = run_id self._context = {key: str(value) for key, value in context.items() if value is not None} self._ensure_client() - def _ensure_client(self) -> bool: - """Build the traces client lazily, so a fork reset is repaired on next write.""" + def _ensure_client(self) -> None: + """Build the client lazily, so a fork reset is repaired on the next write.""" if self._client is not None: - return True + return if self._injected_client: exc = RuntimeError("an injected traces client cannot be reused after a fork") self._disable(str(exc)) @@ -76,20 +69,16 @@ def _ensure_client(self) -> bool: except Exception as exc: # noqa: BLE001 - the run applies its error policy self._disable(f"could not construct the traces client ({exc})") raise - return True def reset_after_fork(self) -> None: """Drop (not close) the inherited client; its socket is the parent's.""" self._client = None - # ------------------------------------------------------------------ write - def write(self, records: Sequence[Any]) -> None: - if not self.enabled or not records or not self._ensure_client(): + if not self.enabled or not records: return - - # The same bytes under a different format are rejected as a conflict, - # so infer from the first record, which is stable within a batch. + self._ensure_client() + # The line format is inferred from the first record; batches are homogeneous. line_format = LineFormat.EPISODE if is_episode(records[0]) else LineFormat.TRACE payload = [self._prepare(record) for record in records] try: @@ -102,18 +91,16 @@ def write(self, records: Sequence[Any]) -> None: ) ) except ForbiddenError as exc: - # No runtime action fixes a 403, so the sink retires either way. - # What differs is whether the batch counts as lost. - if _is_not_enabled(exc): - # Outside the beta: there was never anywhere for these records - # to go, so nothing was lost. Not a warning, not a failure. + # Nothing at runtime fixes a 403, so the sink retires either way; + # what differs is whether the batch counts as lost. + if exc.code == ErrorCode.SERVICE_NOT_ENABLED.value: + # Outside the beta: there was never anywhere for these records to go. self._retire_quietly( f"Prime Traces is not enabled for this account ({exc}); " "continuing with the remaining sinks" ) return - # A credential without the traces scope is something the caller - # can fix; raised so loss accounting and strict callers see it. + # A credential without the traces scope: raised for loss accounting. self._disable(f"this credential cannot write traces ({exc})") raise self.receipts_received += len(receipts) @@ -122,16 +109,13 @@ def write(self, records: Sequence[Any]) -> None: del self.receipts[: -self._receipt_history_size] def _prepare(self, record: Any) -> Any: - """The wire mapping for a record, carrying the run on the envelope and - on every member trace. Producer objects go through their own - ``to_record()`` here rather than inside the transport — same bytes, - but the members are reachable for stamping.""" + """The wire mapping, with the run on the envelope and every member trace.""" if self._run_id is None: return record return stamp_run(to_mapping(record), self._run_id, self._context.get("run_type", RUN_KIND)) def flush(self) -> None: - """Uploads are synchronous; nothing is held back here.""" + """Uploads are synchronous.""" def close(self) -> None: client = self._client @@ -142,8 +126,6 @@ def close(self) -> None: except Exception as exc: # noqa: BLE001 - teardown must not raise logger.debug("Error closing the traces client: %s", exc) - # ----------------------------------------------------------- degradation - def _disable(self, reason: str) -> None: if self.enabled: logger.warning("Traces sink disabled: %s", reason) @@ -153,9 +135,3 @@ def _retire_quietly(self, reason: str) -> None: if self.enabled: logger.info("Traces sink off: %s", reason) self.enabled = False - - -def _is_not_enabled(exc: ForbiddenError) -> bool: - """``service_not_enabled``: the account is outside the private beta. The - other 403, ``forbidden``, means the token lacks the ``traces`` scope.""" - return exc.code == ErrorCode.SERVICE_NOT_ENABLED.value diff --git a/packages/prime-runs/src/prime_runs/sinks/train_samples.py b/packages/prime-runs/src/prime_runs/sinks/train_samples.py index 5f3472b6c..2d3c0f754 100644 --- a/packages/prime-runs/src/prime_runs/sinks/train_samples.py +++ b/packages/prime-runs/src/prime_runs/sinks/train_samples.py @@ -1,26 +1,18 @@ """Training samples sink: the step-keyed Parquet table the training viewer reads. -What prime-rl's ``TrainRun.upload_samples`` did, in the SDK: for each training -step, encode the step's episodes to one Parquet object and send it through the -platform's presign -> PUT to object storage -> confirm flow. Like the eval -samples sink this serves *today's* viewer; Prime Traces is the system of record -and this sink leaves the default list once the viewer reads traces natively. - -Which episodes to upload is read off the records: a verifiers episode carries -``run.work`` (``TrainWorkInfo(step=...)``) once its producer has stamped it, and -prime-rl stamps every dispatched episode. Episodes from eval work, bare traces -and JSON episodes have no row here and go to Prime Traces only. - -Objects are additive per step: every upload mints its own -``step_{step}_{uuid}.parquet`` key and the viewer unions every object under a -step, so a step's episodes may arrive across several ``log_episodes`` calls — -prime-rl hands over a step's batch at ship time, and an off-policy episode -dispatched at step N can land in a later batch. Each call becomes one more -object; ``sample_id`` is numbered from a per-upload offset so (step, -sample_id), which the viewer looks samples up by, stays unique across them. -Log a step in one call to keep it to one object. +Every Nth training step (prime-rl's cadence), the step's episodes are encoded to +one Parquet object and sent through presign -> PUT -> confirm. The step is read +off each episode's ``run.work.step`` (verifiers ``TrainRunInfo``); eval work, +bare traces and JSON episodes have no row here and reach Prime Traces only. + +Objects are additive per step: every upload mints its own key and the viewer +unions every object under a step, so a step logged across several calls (an +off-policy episode landing in a later batch) becomes several objects. Each +numbers its ``sample_id``s in its own range, since the viewer looks samples up +by (step, sample_id). """ +import functools import logging import secrets import time @@ -31,39 +23,25 @@ from .. import _fork from .._http import UPLOAD_TIMEOUT, PlatformClient -from ..exceptions import ( - APIError, - RetryableAPIError, - TransportError, - is_record_rejection, - is_transient, -) -from .base import Sink, SinkWriteError, is_episode +from ..exceptions import APIError, RetryableAPIError, TransportError +from .base import Sink, is_episode, send_each logger = logging.getLogger(__name__) -#: Upload every Nth training step (prime-rl's cadence, chosen so the sample -#: table does not receive every rollout of every step). DEFAULT_STEP_INTERVAL = 10 #: Attempts for the object-storage PUT, which is idempotent. UPLOAD_ATTEMPTS = 3 - -#: ``sample_id = range * SAMPLE_ID_STRIDE + row``. In the process that opened -#: the run the range is the step's upload index: the first object numbers its -#: rows from 0 (what a single-object step always did), a straggler object from -#: 2**32, and so on; no object holds 2**32 rows. A forked child inherits a copy -#: of those counters and would hand out the same ranges as its parent, so it -#: draws its ranges at random from the upper half of the range space instead, -#: which the parent's counter never reaches. Every id stays below 2**53: the -#: viewer parses ``sample_id`` as a JSON number. +#: ``sample_id = range * SAMPLE_ID_STRIDE + row``. The range is the step's upload +#: index, so a single-object step still numbers from 0. A forked child inherits +#: the parent's counters and would reuse its ranges, so it draws ranges at random +#: from the upper half of the space instead. Ids stay below 2**53 for the viewer. SAMPLE_ID_STRIDE = 1 << 32 SAMPLE_ID_RANGES = 1 << 21 class Encoder(Protocol): """``(episodes, run_id, step, sample_id_offset=...) -> parquet bytes``, or - ``None`` when there is nothing to upload for the step. Rows are numbered - from ``sample_id_offset``.""" + ``None`` when there is nothing to upload for the step.""" def __call__( self, episodes: Sequence[Any], run_id: str, step: int, *, sample_id_offset: int = 0 @@ -75,9 +53,7 @@ def _field(obj: Any, name: str) -> Any: def training_step(record: Any) -> Optional[int]: - """The training step an episode was dispatched at, from its ``run.work`` - (verifiers ``TrainRunInfo`` / ``TrainWorkInfo``). ``None`` for anything - else: eval work, a record without provenance, a bare trace.""" + """The step a training-work episode was dispatched at, else ``None``.""" run = _field(record, "run") work = _field(run, "work") if run is not None else None if _field(run, "type") != "train" or _field(work, "type") != "train": @@ -87,8 +63,6 @@ def training_step(record: Any) -> Optional[int]: class RftSamplesSink(Sink): - """Encodes each training step's episodes to Parquet and uploads it.""" - name = "rft_samples" def __init__( @@ -107,28 +81,20 @@ def __init__( self._step_interval = step_interval self._upload_client = upload_client self._owns_upload_client = upload_client is None - # An inherited socket belongs to the parent; see reset_after_fork. self._forked_with_injected_client = False + self._forked = False self._run_id: Optional[str] = None + self._uploads_per_step: Dict[int, int] = {} self.steps_written = 0 #: Records with no row in this table (eval work, bare traces, JSON). self.skipped = 0 #: Training episodes left out by the step cadence. Not a loss. self.sampled_out = 0 - #: Uploads attempted per step; each numbers its rows in its own range. - self._uploads_per_step: Dict[int, int] = {} - #: A forked child cannot share the parent's counters; see SAMPLE_ID_RANGES. - self._forked = False _fork.register(self) - # ------------------------------------------------------------------ setup - def reset_after_fork(self) -> None: - """Drop (not close) the inherited storage client: its socket is the - parent's. An owned client is rebuilt on the next upload; an injected - one cannot be, so the child's uploads fail rather than share it. The - inherited upload counters are the parent's too: from here on ranges - are drawn at random, apart from anything the parent hands out.""" + """Drop (not close) the inherited storage client, and stop sharing the + parent's upload counters: from here on ranges are drawn at random.""" if self._owns_upload_client: self._upload_client = None else: @@ -141,8 +107,6 @@ def start(self, run_id: str, context: Mapping[str, str]) -> None: from ..projection import episodes_to_parquet_bytes, parquet_available if not parquet_available(): - # Not a failure: nothing was lost, the episodes still reach - # Prime Traces. Say once how to turn the table on. logger.warning( "Training samples sink off: pyarrow is not installed " "(pip install 'prime-runs[train]'); episodes reach Prime Traces only" @@ -151,36 +115,17 @@ def start(self, run_id: str, context: Mapping[str, str]) -> None: return self._encoder = episodes_to_parquet_bytes - # ------------------------------------------------------------------ write - def write(self, records: Sequence[Any]) -> None: if not self.enabled or not records: return if self._run_id is None: raise RuntimeError("RftSamplesSink.write called before start()") - assert self._encoder is not None # start() set it or disabled the sink - - steps = self._group(records) - failed = 0 - reported: Optional[Exception] = None - for position, (step, episodes) in enumerate(steps): - try: - uploaded = self._upload_step(step, episodes) - except Exception as exc: # noqa: BLE001 - classified below - failed += len(episodes) - if not (is_record_rejection(exc) or is_transient(exc)): - # Sink-wide: the remaining steps would fail the same way. - failed += sum(len(rest) for _, rest in steps[position + 1 :]) - reported = exc - break - if reported is None or is_transient(exc): - reported = exc - else: - if uploaded: - self.steps_written += 1 - - if reported is not None: - raise SinkWriteError(reported, failed_records=failed) from reported + send_each( + [ + (len(episodes), functools.partial(self._upload_step, step, episodes)) + for step, episodes in self._group(records) + ] + ) def _group(self, records: Sequence[Any]) -> List[Tuple[int, List[Any]]]: """Training episodes by step, on the upload cadence, in step order.""" @@ -188,8 +133,7 @@ def _group(self, records: Sequence[Any]) -> List[Tuple[int, List[Any]]]: skipped = 0 for record in records: step = training_step(record) - # The encoder projects attributes (``build_samples``), so a JSON - # episode has no row here either. + # The encoder projects attributes, so a JSON episode has no row either. if step is None or isinstance(record, Mapping) or not is_episode(record): skipped += 1 continue @@ -207,17 +151,13 @@ def _group(self, records: Sequence[Any]) -> List[Tuple[int, List[Any]]]: self.skipped += skipped return sorted(by_step.items()) - def _upload_step(self, step: int, episodes: List[Any]) -> bool: - """Encode and upload one step. ``False`` when the step had nothing to - show (no trajectories) and no request was made.""" + def _upload_step(self, step: int, episodes: List[Any]) -> None: assert self._encoder is not None and self._run_id is not None - # Every upload of a step numbers its rows in its own range, counted by - # attempt: an upload whose confirm was lost may still have landed. payload = self._encoder( episodes, self._run_id, step, sample_id_offset=self._next_range(step) * SAMPLE_ID_STRIDE ) - if payload is None: - return False + if payload is None: # nothing with a trajectory + return # Replayable: a presign only mints a URL. presign = self._client.post( @@ -232,21 +172,18 @@ def _upload_step(self, step: int, episodes: List[Any]) -> bool: raise APIError( f"POST /rft/samples/presign returned no upload URL/key (keys: {sorted(data)})" ) - self._put_object(url, payload) - - # Records the object under the step. Replayable: confirm checks the - # key and the object, then refreshes the run's progress; a second - # confirm of the same key changes nothing. + # Replayable: confirm validates the key and refreshes the run's progress. self._client.post( "/rft/samples/confirm", json_body={"run_id": self._run_id, "step": step, "s3_key": key}, idempotent=True, ) - return True + self.steps_written += 1 def _next_range(self, step: int) -> int: - """The ``sample_id`` range for the next upload of ``step``.""" + """The ``sample_id`` range for the next upload of ``step``. Counted per + attempt: an upload whose confirm was lost may still have landed.""" if self._forked: half = SAMPLE_ID_RANGES // 2 return half + secrets.randbelow(half) @@ -255,9 +192,8 @@ def _next_range(self, step: int) -> int: return upload_index def _put_object(self, url: str, payload: bytes) -> None: - """PUT the object to storage. Bare client: the presigned URL carries its - own credentials and rejects the platform's auth headers. Idempotent, so - transport failures and 5xx are retried.""" + """PUT to storage with a bare client: the presigned URL carries its own + credentials and rejects the platform's auth headers.""" client = self._upload() for attempt in range(UPLOAD_ATTEMPTS): error: APIError @@ -290,7 +226,7 @@ def _upload(self) -> httpx.Client: return self._upload_client def flush(self) -> None: - """Uploads are synchronous; nothing is held back here.""" + """Uploads are synchronous.""" def close(self) -> None: client = self._upload_client diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py index 3a912d0e3..6e33c96d3 100644 --- a/packages/prime-runs/src/prime_runs/worker.py +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -1,16 +1,11 @@ """Background uploader: one daemon thread draining a bounded queue into sinks. -Backpressure: the queue is bounded, so a producer that outruns the uploader -blocks briefly and then drops (counted) rather than stalling the run. - -Fork safety: a forked child inherits the queue's memory but not the thread. -The child starts over empty; the queued records belong to the parent. - -Containment: record-specific failures drop only their batch. A sink-wide -failure disables the sink, while transient failures get a few consecutive -strikes first — and, with a ``retire_cooldown``, the sink is tried again once -the cooldown has passed, so an outage costs a long run a window of records -rather than the rest of the run. The thread never dies on one bad batch. +Backpressure: a producer that outruns the uploader blocks briefly, then drops +(counted). Fork safety: a child starts over with an empty queue; the queued +records belong to the parent. Containment: record-specific failures drop only +their batch, transient failures get a few consecutive strikes before the sink +is retired (for good, or until ``retire_cooldown`` has passed), and any other +failure retires the sink at once. The thread never dies on one bad batch. """ import logging @@ -39,20 +34,16 @@ class _Flush: event: threading.Event = field(default_factory=threading.Event) -#: "Nothing pulled ahead" — distinct from ``None``, which is the stop sentinel. +#: "Nothing pulled ahead"; ``None`` is the stop sentinel. _NOTHING: Any = object() -def _deadline(timeout: Optional[float]) -> Optional[float]: - if timeout is None: - return None - return time.monotonic() + max(0.0, timeout) +def deadline_after(timeout: Optional[float]) -> Optional[float]: + return None if timeout is None else time.monotonic() + max(0.0, timeout) -def _remaining(deadline: Optional[float]) -> Optional[float]: - if deadline is None: - return None - return max(0.0, deadline - time.monotonic()) +def time_left(deadline: Optional[float]) -> Optional[float]: + return None if deadline is None else max(0.0, deadline - time.monotonic()) class UploadWorker: @@ -80,27 +71,23 @@ def __init__( self._lock = threading.Lock() #: Records never handed to any sink because the queue was full. self.dropped = 0 - #: Records a particular sink could not store, by sink name. Kept apart - #: from ``dropped``: another sink may well have stored them. + #: Records a sink could not store, by sink name. Kept apart from + #: ``dropped``: another sink may well have stored them. self.failed_records: dict = {} self._transient_failures: dict = {} - #: Sinks this worker retired after a failure. Records that skip one of - #: these are lost to it and counted; a sink that switched itself off - #: without raising (nowhere for the records to go) is not in here. + #: Sinks retired by this worker. Records that skip one are counted as + #: lost to it; a sink that switched itself off quietly is not in here. self._retired: set = set() #: Sink name -> monotonic time at which a paused sink is tried again. self._cooldowns: dict = {} _fork.register(self) - # ----------------------------------------------------------------- thread - def start(self) -> None: with self._lock: self._start_locked() def _start_locked(self) -> None: - """Start while ``_lock`` is held; a closed worker is one-shot.""" - if self._stopping.is_set(): + if self._stopping.is_set(): # a closed worker is one-shot return if self._thread is not None and self._thread.is_alive(): return @@ -109,9 +96,7 @@ def _start_locked(self) -> None: def _run(self) -> None: # An item pulled off the queue while coalescing that could not join the - # batch (a control item — the stop sentinel is None, hence the marker — - # or records of the other kind); it is the next thing to process, ahead - # of anything queued after it. + # batch; it is processed next, ahead of anything queued after it. pending: Any = _NOTHING while True: item = pending if pending is not _NOTHING else self._queue.get() @@ -131,21 +116,11 @@ def _run(self) -> None: self._queue.task_done() def _coalesce(self, first: Sequence[Any]) -> "tuple[list, Any]": - """Merge everything already queued behind ``first`` into one batch. - - A producer hands over one episode at a time as rollouts finish, and - each batch costs every sink a request: one ``POST /samples`` per - episode runs into the platform's per-minute limit on any fast eval, - and the retries then back the queue up until records are dropped. - Draining what has accumulated while the previous request was in - flight makes the request rate track upload latency instead of rollout - throughput. Sinks split a large batch by size themselves. - - Stops at a control item (a flush barrier must not be reordered past - the records queued before it) and at a batch of the other record kind - (the traces sink infers the line format from the first record). - Returns the merged batch and the item that stopped it, if any. - """ + """Merge everything already queued behind ``first`` into one batch, so + the request rate tracks upload latency rather than rollout throughput. + Stops at a control item (a flush barrier must not be reordered past the + records before it) and at a batch of the other record kind (the traces + sink infers the line format from the first record).""" batch = list(first) kind = is_episode(batch[0]) if batch else None while True: @@ -178,7 +153,6 @@ def _dispatch(self, records: Sequence[Any]) -> None: self._transient_failures.pop(sink.name, None) def _maybe_revive(self, sink: Sink) -> None: - """Re-enable a sink whose cooldown has passed. Its strikes start over.""" until = self._cooldowns.get(sink.name) if until is None or time.monotonic() < until: return @@ -198,13 +172,6 @@ def _flush_sinks(self) -> None: self._fail_sink(sink, exc) def _fail_sink(self, sink: Sink, exc: Exception, *, dropped: int = 0) -> None: - """Account for a failed operation and decide whether to retire its sink. - - Record-specific failures drop only the current batch. Sink-wide - permanent failures retire immediately; transient failures retire after - ``TRANSIENT_FAILURE_LIMIT`` consecutive strikes — for good, or until - the worker's ``retire_cooldown`` has passed. - """ name = sink.name if dropped: self.failed_records[name] = self.failed_records.get(name, 0) + dropped @@ -271,18 +238,12 @@ def _notify(self, name: str, exc: Exception) -> None: except Exception: # noqa: BLE001 - the handler is the caller's problem logger.debug("Error handler raised while reporting a sink failure", exc_info=True) - # ------------------------------------------------------------------ queue - def submit(self, records: Sequence[Any]) -> bool: - """Hand a batch to the uploader. ``False`` means it was dropped: the - queue stayed full for the whole put timeout.""" - # No sinks (a disabled run): there is nowhere for the records to go, so - # do not copy them into a queue or start a thread to find that out. - # Retired sinks are not the same case — their losses are counted. - if not self.sinks: + """Queue a batch. ``False`` means it was dropped: the queue stayed full + for the whole put timeout.""" + if not self.sinks: # a disabled run: nowhere for the records to go return True - # Serialize acceptance with close(): an accepted batch is queued before - # the stop sentinel, while a submission after close is refused. + # Serialize with close(): an accepted batch is queued before the stop sentinel. with self._lock: if self._stopping.is_set(): return False @@ -307,16 +268,15 @@ def flush(self, timeout: Optional[float] = None) -> bool: if self._thread is None or not self._thread.is_alive(): self._flush_sinks() return True - deadline = _deadline(timeout) + deadline = deadline_after(timeout) barrier = _Flush() try: - # The barrier gets the caller's drain budget, not the short - # producer put timeout: a full queue is when finish() most needs it. - self._queue.put(barrier, timeout=_remaining(deadline)) + # The barrier gets the caller's drain budget, not the producer put timeout. + self._queue.put(barrier, timeout=time_left(deadline)) except queue.Full: logger.warning("Could not enqueue a flush barrier before the drain timeout") return False - return barrier.event.wait(_remaining(deadline)) + return barrier.event.wait(time_left(deadline)) def close(self, timeout: Optional[float] = 30.0) -> None: """Drain, stop the thread, and close every sink.""" @@ -324,15 +284,14 @@ def close(self, timeout: Optional[float] = 30.0) -> None: self._stopping.set() thread = self._thread if thread is not None and thread.is_alive(): - deadline = _deadline(timeout) + deadline = deadline_after(timeout) try: - self._queue.put(None, timeout=_remaining(deadline)) + self._queue.put(None, timeout=time_left(deadline)) except queue.Full: logger.warning("Upload queue remained saturated through the close timeout") - thread.join(_remaining(deadline)) + thread.join(time_left(deadline)) if thread.is_alive(): - # Closing the sinks now would pull a client or file handle out - # from under a request still running on that thread. + # Closing the sinks now would pull a client out from under a request. logger.warning( "Uploader still running after %ss; leaving it and its sinks open. " "Records still in flight may not finish before the process exits.", @@ -346,10 +305,7 @@ def close(self, timeout: Optional[float] = 30.0) -> None: except Exception as exc: # noqa: BLE001 - teardown must not raise logger.debug("Error closing sink %s: %s", sink.name, exc) - # ------------------------------------------------------------------- fork - def reset_after_fork(self) -> None: - """Give the child a clean uploader; what was queued belongs to the parent.""" self._queue = queue.Queue(maxsize=self.max_queue_size) self._thread = None self._stopping = threading.Event() diff --git a/packages/prime-runs/tests/_fakes.py b/packages/prime-runs/tests/_fakes.py index 00522ae49..1a419798a 100644 --- a/packages/prime-runs/tests/_fakes.py +++ b/packages/prime-runs/tests/_fakes.py @@ -1,11 +1,4 @@ -"""Stand-ins for verifiers ``Trace``/``Episode``. - -The SDK never imports a producer package — records are duck-typed — so the -tests must not either, or they would be testing an import rather than the -protocol. These objects implement exactly the surface -:mod:`prime_runs.projection` and :mod:`prime_runs.metrics` touch, which makes -that surface explicit and makes an accidental widening of it fail here. -""" +"""Stand-ins for verifiers ``Trace``/``Episode``.""" from dataclasses import dataclass, field from typing import Any, Dict, List, Optional diff --git a/packages/prime-runs/tests/test_config_source.py b/packages/prime-runs/tests/test_config_source.py index a0120afe1..9daf1f2d8 100644 --- a/packages/prime-runs/tests/test_config_source.py +++ b/packages/prime-runs/tests/test_config_source.py @@ -1,5 +1,4 @@ -"""The config a run is actually configured with: the file someone wrote, kept -byte for byte, rather than a projection that cannot show it.""" +"""The config a run is actually configured with: the file someone wrote, kept""" import json @@ -12,7 +11,6 @@ CONFIG_SOURCE_KEY, MAX_CONFIG_SOURCE_BYTES, ConfigSource, - RunSpec, ) from prime_runs.run import _normalize_config @@ -30,7 +28,6 @@ def test_a_path_is_read_verbatim(tmp_path): - """Comments, ordering and section grouping are the point — a dict loses all three.""" path = tmp_path / "eval.toml" path.write_text(EVAL_TOML) @@ -56,8 +53,6 @@ def test_the_format_is_inferred_from_the_suffix(tmp_path): def test_inline_text_has_to_be_explicit(tmp_path): - """A bare string is a path. Guessing would turn a typo'd filename into a run - whose config tab proudly displays the filename.""" with pytest.raises(ConfigurationError, match="does not exist"): ConfigSource.coerce("environment = 'gsm8k'") @@ -66,7 +61,6 @@ def test_inline_text_has_to_be_explicit(tmp_path): def test_an_oversized_file_is_refused_at_init_not_truncated(tmp_path): - """Silently storing half a config is worse than not storing one.""" path = tmp_path / "eval.toml" path.write_text("x = 1\n" * MAX_CONFIG_SOURCE_BYTES) @@ -95,34 +89,16 @@ def test_a_mapping_round_trips(): def test_a_mapping_without_text_is_not_a_config_source(): assert ConfigSource.from_mapping({"format": "toml"}) is None - with pytest.raises(ValueError, match="must contain a 'text' string"): - ConfigSource.coerce({"format": "toml"}) # ------------------------------------------------------- config normalization def test_a_mapping_is_taken_exactly_as_given(): - """The caller already decided what to say; second-guessing it would be worse.""" assert _normalize_config({"a": 1, "b": None}) == {"a": 1, "b": None} -def test_a_mapping_is_copied_not_aliased(): - original = {"a": 1} - assert _normalize_config(original) is not original - - -def test_no_config_is_an_empty_config(): - assert _normalize_config(None) == {} - - -def test_an_unusable_config_says_what_was_expected(): - with pytest.raises(TypeError, match="config must be a path"): - _normalize_config(object()) - - def test_a_path_becomes_a_stored_source(tmp_path): - """One parameter, three forms — the same shape ``environments=`` already has.""" path = tmp_path / "eval.toml" path.write_text(EVAL_TOML) @@ -134,8 +110,6 @@ def test_a_path_becomes_a_stored_source(tmp_path): def test_a_mapping_that_looks_like_a_source_is_still_just_a_mapping(): - """Form is decided by type, never by inspecting keys — so a config that - happens to have a ``text`` field is not mistaken for a launch file.""" assert _normalize_config({"text": "hello", "format": "toml"}) == { "text": "hello", "format": "toml", @@ -172,8 +146,6 @@ def test_an_online_run_sends_the_source_in_create_metadata(tmp_path, online): def test_extra_values_can_be_merged_onto_a_launch_file(tmp_path, online): - """A run launched from a file that also wants structured values passes a - mapping carrying the source under ``CONFIG_SOURCE_KEY`` — what verifiers does.""" path = tmp_path / "eval.toml" path.write_text(EVAL_TOML) config = { @@ -210,8 +182,6 @@ def test_a_run_without_a_source_reports_none(): def test_the_source_survives_the_failure_fallback(tmp_path, online): - """The fallback rewrites metadata to record a terminal state. It merges into - the whole config, so the source must still be there afterwards.""" path = tmp_path / "eval.toml" path.write_text(EVAL_TOML) @@ -221,7 +191,3 @@ def test_the_source_survives_the_failure_fallback(tmp_path, online): update = handler.bodies_for("/api/v1/evaluations/eval-abc")[-1] assert update["metadata"][CONFIG_SOURCE_KEY]["text"] == EVAL_TOML assert update["metadata"]["prime_runs"]["status"] == "failed" - - -def test_a_spec_defaults_to_no_source(): - assert RunSpec().config == {} diff --git a/packages/prime-runs/tests/test_evals_backend.py b/packages/prime-runs/tests/test_evals_backend.py index 4f9c53ff7..b8ccea974 100644 --- a/packages/prime-runs/tests/test_evals_backend.py +++ b/packages/prime-runs/tests/test_evals_backend.py @@ -85,8 +85,6 @@ def test_a_published_environment_slug_supplies_dataset_and_default_name( def test_an_unresolvable_environment_fails_the_run_rather_than_being_skipped( make_platform_client, eval_routes ): - """Silently dropping it produces a run attached to the wrong environments — - an upload that looks successful and is discovered wrong much later.""" routes = dict(eval_routes) routes["POST /api/v1/environmentshub/resolve"] = lambda request: httpx.Response( 404, json={"detail": "no such environment"} @@ -108,8 +106,6 @@ def test_a_run_with_no_environments_is_rejected_before_any_request( def test_a_run_without_a_name_gets_one(make_platform_client, eval_routes): - """The API requires a name; the alternative to generating one is a 422 at - the worst possible moment.""" backend, handler = make_backend(make_platform_client, eval_routes) backend.create(RunSpec(environments=[EnvironmentRef(name="gsm8k")])) @@ -137,7 +133,6 @@ def test_finalizing_a_completed_run_posts_its_metrics(make_platform_client, eval def test_an_ambiguous_finalize_failure_is_not_replayed(make_platform_client, eval_routes): - """Finalization enqueues asynchronous processing, so a retry can enqueue it twice.""" routes = dict(eval_routes) routes["POST /api/v1/evaluations/eval-abc/finalize"] = lambda request: httpx.Response(502) backend, handler = make_backend(make_platform_client, routes) @@ -149,9 +144,6 @@ def test_an_ambiguous_finalize_failure_is_not_replayed(make_platform_client, eva def test_a_failed_run_is_recorded_in_metadata(make_platform_client, eval_routes, caplog): - """The platform has no producer-facing way to fail an evaluation. The run - cannot leave RUNNING, but the failure is recorded where an operator and the - dashboard can read it, and the SDK says the run will keep showing as running.""" backend, handler = make_backend(make_platform_client, eval_routes) with caplog.at_level("WARNING"): @@ -164,18 +156,7 @@ def test_a_failed_run_is_recorded_in_metadata(make_platform_client, eval_routes, assert "keep showing as running" in caplog.text -def test_update_sends_nothing_when_there_is_nothing_to_send(make_platform_client, eval_routes): - backend, handler = make_backend(make_platform_client, eval_routes) - - backend.update("eval-abc") - - assert handler.requests == [] - - def test_a_pinned_environment_version_reaches_the_api(make_platform_client, eval_routes): - """The API's EnvironmentReference carries version_id. Dropping it attaches - the run to whatever version the hub resolves today — the difference between - a reproducible eval and one that quietly moved.""" backend, handler = make_backend(make_platform_client, eval_routes) backend.create(RunSpec(name="r", environments=[EnvironmentRef(id="env-1", version_id="v-7")])) @@ -196,8 +177,6 @@ def test_a_version_pin_survives_hub_resolution(make_platform_client, eval_routes def test_the_failure_fallback_preserves_the_run_config(make_platform_client, eval_routes): - """The service writes metadata with a document-level $set, so a PUT carrying - only the terminal block would erase everything finish() just wrote.""" backend, handler = make_backend(make_platform_client, eval_routes) backend.finalize( diff --git a/packages/prime-runs/tests/test_http.py b/packages/prime-runs/tests/test_http.py index 6d954f84d..7105a5978 100644 --- a/packages/prime-runs/tests/test_http.py +++ b/packages/prime-runs/tests/test_http.py @@ -3,7 +3,7 @@ import httpx import pytest -from prime_runs._http import PlatformClient, encode_json +from prime_runs._http import PlatformClient from prime_runs.exceptions import ( APIError, ForbiddenError, @@ -37,8 +37,6 @@ def client_for(handler, *, base_url: str = "http://testserver", **kwargs) -> Pla ], ) def test_status_codes_map_to_types_callers_can_branch_on(status, expected): - """The classes are ``prime_traces``' own, so a producer that already - handles the traces client's errors handles these with the same clauses.""" client = client_for(lambda request: httpx.Response(status, json={"detail": "nope"})) with pytest.raises(expected) as caught: @@ -50,8 +48,6 @@ def test_status_codes_map_to_types_callers_can_branch_on(status, expected): def test_a_forbidden_response_is_permanent_so_a_sink_retires_on_it(): - """403 is the gated-account signal. Retrying it for the rest of a run would - log one failure per batch and never succeed.""" client = client_for(lambda request: httpx.Response(403, json={"code": "service_not_enabled"})) with pytest.raises(ForbiddenError) as caught: @@ -60,13 +56,6 @@ def test_a_forbidden_response_is_permanent_so_a_sink_retires_on_it(): assert not is_transient(caught.value) -def test_an_unauthorized_error_says_what_to_do_about_it(): - client = client_for(lambda request: httpx.Response(401, json={"detail": "bad token"})) - - with pytest.raises(UnauthorizedError, match="PRIME_API_KEY"): - client.get("/evaluations/x") - - def test_retryable_statuses_are_retried_then_surface(no_sleep): attempts = [] @@ -84,17 +73,7 @@ def handler(request): assert all(0.0 < delay <= 30.0 for delay in no_sleep) -def test_a_retry_succeeds_without_bothering_the_caller(no_sleep): - responses = [httpx.Response(429), httpx.Response(200, json={"ok": True})] - - client = client_for(lambda request: responses.pop(0)) - - assert client.get("/evaluations/x") == {"ok": True} - - def test_retry_after_is_honoured(no_sleep): - """The schedule itself is ``prime_traces.core.client.retry_delay``; what is - ours is feeding it the server's header.""" responses = [ httpx.Response(429, headers={"Retry-After": "7.5"}), httpx.Response(200, json={"ok": True}), @@ -151,15 +130,7 @@ def test_a_non_json_body_names_the_request_that_produced_it(): client.get("/evaluations/x") -def test_encoding_refuses_values_json_cannot_carry(): - """Bare ``NaN`` is JavaScript, not JSON; it comes back as an opaque 400.""" - with pytest.raises(ValueError): - encode_json({"reward": float("nan")}) - - def test_an_ambiguous_failure_does_not_replay_a_create(no_sleep): - """If the platform created the run and the response was lost, a retry makes - a second one and only the second is tracked — an orphaned duplicate run.""" attempts = [] def handler(request): @@ -174,8 +145,6 @@ def handler(request): def test_a_refusal_is_replayed_even_for_a_create(no_sleep): - """429 is refused before the server does any work, so there is nothing on - the other side to duplicate.""" responses = [httpx.Response(429), httpx.Response(201, json={"evaluation_id": "e1"})] client = client_for(lambda request: responses.pop(0)) @@ -184,7 +153,6 @@ def test_a_refusal_is_replayed_even_for_a_create(no_sleep): def test_a_connection_that_was_never_made_is_replayed_for_a_create(no_sleep): - """Nothing reached the server, so replaying cannot duplicate anything.""" calls = [] def handler(request): @@ -197,7 +165,6 @@ def handler(request): def test_a_read_timeout_does_not_replay_a_create(no_sleep): - """The bytes went out; the platform may have processed them.""" calls = [] def handler(request): @@ -211,7 +178,6 @@ def handler(request): def test_a_post_declared_idempotent_still_retries(no_sleep): - """Get-or-create and terminal-state writes are safe to replay.""" responses = [httpx.Response(502), httpx.Response(200, json={"data": {"id": "env-1"}})] client = client_for(lambda request: responses.pop(0)) @@ -229,8 +195,6 @@ def test_a_post_declared_idempotent_still_retries(no_sleep): ], ) def test_a_base_url_written_with_the_api_prefix_is_not_doubled(given): - """``Config`` strips the suffix, so an explicit ``base_url=`` that does not - would 404 on exactly the value that works through the environment.""" assert client_for(lambda request: httpx.Response(200, json={}), base_url=given).api_prefix == ( "http://testserver/api/v1" ) diff --git a/packages/prime-runs/tests/test_init.py b/packages/prime-runs/tests/test_init.py index ca0debdc0..de301ce7d 100644 --- a/packages/prime-runs/tests/test_init.py +++ b/packages/prime-runs/tests/test_init.py @@ -14,8 +14,6 @@ def test_no_api_key_disables_the_run_and_says_so(caplog): - """Loudly, not silently: a user who forgot ``prime login`` must not believe - the run was tracked.""" with caplog.at_level("WARNING"): run = pr.init(name="local") @@ -39,8 +37,6 @@ def test_an_unknown_mode_is_rejected(): def test_a_disabled_run_still_answers_every_call(tmp_path): - """Same object shape, so producer code needs no branching — and nothing - touches the network or the filesystem.""" run = pr.init(mode="disabled") run.log_traces([{"id": "t1"}]) @@ -147,7 +143,6 @@ def close(self): def test_an_online_run_has_both_transports_by_default( monkeypatch, make_platform_client, eval_routes ): - """Traces is the system of record; the sample table is what the viewer reads.""" handler = RecordingHandler(eval_routes) monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler)) monkeypatch.setattr("prime_traces.TracesClient", lambda **_: object()) @@ -172,17 +167,7 @@ def test_episodes_stream_to_the_sample_table_while_the_run_is_going(online): run.finish() -def test_finishing_an_online_run_finalizes_it_with_its_metrics(online): - run, handler = online() - - run.finish(summary={"avg_reward": 0.9}) - - body = handler.bodies_for("/api/v1/evaluations/eval-abc/finalize")[0] - assert body["metrics"]["avg_reward"] == 0.9 - - def test_the_end_to_end_shape_a_producer_writes(online): - """The whole surface, in the order verifiers calls it.""" from prime_runs import metrics episodes = [make_episode(f"ep-{n}", [make_trace(idx=n, reward=float(n))]) for n in range(3)] @@ -207,25 +192,12 @@ def test_the_end_to_end_shape_a_producer_writes(online): assert run.errors == [] -def test_a_second_init_in_one_process_opens_its_own_run(): - first = pr.init(mode="disabled") - first.finish() - - second = pr.init(mode="disabled") - second.finish() - - assert second.id != first.id - - # --------------------------------------------------------------------- fork @pytest.mark.skipif(not hasattr(os, "fork"), reason="fork is POSIX-only") @pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning") def test_a_forked_child_does_not_duplicate_the_parents_records_or_close_its_run(online): - """At fork time the parent may have records in the upload queue. The child - inherits a copy; writing them would upload every record twice, and the - inherited atexit hook must not finalize the parent's run.""" run, handler = online() run.log_traces([{"sample_id": f"parent-{n}"} for n in range(5)]) @@ -368,8 +340,6 @@ def test_a_failed_training_run_is_marked_failed_with_the_reason(online_train): def test_training_uploaders_pause_rather_than_retire_after_an_outage(online_train): - """A training run lasts days; a sink that struck out on a platform blip - is tried again after a cooldown instead of losing the rest of the run.""" from prime_runs.run import TRAIN_RETIRE_COOLDOWN run, _, _, _ = online_train() @@ -392,8 +362,6 @@ def test_attaching_to_a_managed_run_registers_nothing(online_train): def test_an_attached_run_leaves_failure_marking_to_the_launcher(online_train, caplog): - """A managed launch marks its own run failed when the process goes away; - reporting it from here would race that (prime-rl's contract).""" run, handler, _, _ = online_train(id="run-managed") with caplog.at_level("INFO"): @@ -456,8 +424,6 @@ def test_kind_arguments_are_checked_before_anything_else(kwargs): def test_training_records_are_stamped_with_the_train_run_type(): - """Bare dicts get ``run.type`` from the run's kind, matching what verifiers' - ``TrainRunInfo`` carries, so the traces service files them as training.""" from prime_runs.run import _sink_context from prime_runs.sinks.base import stamp_run diff --git a/packages/prime-runs/tests/test_metrics.py b/packages/prime-runs/tests/test_metrics.py index 3673cc5d9..ca2888248 100644 --- a/packages/prime-runs/tests/test_metrics.py +++ b/packages/prime-runs/tests/test_metrics.py @@ -1,9 +1,4 @@ -"""Run-level aggregates, moved here from verifiers. - -The aggregation rules encode judgements about what a run's headline number -means. They are asserted here because a silent change to any of them changes -every number on the dashboard without changing a line of producer code. -""" +"""Run-level aggregates, moved here from verifiers.""" from _fakes import Reward, make_episode, make_trace @@ -11,11 +6,6 @@ def test_only_trainable_traces_are_scored(): - """A judge scoring 0 must not halve a solver's average. - - Fixed agents often carry no rewards at all, so including them dilutes every - mean with structural zeros. - """ episode = make_episode( "ep-1", [ @@ -28,7 +18,6 @@ def test_only_trainable_traces_are_scored(): def test_all_traces_count_when_none_are_trainable(): - """The dashboard's fallback: an eval of a frozen model still has a score.""" episode = make_episode( "ep-1", [ @@ -53,8 +42,6 @@ def test_sub_rewards_and_env_metrics_average_together(): def test_a_metric_present_on_only_some_traces_averages_over_those_traces(): - """Counts are per-key, not per-run: a metric only some environments emit - must not be divided by traces that never reported it.""" episodes = [ make_episode("ep-1", [make_trace(metrics={"partial": 1.0})]), make_episode("ep-2", [make_trace(metrics={})]), @@ -64,7 +51,6 @@ def test_a_metric_present_on_only_some_traces_averages_over_those_traces(): def test_avg_error_counts_episodes_not_traces(): - """A hook failure counts even when its traces are clean or it left none.""" episodes = [ make_episode("ok", [make_trace(), make_trace(trace_id="t2")], ok=True), make_episode("broken", [], ok=False), diff --git a/packages/prime-runs/tests/test_metrics_sink.py b/packages/prime-runs/tests/test_metrics_sink.py index 624d7dbe0..905714b60 100644 --- a/packages/prime-runs/tests/test_metrics_sink.py +++ b/packages/prime-runs/tests/test_metrics_sink.py @@ -29,8 +29,6 @@ def test_each_metrics_dict_is_one_request(make_platform_client, rft_routes): def test_an_ambiguous_failure_is_replayed(make_platform_client, rft_routes, no_sleep): - """The platform keeps one row per step, so a retry after a lost response - cannot double anything, while a lost row is a hole in the curves for good.""" calls = [] def flaky(request: httpx.Request) -> httpx.Response: @@ -93,10 +91,3 @@ def test_a_record_that_is_not_a_mapping_is_counted_not_sent(make_platform_client assert info.value.failed_records == 1 assert isinstance(info.value.cause, TypeError) assert len(handler.bodies_for("/api/v1/rft/metrics")) == 1 - - -def test_write_before_start_is_a_producer_bug(make_platform_client, rft_routes): - sink = RftMetricsSink(make_platform_client(RecordingHandler(rft_routes))) - - with pytest.raises(RuntimeError, match="before start"): - sink.write([{"step": 1}]) diff --git a/packages/prime-runs/tests/test_projection.py b/packages/prime-runs/tests/test_projection.py index 34646fc1d..06073f7eb 100644 --- a/packages/prime-runs/tests/test_projection.py +++ b/packages/prime-runs/tests/test_projection.py @@ -1,9 +1,4 @@ -"""The v0 eval-sample projection, moved here from verifiers. - -These assertions are the contract the current viewer reads. They exist so the -move is provably a relocation and not a rewrite: a run pushed through this SDK -must produce the same rows verifiers produced. -""" +"""The v0 eval-sample projection, moved here from verifiers.""" import pytest from _fakes import Reward, make_episode, make_trace @@ -46,17 +41,25 @@ def test_trace_to_sample_flattens_sub_rewards_to_top_level(): def test_trace_to_sample_does_not_let_a_sub_reward_clobber_a_real_field(): - """``setdefault`` semantics: a sub-reward named ``reward`` must not win. - - Sub-rewards are flattened into the same namespace as the row's own columns, - so an environment that names a reward function after one of them would - otherwise silently overwrite the value the dashboard reads. - """ trace = make_trace(reward=0.25, rewards={"reward": Reward(9.0)}) assert trace_to_sample(trace)["reward"] == 0.25 +def test_non_finite_numbers_become_null_rather_than_losing_the_row(): + trace = make_trace(reward=float("nan"), metrics={"kl": float("inf")}) + trace.rewards = {"format": Reward(float("-inf"))} + + sample = trace_to_sample(trace) + (row,) = build_samples([make_episode("ep-1", [trace])]) + + assert sample["reward"] is None + assert sample["metrics"] == {"kl": None} + assert sample["format"] is None + assert row["info"]["native_wrapper"]["traces"][0]["reward"] is None + assert batch_samples([row]) # encodes: nothing non-finite is left + + def test_trajectory_keeps_one_entry_per_branch(): trace = make_trace(branches=3) @@ -79,7 +82,6 @@ def test_build_samples_emits_one_row_per_episode_with_the_native_wrapper(): def test_summary_trace_is_the_first_trainable_one(): - """A judge or modeled user must not become the row the dashboard shows.""" episode = make_episode( "ep-1", [ @@ -95,12 +97,6 @@ def test_summary_trace_is_the_first_trainable_one(): def test_rollout_numbers_continue_across_streaming_calls(): - """Streaming uploads must number rollouts the way one final upload did. - - The old code built every sample in a single call, so its per-example counter - lived in a local. Streaming means several calls, and without a carried - counter every batch would restart at rollout 1. - """ counters: dict = {} first = build_samples([make_episode("ep-1", [make_trace(idx=4)])], counters) second = build_samples([make_episode("ep-2", [make_trace(idx=4)])], counters) @@ -124,17 +120,12 @@ def test_batch_samples_splits_on_the_payload_ceiling(): def test_batch_samples_refuses_a_sample_that_cannot_be_sent(): - """Dropping it would report a complete run that is missing rows.""" oversized = {"sample_id": "x", "blob": "a" * (MAX_SAMPLES_PAYLOAD_BYTES + 1)} with pytest.raises(ValueError, match="too large"): batch_samples([oversized]) -def test_batch_samples_returns_nothing_for_no_samples(): - assert batch_samples([]) == [] - - def test_is_episode_distinguishes_episodes_from_traces_in_either_form(): from prime_runs.sinks import is_episode diff --git a/packages/prime-runs/tests/test_rft_backend.py b/packages/prime-runs/tests/test_rft_backend.py index 22c2e9ef8..8238e7171 100644 --- a/packages/prime-runs/tests/test_rft_backend.py +++ b/packages/prime-runs/tests/test_rft_backend.py @@ -61,8 +61,6 @@ def test_create_registers_an_external_run_for_the_team(make_platform_client, rft def test_training_environments_are_passed_through_not_resolved( make_platform_client, rft_routes, environment, expected ): - """The RFT API names environments by hub id and resolves them itself; a - hub round-trip here would be a second opinion.""" backend, handler = make_backend(make_platform_client, rft_routes) backend.create(train_spec(environments=[EnvironmentRef.coerce(environment)])) @@ -88,8 +86,6 @@ def test_create_needs_a_base_model(make_platform_client, rft_routes): def test_create_needs_a_team(make_platform_client, rft_routes): - """External runs are created for a team; the API refuses without one, so - say so before the request rather than surfacing its 403.""" backend, handler = make_backend(make_platform_client, rft_routes, team_id=None) with pytest.raises(ConfigurationError, match="team"): @@ -98,8 +94,6 @@ def test_create_needs_a_team(make_platform_client, rft_routes): def test_a_team_outside_the_allowlist_is_a_forbidden_error(make_platform_client, rft_routes): - """The platform enables external runs per team. The producer decides what - a refusal means; the SDK passes the platform's reason through.""" routes = dict(rft_routes) routes["POST /api/v1/rft/external-runs"] = lambda request: httpx.Response( 403, json={"detail": "External training runs are not enabled for this team"} @@ -112,7 +106,6 @@ def test_a_team_outside_the_allowlist_is_a_forbidden_error(make_platform_client, def test_create_is_not_replayed_after_an_ambiguous_failure(make_platform_client, rft_routes): - """A retry after a lost response would register a second run.""" routes = dict(rft_routes) routes["POST /api/v1/rft/external-runs"] = lambda request: httpx.Response(502) backend, handler = make_backend(make_platform_client, routes) @@ -133,7 +126,6 @@ def test_attach_makes_no_request(make_platform_client, rft_routes): def test_update_is_a_no_op(make_platform_client, rft_routes): - """The RFT API takes the config at creation and has no summary document.""" backend, handler = make_backend(make_platform_client, rft_routes) backend.update("run-abc", config={"a": 1}, summary={"loss": 0.1}) @@ -151,8 +143,6 @@ def test_a_completed_run_is_finalized_with_exit_code_zero(make_platform_client, def test_finalize_is_replayed_after_an_ambiguous_failure(make_platform_client, rft_routes): - """Compare-and-set on the server, and "already finalized" is a success: - a lost response costs nothing to replay, unlike the eval finalize.""" attempts = [] def flaky(request: httpx.Request) -> httpx.Response: @@ -171,7 +161,6 @@ def flaky(request: httpx.Request) -> httpx.Response: def test_finalize_falls_back_to_a_completed_status_update(make_platform_client, rft_routes): - """A finalize-specific failure must not strand an otherwise completed run.""" routes = dict(rft_routes) routes["POST /api/v1/rft/finalize"] = lambda request: httpx.Response( 400, json={"detail": "finalize unavailable"} @@ -202,8 +191,6 @@ def test_finalize_falls_back_to_a_completed_status_update(make_platform_client, def test_other_terminal_states_are_reported_as_failed_with_the_reason( make_platform_client, rft_routes, status, error, expected_message ): - """The RFT vocabulary is completed | failed; the finer distinction rides - in error_message so an operator can still tell a Ctrl-C from a crash.""" backend, handler = make_backend(make_platform_client, rft_routes) backend.finalize("run-abc", status=status, error=error) @@ -216,8 +203,6 @@ def test_other_terminal_states_are_reported_as_failed_with_the_reason( def test_marking_an_already_closed_run_failed_is_not_an_error(make_platform_client, rft_routes): - """409: a managed launcher or an earlier attempt got there first. The - outcome is recorded either way, so teardown must not raise over it.""" routes = dict(rft_routes) routes["PUT /api/v1/rft/external-runs/run-abc/status"] = lambda request: httpx.Response( 409, json={"detail": "Cannot update status of run in 'COMPLETED' state."} @@ -227,14 +212,3 @@ def test_marking_an_already_closed_run_failed_is_not_an_error(make_platform_clie backend.finalize("run-abc", status=RunStatus.FAILED, error="boom") assert handler.paths() == ["PUT /api/v1/rft/external-runs/run-abc/status"] - - -def test_log_metrics_posts_one_step(make_platform_client, rft_routes): - backend, handler = make_backend(make_platform_client, rft_routes) - - backend.log_metrics("run-abc", {"step": 3, "loss": 0.5}) - - assert handler.bodies_for("/api/v1/rft/metrics")[0] == { - "run_id": "run-abc", - "metrics": {"step": 3, "loss": 0.5}, - } diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index 2f9f06b46..2b67b3a14 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -54,15 +54,6 @@ def make_run(backend=None, sinks=None, config=None, **kwargs) -> Run: ) -def test_the_handle_exposes_what_a_producer_prints(): - run = make_run() - - assert run.id == "run-1" - assert run.url == "https://app.example/run-1" - assert run.status is RunStatus.RUNNING - run.finish() - - def test_sinks_are_started_with_the_run_id_and_provenance(): sink = FakeSink() run = make_run(sinks=[sink]) @@ -101,17 +92,6 @@ def test_an_empty_batch_is_not_sent(): run.finish() -def test_episodes_take_the_same_path_as_traces(): - sink = FakeSink() - run = make_run(sinks=[sink]) - - run.log_episodes([{"id": "ep-1", "traces": [{"id": "t1"}]}]) - run.flush() - - assert sink.batches == [[{"id": "ep-1", "traces": [{"id": "t1"}]}]] - run.finish() - - def test_summary_can_be_built_up_before_finish(): backend = FakeBackend() run = make_run(backend) @@ -124,17 +104,7 @@ def test_summary_can_be_built_up_before_finish(): assert backend.finalized[0]["summary"] == {"avg_reward": 0.75, "avg_error": 0.0} -def test_a_finished_run_refuses_more_summary(): - run = make_run() - run.finish() - - with pytest.raises(RunFinishedError): - run.update_summary({"late": 1.0}) - - def test_non_finite_summary_values_are_dropped_rather_than_failing_the_request(): - """A diverged loss serializes as bare ``NaN``, which strict JSON rejects — - the whole request fails on a payload nobody can inspect.""" backend = FakeBackend() run = make_run(backend) @@ -145,8 +115,6 @@ def test_non_finite_summary_values_are_dropped_rather_than_failing_the_request() def test_finish_flushes_records_before_reporting_the_terminal_status(): - """A dashboard reacting to the status must never see a finished run with - samples still landing.""" order: List[str] = [] class OrderedSink(FakeSink): @@ -230,17 +198,25 @@ def test_finish_rejects_a_nonterminal_status_without_closing_the_run(status): assert backend.finalized[0]["status"] is RunStatus.COMPLETED -def test_logging_after_finish_is_a_producer_bug(): - run = make_run() +@pytest.mark.parametrize( + "operation", + [ + lambda run: run.log_traces([{"id": "t1"}]), + lambda run: run.log_episodes([{"id": "ep1", "traces": []}]), + lambda run: run.log_metrics({"loss": 0.5}), + lambda run: run.update_summary({"late": 1.0}), + ], + ids=["traces", "episodes", "metrics", "summary"], +) +def test_writes_after_finish_are_a_producer_bug(operation): + run = make_run(metrics_sinks=[FakeSink("metrics")]) run.finish() with pytest.raises(RunFinishedError): - run.log_traces([{"id": "t1"}]) + operation(run) def test_logging_that_started_first_is_queued_before_concurrent_finish(): - """Finish must not close the worker between the live check and enqueue.""" - class SignalingBackend(FakeBackend): def __init__(self) -> None: super().__init__() @@ -382,8 +358,6 @@ def interrupt_finish(*args, **kwargs): @pytest.mark.parametrize("interrupt", [KeyboardInterrupt, asyncio.CancelledError]) def test_an_interrupt_is_recorded_as_a_decision_not_a_fault(interrupt): - """Ctrl-C must not land in the same bucket as a broken eval, nor in the - one for a process that vanished: the platform has a status for it.""" backend = FakeBackend() with pytest.raises(interrupt): @@ -407,8 +381,6 @@ def test_cancelled_is_terminal_and_finalizes_like_failed(): def test_finish_timeout_is_set_per_run_and_overridden_per_call(): - """An abort path passes a short budget so an interrupted producer is not - pinned behind a slow platform for the full drain.""" run = make_run(finish_timeout=1.5) observed = {} @@ -423,17 +395,7 @@ def record_flush(timeout=None): assert 0.0 < observed["flush"] <= 0.25 -def test_finish_timeout_default_is_the_upload_budget(): - from prime_runs.run import DEFAULT_FINISH_TIMEOUT - - run = make_run() - assert run._finish_timeout == DEFAULT_FINISH_TIMEOUT == 300.0 - run.finish() - - def test_finish_hands_the_full_config_to_finalize(): - """The evaluations API replaces metadata wholesale, so a backend recording - terminal state inside it needs the whole picture to merge into.""" backend = FakeBackend() run = make_run(backend, config={"num_rollouts": 4}) @@ -443,7 +405,6 @@ def test_finish_hands_the_full_config_to_finalize(): def test_finish_warns_when_uploads_do_not_drain(caplog): - """Finalizing over an unfinished upload silently drops records.""" run = make_run(sinks=[FakeSink()]) run._finish_timeout = 0.01 run._worker.flush = lambda timeout=None: False @@ -476,8 +437,6 @@ def record_close(timeout=None): def test_a_process_that_exits_without_finishing_reports_crashed(): - """The producer never said the run failed — it stopped existing. The - distinction tells an operator where to look.""" backend = FakeBackend() run = make_run(backend) @@ -487,7 +446,6 @@ def test_a_process_that_exits_without_finishing_reports_crashed(): def test_a_forked_handle_gets_a_fresh_lock_and_loses_lifecycle_ownership(): - """The child's inherited atexit hook must never finalize the parent's run.""" backend = FakeBackend() run = make_run(backend) inherited_lock = run._finish_lock @@ -501,7 +459,6 @@ def test_a_forked_handle_gets_a_fresh_lock_and_loses_lifecycle_ownership(): def test_a_backend_failure_does_not_escape_into_the_producer_by_default(): - """Six hours of rollouts must not be lost to a 502 on a telemetry call.""" backend = FakeBackend(fail_on="finalize") run = make_run(backend) @@ -564,17 +521,7 @@ def test_a_sink_error_is_recorded_on_the_run(): run.finish() -def test_dropped_records_are_reported_on_the_handle(): - run = make_run() - run._worker.dropped = 3 - - assert run.dropped_records == 3 - run.finish() - - def test_an_upload_failure_reaches_the_caller_in_raise_mode(): - """A sink fails on the uploader thread, where raising reaches nobody. Under - on_error="raise" the failure has to surface where a test is looking.""" run = make_run(sinks=[FakeSink("broken", fail_on_write=True)], on_error="raise") run.log_traces([{"id": "t1"}]) @@ -618,8 +565,6 @@ def close(self) -> None: def test_an_account_outside_the_beta_is_not_a_failed_run(): - """Nothing was lost — the records went to every sink that applies to this - account — so the run finishes clean even under ``on_error="raise"``.""" sink = TracesSink(client=_ForbiddenClient("service_not_enabled")) run = make_run(sinks=[sink], on_error="raise") run.log_traces([{"id": "t1"}]) @@ -701,14 +646,6 @@ def test_log_metrics_without_a_metrics_sink_is_a_no_op(): run.finish() -def test_log_metrics_on_a_finished_run_is_a_producer_bug(): - run = make_run(metrics_sinks=[FakeSink("metrics")]) - run.finish() - - with pytest.raises(RunFinishedError): - run.log_metrics({"loss": 0.5}) - - def test_finish_drains_and_closes_the_metrics_uploader_too(): metrics = FakeSink("metrics") run = make_run(metrics_sinks=[metrics]) diff --git a/packages/prime-runs/tests/test_samples_sink.py b/packages/prime-runs/tests/test_samples_sink.py index 2a46ff2e9..bd5fe3626 100644 --- a/packages/prime-runs/tests/test_samples_sink.py +++ b/packages/prime-runs/tests/test_samples_sink.py @@ -26,7 +26,6 @@ def test_episodes_are_projected_and_posted(make_platform_client, eval_routes): def test_rollout_numbering_is_continuous_across_streamed_batches(make_platform_client, eval_routes): - """Streaming must produce the numbering the old one-shot upload produced.""" sink, handler = make_sink(make_platform_client, eval_routes) sink.write([make_episode("ep-1", [make_trace(idx=0)])]) @@ -48,10 +47,6 @@ def test_a_producer_that_already_speaks_v0_is_passed_through(make_platform_clien def test_records_this_sink_cannot_project_are_skipped_not_fatal( make_platform_client, eval_routes, caplog ): - """A JSON episode or a bare trace has no v0 projection. Raising would have - the worker retire the sink for the run — outside the traces beta that is - an empty viewer — so the record is skipped, warned once, and counted, - while the episode object in the same batch is still stored.""" sink, handler = make_sink(make_platform_client, eval_routes) with caplog.at_level("WARNING"): @@ -86,7 +81,6 @@ def test_an_empty_batch_makes_no_request(make_platform_client, eval_routes): def test_a_rejected_sample_batch_does_not_abort_later_batches( make_platform_client, eval_routes, monkeypatch ): - """One bad request must not hide valid rows later in the same worker batch.""" statuses = iter([200, 422, 200, 200]) def respond(request): diff --git a/packages/prime-runs/tests/test_traces_sink.py b/packages/prime-runs/tests/test_traces_sink.py index 48c59fc3e..aea3a7cd0 100644 --- a/packages/prime-runs/tests/test_traces_sink.py +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -31,8 +31,6 @@ def make_sink(client=None, **kwargs) -> TracesSink: def test_records_go_out_with_provenance_but_not_the_join_key(): - """``run.id`` inside the document is the indexed column; ``context`` is an - upload-scoped map that answers a different question.""" client = FakeTracesClient() sink = make_sink(client) @@ -61,8 +59,6 @@ def test_the_line_format_is_inferred_from_the_records(): def test_a_bare_mapping_gets_the_run_stamped_onto_a_copy(): - """A dict has no stamping convention, and an upload with no ``run.id`` is - orphaned — unqueryable and undeletable by run.""" client = FakeTracesClient() sink = make_sink(client) original = {"id": "t1"} @@ -74,9 +70,6 @@ def test_a_bare_mapping_gets_the_run_stamped_onto_a_copy(): def test_a_record_that_names_its_run_keeps_it(): - """Producers stamp the run at rollout time; the sink never overrides a - ``run`` that is already there, which is how a second source of truth - would appear.""" client = FakeTracesClient() sink = make_sink(client) @@ -86,10 +79,6 @@ def test_a_record_that_names_its_run_keeps_it(): def test_an_episode_s_run_reaches_every_member_trace(): - """The traces service derives ``run_id`` from ``trace.run.id`` only and never - reads the envelope's ``run``, while verifiers records the run on the episode - and its ``Trace`` has no ``run`` field at all. Without this, every row of an - episode upload lands with an empty ``run_id`` (seen live, 2026-08-21).""" client = FakeTracesClient() sink = make_sink(client) episode = make_episode("ep-1", [make_trace(trace_id="a"), make_trace(trace_id="b")]) @@ -120,23 +109,7 @@ def test_a_member_that_names_its_own_run_is_left_alone(): assert episode["traces"][0] == {"id": "a"}, "the caller's members were not mutated" -def test_producer_objects_are_serialized_once_through_to_record(): - """The sink calls ``to_record()`` itself so members are reachable; the bytes - the transport sees are the same ones it would have produced.""" - client = FakeTracesClient() - sink = make_sink(client) - trace = make_trace() - - sink.write([trace]) - - sent = client.calls[0][0][0] - assert sent == {**trace.to_record(), "run": {"id": "run-1", "type": "eval"}} - - def test_an_account_outside_the_beta_retires_the_sink_without_a_failure(caplog): - """Prime Traces is in closed beta. For everyone outside it there was never - anywhere for these records to go, so the sink turns itself off quietly: - no exception for the worker to count, nothing above INFO in the log.""" client = FakeTracesClient( raises=ForbiddenError("not in beta", status_code=403, code="service_not_enabled") ) @@ -153,8 +126,6 @@ def test_an_account_outside_the_beta_retires_the_sink_without_a_failure(caplog): def test_a_credential_without_the_traces_scope_is_still_a_failure(caplog): - """The other 403 is fixable — mint a token with the scope — so it is raised - for loss accounting and strict callers, and the sink still retires.""" client = FakeTracesClient( raises=ForbiddenError("missing scope: traces", status_code=403, code="forbidden") ) @@ -169,7 +140,6 @@ def test_a_credential_without_the_traces_scope_is_still_a_failure(caplog): def test_a_transient_failure_is_raised_so_the_worker_can_report_it(): - """Unlike a 403, this one is about the moment, not the account.""" client = FakeTracesClient(raises=RetryableAPIError("busy", status_code=503)) sink = make_sink(client) @@ -178,16 +148,6 @@ def test_a_transient_failure_is_raised_so_the_worker_can_report_it(): assert sink.enabled is True -def test_a_disabled_sink_stops_calling_the_service(): - client = FakeTracesClient() - sink = make_sink(client) - sink.enabled = False - - sink.write([{"id": "t1"}]) - - assert client.calls == [] - - def test_receipt_history_is_bounded_while_the_total_is_retained(): client = FakeTracesClient() sink = make_sink(client, receipt_history_size=2) @@ -199,18 +159,7 @@ def test_receipt_history_is_bounded_while_the_total_is_retained(): assert len(sink.receipts) == 2 -def test_closing_the_sink_closes_the_client(): - client = FakeTracesClient() - sink = make_sink(client) - - sink.close() - - assert client.closed is True - - def test_a_missing_traces_client_disables_the_sink_and_reports_the_failure(monkeypatch, caplog): - """The run applies warn/raise policy, so the sink must surface this failure.""" - def explode(**kwargs): raise RuntimeError("no credentials") diff --git a/packages/prime-runs/tests/test_train_samples_sink.py b/packages/prime-runs/tests/test_train_samples_sink.py index 6c5aa8946..23401e54b 100644 --- a/packages/prime-runs/tests/test_train_samples_sink.py +++ b/packages/prime-runs/tests/test_train_samples_sink.py @@ -5,7 +5,7 @@ from _fakes import make_episode, make_trace, make_train_episode from conftest import RecordingHandler, StorageHandler -from prime_runs.exceptions import APIError, RetryableAPIError, TransportError +from prime_runs.exceptions import APIError, RetryableAPIError from prime_runs.sinks import RftSamplesSink, training_step from prime_runs.sinks.base import SinkWriteError @@ -79,8 +79,6 @@ def test_the_cadence_is_configurable(make_platform_client, rft_routes): def test_records_without_training_work_have_no_row_here(make_platform_client, rft_routes): - """Eval-work episodes, bare traces and JSON episodes reach Prime Traces - only; here they are counted, not uploaded, and not a loss.""" sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes) sink.write( @@ -113,8 +111,6 @@ def test_nothing_to_encode_makes_no_request(make_platform_client, rft_routes): def test_a_failed_presign_loses_that_step_only(make_platform_client, rft_routes): - """Presign is replayable (it only mints a URL), so a blip is retried; a - step whose presign keeps failing is lost on its own, not the batch.""" import json def flaky(request: httpx.Request) -> httpx.Response: @@ -178,23 +174,6 @@ def test_a_storage_rejection_is_not_retried(make_platform_client, rft_routes): assert not isinstance(info.value.cause, RetryableAPIError) -def test_a_dropped_connection_to_storage_is_transient(make_platform_client, rft_routes, no_sleep): - def refuse(request: httpx.Request) -> httpx.Response: - raise httpx.ConnectError("connection refused") - - storage = StorageHandler() - storage.__call__ = refuse # type: ignore[method-assign] - upload = httpx.Client(transport=httpx.MockTransport(refuse)) - handler = RecordingHandler(rft_routes) - sink = RftSamplesSink(make_platform_client(handler), encoder=Encoder(), upload_client=upload) - sink.start("run-abc", {}) - - with pytest.raises(SinkWriteError) as info: - sink.write([make_train_episode("e1", step=10)]) - - assert isinstance(info.value.cause, TransportError) - - def test_an_encoder_failure_is_a_record_rejection(make_platform_client, rft_routes): def broken(episodes, run_id, step, sample_id_offset=0): raise ValueError("cannot serialize") @@ -221,9 +200,6 @@ def test_a_presign_without_a_url_is_an_api_error(make_platform_client, rft_route def test_a_step_logged_again_gets_its_own_sample_id_range(make_platform_client, rft_routes): - """Objects are additive per step on the platform (each presign mints a new - key and the viewer unions them), and the viewer looks samples up by - (step, sample_id): a straggler batch must not reuse the first object's ids.""" from prime_runs.sinks.train_samples import SAMPLE_ID_STRIDE sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes) @@ -239,8 +215,6 @@ def test_a_step_logged_again_gets_its_own_sample_id_range(make_platform_client, def test_an_upload_that_failed_still_advances_the_sample_id_range(make_platform_client, rft_routes): - """An attempt that failed late may still have landed an object; the next - one for the step must not overlap it.""" from prime_runs.sinks.train_samples import SAMPLE_ID_STRIDE storage = StorageHandler(status_codes=[403]) @@ -254,10 +228,6 @@ def test_an_upload_that_failed_still_advances_the_sample_id_range(make_platform_ def test_a_forked_child_numbers_its_rows_apart_from_the_parent(make_platform_client, rft_routes): - """The child inherits a copy of the parent's upload counters; continuing - them would hand out the parent's next range. The child's ranges come from - a band the parent's counter never reaches, and every id stays a JSON-safe - integer for the viewer.""" from prime_runs.sinks.train_samples import SAMPLE_ID_RANGES, SAMPLE_ID_STRIDE assert SAMPLE_ID_RANGES * SAMPLE_ID_STRIDE <= 2**53 @@ -279,9 +249,6 @@ def test_a_forked_child_numbers_its_rows_apart_from_the_parent(make_platform_cli def test_a_confirm_whose_response_was_lost_is_replayed(make_platform_client, rft_routes, no_sleep): - """Confirm validates the key and refreshes progress; a second confirm of - the same key changes nothing, so an ambiguous failure is retried rather - than costing the step.""" attempts = [] def flaky(request: httpx.Request) -> httpx.Response: @@ -304,7 +271,6 @@ def flaky(request: httpx.Request) -> httpx.Response: def test_without_pyarrow_the_sink_turns_itself_off_quietly( make_platform_client, rft_routes, monkeypatch, caplog ): - """No loss: the episodes still reach Prime Traces. Say how to turn it on.""" monkeypatch.setattr("prime_runs.projection.parquet_available", lambda: False) handler = RecordingHandler(rft_routes) sink = RftSamplesSink(make_platform_client(handler)) @@ -350,8 +316,6 @@ def test_training_step_reads_the_episode_provenance(record, expected): def test_the_default_encoder_writes_the_viewer_table(): - """The real Parquet encoder, when pyarrow is installed: one row per episode - with the RFT-only columns layered on the shared sample projection.""" pa = pytest.importorskip("pyarrow") pq = pytest.importorskip("pyarrow.parquet") import io @@ -385,6 +349,28 @@ def test_the_default_encoder_writes_the_viewer_table(): assert isinstance(pa.Table, type) +def test_the_default_encoder_writes_non_finite_numbers_as_null(): + pytest.importorskip("pyarrow") + import io + import json + import math + + import pyarrow.parquet as pq + + from prime_runs.projection import episodes_to_parquet_bytes + + episode = make_train_episode("e1", step=10, advantage=float("inf")) + episode.traces[0].metrics["kl"] = float("nan") + + payload = episodes_to_parquet_bytes([episode], "run-abc", 10) + + (row,) = pq.read_table(io.BytesIO(payload)).to_pylist() + strict = dict(parse_constant=lambda constant: pytest.fail(f"bare {constant} in a column")) + assert json.loads(row["metrics"], **strict) == {"kl": None} + assert json.loads(row["trajectory"], **strict)[-1]["advantage"] is None + assert math.isinf(row["advantage"]) # the float column keeps the value + + def test_the_default_encoder_numbers_rows_from_the_offset(): pytest.importorskip("pyarrow") import io @@ -414,8 +400,6 @@ def test_the_default_encoder_skips_episodes_without_a_trajectory(): def test_a_forked_child_drops_the_inherited_storage_client(make_platform_client, rft_routes): - """The socket belongs to the parent: the child must neither reuse it nor - close it (a close would send close_notify on the parent's connection).""" handler = RecordingHandler(rft_routes) owned = RftSamplesSink(make_platform_client(handler), encoder=Encoder()) owned.start("run-abc", {}) @@ -437,12 +421,3 @@ def test_a_forked_child_drops_the_inherited_storage_client(make_platform_client, injected.write([make_train_episode("e1", step=10)]) assert isinstance(info.value.cause, RuntimeError) assert "after a fork" in str(info.value.cause) - - -def test_the_sink_is_registered_for_fork_resets(): - from prime_runs import _fork - from prime_runs.sinks import RftSamplesSink - - sink = RftSamplesSink(object(), encoder=Encoder()) # type: ignore[arg-type] - - assert sink in _fork._registry diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py index 0b17e8e8e..fbf6bf620 100644 --- a/packages/prime-runs/tests/test_worker.py +++ b/packages/prime-runs/tests/test_worker.py @@ -38,18 +38,6 @@ def test_records_reach_every_enabled_sink(): worker.close() -def test_a_disabled_sink_is_skipped(): - live, dead = FakeSink("live"), FakeSink("dead") - dead.enabled = False - worker = UploadWorker([live, dead]) - - worker.submit([{"id": 1}]) - drain(worker) - - assert live.batches and not dead.batches - worker.close() - - def test_one_sink_failing_does_not_stop_the_others(): broken, healthy = FakeSink("broken", fail_on_write=True), FakeSink("healthy") reported = [] @@ -64,24 +52,7 @@ def test_one_sink_failing_does_not_stop_the_others(): worker.close() -def test_a_failed_sink_is_not_called_again(): - """One log line per batch for the rest of a run hides whatever failed first.""" - broken = FakeSink("broken", fail_on_write=True) - reported = [] - worker = UploadWorker([broken], on_error=lambda name, exc: reported.append(name)) - - for _ in range(3): - worker.submit([{"id": 1}]) - drain(worker) - - assert reported == ["broken"] - worker.close() - - def test_records_that_skip_a_sink_retired_by_error_are_counted_as_lost_to_it(): - """After the first failed batch the sink is off, but the producer keeps - logging; every later batch is just as lost to that sink as the first. Live - run 2026-08-21: five episodes, footer said "1 failed via traces".""" broken = FakeSink("broken", fail_on_write=True) worker = UploadWorker([broken]) @@ -94,9 +65,6 @@ def test_records_that_skip_a_sink_retired_by_error_are_counted_as_lost_to_it(): def test_records_that_skip_a_sink_which_switched_itself_off_are_not_counted(): - """A sink that retires quietly (nowhere for the records to go, e.g. outside - the traces beta) lost nothing, and must not start a failure count.""" - class QuietSink(FakeSink): def write(self, records) -> None: self.enabled = False # retires without raising, like service_not_enabled @@ -113,7 +81,6 @@ def write(self, records) -> None: def test_a_full_queue_drops_rather_than_blocking_the_producer(): - """Stalling a training run to protect telemetry is the wrong trade.""" sink = BlockingSink() worker = UploadWorker([sink], max_queue_size=1, put_timeout=0.05) @@ -191,20 +158,7 @@ def test_submitting_after_close_is_refused(): assert sink.batches == [] -def test_flush_without_a_running_thread_still_flushes_the_sinks(): - sink = FakeSink() - worker = UploadWorker([sink]) - - assert worker.flush(timeout=1.0) is True - assert sink.flushes == 1 - - def test_a_forked_child_starts_over_instead_of_re_uploading_the_parents_queue(): - """The queued records belong to the parent, which still has a live thread. - - Inheriting them would upload each record twice; inheriting the lock could - deadlock the child on its first write. - """ sink = FakeSink() worker = UploadWorker([sink], max_queue_size=4) worker._queue.put([{"id": "parents"}]) @@ -220,9 +174,6 @@ def test_a_forked_child_starts_over_instead_of_re_uploading_the_parents_queue(): def test_close_leaves_sinks_open_when_the_uploader_will_not_stop(caplog): - """Closing them would pull an httpx client or a file handle out from under - a request still running on that thread.""" - class WedgedSink(FakeSink): def __init__(self) -> None: super().__init__("wedged") @@ -244,9 +195,6 @@ def write(self, records) -> None: def test_a_forked_child_resets_every_registered_holder_of_a_connection(): - """One process-wide hook, not one per object: a per-instance - register_at_fork can never be undone, so it would pin every run the process - ever opened and re-run hooks for runs that finished hours ago.""" from prime_runs import _fork class Holder: @@ -264,32 +212,7 @@ def reset_after_fork(self) -> None: assert holder.reset == 1 -def test_one_registered_object_raising_does_not_block_the_others(): - from prime_runs import _fork - - class Boom: - def reset_after_fork(self) -> None: - raise RuntimeError("nope") - - class Fine: - def __init__(self) -> None: - self.reset = 0 - - def reset_after_fork(self) -> None: - self.reset += 1 - - fine = Fine() - _fork.register(Boom()) - _fork.register(fine) - - _fork._reset_all() - - assert fine.reset == 1 - - def test_a_transient_failure_drops_the_batch_but_keeps_the_sink(): - """One gateway blip must not empty the rest of the run's dashboard. The - batch is already lost; retiring the sink loses every batch after it too.""" from prime_runs.exceptions import RetryableAPIError class BlipSink(FakeSink): @@ -320,8 +243,6 @@ def write(self, records) -> None: ids=["local-encoding", "traces-validation", "samples-validation"], ) def test_a_record_rejection_drops_only_its_batch(error): - """Malformed content must not prevent later valid records from uploading.""" - class RejectOnceSink(FakeSink): def __init__(self) -> None: super().__init__("reject-once") @@ -353,7 +274,6 @@ def write(self, records) -> None: def test_a_sustained_outage_eventually_retires_the_sink(): - """A blip is forgiven; hours of re-attempting every batch is not useful.""" from prime_runs.exceptions import TransportError from prime_runs.worker import TRANSIENT_FAILURE_LIMIT @@ -374,9 +294,6 @@ def write(self, records) -> None: def test_a_cooldown_gives_a_retired_sink_another_chance(): - """A training run outlives any platform deploy: after the cooldown the - sink is tried again, and records that arrived meanwhile are counted as - lost to it, not silently forgotten.""" from prime_runs.exceptions import TransportError from prime_runs.worker import TRANSIENT_FAILURE_LIMIT @@ -435,8 +352,6 @@ def write(self, records) -> None: def test_a_success_forgives_earlier_blips(): - """Strikes are consecutive: an intermittent gateway must never accumulate - its way to a retirement across an otherwise healthy run.""" from prime_runs.exceptions import RetryableAPIError from prime_runs.worker import TRANSIENT_FAILURE_LIMIT @@ -469,7 +384,6 @@ def write(self, records) -> None: def test_a_permanent_failure_retires_the_sink_immediately(): - """A gated account or a rejected credential fails identically forever.""" from prime_runs.exceptions import UnauthorizedError class DeniedSink(FakeSink): @@ -487,9 +401,6 @@ def write(self, records) -> None: def test_a_failed_batch_is_counted_once_per_sink_not_once_per_run(): - """Default online runs write to two sinks. Adding both to one total would - report twice the loss, and would report loss at all when the other sink - stored the records.""" from prime_runs.exceptions import RetryableAPIError class BlipSink(FakeSink): @@ -509,9 +420,6 @@ def write(self, records) -> None: def test_records_queued_during_an_upload_go_out_as_one_batch(): - """One request per episode runs into the platform's per-minute limit on a - fast eval; whatever accumulates while a request is in flight must go out - together, so the request rate tracks upload latency instead.""" sink = BlockingSink() worker = UploadWorker([sink]) @@ -527,8 +435,6 @@ def test_records_queued_during_an_upload_go_out_as_one_batch(): def test_coalescing_keeps_traces_and_episodes_apart(): - """The traces sink infers the line format from a batch's first record, so - mixing episodes into a batch of bare traces would misfile every record.""" sink = BlockingSink() worker = UploadWorker([sink]) @@ -551,8 +457,6 @@ def test_coalescing_keeps_traces_and_episodes_apart(): def test_a_flush_barrier_is_not_reordered_past_the_records_before_it(): - """Coalescing must stop at the barrier: flush() promises that everything - submitted before it has been written when it returns.""" sink = BlockingSink() worker = UploadWorker([sink]) @@ -574,8 +478,6 @@ def test_a_flush_barrier_is_not_reordered_past_the_records_before_it(): def test_a_worker_with_no_sinks_never_starts_a_thread(): - """A disabled run submits every batch; copying them into a queue and - starting a thread to find there is nowhere to put them is pure cost.""" worker = UploadWorker([]) assert worker.submit([{"id": 1}]) is True From aba26790514749fe484e995f99591fd0d35e2350 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Tue, 1 Sep 2026 17:16:51 -0700 Subject: [PATCH 38/39] fix(runs): keep metric cleaning array-safe and count a split record once - _finite compared every cleaned value with `!= {}` to drop empty nested mappings; an array-like metric value makes that comparison ambiguous and would raise on the producer thread. Only identity and mapping checks now. - EvalSamplesSink attributes a record whose rows straddle two HTTP batches to the first batch, so a lost record is counted once, as before the send_each refactor. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Fk5fBAAaTiPfJ5UpNyqxXR --- packages/prime-runs/src/prime_runs/run.py | 10 ++++++++-- .../src/prime_runs/sinks/samples.py | 7 +++++-- packages/prime-runs/tests/test_run.py | 20 +++++++++++++++++++ .../prime-runs/tests/test_samples_sink.py | 20 +++++++++++++++++++ 4 files changed, 53 insertions(+), 4 deletions(-) diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index aaf425f49..bc89b98d2 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -570,8 +570,14 @@ def _finite(value: Any) -> Any: if isinstance(value, float) and not math.isfinite(value): return _DROP if isinstance(value, Mapping): - items = ((key, _finite(item)) for key, item in value.items()) - return {key: item for key, item in items if item is not _DROP and item != {}} + cleaned = {} + for key, item in value.items(): + kept = _finite(item) + # Identity checks only: an array-like value would make `==` ambiguous. + if kept is _DROP or (isinstance(item, Mapping) and not kept): + continue + cleaned[key] = kept + return cleaned if isinstance(value, (list, tuple)): return [item for item in (_finite(item) for item in value) if item is not _DROP] return value diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py index b9f384481..3bec0d680 100644 --- a/packages/prime-runs/src/prime_runs/sinks/samples.py +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -5,7 +5,7 @@ """ import logging -from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple +from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Set, Tuple from .._http import UPLOAD_TIMEOUT, PlatformClient, encode_json from ..projection import batch_samples, build_samples @@ -42,8 +42,11 @@ def write(self, records: Sequence[Any]) -> None: samples, owners = self._to_samples(records) units: List[Unit] = [] offset = 0 + seen: Set[int] = set() for batch in batch_samples(samples): - batch_owners = set(owners[offset : offset + len(batch)]) + # A record whose rows straddle two batches is counted once, with the first. + batch_owners = set(owners[offset : offset + len(batch)]) - seen + seen |= batch_owners offset += len(batch) units.append((len(batch_owners), self._sender(batch))) send_each(units) diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py index 2b67b3a14..f6ba0a051 100644 --- a/packages/prime-runs/tests/test_run.py +++ b/packages/prime-runs/tests/test_run.py @@ -625,6 +625,26 @@ def test_log_metrics_drops_non_finite_values_inside_sequences(): run.finish() +def test_log_metrics_leaves_array_like_values_alone(): + class Arrayish: # numpy-style: comparisons are element-wise, truthiness ambiguous + def __eq__(self, other): + raise ValueError("ambiguous") + + __ne__ = __eq__ + __bool__ = __eq__ + + metrics = FakeSink("metrics") + run = make_run(metrics_sinks=[metrics]) + histogram = Arrayish() + + run.log_metrics({"hist": histogram, "nested": {"hist": histogram, "nan": float("nan")}}) + run.flush() + + row = metrics.batches[0][0] + assert row["hist"] is histogram and row["nested"] == {"hist": histogram} + run.finish() + + def test_log_metrics_keeps_a_step_the_producer_already_set(): metrics = FakeSink("metrics") run = make_run(metrics_sinks=[metrics]) diff --git a/packages/prime-runs/tests/test_samples_sink.py b/packages/prime-runs/tests/test_samples_sink.py index bd5fe3626..8349ba0ee 100644 --- a/packages/prime-runs/tests/test_samples_sink.py +++ b/packages/prime-runs/tests/test_samples_sink.py @@ -1,10 +1,12 @@ """The legacy sample sink that keeps today's viewer working.""" import httpx +import pytest from _fakes import make_episode, make_trace from conftest import RecordingHandler from prime_runs.sinks import EvalSamplesSink +from prime_runs.sinks.base import SinkWriteError from prime_runs.worker import UploadWorker @@ -70,6 +72,24 @@ def test_records_this_sink_cannot_project_are_skipped_not_fatal( assert "(dict)" in warnings[0] +def test_a_record_split_across_batches_is_counted_once( + make_platform_client, eval_routes, monkeypatch +): + eval_routes["POST /api/v1/evaluations/eval-abc/samples"] = lambda request: httpx.Response(502) + # Every episode falls back to one row per trace, and every row is its own batch. + monkeypatch.setattr("prime_runs.projection.MAX_SAMPLES_PAYLOAD_BYTES", 1) + monkeypatch.setattr( + "prime_runs.sinks.samples.batch_samples", lambda samples: [[sample] for sample in samples] + ) + sink, handler = make_sink(make_platform_client, eval_routes) + + with pytest.raises(SinkWriteError) as info: + sink.write([make_episode("ep-1", [make_trace(trace_id="a"), make_trace(trace_id="b")])]) + + assert len(handler.bodies_for("/api/v1/evaluations/eval-abc/samples")) == 2 + assert info.value.failed_records == 1 + + def test_an_empty_batch_makes_no_request(make_platform_client, eval_routes): sink, handler = make_sink(make_platform_client, eval_routes) From 1b65b9a1dd22c41ec99e670646747ef912b68f51 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 2 Sep 2026 15:02:13 -0700 Subject: [PATCH 39/39] fix(runs): key every uploaded record to the run in the traces sink `stamp_run` kept a record's own `run` block, and verifiers and prime-rl always record one with their local run id, so the only way to get the platform's id into the trace document was for the producers to adopt it. The sink now sets `run.id`/`run.type` on the uploaded copy (envelope and every member trace) and keeps the rest of the block (`name`, `work`); producers keep their own ids and never learn the platform's. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01FHMjWstGWd1oGKjoj5RgeT --- packages/prime-runs/README.md | 15 ++++---- packages/prime-runs/src/prime_runs/run.py | 4 +- .../prime-runs/src/prime_runs/sinks/base.py | 25 ++++++++----- packages/prime-runs/tests/test_traces_sink.py | 37 ++++++++++++++----- 4 files changed, 54 insertions(+), 27 deletions(-) diff --git a/packages/prime-runs/README.md b/packages/prime-runs/README.md index 1e7a3bb9f..47ce21b67 100644 --- a/packages/prime-runs/README.md +++ b/packages/prime-runs/README.md @@ -15,7 +15,6 @@ uv add 'prime-runs[train]' # training runs: adds pyarrow for the sample table ```python import prime_runs as pr -from verifiers.v1 import EvalRunInfo run = pr.init( name="gsm8k-qwen3-8b", @@ -27,16 +26,18 @@ run = pr.init( print(run.url) # https://app.primeintellect.ai/dashboard/evaluations/... for episode in rollouts: - episode.record_run(EvalRunInfo(id=run.id)) # every trace carries the run id - run.log_episodes([episode]) # a queue put; bare traces: log_traces() + run.log_episodes([episode]) # a queue put; bare traces: log_traces() run.finish(summary=pr.metrics.from_episodes(episodes)) ``` -`init()` is called before the first rollout, and the id it returns is the run -id everywhere, including inside every trace document. A `with run:` block -finishes for you: an exception marks the run `failed`, Ctrl-C `cancelled`, and -a process that exits without finishing is reported `crashed` by an atexit hook. +`init()` is called before the first rollout. Every record the run uploads is +keyed to it: the SDK sets `run.id` and `run.type` on the uploaded copy of each +trace and episode, over whatever run id the producer recorded locally, and +keeps the rest of that block (`name`, `work`). A producer never needs to know +the platform's id; `run.url` is the handle. A `with run:` block finishes for +you: an exception marks the run `failed`, Ctrl-C `cancelled`, and a process +that exits without finishing is reported `crashed` by an atexit hook. `config=` takes the path to the launched file (kept verbatim under `config_source`, comments and all) or a mapping stored as given; put a file diff --git a/packages/prime-runs/src/prime_runs/run.py b/packages/prime-runs/src/prime_runs/run.py index bc89b98d2..384082736 100644 --- a/packages/prime-runs/src/prime_runs/run.py +++ b/packages/prime-runs/src/prime_runs/run.py @@ -1,7 +1,7 @@ """The run handle, and ``init()`` that opens one. -``init()`` is called before the first rollout; the id it returns is the run id -everywhere, including inside every trace document the producer writes. +``init()`` is called before the first rollout; the id it returns keys every +record the run uploads, whatever run id the producer recorded locally. """ import asyncio diff --git a/packages/prime-runs/src/prime_runs/sinks/base.py b/packages/prime-runs/src/prime_runs/sinks/base.py index 5418e16d0..70185ab36 100644 --- a/packages/prime-runs/src/prime_runs/sinks/base.py +++ b/packages/prime-runs/src/prime_runs/sinks/base.py @@ -81,20 +81,27 @@ def is_episode(record: Any) -> bool: def stamp_run(mapping: Mapping[str, Any], run_id: str, run_type: str = RUN_KIND) -> Dict[str, Any]: - """A copy of ``mapping`` with ``run`` on the envelope and on every member - trace that lacks one. A record that already names a run keeps it. Members - matter because the traces service reads ``run_id`` from ``trace.run.id`` - only, while verifiers records the run on the episode.""" + """A copy of ``mapping`` keyed to this run: ``run.id`` and ``run.type`` are + set on the envelope and on every member trace, over whatever run the + producer recorded there (its own local id), and the rest of that block + (``name``, ``work``) is kept. Members matter because the traces service + reads ``run_id`` from ``trace.run.id`` only, while verifiers records the + run on the episode.""" stamped = dict(mapping) - if not stamped.get("run"): - stamped["run"] = {"id": run_id, "type": run_type} + stamped["run"] = _rekey(stamped.get("run"), run_id, run_type) members = stamped.get("traces") if isinstance(members, list): - run = stamped["run"] stamped["traces"] = [ - {**member, "run": run} - if isinstance(member, Mapping) and not member.get("run") + {**member, "run": _rekey(member.get("run") or stamped["run"], run_id, run_type)} + if isinstance(member, Mapping) else member for member in members ] return stamped + + +def _rekey(run: Any, run_id: str, run_type: str) -> Dict[str, Any]: + keyed = dict(run) if isinstance(run, Mapping) else {} + keyed["id"] = run_id + keyed["type"] = run_type + return keyed diff --git a/packages/prime-runs/tests/test_traces_sink.py b/packages/prime-runs/tests/test_traces_sink.py index aea3a7cd0..9d59c8f10 100644 --- a/packages/prime-runs/tests/test_traces_sink.py +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -69,13 +69,30 @@ def test_a_bare_mapping_gets_the_run_stamped_onto_a_copy(): assert original == {"id": "t1"}, "the caller's dict was not mutated" -def test_a_record_that_names_its_run_keeps_it(): +def test_a_record_that_names_its_own_run_is_rekeyed_to_this_one(): client = FakeTracesClient() sink = make_sink(client) + original = {"id": "t1", "run": {"id": "theirs", "type": "train", "name": "local"}} - sink.write([{"id": "t1", "run": {"id": "theirs", "type": "eval"}}]) + sink.write([original]) + + assert client.calls[0][0][0]["run"] == {"id": "run-1", "type": "eval", "name": "local"} + assert original["run"] == {"id": "theirs", "type": "train", "name": "local"}, "not mutated" + + +def test_a_training_episode_keeps_its_dispatch_step_while_being_rekeyed(): + """prime-rl stamps ``TrainRunInfo(id=$PRL_RUN_ID, work=...)`` on every + dispatched episode; the platform's id replaces the launcher's, the work + (the dispatch step) stays.""" + client = FakeTracesClient() + sink = TracesSink(client=client) + sink.start("run-1", {"run_type": "train"}) + work = {"type": "train", "step": 10} + episode = {"id": "ep-1", "run": {"id": "prl-run", "type": "train", "work": work}, "traces": []} + + sink.write([episode]) - assert client.calls[0][0][0]["run"] == {"id": "theirs", "type": "eval"} + assert client.calls[0][0][0]["run"] == {"id": "run-1", "type": "train", "work": work} def test_an_episode_s_run_reaches_every_member_trace(): @@ -91,22 +108,24 @@ def test_an_episode_s_run_reaches_every_member_trace(): assert [member["id"] for member in sent["traces"]] == ["a", "b"] -def test_a_member_that_names_its_own_run_is_left_alone(): +def test_every_member_trace_is_rekeyed_too(): client = FakeTracesClient() sink = make_sink(client) - theirs = {"id": "other-run", "type": "eval"} + theirs = {"id": "other-run", "type": "eval", "name": "other"} episode = { "id": "ep-1", - "run": {"id": "env-run", "type": "eval"}, + "run": {"id": "env-run", "type": "eval", "name": "local"}, "traces": [{"id": "a"}, {"id": "b", "run": theirs}], } sink.write([episode]) - members = client.calls[0][0][0]["traces"] - assert members[0]["run"] == {"id": "env-run", "type": "eval"} - assert members[1]["run"] == theirs + sent = client.calls[0][0][0] + assert sent["run"] == {"id": "run-1", "type": "eval", "name": "local"} + assert sent["traces"][0]["run"] == sent["run"] + assert sent["traces"][1]["run"] == {"id": "run-1", "type": "eval", "name": "other"} assert episode["traces"][0] == {"id": "a"}, "the caller's members were not mutated" + assert theirs == {"id": "other-run", "type": "eval", "name": "other"} def test_an_account_outside_the_beta_retires_the_sink_without_a_failure(caplog):