diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index dccac8331..f572c8147 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="" @@ -168,6 +169,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..da263298a --- /dev/null +++ b/.github/workflows/release-runs.yml @@ -0,0 +1,85 @@ +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 + + # 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 + 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: 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: | + 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" 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 new file mode 100644 index 000000000..5824f4117 --- /dev/null +++ b/packages/prime-runs/README.md @@ -0,0 +1,146 @@ +# Prime Runs SDK + +Track eval 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="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_episodes([episode]) # bare traces go through log_traces() + +run.finish(summary=pr.metrics.from_episodes(episodes)) +``` + +`init()` opens the run and returns a handle carrying its ID and dashboard URL. +Records stream out on a background thread while the run proceeds, so the +dashboard fills in as rollouts land. `finish()` closes the run out; a `with` +block does that for you, including when the body raises. Run-level outputs go +in `finish(summary=...)`, or incrementally through `update_summary()`; +`metrics.from_episodes()` returns them in the shape the dashboard reads +(`metrics.RunSummary`). + +## Identity + +`init()` is called **before** the first rollout, and the ID it returns is *the* +run ID everywhere — including inside every trace document you write: + +```python +run = pr.init(...) +trace.record_run(EvalRunInfo(id=run.id)) # verifiers +``` + +The ingestion service indexes `run.id` from the trace body, so "every trace for +this run" is a fast query. Bare dicts without a `run` key are stamped for you; +producer objects are passed through untouched. + +## Config + +`config=` takes the path to the file the run was launched from, or a mapping. + +| you pass | what is stored | +| --- | --- | +| a path | the file, byte for byte, under `config_source` | +| a mapping | exactly as given | + +The file is the run's real configuration — comments, key order and section +grouping included — so it is stored verbatim, not parsed. A `str` or `Path` is +always a path; use `pr.ConfigSource(text=...)` for contents already in memory. +To send structured values *and* the file, put the file under +`pr.CONFIG_SOURCE_KEY` in the mapping: + +```python +config = {**cfg.model_dump(exclude_unset=True), + pr.CONFIG_SOURCE_KEY: pr.ConfigSource.from_file("eval.toml").to_dict()} +``` + +Nothing is redacted — keep credentials in the environment, not in the file. + +## Modes + +| mode | what happens | +| --- | --- | +| `online` | the run lives on the platform (default when an API key is present) | +| `disabled` | every call is a no-op, with the same object shape | + +Set the mode explicitly, or through `$PRIME_RUNS_MODE`. A missing API key +disables the run with a warning — it never silently writes somewhere else. + +## What the run handle does for you + +- **Streams instead of buffering.** Records go out as they are produced. +- **Contains its own errors.** With the default `on_error="warn"`, nothing the + platform raises escapes into your loop. Use `on_error="raise"` in tests and + CI, where a silent upload failure is the bug; failures surface from `flush()` + and `finish()`. Platform errors are the `prime_traces` exception family + (`pr.APIError` and friends), so one set of `except` clauses covers both SDKs. +- **Applies backpressure.** The upload queue is bounded; a producer that + durably outruns the uploader has records dropped and counted + (`run.dropped_records`) rather than stalled. Per-sink losses are in + `run.failed_records`. +- **Waits for its own uploads.** `finish()` gives queued records the same + budget a single upload gets (300s) and warns if they do not drain. +- **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()` / `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. + +## Configuration + +Resolved from environment variables first, then `~/.prime/config.json`: + +| setting | env var | default | +| --- | --- | --- | +| API key | `PRIME_API_KEY` | — | +| team | `PRIME_TEAM_ID` | — | +| platform API | `PRIME_API_BASE_URL` | `https://api.primeintellect.ai` | +| dashboard | `PRIME_FRONTEND_URL` | `https://app.primeintellect.ai` | +| traces service | `PRIME_TRACES_URL` | resolved by `prime-traces` | + +Or pass `api_key=`, `base_url=`, `team_id=` to `init()`. + +## Transports + +An online run writes every record to two sinks: Prime Traces (the system of +record — streaming, episode-aware, content-addressed and therefore idempotent +on retry) and the flat v0 sample table today's viewer reads. Both run because +Prime Traces is gated to an allowlist; an account outside it has the traces sink +turn itself off at the first upload — not counted as a failure, since nothing was +lost. When the viewer reads traces natively the default sink list drops one entry +and no producer changes. + +`prime_runs.projection` holds `trace_to_sample` / `build_samples`, the v0 +projection moved here from verifiers. `prime_runs.metrics.from_episodes` is the +run-level aggregation the eval dashboard reads. Both are duck-typed; no producer +package is imported — this is a leaf package, because the `prime` CLI depends on +`verifiers` and verifiers depends on this. + +## Status + +Eval runs only. Training runs will arrive with a backend over +`/api/v1/rft/external-runs`, designed against prime-rl's actual needs. + +There is currently no producer-facing way to mark an evaluation **failed**. The +SDK records the terminal state under `metadata.prime_runs` and warns that the +run will keep showing as running on the dashboard. diff --git a/packages/prime-runs/pyproject.toml b/packages/prime-runs/pyproject.toml new file mode 100644 index 000000000..3f895ac21 --- /dev/null +++ b/packages/prime-runs/pyproject.toml @@ -0,0 +1,74 @@ +[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. +# +# 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", +] +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", + "Programming Language :: Python :: 3.13", + "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..fe66244e2 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/__init__.py @@ -0,0 +1,60 @@ +"""Prime Intellect Runs SDK. + +Track eval runs on the Prime platform:: + + import prime_runs as pr + + run = pr.init(name="gsm8k-qwen3-8b", environments=["gsm8k"], model=model) + print(run.url) + + for episode in rollouts: # stamp run.id onto the traces + run.log_traces([episode]) + + run.finish(summary=pr.metrics.from_episodes(episodes)) + +This is a leaf package: the ``prime`` CLI depends on ``verifiers``, so verifiers +can never depend on ``prime``. Nothing here imports a producer package; records +are duck-typed through ``to_record()``. +""" + +from . import metrics, projection +from .exceptions import ( + APIError, + ConfigurationError, + EnvironmentResolutionError, + ForbiddenError, + NotFoundError, + PaymentRequiredError, + PrimeRunsError, + RetryableAPIError, + RunFinishedError, + TransportError, + UnauthorizedError, +) +from .models import CONFIG_SOURCE_KEY, ConfigSource, EnvironmentRef, RunStatus +from .run import MODE_ENV, Run, init + +__version__ = "0.1.0" + +__all__ = [ + "init", + "Run", + "RunStatus", + "ConfigSource", + "CONFIG_SOURCE_KEY", + "EnvironmentRef", + "MODE_ENV", + "metrics", + "projection", + "PrimeRunsError", + "ConfigurationError", + "EnvironmentResolutionError", + "APIError", + "RunFinishedError", + "ForbiddenError", + "NotFoundError", + "PaymentRequiredError", + "RetryableAPIError", + "TransportError", + "UnauthorizedError", +] 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..ff2eefb2f --- /dev/null +++ b/packages/prime-runs/src/prime_runs/_fork.py @@ -0,0 +1,49 @@ +"""One process-wide ``os.register_at_fork`` hook, shared by everything stateful. + +A forked child inherits copies of every open socket and buffered file handle. +Anything holding a connection or a file registers here and gets told to start +over in the child. Two rules for ``reset_after_fork``: + +- **Drop, do not close.** Closing an inherited transport can send bytes down a + socket the parent is still using. +- **Do not flush.** An inherited buffer holds records the parent will write. + +Registration is weak and the hook is installed once: ``register_at_fork`` +cannot be undone, so per-instance registration would pin every run forever. +""" + +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 new file mode 100644 index 000000000..713ebed65 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/_http.py @@ -0,0 +1,177 @@ +"""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 +import sys +import time +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 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) +#: 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 + + +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 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.""" + 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: + # Strip a trailing /api/v1 — the same normalization prime_traces applies — + # so an explicit base_url written with the suffix does not double it. + self.base_url = base_url.rstrip("/").removesuffix("/api/v1") + self.api_prefix = f"{self.base_url}/api/v1" + self.max_attempts = max(1, max_attempts) + self._owns_client = client is None + self._headers = { + "Authorization": f"Bearer {api_key}", + "User-Agent": _user_agent(), + } + self._timeout = timeout + self._client = client or self._new_client() + if self._owns_client: + # Only a pool we opened ourselves is ours to rebuild after a fork. + _fork.register(self) + + def _new_client(self) -> httpx.Client: + return httpx.Client( + headers=dict(self._headers), + follow_redirects=True, + timeout=self._timeout, + ) + + def reset_after_fork(self) -> None: + """Rebuild the pool in a forked child. The old one is dropped, not + closed: closing could send ``close_notify`` on the parent's socket.""" + self._client = self._new_client() + + 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, + idempotent: Optional[bool] = None, + ) -> Dict[str, Any]: + """Send one request, retrying transient failures. Returns the JSON body. + + ``idempotent`` defaults to ``method != "POST"``; a POST that is safe to + replay (get-or-create) passes ``idempotent=True`` explicitly. + """ + url = f"{self.api_prefix}{path}" + body = content if content is not None else (encode_json(json_body) if json_body else None) + request_kwargs: Dict[str, Any] = { + "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" + + for attempt in range(attempts): + error: APIError + try: + response = self._client.request(method, url, **request_kwargs) + except httpx.TimeoutException as exc: + error = APITimeoutError(f"{method} {path} timed out: {exc}") + ambiguous = isinstance(exc, AMBIGUOUS_TRANSPORT_ERRORS) + except httpx.RequestError as exc: + error = TransportError(f"{method} {path} failed: {type(exc).__name__}: {exc}") + ambiguous = isinstance(exc, AMBIGUOUS_TRANSPORT_ERRORS) + else: + try: + raise_for_response(response) + except RetryableAPIError as exc: + error = exc + ambiguous = exc.status_code not in UNAMBIGUOUS_RETRY_STATUS + else: + return _decode(response) + + last = attempt == attempts - 1 + if last or (ambiguous and not replayable): + # Possibly processed already; a duplicate cannot be undone. + 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) + + 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 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} diff --git a/packages/prime-runs/src/prime_runs/backend.py b/packages/prime-runs/src/prime_runs/backend.py new file mode 100644 index 000000000..68e226c55 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/backend.py @@ -0,0 +1,269 @@ +"""Run backends: the contract, the evaluations backend, and the disabled no-op. + +A backend owns the *lifecycle* of a run — creating it, updating what is known +about it, closing it out with a terminal status. It does not move records; +that is a sink's job (see :mod:`prime_runs.sinks`). + +:class:`EvalsBackend` works over ``/api/v1/evaluations/*``, resolving +environments through the hub's get-or-create so a local run uploads without +``prime env push``. The eval API has no producer-facing way to mark a run +failed: ``finalize`` moves a run to COMPLETED and ``UpdateEvaluationRequest`` +carries no status. A failed or crashed run therefore keeps showing as running; +the terminal state is recorded in ``metadata.prime_runs`` so it is at least +visible. +""" + +import logging +import uuid +from datetime import datetime, timezone +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 + +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.""" + + def __init__( + self, + client: PlatformClient, + *, + frontend_url: str, + team_id: Optional[str] = None, + ) -> None: + self._client = client + self._frontend_url = frontend_url.rstrip("/") + self._team_id = team_id + + # ------------------------------------------------------------------ create + + def create(self, spec: RunSpec) -> RunHandle: + 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().' + ) + + run_name: str = spec.name or _default_name(spec) + payload: Dict[str, Any] = { + "name": run_name, + "environments": environments, + "tags": list(spec.tags), + } + _set_if(payload, "model_name", spec.model) + # The API's `dataset` column is always the environment under another name. + _set_if(payload, "dataset", _first_environment_name(spec)) + _set_if(payload, "framework", spec.framework) + _set_if(payload, "description", spec.description) + _set_if(payload, "metadata", spec.config or None) + _set_if(payload, "team_id", spec.team_id or self._team_id) + + # Not replayable (the POST default): a retry after a lost response would + # create a second run and only the second would be tracked. + response = self._client.post("/evaluations/", json_body=payload) + run_id = response.get("evaluation_id") + if not run_id: + raise APIError( + 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), + ) + + 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: + """``config`` must be the whole config: the service writes metadata with + a document-level ``$set``.""" + payload: Dict[str, Any] = {} + _set_if(payload, "metadata", config or None) + _set_if(payload, "metrics", summary or None) + if not payload: + return + self._client.put(f"/evaluations/{run_id}", json_body=payload) + + # ---------------------------------------------------------------- finalize + + def finalize( + self, + run_id: str, + *, + status: RunStatus, + summary: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, + ) -> None: + if status is RunStatus.COMPLETED: + body: Dict[str, Any] = {} + _set_if(body, "metrics", summary or None) + # Finalization enqueues the platform's statistics task; replaying an + # ambiguous failure could enqueue it twice, so this stays non-idempotent. + self._client.post( + f"/evaluations/{run_id}/finalize", + json_body=body or {"metrics": {}}, + ) + return + + # No status endpoint exists yet. Record the terminal state in metadata, + # merged into the full config because this PUT replaces the document. + terminal = { + "status": status.value, + "finished_at": datetime.now(timezone.utc).isoformat(), + } + if error: + terminal["error"] = error + self.update(run_id, config={**(config or {}), "prime_runs": terminal}, summary=summary) + logger.warning( + "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, + ) + + def close(self) -> None: + self._client.close() + + # ----------------------------------------------------------- environments + + def _resolve_environments(self, refs: List[EnvironmentRef]) -> List[Dict[str, Any]]: + """Hub references as the API's ``EnvironmentReference`` objects, carrying + a pinned ``version_id`` through. A reference that cannot be resolved + raises rather than being skipped: a run silently attached to the wrong + environments looks like a successful upload.""" + resolved: List[Dict[str, Any]] = [] + for ref in refs: + 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: + if ref.slug: + owner_slug, name = ref.slug.split("/", 1) + try: + response = self._client.get(f"/environmentshub/{owner_slug}/{name}/@latest") + except APIError 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: + # Get-or-create: a replay returns the same environment. + response = self._client.post( + "/environmentshub/resolve", json_body=body, idempotent=True + ) + except APIError 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: + payload[key] = value + + +def _first_environment_name(spec: RunSpec) -> Optional[str]: + for ref in spec.environments: + if ref.name: + return ref.name + if ref.slug: + return ref.slug.split("/", 1)[1] + return None + + +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/config.py b/packages/prime-runs/src/prime_runs/config.py new file mode 100644 index 000000000..c25440a1f --- /dev/null +++ b/packages/prime-runs/src/prime_runs/config.py @@ -0,0 +1,18 @@ +"""Configuration: ``prime_traces.core.Config`` (``~/.prime/config.json`` plus +environment variables, env taking precedence) with the dashboard URL added.""" + +import os + +from prime_traces.core import Config as _TracesConfig + + +class Config(_TracesConfig): + DEFAULT_FRONTEND_URL: str = "https://app.primeintellect.ai" + + @property + def frontend_url(self) -> str: + """Dashboard base URL; fallback when a create response omits ``viewer_url``.""" + env_val = os.getenv("PRIME_FRONTEND_URL") + if env_val: + return env_val.rstrip("/") + 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..adc5df8b9 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/exceptions.py @@ -0,0 +1,62 @@ +"""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 for errors raised by the SDK itself, before any request is made.""" + + +class ConfigurationError(PrimeRunsError): + """Missing API key, unreadable config file, unknown mode.""" + + +class EnvironmentResolutionError(PrimeRunsError): + """An environment named in ``init()`` could not be resolved to a hub ID — + usually a typo or a permissions problem, not an outage.""" + + +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/metrics.py b/packages/prime-runs/src/prime_runs/metrics.py new file mode 100644 index 000000000..ee0f12670 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/metrics.py @@ -0,0 +1,49 @@ +"""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, TypedDict + + +class RunSummary(TypedDict): + """The run-level aggregates the eval dashboard reads. A summary may carry + more than this — ``finish(summary=...)`` stores whatever it is given — but + these three are what the dashboard renders.""" + + avg_reward: float + avg_metrics: Dict[str, float] + avg_error: float + + +def from_episodes(episodes: Sequence[Any], traces: Optional[Sequence[Any]] = None) -> RunSummary: + """Run-level aggregates in the shape the eval dashboard reads. + + Rewards and metrics aggregate over the trainable traces only — fixed agents + (a judge, a modeled user) often carry no rewards and would dilute every mean + with structural zeros — falling back to all traces when none are trainable, + the same rule the dashboard applies. ``avg_error`` is the share of EPISODES + that aren't ok: a hook failure counts even when its traces are clean or it + left none behind. + """ + 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..cf27c838c --- /dev/null +++ b/packages/prime-runs/src/prime_runs/models.py @@ -0,0 +1,213 @@ +"""Types shared across backends, sinks and the ``Run`` handle. + +Response bodies are deliberately not modeled: backends pull the two or three +fields they need and hand back a ``RunHandle``. +""" + +import os +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from typing import Any, Dict, List, Literal, Mapping, Optional, Union + +from .exceptions import ConfigurationError + +Mode = Literal["online", "disabled"] +"""``online`` talks to the platform; ``disabled`` makes every call a no-op with +the same object shape.""" + +OnError = Literal["warn", "raise"] + +RUN_KIND = "eval" +"""Stamped as ``run.type`` on records and sent as upload provenance.""" + + +class RunStatus(str, Enum): + """Terminal state a producer can report. + + ``failed`` means the producer said the run failed; ``crashed`` means the + process exited without saying (only the SDK's atexit hook reports it). + """ + + 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; ``slug`` looks up a 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(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"), + ) + 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.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") + + +CONFIG_SOURCE_KEY = "config_source" +"""Where a config file lands inside a run's config, as a :class:`ConfigSource` +dict. Present means "render this verbatim"; absent means "render the structure".""" + +MAX_CONFIG_SOURCE_BYTES = 256 * 1024 +"""A hand-written run config is kilobytes; anything past this is refused at +``init()`` rather than truncated.""" + +_CONFIG_SOURCE_FORMATS = { + ".toml": "toml", + ".json": "json", + ".yaml": "yaml", + ".yml": "yaml", +} + + +@dataclass +class ConfigSource: + """The config file a run was started from, kept byte-for-byte. + + That file *is* the run's configuration — comments, key order and section + grouping included — where a resolved model dump is a different artifact. + Nothing here is redacted: keep credentials in the environment, not the file. + """ + + 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={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={str(path)!r} could not be read: {exc}") from exc + if len(raw) > MAX_CONFIG_SOURCE_BYTES: + raise ConfigurationError( + 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={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 the config-file form of ``init(config=...)``. + + 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 + if isinstance(value, Mapping): + source = cls.from_mapping(value) + if source is None: + raise ValueError("a config-source mapping must contain a 'text' string") + return source + if isinstance(value, (str, os.PathLike)): + return cls.from_file(value) + raise TypeError( + "a config source must be a path, a ConfigSource or a mapping, " + f"got {type(value).__name__}" + ) + + 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"the 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: ``init()``'s arguments after + normalization. ``config`` is the run's inputs; outputs accumulate on the + handle as ``summary``.""" + + name: Optional[str] = None + 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 + config: 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 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..1474bb3ae --- /dev/null +++ b/packages/prime-runs/src/prime_runs/projection.py @@ -0,0 +1,181 @@ +"""Projection of native traces onto the platform's v0 eval-sample format. + +Moved here from ``verifiers.v1.utils.platform``: this is knowledge about a +platform wire format, so it lives in the client for that wire. Duck-typed — +verifiers ``Trace``/``Episode`` satisfy it structurally and are not imported. + +The projection serves the *current* viewer, which reads the flat sample table; +it leaves the default path once the Viewer API reads traces natively. +""" + +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 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..758d0d3c1 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/run.py @@ -0,0 +1,526 @@ +"""The run handle, and ``init()`` that produces one. + +A run is a long-lived thing with a status, so it is an object rather than three +stateless calls: records stream instead of buffering, errors are contained, and +a process that dies still reports a terminal status. + +``init()`` is called *before* rollouts start, and the ID it returns is *the* run +ID everywhere — including inside every trace document the producer writes. +Nothing is re-stamped afterwards. +""" + +import atexit +import logging +import math +import os +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 .backend import Backend, DisabledBackend, EvalsBackend, disabled_run_id +from .config import Config +from .exceptions import ConfigurationError, RunFinishedError +from .models import ( + CONFIG_SOURCE_KEY, + RUN_KIND, + ConfigSource, + EnvironmentRef, + Mode, + OnError, + RunHandle, + RunSpec, + RunStatus, +) +from .sinks import EvalSamplesSink, Sink, TracesSink +from .worker import UploadWorker + +logger = logging.getLogger(__name__) + +MODE_ENV = "PRIME_RUNS_MODE" +#: How long ``finish()`` gives queued uploads to drain before finalizing anyway. +#: Derived from the upload timeout: a single in-flight sample POST may take +#: this long, and a shorter budget would abandon an upload about to succeed. +DEFAULT_FINISH_TIMEOUT = float(UPLOAD_TIMEOUT.read or 300.0) + + +class Run: + """A live run: an ID, a URL, somewhere to put traces, a summary. + + With the default ``on_error="warn"`` nothing raised by the platform escapes + into a producer's loop; ``on_error="raise"`` surfaces the first failure + from :meth:`flush` or :meth:`finish`, for tests and CI. + """ + + def __init__( + self, + *, + backend: Backend, + handle: RunHandle, + spec: RunSpec, + sinks: Optional[List[Sink]] = None, + mode: Mode = "online", + on_error: OnError = "warn", + ) -> None: + self._backend = backend + self._handle = handle + self._spec = spec + self._mode: Mode = mode + self._on_error: OnError = on_error + self._status = RunStatus.RUNNING + # A forked child inherits this handle but must not close the parent's run. + self._owns_lifecycle = True + + self.config: Dict[str, Any] = dict(spec.config) + self.summary: Dict[str, Any] = {} + self.errors: List[str] = [] + # Under on_error="raise", a failure on the uploader thread is held here + # and re-raised from flush() or finish(), where the caller is looking. + self._deferred_error: Optional[BaseException] = None + + self._finish_timeout = DEFAULT_FINISH_TIMEOUT + self._finish_lock = threading.Lock() + self._finishing = False + self._finished = False + self._atexit_hook = self._on_process_exit + _fork.register(self) + + sinks = sinks or [] + 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 + self._note(f"starting sink {sink.name}", exc) + if self._on_error == "raise": + # The backend may already have created a remote run; close + # it out before the failure reaches the caller. finish() + # re-raises the error noted above. + self.finish(status=RunStatus.FAILED, error=_describe(exc)) + raise exc + + atexit.register(self._atexit_hook) + + # -------------------------------------------------------------- identity + + @property + def id(self) -> str: + """The run ID. Stamp this onto every trace the run produces.""" + return self._handle.id + + @property + def name(self) -> Optional[str]: + return self._handle.name + + @property + def url(self) -> Optional[str]: + """The dashboard URL; ``None`` when disabled.""" + return self._handle.url + + @property + def config_source(self) -> Optional[ConfigSource]: + """The config file this run was launched from, if one was given.""" + raw = self.config.get(CONFIG_SOURCE_KEY) + return ConfigSource.from_mapping(raw) if isinstance(raw, Mapping) else None + + @property + def mode(self) -> Mode: + return self._mode + + @property + def status(self) -> RunStatus: + return self._status + + @property + def finished(self) -> bool: + return self._finished + + @property + def dropped_records(self) -> int: + """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. Per sink because + another sink may still hold them.""" + return dict(self._worker.failed_records) + + def __repr__(self) -> str: + return f"" + + # ------------------------------------------------------------------- log + + def log_traces(self, traces: Iterable[Any]) -> None: + """Hand bare traces to the sinks. Returns immediately. + + Accepts verifiers ``Trace`` objects or plain JSON mappings. 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.""" + flushed = self._worker.flush(timeout=timeout) + self._raise_deferred() + return flushed + + # ---------------------------------------------------------------- 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: 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 + try: + self._finish_once(summary, resolved, error) + finally: + self._finishing = False + self._finished = True + + def _finish_once( + self, + summary: Optional[Mapping[str, Any]], + resolved: RunStatus, + error: Optional[str], + ) -> None: + if summary: + self.summary.update(_clean_metrics(summary)) + self._status = resolved + + deadline = time.monotonic() + max(0.0, self._finish_timeout) + + def remaining() -> float: + return max(0.0, deadline - time.monotonic()) + + # Records first, so a dashboard reacting to the terminal status never + # sees a finished run with samples still landing. + 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()) + + if self._owns_lifecycle: + self._teardown_step( + "updating the run", + lambda: self._backend.update( + self.id, config=self.config or None, summary=self.summary or None + ), + ) + self._teardown_step( + "finalizing the run", + lambda: self._backend.finalize( + self.id, + status=resolved, + summary=self.summary or None, + error=error or (self.errors[0] if self.errors else None), + config=self.config or None, + ), + ) + self._teardown_step("closing the backend", self._backend.close) + atexit.unregister(self._atexit_hook) + + if self._worker.dropped: + logger.warning( + "Run %s finished with %d record(s) that reached no sink; the producer " + "outran the uploader.", + self.id, + self._worker.dropped, + ) + for sink_name, count in self._worker.failed_records.items(): + logger.warning( + "Run %s: the %s sink could not store %d record(s)", self.id, sink_name, count + ) + # Last, so a run whose teardown failed is still closed out first. + self._raise_deferred() + + def fail(self, error: Union[str, BaseException]) -> None: + """Close the run out as failed.""" + self.finish(status=RunStatus.FAILED, error=_describe(error)) + + # -------------------------------------------------------- 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() + return False + + if isinstance(exc, KeyboardInterrupt): + # An interrupt is a decision, not a fault. + status, error = RunStatus.CRASHED, "interrupted" + else: + status, error = RunStatus.FAILED, _describe(exc) + + try: + self.finish(status=status, error=error) + except Exception as finish_error: + # The producer's exception is why this block is unwinding; a + # telemetry teardown error must not replace it. Control-flow + # exceptions (KeyboardInterrupt, SystemExit) deliberately pass. + 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 + + def reset_after_fork(self) -> None: + """Make an inherited handle safe in a forked child: fresh lock, and the + parent keeps ownership of the lifecycle so the child's atexit hook + cannot finalize a still-running run.""" + self._finish_lock = threading.Lock() + self._finishing = False + self._owns_lifecycle = False + self._deferred_error = None + + def _on_process_exit(self) -> None: + """The process is exiting and nobody called ``finish()``: report CRASHED + rather than FAILED, since the producer never said the run failed.""" + if self._finished: + return + logger.warning("Run %s was never finished; reporting it as crashed", self.id) + try: + self.finish(status=RunStatus.CRASHED, error="process exited without finishing the run") + except Exception as exc: # noqa: BLE001 - never a traceback from atexit + logger.warning("Run %s: reporting the crash failed: %s", self.id, exc) + + def _require_live(self, operation: str) -> None: + 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." + ) + + def _record_sink_error(self, sink_name: str, exc: Exception) -> None: + """Called on the uploader thread when a sink gives up on a batch.""" + self._note(f"writing to the {sink_name} sink", exc) + + def _teardown_step(self, what: str, call: Any) -> None: + """Run one teardown step without letting it skip the steps after it.""" + try: + call() + except Exception as exc: # noqa: BLE001 - teardown must continue + self._note(what, 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 + else: + logger.warning("Run %s: %s", self.id, message) + + def _raise_deferred(self) -> None: + """Re-raise the first held failure, once.""" + exc = self._deferred_error + if exc is None: + return + self._deferred_error = None + raise exc + + +# --------------------------------------------------------------------- init + + +def init( + *, + name: Optional[str] = None, + 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, + mode: Optional[Mode] = None, + team_id: Optional[str] = None, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + on_error: OnError = "warn", +) -> Run: + """Start a run and return a handle to it. + + Call this *before* the first rollout: the ID it returns is what every trace + in the run should carry, and the URL is what a producer prints. + + ``mode`` defaults to ``$PRIME_RUNS_MODE``, else online when there is an API + key and disabled (with a warning) when there is not. ``config`` is what the + run was configured with: the path to the file it was launched from (stored + byte for byte under ``config_source``), or a mapping taken as given. + """ + 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, + environments=[EnvironmentRef.coerce(entry) for entry in (environments or [])], + model=model, + framework=framework, + description=description, + tags=list(tags or []), + team_id=team_id, + config=_normalize_config(config), + ) + resolved_mode = _resolve_mode(mode, api_key=api_key) + + backend: Backend + sinks: List[Sink] + if resolved_mode == "disabled": + backend = DisabledBackend() + handle = RunHandle(id=disabled_run_id(), name=name) + sinks = [] + else: + if not api_key: + raise ConfigurationError( + 'mode="online" needs an API key. Set PRIME_API_KEY, run `prime login`, ' + '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) + 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. + sinks = [TracesSink(api_key=api_key, team_id=team_id), EvalSamplesSink(client)] + + run = Run( + backend=backend, + handle=handle, + spec=spec, + sinks=sinks, + mode=resolved_mode, + on_error=on_error, + ) + if run.url: + logger.info("Run %s: %s", run.id, run.url) + return run + + +def _resolve_mode(mode: Optional[Mode], *, api_key: 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", "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`); " + "this run will not be tracked." + ) + mode = "disabled" + return mode # type: ignore[return-value] + + +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: + context["model"] = spec.model + return context + + +def _normalize_config(value: Any) -> Dict[str, Any]: + """A producer's config as a plain dict. + + A path is the file the run was launched from, kept byte for byte under + ``config_source`` (stored, not parsed — the platform can read TOML). A + mapping is taken exactly as given. + """ + if value is None: + return {} + if isinstance(value, Mapping): + 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()} + raise TypeError( + f"config must be a path to the run's config file or a mapping, got {type(value).__name__}" + ) + + +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 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): + 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", "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..6f5a7ffdf --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/__init__.py @@ -0,0 +1,14 @@ +"""Record transports. Independent of backends, and of each other.""" + +from .base import Sink, is_episode, stamp_run, to_mapping +from .samples import EvalSamplesSink +from .traces import TracesSink + +__all__ = [ + "Sink", + "is_episode", + "stamp_run", + "to_mapping", + "EvalSamplesSink", + "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..fee20ccbd --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/base.py @@ -0,0 +1,83 @@ +"""The contract a record sink implements, plus the record helpers sinks share. + +A sink moves records — traces, episodes — to wherever they are stored. It knows +nothing about run lifecycle. Sinks are independent of each other: during the +transition the traces sink and the legacy eval-samples sink both run, and +retiring one is a change to the default sink list, not to any producer. + +Every sink must be degradable: one that cannot write sets ``enabled = False`` +and says why, once. +""" + +from typing import Any, Dict, Mapping, Protocol, Sequence + +from ..models import RUN_KIND + + +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]) -> 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: 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) + 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()") + + +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`` 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/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py new file mode 100644 index 000000000..d99b0d719 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -0,0 +1,76 @@ +"""Legacy sink: the flat eval-sample table behind today's viewer. + +Prime Traces is gated to an allowlist and the viewer reads the v0 sample table, +so both sinks run until the viewer reads traces natively; retiring this one is +a change to the default sink list. + +``POST /samples`` appends, so a request whose response was lost is not +replayed: a lost batch is recoverable, duplicated rows skew every average. +""" + +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, is_episode + +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 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]) -> 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, + 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]]: + """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) 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" + ) + return samples + + def flush(self) -> None: + """Writes are synchronous; the uploader thread owns the asynchrony.""" + + def close(self) -> None: + """The client is shared with the backend, which closes it.""" diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py new file mode 100644 index 000000000..79b2831f9 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -0,0 +1,160 @@ +"""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 prime_traces import ErrorCode, LineFormat, TracesClient + +from .. import _fork +from ..exceptions import ForbiddenError +from .base import Sink, is_episode, stamp_run, to_mapping + +logger = logging.getLogger(__name__) + +DEFAULT_RECEIPT_HISTORY_SIZE = 100 + + +class TracesSink(Sink): + """Uploads records through the Prime Traces service.""" + + name = "traces" + + def __init__( + self, + *, + client: Optional[Any] = None, + api_key: Optional[str] = None, + team_id: Optional[str] = None, + 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 + # 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 team_id is not None: + self._client_kwargs["team_id"] = team_id + self._compress = compress + self._run_id: 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 + + def start(self, run_id: str, context: Mapping[str, str]) -> None: + self._run_id = run_id + self._context = {key: str(value) for key, value in context.items() if value is not None} + self._ensure_client() + + def _ensure_client(self) -> bool: + """Build the traces client lazily, so a fork reset is repaired on next write.""" + if self._client is not None: + return True + if self._injected_client: + exc = RuntimeError("an injected traces client cannot be reused after a fork") + self._disable(str(exc)) + raise exc + try: + self._client = TracesClient(**self._client_kwargs) + except Exception as exc: # noqa: BLE001 - the run applies its error policy + self._disable(f"could not construct the traces client ({exc})") + raise + return True + + def reset_after_fork(self) -> None: + """Drop (not close) the inherited client; its socket is the parent's.""" + self._client = None + + # ------------------------------------------------------------------ write + + def write(self, records: Sequence[Any]) -> None: + if not self.enabled or not records or not self._ensure_client(): + return + + # The same bytes under a different format are rejected as a conflict, + # so infer from the first record, which is stable within a batch. + 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=line_format, + context=dict(self._context) or None, + compress=self._compress, + ) + ) + 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}); " + "continuing with the remaining sinks" + ) + return + # A credential without the traces scope is something the caller + # can fix; raised so loss accounting and strict callers see it. + self._disable(f"this credential cannot write traces ({exc})") + raise + 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: + """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(to_mapping(record), self._run_id) + + def flush(self) -> None: + """Uploads are synchronous; nothing is held back here.""" + + def close(self) -> None: + client = self._client + self._client = None + if client is not None: + 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 + + def _retire_quietly(self, reason: str) -> None: + if self.enabled: + logger.info("Traces sink off: %s", reason) + self.enabled = False + + +def _is_not_enabled(exc: ForbiddenError) -> bool: + """``service_not_enabled``: the account is outside the private beta. The + other 403, ``forbidden``, means the token lacks the ``traces`` scope.""" + return exc.code == ErrorCode.SERVICE_NOT_ENABLED.value diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py new file mode 100644 index 000000000..3391e66a9 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -0,0 +1,255 @@ +"""Background uploader: one daemon thread draining a bounded queue into sinks. + +Backpressure: the queue is bounded, so a producer that outruns the uploader +blocks briefly and then drops (counted) rather than stalling the run. + +Fork safety: a forked child inherits the queue's memory but not the thread. +The child starts over empty; the queued records belong to the parent. + +Containment: 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 queue +import threading +import time +from dataclasses import dataclass, field +from typing import Any, Callable, Optional, Sequence + +from . import _fork +from .exceptions import is_transient +from .sinks.base import Sink + +logger = logging.getLogger(__name__) + +DEFAULT_QUEUE_SIZE = 256 +DEFAULT_PUT_TIMEOUT = 5.0 +#: Consecutive transient failures before a sink is retired. +TRANSIENT_FAILURE_LIMIT = 3 + + +@dataclass +class _Flush: + """A barrier the caller waits on.""" + + 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 of record batches into a list of sinks.""" + + def __init__( + self, + sinks: Sequence[Sink], + *, + max_queue_size: int = DEFAULT_QUEUE_SIZE, + put_timeout: float = DEFAULT_PUT_TIMEOUT, + on_error: Optional[Callable[[str, Exception], None]] = None, + ) -> None: + self.sinks = sinks + self.max_queue_size = max_queue_size + self.put_timeout = put_timeout + self._on_error = on_error + 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() + #: 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 = {} + #: 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 + + 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() + 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, records: Sequence[Any]) -> None: + for sink in self.sinks: + 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, None) + + def _flush_sinks(self) -> None: + for sink in self.sinks: + if not sink.enabled: + continue + try: + sink.flush() + except Exception as exc: # noqa: BLE001 + self._fail_sink(sink, exc) + + def _fail_sink(self, sink: Sink, exc: Exception, *, dropped: int = 0) -> None: + """The batch is gone (the transports already retried). Decide whether + the sink is too: permanent failures retire it at once, transient ones + after ``TRANSIENT_FAILURE_LIMIT`` consecutive strikes.""" + name = sink.name + 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 + 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 + self._retired.add(name) + logger.warning( + "Sink %s disabled after %d consecutive transient failures: %s: %s", + name, + strikes, + type(exc).__name__, + exc, + ) + 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) + + 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 + + def submit(self, records: Sequence[Any]) -> bool: + """Hand a batch to the uploader. ``False`` means it was dropped: the + queue stayed full for the whole put timeout.""" + if self._stopping.is_set(): + return False + thread = self._thread + if thread is None or not thread.is_alive(): + self.start() + try: + self._queue.put(records, timeout=self.put_timeout) + return True + except queue.Full: + count = len(records) + self.dropped += count + logger.warning( + "Upload queue full after %.1fs; dropped %d item(s) (%d total). " + "The producer is outrunning the uploader.", + self.put_timeout, + count, + self.dropped, + ) + return False + + 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 + deadline = _deadline(timeout) + barrier = _Flush() + try: + # The barrier gets the caller's drain budget, not the short + # producer put timeout: a full queue is when finish() most needs it. + self._queue.put(barrier, timeout=_remaining(deadline)) + except queue.Full: + logger.warning("Could not enqueue a flush barrier before the drain timeout") + return False + 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=_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 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.", + timeout, + ) + return + 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", sink.name, exc) + + # ------------------------------------------------------------------- fork + + def reset_after_fork(self) -> None: + """Give the child a clean uploader; what was queued belongs to the parent.""" + self._queue = queue.Queue(maxsize=self.max_queue_size) + self._thread = None + self._stopping = threading.Event() + 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..9c7166c9c --- /dev/null +++ b/packages/prime-runs/tests/conftest.py @@ -0,0 +1,143 @@ +"""Shared fixtures. Every test is hermetic: no network, no real ~/.prime.""" + +from pathlib import Path +from typing import Any, Callable, Dict, List, Mapping, Sequence + +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_RUNS_MODE", +) + + +@pytest.fixture(autouse=True) +def isolated_prime_config(monkeypatch, tmp_path): + """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) + 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: Mapping[str, str]) -> None: + self.started.append((run_id, dict(context))) + + def write(self, records: Sequence[Any]) -> None: + if self.fail_on_write: + raise RuntimeError("sink is broken") + self.batches.append(list(records)) + + 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_config_source.py b/packages/prime-runs/tests/test_config_source.py new file mode 100644 index 000000000..a0120afe1 --- /dev/null +++ b/packages/prime-runs/tests/test_config_source.py @@ -0,0 +1,227 @@ +"""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 + +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 +""" + + +# ------------------------------------------------------ the config-file form + + +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 + + +def test_a_mapping_is_taken_exactly_as_given(): + """The caller already decided what to say; second-guessing it would be worse.""" + assert _normalize_config({"a": 1, "b": None}) == {"a": 1, "b": None} + + +def test_a_mapping_is_copied_not_aliased(): + original = {"a": 1} + assert _normalize_config(original) is not original + + +def test_no_config_is_an_empty_config(): + assert _normalize_config(None) == {} + + +def test_an_unusable_config_says_what_was_expected(): + with pytest.raises(TypeError, match="config must be a path"): + _normalize_config(object()) + + +def test_a_path_becomes_a_stored_source(tmp_path): + """One parameter, three forms — the same shape ``environments=`` already has.""" + path = tmp_path / "eval.toml" + path.write_text(EVAL_TOML) + + 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 + + +@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, handler = online(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" + assert metadata[CONFIG_SOURCE_KEY]["filename"] == "eval.toml" + + +def test_extra_values_can_be_merged_onto_a_launch_file(tmp_path, online): + """A run launched from a file that also wants structured values passes a + mapping carrying the source under ``CONFIG_SOURCE_KEY`` — what verifiers does.""" + path = tmp_path / "eval.toml" + path.write_text(EVAL_TOML) + config = { + "model": "deepseek/deepseek-v4-flash", + CONFIG_SOURCE_KEY: ConfigSource.from_file(path).to_dict(), + } + + run, handler = online(config=config) + 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["model"] == "deepseek/deepseek-v4-flash" + assert metadata[CONFIG_SOURCE_KEY]["text"] == EVAL_TOML + assert run.config_source.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="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(): + run = pr.init(environments=["gsm8k"], mode="disabled") + + assert run.config_source is None + run.finish() + + +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, handler = online(config=path) + run.fail("something broke") + + update = handler.bodies_for("/api/v1/evaluations/eval-abc")[-1] + assert update["metadata"][CONFIG_SOURCE_KEY]["text"] == EVAL_TOML + assert update["metadata"]["prime_runs"]["status"] == "failed" + + +def test_a_spec_defaults_to_no_source(): + assert RunSpec().config == {} diff --git a/packages/prime-runs/tests/test_evals_backend.py b/packages/prime-runs/tests/test_evals_backend.py new file mode 100644 index 000000000..4f9c53ff7 --- /dev/null +++ b/packages/prime-runs/tests/test_evals_backend.py @@ -0,0 +1,213 @@ +"""Eval run lifecycle against ``/api/v1/evaluations/*``.""" + +import httpx +import pytest +from conftest import RecordingHandler + +from prime_runs.backend import EvalsBackend +from prime_runs.exceptions import ( + ConfigurationError, + EnvironmentResolutionError, + RetryableAPIError, +) +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"}] + + +@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_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 +): + """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_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_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/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_update_sends_nothing_when_there_is_nothing_to_send(make_platform_client, eval_routes): + backend, handler = make_backend(make_platform_client, eval_routes) + + backend.update("eval-abc") + + assert handler.requests == [] + + +def test_a_pinned_environment_version_reaches_the_api(make_platform_client, eval_routes): + """The API's EnvironmentReference carries version_id. Dropping it attaches + the run to whatever version the hub resolves today — the difference between + a reproducible eval and one that quietly moved.""" + backend, handler = make_backend(make_platform_client, eval_routes) + + backend.create(RunSpec(name="r", environments=[EnvironmentRef(id="env-1", version_id="v-7")])) + + 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_http.py b/packages/prime-runs/tests/test_http.py new file mode 100644 index 000000000..6d954f84d --- /dev/null +++ b/packages/prime-runs/tests/test_http.py @@ -0,0 +1,244 @@ +"""Transport behaviour: error mapping and retry.""" + +import httpx +import pytest + +from prime_runs._http import PlatformClient, encode_json +from prime_runs.exceptions import ( + APIError, + ForbiddenError, + NotFoundError, + PaymentRequiredError, + RetryableAPIError, + TransportError, + UnauthorizedError, + is_transient, +) + + +def client_for(handler, *, base_url: str = "http://testserver", **kwargs) -> PlatformClient: + return PlatformClient( + api_key="test-key", + base_url=base_url, + client=httpx.Client(transport=httpx.MockTransport(handler)), + **kwargs, + ) + + +@pytest.mark.parametrize( + "status,expected", + [ + (401, UnauthorizedError), + (402, PaymentRequiredError), + (403, ForbiddenError), + (404, NotFoundError), + (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 + 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(): + """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"})) + + 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={"detail": "overloaded"}) + + with pytest.raises(RetryableAPIError) as caught: + client_for(handler, max_attempts=3).get("/evaluations/x") + + assert len(attempts) == 3 + 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): + responses = [httpx.Response(429), httpx.Response(200, json={"ok": True})] + + client = client_for(lambda request: responses.pop(0)) + + assert client.get("/evaluations/x") == {"ok": True} + + +def test_retry_after_is_honoured(no_sleep): + """The schedule itself is ``prime_traces.core.client.retry_delay``; what is + ours is feeding it the server's header.""" + responses = [ + httpx.Response(429, headers={"Retry-After": "7.5"}), + httpx.Response(200, json={"ok": True}), + ] + + 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): + 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_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)) + + 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(APIError, 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")}) + + +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) + + +@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})] + + 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_init.py b/packages/prime-runs/tests/test_init.py new file mode 100644 index 000000000..c42165279 --- /dev/null +++ b/packages/prime-runs/tests/test_init.py @@ -0,0 +1,225 @@ +"""``init()``: mode resolution, disabled runs, online runs.""" + +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 + +# -------------------------------------------------------------------- modes + + +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") + + assert run.mode == "disabled" + assert "not be tracked" in caplog.text + run.finish() + + +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") + + assert run.mode == "disabled" + run.finish() + + +def test_an_unknown_mode_is_rejected(): + with pytest.raises(ConfigurationError, match="not one of"): + pr.init(mode="sideways") + + +def test_a_disabled_run_still_answers_every_call(tmp_path): + """Same object shape, so producer code needs no branching — and nothing + touches the network or the filesystem.""" + run = pr.init(mode="disabled") + + run.log_traces([{"id": "t1"}]) + run.finish(summary={"avg_reward": 1.0}) + + assert run.id.startswith("disabled-") + assert run.url is None + assert run.status is RunStatus.COMPLETED + assert not list(tmp_path.iterdir()) # tmp_path is $HOME here + + +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 + + +@pytest.fixture +def online(monkeypatch, make_platform_client, eval_routes): + """``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", + **kwargs, + ) + return run, handler + + return _init + + +class _NullSink: + name = "traces" + enabled = True + + def start(self, run_id, context): + pass + + def write(self, records): + pass + + def flush(self): + pass + + def close(self): + pass + + +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_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() + + 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 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)] + + 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 == [] + + +def test_a_second_init_in_one_process_opens_its_own_run(): + first = pr.init(mode="disabled") + first.finish() + + second = pr.init(mode="disabled") + second.finish() + + assert second.id != first.id + + +# --------------------------------------------------------------------- fork + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork is POSIX-only") +@pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning") +def test_a_forked_child_does_not_duplicate_the_parents_records_or_close_its_run(online): + """At fork time the parent may have records in the upload queue. The child + inherits a copy; writing them would upload every record twice, and the + inherited atexit hook must not finalize the parent's run.""" + run, handler = online() + run.log_traces([{"sample_id": f"parent-{n}"} for n in range(5)]) + + pid = os.fork() + if pid == 0: # pragma: no cover - asserted through the child's exit code + code = 0 + try: + # ``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() + 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 + finally: + os._exit(code) + + _, status = os.waitpid(pid, 0) + assert os.waitstatus_to_exitcode(status) == 0 + run.finish() + + 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_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() + run.finish() 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..34646fc1d --- /dev/null +++ b/packages/prime-runs/tests/test_projection.py @@ -0,0 +1,144 @@ +"""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, + 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_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 new file mode 100644 index 000000000..9eb6df6d3 --- /dev/null +++ b/packages/prime-runs/tests/test_run.py @@ -0,0 +1,505 @@ +"""The run handle: lifecycle, containment, terminal status.""" + +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: + def __init__(self, fail_on: Optional[str] = None) -> None: + self.fail_on = fail_on + self.updates: List[Dict[str, 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 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 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, "config": config} + ) + + def close(self) -> None: + self.closed = True + + +def make_run(backend=None, sinks=None, config=None, **kwargs) -> Run: + backend = backend or FakeBackend() + spec = RunSpec(name="test-run", framework="verifiers", model="Qwen3-8B", config=config or {}) + 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 + 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"}]) + run.flush() + + assert sink.batches == [[{"id": "t1"}]] + 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_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.""" + backend = FakeBackend() + run = make_run(backend) + + run.finish(summary={"loss": float("nan"), "grad": float("inf"), "reward": 0.5}) + + assert run.summary == {"reward": 0.5} + assert backend.finalized[0]["summary"] == {"reward": 0.5} + + +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) -> None: + order.append("write") + super().write(records) + + class OrderedBackend(FakeBackend): + 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, config=config) + + 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_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() + 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() + + 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_traces([{"id": "t1"}]) + + 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_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 + + +@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.""" + 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_finish_hands_the_full_config_to_finalize(): + """The evaluations API replaces metadata wholesale, so a backend recording + terminal state inside it needs the whole picture to merge into.""" + backend = FakeBackend() + run = make_run(backend, config={"num_rollouts": 4}) + + 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_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.""" + backend = FakeBackend() + run = make_run(backend) + + run._on_process_exit() + + assert backend.finalized[0]["status"] is RunStatus.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 + + run.reset_after_fork() + + assert run._finish_lock is not inherited_lock + run.finish() + assert backend.finalized == [] + assert backend.closed is True + + +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() + + 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") + + 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)]) + + run.log_traces([{"id": "t1"}]) + run.flush() + + assert any("broken" in error for error in run.errors) + run.finish() + + +def test_dropped_records_are_reported_on_the_handle(): + run = make_run() + run._worker.dropped = 3 + + assert run.dropped_records == 3 + run.finish() + + +def test_an_upload_failure_reaches_the_caller_in_raise_mode(): + """A sink fails on the uploader thread, where raising reaches nobody. Under + on_error="raise" the failure has to surface where a test is looking.""" + run = make_run(sinks=[FakeSink("broken", fail_on_write=True)], on_error="raise") + + run.log_traces([{"id": "t1"}]) + 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() + + +class _ForbiddenClient: + def __init__(self, code: str) -> None: + self.code = code + + def upload_records(self, records, **kwargs): + raise ForbiddenError("403", status_code=403, code=self.code) + + 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): + run.finish() + + assert run.failed_records == {"traces": 1} 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..6ea84556a --- /dev/null +++ b/packages/prime-runs/tests/test_samples_sink.py @@ -0,0 +1,63 @@ +"""The legacy sample sink that keeps today's viewer working.""" + +import pytest +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_fail_explicitly(make_platform_client, eval_routes): + 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"): + sink.write([make_trace()]) + + 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..d3d15c7c1 --- /dev/null +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -0,0 +1,224 @@ +"""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"}}]) + + _, kwargs = client.calls[0] + assert kwargs["context"] == { + "source": "prime-runs", + "run_kind": "eval", + "framework": "verifiers", + } + 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()]) + 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 + assert client.calls[2][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_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]) + + sent = client.calls[0][0][0] + assert sent == {**trace.to_record(), "run": {"id": "run-1", "type": "eval"}} + + +def test_an_account_outside_the_beta_retires_the_sink_without_a_failure(caplog): + """Prime Traces is in closed beta. For everyone outside it there was never + anywhere for these records to go, so the sink turns itself off quietly: + no exception for the worker to count, nothing above INFO in the log.""" + client = FakeTracesClient( + raises=ForbiddenError("not in beta", status_code=403, code="service_not_enabled") + ) + 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="missing scope"): + sink.write([{"id": "t1"}]) + + assert sink.enabled is False + assert "cannot write traces" 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_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) + + sink.close() + + assert client.closed is True + + +def test_a_missing_traces_client_disables_the_sink_and_reports_the_failure(monkeypatch, caplog): + """The run applies warn/raise policy, so the sink must surface this failure.""" + + def explode(**kwargs): + raise RuntimeError("no credentials") + + monkeypatch.setattr("prime_runs.sinks.traces.TracesClient", explode) + sink = TracesSink() + + with caplog.at_level("WARNING"): + with pytest.raises(RuntimeError, match="no credentials"): + 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..d7c251708 --- /dev/null +++ b/packages/prime-runs/tests/test_worker.py @@ -0,0 +1,397 @@ +"""The background uploader: backpressure, containment, fork safety.""" + +import queue +import threading + +from conftest import FakeSink + +from prime_runs.worker import UploadWorker + + +class BlockingSink(FakeSink): + def __init__(self) -> None: + super().__init__("blocking") + self.entered = threading.Event() + self.released = threading.Event() + + def write(self, records) -> None: + self.entered.set() + self.released.wait(5.0) + super().write(records) + + +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([{"id": 1}]) + drain(worker) + + for sink in sinks: + assert sink.batches == [[{"id": 1}]] + worker.close() + + +def test_a_disabled_sink_is_skipped(): + live, dead = FakeSink("live"), FakeSink("dead") + dead.enabled = False + worker = UploadWorker([live, dead]) + + worker.submit([{"id": 1}]) + drain(worker) + + assert live.batches and not dead.batches + worker.close() + + +def test_one_sink_failing_does_not_stop_the_others(): + broken, healthy = FakeSink("broken", fail_on_write=True), FakeSink("healthy") + reported = [] + worker = UploadWorker([broken, healthy], on_error=lambda name, exc: reported.append(name)) + + worker.submit([{"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([{"id": 1}]) + drain(worker) + + assert reported == ["broken"] + worker.close() + + +def test_records_that_skip_a_sink_retired_by_error_are_counted_as_lost_to_it(): + """After the first failed batch the sink is off, but the producer keeps + logging; every later batch is just as lost to that sink as the first. Live + run 2026-08-21: five episodes, footer said "1 failed via traces".""" + broken = FakeSink("broken", fail_on_write=True) + worker = UploadWorker([broken]) + + 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() + 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([{"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([{"id": 1}]) + accepted = worker.submit([{"id": 2}, {"id": 3}]) + + assert accepted is False + assert worker.dropped == 2 + + sink.released.set() + 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([{"id": 0}]) + assert sink.entered.wait(1.0) + assert worker.submit([{"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]["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([{"id": 0}]) + assert sink.entered.wait(1.0) + assert worker.submit([{"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]["id"] for batch in sink.batches] == [0, 1] + assert sink.closed is True + + +def test_close_drains_then_closes_every_sink(): + sink = FakeSink() + worker = UploadWorker([sink]) + + worker.submit([{"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([{"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([{"id": "parents"}]) + old_queue = worker._queue + + 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) -> None: + self.released.wait(10.0) + + sink = WedgedSink() + worker = UploadWorker([sink]) + worker.submit([{"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 + + +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) -> None: + raise RetryableAPIError("bad gateway", status_code=502) + + sink = BlipSink("blippy") + worker = UploadWorker([sink]) + + worker.submit([{"id": 1}, {"id": 2}]) + drain(worker) + + assert sink.enabled is True + # 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() + + +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) -> None: + raise TransportError("connection refused") + + sink = DeadSink("dead") + worker = UploadWorker([sink]) + + for _ in range(TRANSIENT_FAILURE_LIMIT): + worker.submit([{"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) -> None: + self.calls += 1 + if self.calls % 2 == 1: + raise RetryableAPIError("bad gateway", status_code=502) + 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([{"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) -> None: + raise UnauthorizedError("nope", status_code=401) + + sink = DeniedSink("denied") + worker = UploadWorker([sink]) + + worker.submit([{"id": 1}]) + drain(worker) + + 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) -> None: + raise RetryableAPIError("bad gateway", status_code=502) + + broken, healthy = BlipSink("broken"), FakeSink("healthy") + worker = UploadWorker([broken, healthy]) + + worker.submit([{"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() 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 cf8552cf5..cd3ce2368 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,31 @@ toml = [ { name = "tomli" }, ] +[[package]] +name = "prime-runs" +source = { editable = "packages/prime-runs" } +dependencies = [ + { name = "httpx" }, + { name = "prime-traces" }, +] + +[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 = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, + { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.23.0" }, + { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.13.1" }, +] +provides-extras = ["dev"] + [[package]] name = "prime-sandboxes" source = { editable = "packages/prime-sandboxes" }