From 0dcb62a8167f98c71180bb0a7925d6780e1eb829 Mon Sep 17 00:00:00 2001 From: Cooper Miller Date: Wed, 19 Aug 2026 13:52:53 -0700 Subject: [PATCH 01/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] 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/27] =?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/27] 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):