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..05f989796 --- /dev/null +++ b/.github/workflows/release-runs.yml @@ -0,0 +1,91 @@ +name: Release prime-runs + +on: + push: + branches: [ "main" ] + paths: + - 'packages/prime-runs/**' + - '.github/workflows/release-runs.yml' + workflow_dispatch: + inputs: + force_release: + description: 'Force release even if tag exists' + required: false + default: 'false' + +jobs: + release: + runs-on: ubuntu-latest + environment: pypi-prod + permissions: + contents: write + packages: write + id-token: write + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Get version + id: version + working-directory: packages/prime-runs + run: | + VERSION=$(grep -E "^__version__\s*=\s*" src/prime_runs/__init__.py | sed -E 's/^__version__[[:space:]]*=[[:space:]]*"([^"]+)".*$/\1/') + echo "version=$VERSION" >> $GITHUB_OUTPUT + echo "Version: $VERSION" + + - name: Check for existing tag + id: check_tag + run: | + TAG="prime-runs-v${{ steps.version.outputs.version }}" + if git tag -l "$TAG" | grep -q .; then + echo "exists=true" >> $GITHUB_OUTPUT + echo "Tag $TAG already exists" + else + echo "exists=false" >> $GITHUB_OUTPUT + echo "Tag $TAG does not exist" + fi + + - name: Refuse untagged releases outside main + if: github.event_name == 'workflow_dispatch' && steps.check_tag.outputs.exists != 'true' && github.ref != 'refs/heads/main' + run: | + echo "::error title=Invalid release ref::New prime-runs releases must be dispatched from main, not ${{ github.ref }}." + exit 1 + + # A forced release must build the tagged source, not whatever the + # dispatch ref currently points at. + - name: Check out existing tag + 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..47ce21b67 --- /dev/null +++ b/packages/prime-runs/README.md @@ -0,0 +1,124 @@ +# Prime Runs SDK + +Track evaluation and training runs on the Prime Intellect platform: `init()` +opens the run, records stream out while it proceeds, and `finish()` closes it +out with a terminal status. + +## Install + +```bash +uv add prime-runs # or: pip install prime-runs +uv add 'prime-runs[train]' # training runs: adds pyarrow for the sample table +``` + +## Eval runs + +```python +import prime_runs as pr + +run = pr.init( + name="gsm8k-qwen3-8b", + environments=["gsm8k"], # hub names (get-or-create) or owner/name slugs + model="Qwen/Qwen3-8B", + framework="verifiers", + config="eval.toml", # the launched file, stored byte for byte +) +print(run.url) # https://app.primeintellect.ai/dashboard/evaluations/... + +for episode in rollouts: + run.log_episodes([episode]) # a queue put; bare traces: log_traces() + +run.finish(summary=pr.metrics.from_episodes(episodes)) +``` + +`init()` is called before the first rollout. Every record the run uploads is +keyed to it: the SDK sets `run.id` and `run.type` on the uploaded copy of each +trace and episode, over whatever run id the producer recorded locally, and +keeps the rest of that block (`name`, `work`). A producer never needs to know +the platform's id; `run.url` is the handle. A `with run:` block finishes for +you: an exception marks the run `failed`, Ctrl-C `cancelled`, and a process +that exits without finishing is reported `crashed` by an atexit hook. + +`config=` takes the path to the launched file (kept verbatim under +`config_source`, comments and all) or a mapping stored as given; put a file +under `pr.CONFIG_SOURCE_KEY` in the mapping to send both. Nothing is redacted. + +## Training runs + +```python +run = pr.init( + kind="train", + name="qwen3-8b-gsm8k-rl", + model="Qwen/Qwen3-8B", # the base model + environments=["primeintellect/gsm8k"], # hub ids, passed through + training=pr.TrainingSpec(max_steps=1000, batch_size=64, rollouts_per_example=8), + config=train_config.model_dump(), + team_id="team_...", # external runs belong to a team +) + +for step, (episodes, metrics) in enumerate(training_loop): + run.log_episodes(episodes) # episodes carry run.work.step (TrainRunInfo) + run.log_metrics(metrics, step=step) + +run.finish() +``` + +- The platform enables external runs per team; a team outside the allowlist + gets a `ForbiddenError` from `init()`. +- `init(kind="train", id=os.environ["RUN_ID"])` attaches to an external run a + launcher already created: nothing is registered, the platform keeps the run's + failure marking, and a clean `finish()` still completes it. +- Metrics are one row per `log_metrics` call, on their own uploader. The sample + table gets one Parquet object per upload, every 10th step, keyed by the step + an episode was dispatched at; a step logged in several calls gets several + objects, and the viewer shows their union. +- The status vocabulary is `completed | failed`; `cancelled` and `crashed` + arrive as `failed` with the reason in `error_message`. + +## How it behaves + +- **Streams.** Records go out on a background thread as they are logged; + whatever queues up during one request goes out as the next. +- **Contains its errors.** With the default `on_error="warn"` nothing the + platform raises escapes into your loop; `on_error="raise"` surfaces the first + failure from `flush()` or `finish()`, for tests and CI. Platform errors are + the `prime_traces` exception family. +- **Degrades.** A transient failure costs its batch, three in a row retire the + sink (a training run pauses it for five minutes instead), and a full queue + drops records rather than stalling the run. Losses are counted in + `run.dropped_records` and `run.failed_records`. +- **Drains on exit.** `finish()` gives queued uploads up to `finish_timeout` + (300 s) before closing the run out; an abort path can pass + `finish(timeout=...)`. + +An online run writes to Prime Traces (the system of record, gated to an +allowlist; outside it that sink turns itself off quietly) and to the sample +table today's viewer reads. `log_*()` are queue puts, safe inside a coroutine; +`init()` and `finish()` do network I/O. + +## Configuration + +| Source | Meaning | +| ---------------------- | ---------------------------------------------------------------------- | +| `PRIME_API_KEY` | Platform API token | +| `PRIME_TEAM_ID` | Team context; required for training runs | +| `PRIME_API_BASE_URL` | Platform API; defaults to `https://api.primeintellect.ai` | +| `PRIME_FRONTEND_URL` | Dashboard; defaults to `https://app.primeintellect.ai` | +| `PRIME_TRACES_URL` | Prime Traces service, resolved by `prime-traces` | +| `PRIME_RUNS_MODE` | `online` or `disabled`; unset means online when there is an API key | +| `~/.prime/config.json` | Shared prime CLI config (`api_key`, `team_id`, `base_url`) | + +Precedence is `init()` argument → environment variable → config file. A missing +API key disables the run with a warning. + +## Not yet available + +- **Failed or cancelled evaluations on the dashboard.** The evaluations API has + no producer-facing status endpoint yet, so the terminal state is recorded + under `metadata.prime_runs` and the run keeps showing as running. + +## Related packages + +- [prime-traces](https://github.com/PrimeIntellect-ai/prime/tree/main/packages/prime-traces) — Prime Traces SDK +- [prime](https://github.com/PrimeIntellect-ai/prime/tree/main/packages/prime) — Prime CLI +- [prime-evals](https://github.com/PrimeIntellect-ai/prime/tree/main/packages/prime-evals) — Evals SDK diff --git a/packages/prime-runs/pyproject.toml b/packages/prime-runs/pyproject.toml new file mode 100644 index 000000000..0c6aa5fca --- /dev/null +++ b/packages/prime-runs/pyproject.toml @@ -0,0 +1,79 @@ +[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.3", +] +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] +# The training sample table is Parquet. Optional: an eval producer never needs +# it, and prime-rl already carries pyarrow. +train = [ + "pyarrow>=15.0.0", +] +dev = [ + "pytest>=7.0.0", + "pytest-asyncio>=0.23.0", + "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..ac86cfd66 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/__init__.py @@ -0,0 +1,56 @@ +"""Prime Intellect Runs SDK: track eval and training runs on the Prime platform. + +A leaf package: the ``prime`` CLI depends on ``verifiers`` and verifiers depends +on this, so nothing here imports a producer package; records are duck-typed. +""" + +from . import metrics, projection +from .exceptions import ( + APIError, + ConfigurationError, + EnvironmentResolutionError, + ForbiddenError, + NotFoundError, + PaymentRequiredError, + PrimeRunsError, + RetryableAPIError, + RunFinishedError, + TransportError, + UnauthorizedError, +) +from .models import ( + CONFIG_SOURCE_KEY, + ConfigSource, + EnvironmentRef, + RunKind, + RunStatus, + TrainingSpec, +) +from .run import MODE_ENV, Run, init + +__version__ = "0.1.0" + +__all__ = [ + "init", + "Run", + "RunStatus", + "RunKind", + "TrainingSpec", + "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..f44d29c79 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/_fork.py @@ -0,0 +1,38 @@ +"""One process-wide ``os.register_at_fork`` hook for everything holding a +connection. A child must *drop* inherited transports (closing one can write to +the parent's socket) and never flush inherited buffers (those records are the +parent's). Registration is weak and the hook is installed once, since +``register_at_fork`` cannot be undone.""" + +import logging +import os +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: + global _lock + _lock = threading.Lock() # the inherited one may be held by a thread that is not here + for obj in list(_registry): + try: + obj.reset_after_fork() + 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..9e0fd725b --- /dev/null +++ b/packages/prime-runs/src/prime_runs/_http.py @@ -0,0 +1,147 @@ +"""Authenticated client for ``{base_url}/api/v1`` with a per-call retry policy. + +A failure is *ambiguous* when the server may already have processed the request +(a 502/504, a read timeout). Replaying one is safe for a GET or PUT and not for +``POST /evaluations/``, which would create a second run; callers declare intent +with ``idempotent=``. Connect errors and 429 are replayed for every method. +Error mapping and backoff come from ``prime_traces.core.client``. +""" + +import json +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. +UPLOAD_TIMEOUT = httpx.Timeout(300.0, connect=10.0) +#: Refused before any work was done. 503 is excluded: it may come from an +#: intermediary after forwarding. +UNAMBIGUOUS_RETRY_STATUS = frozenset({429}) +DEFAULT_MAX_ATTEMPTS = 5 + + +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; bare ``NaN`` is refused here rather than as a 400.""" + return json.dumps(value, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode( + "utf-8" + ) + + +class PlatformClient: + def __init__( + self, + *, + api_key: str, + base_url: str, + timeout: httpx.Timeout = DEFAULT_TIMEOUT, + max_attempts: int = DEFAULT_MAX_ATTEMPTS, + client: Optional[httpx.Client] = None, + ) -> None: + self.base_url = base_url.rstrip("/").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 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; the old one is dropped, not closed.""" + 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"``.""" + 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, + } + if timeout is not None: # None would disable httpx timeouts, not restore the default + request_kwargs["timeout"] = timeout + attempts = max_attempts or self.max_attempts + replayable = idempotent if idempotent is not None else method.upper() != "POST" + + 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) + + if attempt == attempts - 1 or (ambiguous and not replayable): + 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 _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..f10acaed4 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/backend.py @@ -0,0 +1,363 @@ +"""Run lifecycles: create, update, finalize. Sinks, not backends, move records. + +``EvalsBackend`` speaks ``/api/v1/evaluations``; ``RftBackend`` speaks +``/api/v1/rft/external-runs``; ``DisabledBackend`` is the no-op behind +``mode="disabled"``. +""" + +import logging +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, TrainingSpec + +logger = logging.getLogger(__name__) + + +class Backend(Protocol): + def create(self, spec: RunSpec) -> RunHandle: ... + + def update( + self, + run_id: str, + *, + config: Optional[Dict[str, Any]] = None, + summary: Optional[Dict[str, Any]] = None, + ) -> None: + """Persist the run's whole config (not a patch) and/or summary.""" + ... + + 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, once per run.""" + ... + + def close(self) -> None: ... + + +class EvalsBackend: + """Environments are resolved through the hub's get-or-create, so a local run + uploads without ``prime env push``. The API has no producer-facing way to mark + a run failed, so a non-completed terminal state is recorded in + ``metadata.prime_runs`` and the run keeps showing as running.""" + + def __init__( + self, + 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 + + 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) + _set_if(payload, "dataset", _first_environment_name(spec)) # the env under another name + _set_if(payload, "framework", spec.framework) + _set_if(payload, "description", spec.description) + _set_if(payload, "metadata", spec.config or None) + _set_if(payload, "team_id", spec.team_id or self._team_id) + + # Not replayable: a retry after a lost response would create a second run. + response = self._client.post("/evaluations/", json_body=payload) + run_id = response.get("evaluation_id") + if not run_id: + 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}" + + def update( + self, + run_id: str, + *, + config: Optional[Dict[str, Any]] = None, + summary: Optional[Dict[str, Any]] = None, + ) -> None: + payload: Dict[str, Any] = {} + _set_if(payload, "metadata", config or None) + _set_if(payload, "metrics", summary or None) + if payload: + self._client.put(f"/evaluations/{run_id}", json_body=payload) + + 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) + # Not replayable: finalization enqueues the platform's statistics task. + self._client.post(f"/evaluations/{run_id}/finalize", json_body=body or {"metrics": {}}) + return + + # Merged into the whole config: the PUT replaces the metadata document. + terminal = {"status": status.value, "finished_at": datetime.now(timezone.utc).isoformat()} + if error: + terminal["error"] = error + self.update(run_id, config={**(config or {}), "prime_runs": terminal}, summary=summary) + 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() + + def _resolve_environments(self, refs: List[EnvironmentRef]) -> List[Dict[str, Any]]: + """Hub references as the API's ``EnvironmentReference`` objects. One that + cannot be resolved raises rather than being skipped: a run silently + attached to the wrong environments looks like a successful upload.""" + resolved: List[Dict[str, Any]] = [] + for ref in refs: + entry: Dict[str, Any] = {"id": ref.id or self._lookup_environment(ref)} + _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 + environment_id = (response.get("data") or response).get("id") + if not environment_id: + raise EnvironmentResolutionError(f"Hub returned no id for environment {ref.slug!r}") + return str(environment_id) + + 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: + stem = _first_environment_name(spec) or spec.framework or "eval" + return f"{stem}-{uuid.uuid4().hex[:8]}" + + +def disabled_run_id() -> str: + return f"disabled-{uuid.uuid4().hex[:16]}" + + +class DisabledBackend: + 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 + + +#: The RFT API's vocabulary is ``completed | failed``; the finer distinctions +#: travel in ``error_message``. +RFT_TERMINAL_STATUS = { + RunStatus.COMPLETED: "completed", + RunStatus.FAILED: "failed", + RunStatus.CANCELLED: "failed", + RunStatus.CRASHED: "failed", +} + + +class RftBackend: + """External training runs are created for a team, and the platform enables + them per team: ``team_id`` is required and a team outside the allowlist gets + a 403 from :meth:`create`. The API takes the config at creation and has no + summary document, so :meth:`update` is a no-op.""" + + def __init__( + self, + 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 + + def create(self, spec: RunSpec) -> RunHandle: + if not spec.model: + raise ConfigurationError("A training run needs model= — the base model being trained.") + team_id = spec.team_id or self._team_id + if not team_id: + raise ConfigurationError( + "A training run needs a team: pass team_id= or set PRIME_TEAM_ID. " + "External training runs are created for a team, and the platform " + "enables them per team." + ) + training = spec.training or TrainingSpec() + payload: Dict[str, Any] = { + "base_model": spec.model, + "max_steps": int(training.max_steps), + # Passed through: the RFT API resolves hub ids itself. + "environments": [_training_environment(ref) for ref in spec.environments], + "team_id": team_id, + } + _set_if(payload, "name", spec.name) + _set_if(payload, "batch_size", training.batch_size) + _set_if(payload, "rollouts_per_example", training.rollouts_per_example) + _set_if(payload, "seq_len", training.seq_len) + _set_if(payload, "run_config", spec.config or None) + _set_if(payload, "wandb_project", training.wandb_project) + _set_if(payload, "wandb_entity", training.wandb_entity) + _set_if(payload, "wandb_run_name", training.wandb_run_name) + + # Not replayable: a retry after a lost response would register a second run. + response = self._client.post("/rft/external-runs", json_body=payload) + run = response.get("run") or {} + run_id = run.get("id") + if not run_id: + raise APIError(f"POST /rft/external-runs returned no run id (keys: {sorted(response)})") + return RunHandle( + id=str(run_id), + name=str(run.get("name") or spec.name or "") or None, + url=self.url_for(str(run_id)), + ) + + def attach(self, run_id: str) -> RunHandle: + """A handle for a run a launcher already created. No request: the first + write proves access. The run must be an external one; the monitoring + endpoints answer 400 for a hosted run's id.""" + return RunHandle(id=run_id, url=self.url_for(run_id)) + + def url_for(self, run_id: str) -> str: + return f"{self._frontend_url}/dashboard/training/{run_id}" + + def update( + self, + run_id: str, + *, + config: Optional[Dict[str, Any]] = None, + summary: Optional[Dict[str, Any]] = None, + ) -> None: + if config or summary: + logger.debug("Run %s: the RFT API has no config/summary update; nothing sent", run_id) + + def finalize( + self, + run_id: str, + *, + status: RunStatus, + summary: Optional[Dict[str, Any]] = None, + error: Optional[str] = None, + config: Optional[Dict[str, Any]] = None, + ) -> None: + if status is RunStatus.COMPLETED: + # Compare-and-set server-side, and "already finalized" is a success. + try: + self._client.post( + "/rft/finalize", json_body={"run_id": run_id, "exit_code": 0}, idempotent=True + ) + except APIError as exc: + # The status endpoint may still close the run out. + logger.warning( + "Run %s could not be finalized (%s); falling back to a status update", + run_id, + exc, + ) + self._set_terminal_status(run_id, status=status) + return + + message = error or status.value + if status is not RunStatus.FAILED: + message = f"{status.value}: {message}" + self._set_terminal_status(run_id, status=status, error_message=message) + + def _set_terminal_status( + self, run_id: str, *, status: RunStatus, error_message: Optional[str] = None + ) -> None: + payload: Dict[str, Any] = {"status": RFT_TERMINAL_STATUS[status]} + _set_if(payload, "error_message", error_message) + try: + self._client.put(f"/rft/external-runs/{run_id}/status", json_body=payload) + except APIError as exc: + if exc.status_code == 409: # already closed: the outcome is recorded either way + logger.info( + "Run %s is already closed on the platform; not marking it %s", + run_id, + status.value, + ) + return + raise + + def close(self) -> None: + self._client.close() + + +def _training_environment(ref: EnvironmentRef) -> Dict[str, Any]: + entry: Dict[str, Any] = {"id": ref.id or ref.slug or ref.name} + _set_if(entry, "version_id", ref.version_id) + return entry diff --git a/packages/prime-runs/src/prime_runs/config.py b/packages/prime-runs/src/prime_runs/config.py new file mode 100644 index 000000000..2a4caba54 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/config.py @@ -0,0 +1,17 @@ +"""``prime_traces.core.Config`` (env vars, then ``~/.prime/config.json``) plus the +dashboard URL.""" + +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: + 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..37658eeaf --- /dev/null +++ b/packages/prime-runs/src/prime_runs/exceptions.py @@ -0,0 +1,76 @@ +"""Errors raised before any request derive from ``PrimeRunsError``; errors from +the platform are ``prime_traces``' own, so one set of ``except`` clauses covers +both SDKs. Under ``on_error="warn"`` none of them escape into a producer loop.""" + +from prime_traces.exceptions import ( + APIError, + APITimeoutError, + ForbiddenError, + LineFormatConflictError, + NotFoundError, + PaymentRequiredError, + RetryableAPIError, + TraceTooLargeError, + TransportError, + UnauthorizedError, + ValidationRejectedError, +) + +__all__ = [ + "APIError", + "APITimeoutError", + "ConfigurationError", + "EnvironmentResolutionError", + "ForbiddenError", + "LineFormatConflictError", + "NotFoundError", + "PaymentRequiredError", + "PrimeRunsError", + "RetryableAPIError", + "RunFinishedError", + "TraceTooLargeError", + "TransportError", + "UnauthorizedError", + "ValidationRejectedError", + "is_record_rejection", + "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.""" + + +class RunFinishedError(PrimeRunsError): + """A finished run was written to again: a producer bug.""" + + +def is_record_rejection(exc: BaseException) -> bool: + """A failure specific to the submitted batch: a later batch can still succeed.""" + if isinstance( + exc, + ( + TypeError, + ValueError, + TraceTooLargeError, + ValidationRejectedError, + LineFormatConflictError, + ), + ): + return True + # The legacy samples API returns framework-generated 413/422 responses, + # which the shared traces mapper leaves as plain APIError. + return isinstance(exc, APIError) and exc.status_code in (413, 422) + + +def is_transient(exc: BaseException) -> bool: + """A failure about this moment (retry later) rather than this run (stop).""" + return isinstance(exc, (RetryableAPIError, TransportError)) diff --git a/packages/prime-runs/src/prime_runs/metrics.py b/packages/prime-runs/src/prime_runs/metrics.py new file mode 100644 index 000000000..de5f31f0c --- /dev/null +++ b/packages/prime-runs/src/prime_runs/metrics.py @@ -0,0 +1,40 @@ +"""Run-level aggregates over native episodes, in the shape the eval dashboard +reads; the aggregation verifiers used, so migrated runs keep their numbers. +Pass the result to ``run.finish(summary=...)``.""" + +from typing import Any, Dict, Optional, Sequence, TypedDict + + +class RunSummary(TypedDict): + avg_reward: float + avg_metrics: Dict[str, float] + avg_error: float + + +def from_episodes(episodes: Sequence[Any], traces: Optional[Sequence[Any]] = None) -> RunSummary: + """Rewards and metrics aggregate over the trainable traces only (a judge or + a modeled user would dilute every mean with structural zeros), falling back + to all traces when none are trainable, as the dashboard does. ``avg_error`` + is the share of episodes that are not ok.""" + if traces is None: + traces = [trace for episode in episodes for trace in episode.traces] + scored = [trace for trace in traces if trace.agent.trainable] or list(traces) + + 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..c2419d3f9 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/models.py @@ -0,0 +1,199 @@ +"""Types shared across backends, sinks and the ``Run`` handle. Response bodies +are not modeled: backends pull the fields they need into a ``RunHandle``.""" + +import os +from dataclasses import dataclass, field +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"] +OnError = Literal["warn", "raise"] + +RunKind = Literal["eval", "train"] +"""Stamped as ``run.type`` on records and sent as upload provenance; the traces +service's vocabulary, matching verifiers' ``EvalRunInfo`` / ``TrainRunInfo``.""" + +RUN_KIND: RunKind = "eval" + + +class RunStatus(str, Enum): + """``failed``: the producer said so. ``cancelled``: stopped on purpose (an + interrupt). ``crashed``: the process exited without saying (only the + atexit hook reports it).""" + + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + CRASHED = "crashed" + + def is_terminal(self) -> bool: + return self is not RunStatus.RUNNING + + +@dataclass +class EnvironmentRef: + """``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 ``ConfigSource`` dict.""" + +MAX_CONFIG_SOURCE_BYTES = 256 * 1024 + +_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. Nothing 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": + 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"]: + """A ``str`` or ``PathLike`` is a *path*, never inline text.""" + if value is None or isinstance(value, cls): + return value + if isinstance(value, (str, os.PathLike)): + return cls.from_file(value) + raise TypeError( + f"a config source must be a path or a ConfigSource, 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 TrainingSpec: + """Display fields a training run is registered with. ``max_steps`` is + required by the platform; ``0`` means unknown.""" + + max_steps: int = 0 + batch_size: Optional[int] = None + rollouts_per_example: Optional[int] = None + seq_len: Optional[int] = None + wandb_project: Optional[str] = None + wandb_entity: Optional[str] = None + wandb_run_name: Optional[str] = None + + def __post_init__(self) -> None: + if self.max_steps < 0: + raise ValueError("max_steps must be non-negative") + + +@dataclass +class RunSpec: + """``init()``'s arguments after normalization.""" + + 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) + kind: RunKind = RUN_KIND + training: Optional[TrainingSpec] = None + + +@dataclass +class RunHandle: + 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..958b751f5 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/projection.py @@ -0,0 +1,267 @@ +"""Projections of native episodes onto the platform's sample tables: the v0 +eval-sample format (moved here from ``verifiers.v1.utils.platform``) and the +training viewer's Parquet table (moved here from prime-rl). Duck-typed: verifiers +``Trace``/``Episode`` satisfy it structurally and are not imported.""" + +import json +import logging +import math +from typing import Any, Dict, Iterable, List, Optional, Sequence + +from ._http import encode_json + +logger = logging.getLogger(__name__) + +#: The platform's per-request ceiling on ``POST /samples``. +MAX_SAMPLES_PAYLOAD_BYTES = 25 * 1024 * 1024 +ENVELOPE_BYTES = len(b'{"samples":[]}') + + +def json_bytes(value: Any) -> int: + 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 _nullify_non_finite(value: Any) -> Any: + """NaN and infinity are not JSON. The tables show them as ``null`` + rather than losing the row (or the batch) they sit in.""" + if isinstance(value, float) and not math.isfinite(value): + return None + if isinstance(value, dict): + return {key: _nullify_non_finite(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_nullify_non_finite(item) for item in value] + return value + + +def _json_column(value: Any) -> str: + return json.dumps(_nullify_non_finite(value), allow_nan=False) + + +def summary_trace_index(episode: Any) -> int: + """The first trainable trace, else the first trace: the one whose flat + projection (and ``task.data.idx``) represents the episode.""" + return next( + (index for index, trace in enumerate(episode.traces) if trace.agent.trainable), + 0, + ) + + +def trace_to_sample( + trace: Any, rollout_number: int = 1, episode_id: Optional[str] = None +) -> Dict[str, Any]: + """One trace as a v0 sample row. ``completion`` is the final branch's + messages, ``trajectory`` one message list per branch.""" + task = trace.task.data.model_dump(mode="json", exclude_none=True) + branches = trace.branches + sample: Dict[str, Any] = { + "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"), + "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, + } + # Sub-rewards are flattened to top-level keys the way v0 did; env metrics stay nested. + for name, reward in trace.rewards.items(): + if reward is not None: + sample.setdefault(name, reward.score) + return _nullify_non_finite(sample) + + +def episode_to_samples(episode: Any, rollout_number: int) -> List[Dict[str, Any]]: + """Normally one row: the summary trace's projection, with the whole native + episode in ``info.native_wrapper``. An episode too large for one request + falls back to one projected row per trace.""" + if not episode.traces: + return [] + summary_index = summary_trace_index(episode) + sample = trace_to_sample(episode.traces[summary_index], rollout_number, episode.id) + sample["sample_id"] = episode.id + sample["info"] = _nullify_non_finite( + { + **(sample["info"] or {}), + "native_wrapper": episode.to_record(), + "native_trace_index": summary_index, + } + ) + if ENVELOPE_BYTES + json_bytes(sample) <= MAX_SAMPLES_PAYLOAD_BYTES: + return [sample] + logger.warning( + "Episode %s exceeds the platform sample limit; uploading projected traces", episode.id + ) + 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 v0 samples. A streaming producer passes the same + ``rollout_numbers`` dict every call so numbering stays consistent.""" + counts = rollout_numbers if rollout_numbers is not None else {} + samples: List[Dict[str, Any]] = [] + for episode in episodes: + 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; a sample too large to send + alone raises rather than being silently dropped.""" + batches: List[List[Dict[str, Any]]] = [] + batch: List[Dict[str, Any]] = [] + payload_bytes = ENVELOPE_BYTES + 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)" + ) + next_bytes = payload_bytes + (1 if batch else 0) + sample_bytes # +1: the joining comma + if batch and next_bytes > MAX_SAMPLES_PAYLOAD_BYTES: + batches.append(batch) + batch = [] + next_bytes = ENVELOPE_BYTES + sample_bytes + batch.append(sample) + payload_bytes = next_bytes + if batch: + batches.append(batch) + return batches + + +def parquet_available() -> bool: + try: + import pyarrow # noqa: F401 + except ImportError: + return False + return True + + +def train_sample_schema() -> Any: + import pyarrow as pa + + return pa.schema( + [ + ("run_id", pa.string()), + ("step", pa.int64()), + ("tag", pa.string()), + ("problem_id", pa.int64()), + ("sample_id", pa.int64()), + ("prompt", pa.string()), + ("completion", pa.string()), + ("trajectory", pa.string()), + ("answer", pa.string()), + ("env_name", pa.string()), + ("task", pa.string()), + ("info", pa.string()), + ("reward", pa.float64()), + ("advantage", pa.float64()), + ("metrics", pa.string()), + ("timing", pa.string()), + ("num_input_tokens", pa.int64()), + ("num_output_tokens", pa.int64()), + ("created_at", pa.timestamp("us", tz="UTC")), + ] + ) + + +def episodes_to_parquet_bytes( + episodes: Sequence[Any], run_id: str, step: int, *, sample_id_offset: int = 0 +) -> Optional[bytes]: + """One training step's episodes as the viewer's Parquet table, one row per + episode (the v0 projection plus the RFT-only columns), numbered from + ``sample_id_offset``. ``None`` when no episode has a trajectory.""" + import io + from datetime import datetime, timezone + + import pyarrow as pa + import pyarrow.parquet as pq + + advantages: Dict[Any, Optional[float]] = {} + env_names: Dict[Any, str] = {} + for episode in episodes: + if not episode.traces: + continue + summary_trace = episode.traces[summary_trace_index(episode)] + advantages[episode.id] = (summary_trace.info or {}).get("advantage") + env = getattr(episode, "env", None) + env_names[episode.id] = str(getattr(env, "id", None) or "") + + now = datetime.now(timezone.utc) + rows: List[Dict[str, Any]] = [] + for sample_id, sample in enumerate(build_samples(episodes), start=sample_id_offset): + trajectory = sample["trajectory"] + if not trajectory: + continue + advantage = advantages.get(sample["episode_id"]) + trajectory = [{**branch, "advantage": advantage} for branch in trajectory] + try: + problem_id = ( + int(sample["example_id"]) if sample["example_id"] is not None else sample_id + ) + except (TypeError, ValueError): + problem_id = sample_id + rows.append( + { + "run_id": run_id, + "step": step, + "tag": "", + "problem_id": problem_id, + "sample_id": sample_id, + "prompt": "", + "completion": _json_column(sample["completion"]), + "trajectory": _json_column(trajectory), + "answer": "", + "env_name": env_names.get(sample["episode_id"], ""), + "task": _json_column(sample["task"]), + "info": _json_column(sample["info"]), + "reward": sample["reward"], + "advantage": advantage, + "metrics": _json_column(sample["metrics"]), + "timing": _json_column(sample["timing"]), + "num_input_tokens": trajectory[-1]["num_input_tokens"], + "num_output_tokens": trajectory[-1]["num_output_tokens"], + "created_at": now, + } + ) + if not rows: + return None + + table = pa.Table.from_pylist(rows, schema=train_sample_schema()) + buffer = io.BytesIO() + pq.write_table(table, buffer, compression="snappy", use_dictionary=True, write_statistics=True) + return buffer.getvalue() diff --git a/packages/prime-runs/src/prime_runs/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..384082736 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/run.py @@ -0,0 +1,587 @@ +"""The run handle, and ``init()`` that opens one. + +``init()`` is called before the first rollout; the id it returns keys every +record the run uploads, whatever run id the producer recorded locally. +""" + +import asyncio +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, RftBackend, disabled_run_id +from .config import Config +from .exceptions import ConfigurationError, RunFinishedError +from .models import ( + CONFIG_SOURCE_KEY, + ConfigSource, + EnvironmentRef, + Mode, + OnError, + RunHandle, + RunKind, + RunSpec, + RunStatus, + TrainingSpec, +) +from .sinks import EvalSamplesSink, RftMetricsSink, RftSamplesSink, Sink, TracesSink +from .worker import UploadWorker, deadline_after, time_left + +logger = logging.getLogger(__name__) + +MODE_ENV = "PRIME_RUNS_MODE" +#: How long ``finish()`` lets queued uploads drain: one in-flight sample POST +#: may take this long, and a shorter budget would abandon it about to succeed. +DEFAULT_FINISH_TIMEOUT = float(UPLOAD_TIMEOUT.read or 300.0) +#: Training runs last days: a training sink that struck out on transient +#: failures is tried again after this long instead of retired for the run. +TRAIN_RETIRE_COOLDOWN = 300.0 + + +class Run: + """A live run: an id, a URL, somewhere to put records, a summary. + + With ``on_error="warn"`` (the default) nothing the platform raises escapes + into the producer's loop; ``on_error="raise"`` surfaces the first failure + from :meth:`flush` or :meth:`finish`. + """ + + def __init__( + self, + *, + backend: Backend, + handle: RunHandle, + spec: RunSpec, + sinks: Optional[List[Sink]] = None, + mode: Mode = "online", + on_error: OnError = "warn", + finish_timeout: Optional[float] = None, + metrics_sinks: Optional[List[Sink]] = None, + attached: bool = False, + retire_cooldown: Optional[float] = None, + ) -> None: + self._backend = backend + self._handle = handle + self._spec = spec + self._mode: Mode = mode + self._on_error: OnError = on_error + self._status = RunStatus.RUNNING + # A forked child inherits the handle but must not close the parent's run. + self._owns_lifecycle = True + # A run a launcher created and handed over: the platform owns its + # failure marking, this process only completes it. + self._attached = attached + + self.config: Dict[str, Any] = dict(spec.config) + self.summary: Dict[str, Any] = {} + self.errors: List[str] = [] + # on_error="raise": the first uploader-thread failure, re-raised from + # flush() or finish() where the caller is looking. + self._deferred_error: Optional[BaseException] = None + + self._finish_timeout = ( + DEFAULT_FINISH_TIMEOUT if finish_timeout is None else max(0.0, float(finish_timeout)) + ) + # _state_lock guards admission only; _finish_lock is held through the + # slow teardown so concurrent finish() callers wait for it. + self._state_lock = threading.Lock() + self._finish_lock = threading.Lock() + self._finishing = False + self._finished = False + self._atexit_hook = self._on_process_exit + _fork.register(self) + + sinks = sinks or [] + metrics_sinks = metrics_sinks or [] + # Separate uploaders: a slow sample upload must not hold metrics back. + self._worker = UploadWorker( + sinks, on_error=self._record_sink_error, retire_cooldown=retire_cooldown + ) + self._metrics_worker = UploadWorker( + metrics_sinks, on_error=self._record_sink_error, retire_cooldown=retire_cooldown + ) + context = _sink_context(spec) + for sink in [*sinks, *metrics_sinks]: + try: + sink.start(handle.id, context) + except Exception as exc: # noqa: BLE001 - a bad sink is not a bad run + sink.enabled = False + self._note(f"starting sink {sink.name}", exc) + if self._on_error == "raise": + # Close the remote run before the failure reaches the caller. + self.finish(status=RunStatus.FAILED, error=_describe(exc)) + raise exc + + atexit.register(self._atexit_hook) + + @property + def id(self) -> str: + return self._handle.id + + @property + def name(self) -> Optional[str]: + return self._handle.name + + @property + def url(self) -> Optional[str]: + return self._handle.url + + @property + def config_source(self) -> Optional[ConfigSource]: + 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 kind(self) -> RunKind: + return self._spec.kind + + @property + def attached(self) -> bool: + return self._attached + + @property + def dropped_records(self) -> int: + """Records that reached no sink because a queue was full.""" + return self._worker.dropped + self._metrics_worker.dropped + + @property + def failed_records(self) -> Dict[str, int]: + """Records each sink could not store, by sink name.""" + return {**self._worker.failed_records, **self._metrics_worker.failed_records} + + def __repr__(self) -> str: + return f"" + + def log_traces(self, traces: Iterable[Any]) -> None: + """Queue bare traces: verifiers ``Trace`` objects or JSON mappings.""" + self._submit("log_traces", traces) + + def log_episodes(self, episodes: Iterable[Any]) -> None: + """Queue episodes: verifiers ``Episode`` objects or JSON mappings with a + ``traces`` list. The sample tables are projected from episode objects + only; a JSON episode reaches Prime Traces alone.""" + self._submit("log_episodes", episodes) + + def log_metrics(self, values: Mapping[str, Any], *, step: Optional[int] = None) -> None: + """Queue one step's metrics (training runs; a no-op for evals). ``step`` + becomes ``values["step"]`` unless already set, every row gets a + ``_timestamp``, and non-finite numbers are dropped.""" + self._require_live("log_metrics") + record = _clean_metrics(values) + if step is not None: + record.setdefault("step", step) + record.setdefault("_timestamp", time.time()) + with self._state_lock: + self._require_live("log_metrics") + self._metrics_worker.submit([record]) + + def _submit(self, operation: str, records: Iterable[Any]) -> None: + self._require_live(operation) # before consuming a lazy iterable + batch = list(records) + with self._state_lock: + self._require_live(operation) + if batch: + self._worker.submit(batch) + + def update_summary(self, values: Mapping[str, Any]) -> None: + """Merge run-level outputs into :attr:`summary` ahead of :meth:`finish`.""" + with self._state_lock: + self._require_live("update_summary") + self.summary.update(_clean_metrics(values)) + + def flush(self, timeout: Optional[float] = 30.0) -> bool: + """Block until queued records have been written. Under + ``on_error="raise"`` this is the first place an upload failure surfaces.""" + deadline = deadline_after(timeout) + flushed = self._worker.flush(timeout=time_left(deadline)) + flushed = self._metrics_worker.flush(timeout=time_left(deadline)) and flushed + self._raise_deferred() + return flushed + + def finish( + self, + summary: Optional[Mapping[str, Any]] = None, + *, + status: Union[RunStatus, str] = RunStatus.COMPLETED, + error: Optional[str] = None, + timeout: Optional[float] = None, + ) -> None: + """Drain the uploads and close the run out. Idempotent: the first caller + reports the status, concurrent callers wait for that teardown. + ``timeout`` overrides the run's ``finish_timeout`` for this call.""" + 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}") + with self._state_lock: + self._finishing = True + try: + self._finish_once(summary, resolved, error, timeout) + finally: + with self._state_lock: + self._finishing = False + self._finished = True + + def _finish_once( + self, + summary: Optional[Mapping[str, Any]], + resolved: RunStatus, + error: Optional[str], + timeout: Optional[float], + ) -> None: + if summary: + self.summary.update(_clean_metrics(summary)) + self._status = resolved + + budget = self._finish_timeout if timeout is None else max(0.0, float(timeout)) + deadline = deadline_after(budget) + # Records first, so the dashboard never sees a finished run with samples landing. + drained = self._worker.flush(timeout=time_left(deadline)) + drained = self._metrics_worker.flush(timeout=time_left(deadline)) and drained + if not drained: + logger.warning( + "Run %s: uploads did not drain within %ss; finalizing anyway. " + "Some records may be missing from this run.", + self.id, + budget, + ) + self._worker.close(timeout=time_left(deadline)) + self._metrics_worker.close(timeout=time_left(deadline)) + + 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 + ), + ) + if self._attached and resolved is not RunStatus.COMPLETED: + # The launcher marks its own run failed; reporting it here would race that. + logger.info( + "Run %s: %s, but the platform owns the status of an attached run; " + "leaving it to the launcher", + self.id, + resolved.value, + ) + else: + self._teardown_step( + "finalizing the run", + lambda: self._backend.finalize( + self.id, + status=resolved, + summary=self.summary or None, + error=error or (self.errors[0] if self.errors else None), + config=self.config or None, + ), + ) + self._teardown_step("closing the backend", self._backend.close) + atexit.unregister(self._atexit_hook) + + if self.dropped_records: + logger.warning( + "Run %s finished with %d record(s) that reached no sink; the producer " + "outran the uploader.", + self.id, + self.dropped_records, + ) + for sink_name, count in self.failed_records.items(): + logger.warning( + "Run %s: the %s sink could not store %d record(s)", self.id, sink_name, count + ) + self._raise_deferred() # last, so the run is closed out even when this raises + + def fail(self, error: Union[str, BaseException], *, timeout: Optional[float] = None) -> None: + self.finish(status=RunStatus.FAILED, error=_describe(error), timeout=timeout) + + 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, asyncio.CancelledError)): + status, error = RunStatus.CANCELLED, "interrupted" # a decision, not a fault + else: + status, error = RunStatus.FAILED, _describe(exc) + try: + self.finish(status=status, error=error) + except Exception as finish_error: + # The producer's exception is why this block is unwinding; a teardown + # error must not replace it. Control-flow exceptions pass through. + 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 + + def reset_after_fork(self) -> None: + """Fresh locks, and the parent keeps the lifecycle: the child's atexit + hook must not finalize a still-running run.""" + self._state_lock = threading.Lock() + self._finish_lock = threading.Lock() + self._finishing = False + self._owns_lifecycle = False + self._deferred_error = None + + def _on_process_exit(self) -> None: + """atexit: nobody called ``finish()``, so report CRASHED rather than FAILED.""" + if self._finished: + return + logger.warning("Run %s was never finished; reporting it as crashed", self.id) + 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: + self._note(f"writing to the {sink_name} sink", exc) + + def _teardown_step(self, what: str, call: Any) -> None: + 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: + exc = self._deferred_error + if exc is None: + return + self._deferred_error = None + raise exc + + +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", + finish_timeout: Optional[float] = None, + kind: RunKind = "eval", + id: Optional[str] = None, + training: Optional[TrainingSpec] = None, +) -> Run: + """Open a run and return its handle. Call it before the first rollout. + + ``mode`` defaults to ``$PRIME_RUNS_MODE``, else online when there is an API + key and disabled (with a warning) when there is not. ``config`` is the path + to the file the run was launched from (stored byte for byte under + ``config_source``) or a mapping taken as given. ``finish_timeout`` bounds + the drain in :meth:`Run.finish`. + + ``kind="train"`` opens an external training run: ``model`` is the base + model, ``environments`` the hub ids, ``training`` the display fields, and a + team is required. ``id`` attaches to an external run a launcher already + created (``$RUN_ID``): nothing is registered, the platform keeps the run's + failure marking, and a clean finish still completes it. + """ + settings = Config() + api_key = api_key if api_key is not None else settings.api_key + base_url = base_url or settings.base_url + team_id = team_id if team_id is not None else settings.team_id + + if kind not in ("eval", "train"): + raise ConfigurationError(f"kind={kind!r} is not one of 'eval' or 'train'") + if id is not None and kind != "train": + raise ConfigurationError( + "id= attaches to an existing training run; an eval run is always created here." + ) + if training is not None and kind != "train": + raise ConfigurationError("training= only applies to kind='train'") + + spec = RunSpec( + name=name, + environments=[EnvironmentRef.coerce(entry) for entry in (environments or [])], + model=model, + framework=framework, + description=description, + tags=list(tags or []), + team_id=team_id, + config=_normalize_config(config), + kind=kind, + training=training, + ) + resolved_mode = _resolve_mode(mode, api_key=api_key) + + backend: Backend + sinks: List[Sink] + metrics_sinks: List[Sink] = [] + attached = False + if resolved_mode == "disabled": + backend = DisabledBackend() + handle = RunHandle(id=id or 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) + if kind == "train": + rft = RftBackend(client, frontend_url=settings.frontend_url, team_id=team_id) + backend = rft + else: + backend = EvalsBackend(client, frontend_url=settings.frontend_url, team_id=team_id) + try: + if kind == "train" and id is not None: + handle = rft.attach(id) + attached = True + else: + handle = backend.create(spec) + except BaseException: + # Nothing else can release the connection pool yet. + try: + backend.close() + except Exception as close_error: # noqa: BLE001 - preserve the create failure + logger.debug("Error closing the platform client: %s", close_error) + raise + # Both transports run during the transition: traces is the system of + # record, the sample table is what today's viewer reads. + traces = TracesSink(api_key=api_key, team_id=team_id) + if kind == "train": + sinks = [traces, RftSamplesSink(client)] + metrics_sinks = [RftMetricsSink(client)] + else: + sinks = [traces, EvalSamplesSink(client)] + + run = Run( + backend=backend, + handle=handle, + spec=spec, + sinks=sinks, + mode=resolved_mode, + on_error=on_error, + finish_timeout=finish_timeout, + metrics_sinks=metrics_sinks, + attached=attached, + retire_cooldown=TRAIN_RETIRE_COOLDOWN if kind == "train" else None, + ) + if run.url: + logger.info("Run %s: %s", run.id, run.url) + 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; the join key is ``run.id`` inside the document.""" + context = {"source": "prime-runs", "run_type": spec.kind} + if spec.framework: + context["framework"] = spec.framework + if spec.model: + context["model"] = spec.model + return context + + +def _normalize_config(value: Any) -> Dict[str, Any]: + 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) + + +_DROP: Any = object() + + +def _finite(value: Any) -> Any: + """``value`` without NaN/infinity anywhere inside it (strict JSON rejects + them as an opaque 400), or ``_DROP``. Empty nested mappings are dropped.""" + if isinstance(value, float) and not math.isfinite(value): + return _DROP + if isinstance(value, Mapping): + cleaned = {} + for key, item in value.items(): + kept = _finite(item) + # Identity checks only: an array-like value would make `==` ambiguous. + if kept is _DROP or (isinstance(item, Mapping) and not kept): + continue + cleaned[key] = kept + return cleaned + if isinstance(value, (list, tuple)): + return [item for item in (_finite(item) for item in value) if item is not _DROP] + return value + + +def _clean_metrics(metrics: Mapping[str, Any]) -> Dict[str, Any]: + return _finite(dict(metrics)) diff --git a/packages/prime-runs/src/prime_runs/sinks/__init__.py b/packages/prime-runs/src/prime_runs/sinks/__init__.py new file mode 100644 index 000000000..ea393e379 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/__init__.py @@ -0,0 +1,19 @@ +"""Record transports, independent of backends and of each other.""" + +from .base import Sink, is_episode, stamp_run, to_mapping +from .metrics import RftMetricsSink +from .samples import EvalSamplesSink +from .traces import TracesSink +from .train_samples import RftSamplesSink, training_step + +__all__ = [ + "Sink", + "is_episode", + "stamp_run", + "to_mapping", + "EvalSamplesSink", + "RftMetricsSink", + "RftSamplesSink", + "TracesSink", + "training_step", +] diff --git a/packages/prime-runs/src/prime_runs/sinks/base.py b/packages/prime-runs/src/prime_runs/sinks/base.py new file mode 100644 index 000000000..70185ab36 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/base.py @@ -0,0 +1,107 @@ +"""The sink contract, the shared write loop, and the record helpers.""" + +from typing import Any, Callable, Dict, Mapping, Optional, Protocol, Sequence, Tuple + +from ..exceptions import is_record_rejection, is_transient +from ..models import RUN_KIND + + +class SinkWriteError(Exception): + """A sink stored only part of a batch. ``cause`` drives the worker's retry + and retirement policy, ``failed_records`` its loss accounting.""" + + def __init__(self, cause: Exception, *, failed_records: int) -> None: + self.cause = cause + self.failed_records = failed_records + super().__init__( + f"{failed_records} record(s) were not fully stored: {type(cause).__name__}: {cause}" + ) + + +class Sink(Protocol): + """A destination for run records. ``write`` runs on the uploader thread. A + sink that cannot write sets ``enabled = False`` and says why, once.""" + + name: str + enabled: bool + + def start(self, run_id: str, context: Mapping[str, str]) -> None: ... + + def write(self, records: Sequence[Any]) -> None: ... + + def flush(self) -> None: ... + + def close(self) -> None: ... + + +Unit = Tuple[int, Callable[[], Any]] +"""``(record count, send)``: one request's worth of a batch.""" + + +def send_each(units: Sequence[Unit]) -> None: + """Send a batch one unit at a time, then raise ``SinkWriteError`` if any + failed. A record rejection or a transient failure costs its own unit; a + sink-wide failure stops the batch and costs the rest too. A transient + error is what gets reported, so the worker's strike accounting is not + fooled by a record-local one in the same batch.""" + failed = 0 + reported: Optional[Exception] = None + for position, (count, send) in enumerate(units): + try: + send() + except Exception as exc: # noqa: BLE001 - classified below + failed += count + if not (is_record_rejection(exc) or is_transient(exc)): + failed += sum(rest for rest, _ in units[position + 1 :]) + reported = exc + break + if reported is None or is_transient(exc): + reported = exc + if reported is not None: + raise SinkWriteError(reported, failed_records=failed) from reported + + +def to_mapping(record: Any) -> Mapping[str, Any]: + """A dict passes through; anything else must implement ``to_record()``.""" + 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: + if isinstance(record, Mapping): + return "traces" in record + return hasattr(record, "traces") + + +def stamp_run(mapping: Mapping[str, Any], run_id: str, run_type: str = RUN_KIND) -> Dict[str, Any]: + """A copy of ``mapping`` keyed to this run: ``run.id`` and ``run.type`` are + set on the envelope and on every member trace, over whatever run the + producer recorded there (its own local id), and the rest of that block + (``name``, ``work``) is kept. Members matter because the traces service + reads ``run_id`` from ``trace.run.id`` only, while verifiers records the + run on the episode.""" + stamped = dict(mapping) + stamped["run"] = _rekey(stamped.get("run"), run_id, run_type) + members = stamped.get("traces") + if isinstance(members, list): + stamped["traces"] = [ + {**member, "run": _rekey(member.get("run") or stamped["run"], run_id, run_type)} + if isinstance(member, Mapping) + else member + for member in members + ] + return stamped + + +def _rekey(run: Any, run_id: str, run_type: str) -> Dict[str, Any]: + keyed = dict(run) if isinstance(run, Mapping) else {} + keyed["id"] = run_id + keyed["type"] = run_type + return keyed diff --git a/packages/prime-runs/src/prime_runs/sinks/metrics.py b/packages/prime-runs/src/prime_runs/sinks/metrics.py new file mode 100644 index 000000000..f71c675e9 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/metrics.py @@ -0,0 +1,50 @@ +"""Training metrics sink: one ``POST /rft/metrics`` per logged step. + +Requests are replayed after an ambiguous failure: the platform keeps one row per +step, so a duplicate collapses server-side, while a lost row is a hole in the +curves for good. +""" + +from typing import Any, Callable, Mapping, Optional, Sequence + +from .._http import PlatformClient +from .base import Sink, send_each + + +class RftMetricsSink(Sink): + name = "rft_metrics" + + def __init__(self, client: PlatformClient) -> None: + self.enabled = True + self._client = client + self._run_id: Optional[str] = None + self.steps_written = 0 + + def start(self, run_id: str, context: Mapping[str, str]) -> None: + self._run_id = run_id + + def write(self, records: Sequence[Any]) -> None: + if not self.enabled or not records: + return + if self._run_id is None: + raise RuntimeError("RftMetricsSink.write called before start()") + send_each([(1, self._sender(record)) for record in records]) + + def _sender(self, record: Any) -> Callable[[], None]: + def send() -> None: + if not isinstance(record, Mapping): + raise TypeError(f"metrics must be a mapping, got {type(record).__name__}") + self._client.post( + "/rft/metrics", + json_body={"run_id": self._run_id, "metrics": dict(record)}, + idempotent=True, + ) + self.steps_written += 1 + + return send + + def flush(self) -> None: + """Writes are synchronous.""" + + def close(self) -> None: + """The client belongs to the backend.""" diff --git a/packages/prime-runs/src/prime_runs/sinks/samples.py b/packages/prime-runs/src/prime_runs/sinks/samples.py new file mode 100644 index 000000000..3bec0d680 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/samples.py @@ -0,0 +1,102 @@ +"""Legacy sink: the flat eval-sample table today's viewer reads. + +``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, Callable, Dict, List, Mapping, Optional, Sequence, Set, Tuple + +from .._http import UPLOAD_TIMEOUT, PlatformClient, encode_json +from ..projection import batch_samples, build_samples +from .base import Sink, Unit, is_episode, send_each + +logger = logging.getLogger(__name__) + + +class EvalSamplesSink(Sink): + """Projects episodes to v0 samples and posts 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 + # Per-example rollout counter, carried across calls so streamed batches + # number rollouts the way a one-shot upload did. + self._rollout_numbers: Dict[Any, int] = {} + self.samples_written = 0 + #: Records with no v0 projection (JSON episodes, bare traces). They + #: reach Prime Traces only and are not a loss. + self.skipped = 0 + + def start(self, run_id: str, context: Mapping[str, str]) -> None: + 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, owners = self._to_samples(records) + units: List[Unit] = [] + offset = 0 + seen: Set[int] = set() + for batch in batch_samples(samples): + # A record whose rows straddle two batches is counted once, with the first. + batch_owners = set(owners[offset : offset + len(batch)]) - seen + seen |= batch_owners + offset += len(batch) + units.append((len(batch_owners), self._sender(batch))) + send_each(units) + + def _sender(self, batch: List[Dict[str, Any]]) -> Callable[[], None]: + def send() -> None: + self._client.post( + f"/evaluations/{self._run_id}/samples", + content=encode_json({"samples": batch}), + timeout=UPLOAD_TIMEOUT, + idempotent=False, + ) + self.samples_written += len(batch) + + return send + + def _to_samples(self, records: Sequence[Any]) -> Tuple[List[Dict[str, Any]], List[int]]: + """Project episode objects; pass v0 sample dicts through; skip the rest, + warned once. ``owners`` maps each sample back to its input record so + loss accounting survives HTTP batching.""" + samples: List[Dict[str, Any]] = [] + owners: List[int] = [] + skipped = 0 + first_skipped_type: Optional[str] = None + for record_index, record in enumerate(records): + if isinstance(record, Mapping) and "sample_id" in record: + samples.append(dict(record)) + owners.append(record_index) + elif not isinstance(record, Mapping) and is_episode(record): + projected = build_samples([record], self._rollout_numbers) + samples.extend(projected) + owners.extend([record_index] * len(projected)) + else: + skipped += 1 + if first_skipped_type is None: + first_skipped_type = type(record).__name__ + if skipped: + if not self.skipped: + logger.warning( + "The v0 sample table is projected from episode objects; %d record(s) " + "in this batch (%s) have no projection and reach Prime Traces only. " + "Further skips are counted, not logged.", + skipped, + first_skipped_type, + ) + self.skipped += skipped + return samples, owners + + def flush(self) -> None: + """Writes are synchronous.""" + + def close(self) -> None: + """The client belongs to the backend.""" diff --git a/packages/prime-runs/src/prime_runs/sinks/traces.py b/packages/prime-runs/src/prime_runs/sinks/traces.py new file mode 100644 index 000000000..90858f9ac --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/traces.py @@ -0,0 +1,137 @@ +"""Streaming sink over Prime Traces, the system of record. + +Uploads are content-addressed (a retry after a lost response replays rather +than duplicates) and episode-aware. The join key is ``run.id`` inside the trace +document; ``context`` is upload-scoped provenance only. +""" + +import logging +from typing import Any, Dict, Mapping, Optional, Sequence + +from prime_traces import ErrorCode, LineFormat, TracesClient + +from .. import _fork +from ..exceptions import ForbiddenError +from ..models import RUN_KIND +from .base import Sink, is_episode, stamp_run, to_mapping + +logger = logging.getLogger(__name__) + +DEFAULT_RECEIPT_HISTORY_SIZE = 100 + + +class TracesSink(Sink): + 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 URL (PRIME_TRACES_URL) is resolved by prime-traces. + self._client_kwargs: Dict[str, Any] = {} + if api_key is not None: + self._client_kwargs["api_key"] = api_key + 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) + + 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) -> None: + """Build the client lazily, so a fork reset is repaired on the next write.""" + if self._client is not None: + return + 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 + + def reset_after_fork(self) -> None: + """Drop (not close) the inherited client; its socket is the parent's.""" + self._client = None + + def write(self, records: Sequence[Any]) -> None: + if not self.enabled or not records: + return + self._ensure_client() + # The line format is inferred from the first record; batches are homogeneous. + line_format = LineFormat.EPISODE if is_episode(records[0]) else LineFormat.TRACE + payload = [self._prepare(record) for record in records] + try: + receipts = list( + self._client.upload_records( + payload, + line_format=line_format, + context=dict(self._context) or None, + compress=self._compress, + ) + ) + except ForbiddenError as exc: + # Nothing at runtime fixes a 403, so the sink retires either way; + # what differs is whether the batch counts as lost. + if exc.code == ErrorCode.SERVICE_NOT_ENABLED.value: + # Outside the beta: there was never anywhere for these records to go. + self._retire_quietly( + f"Prime Traces is not enabled for this account ({exc}); " + "continuing with the remaining sinks" + ) + return + # A credential without the traces scope: raised for loss accounting. + 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, with the run on the envelope and every member trace.""" + if self._run_id is None: + return record + return stamp_run(to_mapping(record), self._run_id, self._context.get("run_type", RUN_KIND)) + + def flush(self) -> None: + """Uploads are synchronous.""" + + 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) + + 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 diff --git a/packages/prime-runs/src/prime_runs/sinks/train_samples.py b/packages/prime-runs/src/prime_runs/sinks/train_samples.py new file mode 100644 index 000000000..2d3c0f754 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/sinks/train_samples.py @@ -0,0 +1,238 @@ +"""Training samples sink: the step-keyed Parquet table the training viewer reads. + +Every Nth training step (prime-rl's cadence), the step's episodes are encoded to +one Parquet object and sent through presign -> PUT -> confirm. The step is read +off each episode's ``run.work.step`` (verifiers ``TrainRunInfo``); eval work, +bare traces and JSON episodes have no row here and reach Prime Traces only. + +Objects are additive per step: every upload mints its own key and the viewer +unions every object under a step, so a step logged across several calls (an +off-policy episode landing in a later batch) becomes several objects. Each +numbers its ``sample_id``s in its own range, since the viewer looks samples up +by (step, sample_id). +""" + +import functools +import logging +import secrets +import time +from typing import Any, Dict, List, Mapping, Optional, Protocol, Sequence, Tuple + +import httpx +from prime_traces.core.client import retry_delay + +from .. import _fork +from .._http import UPLOAD_TIMEOUT, PlatformClient +from ..exceptions import APIError, RetryableAPIError, TransportError +from .base import Sink, is_episode, send_each + +logger = logging.getLogger(__name__) + +DEFAULT_STEP_INTERVAL = 10 +#: Attempts for the object-storage PUT, which is idempotent. +UPLOAD_ATTEMPTS = 3 +#: ``sample_id = range * SAMPLE_ID_STRIDE + row``. The range is the step's upload +#: index, so a single-object step still numbers from 0. A forked child inherits +#: the parent's counters and would reuse its ranges, so it draws ranges at random +#: from the upper half of the space instead. Ids stay below 2**53 for the viewer. +SAMPLE_ID_STRIDE = 1 << 32 +SAMPLE_ID_RANGES = 1 << 21 + + +class Encoder(Protocol): + """``(episodes, run_id, step, sample_id_offset=...) -> parquet bytes``, or + ``None`` when there is nothing to upload for the step.""" + + def __call__( + self, episodes: Sequence[Any], run_id: str, step: int, *, sample_id_offset: int = 0 + ) -> Optional[bytes]: ... + + +def _field(obj: Any, name: str) -> Any: + return obj.get(name) if isinstance(obj, Mapping) else getattr(obj, name, None) + + +def training_step(record: Any) -> Optional[int]: + """The step a training-work episode was dispatched at, else ``None``.""" + run = _field(record, "run") + work = _field(run, "work") if run is not None else None + if _field(run, "type") != "train" or _field(work, "type") != "train": + return None + step = _field(work, "step") + return step if isinstance(step, int) and not isinstance(step, bool) else None + + +class RftSamplesSink(Sink): + name = "rft_samples" + + def __init__( + self, + client: PlatformClient, + *, + encoder: Optional[Encoder] = None, + step_interval: int = DEFAULT_STEP_INTERVAL, + upload_client: Optional[httpx.Client] = None, + ) -> None: + if step_interval < 1: + raise ValueError("step_interval must be at least 1") + self.enabled = True + self._client = client + self._encoder = encoder + self._step_interval = step_interval + self._upload_client = upload_client + self._owns_upload_client = upload_client is None + self._forked_with_injected_client = False + self._forked = False + self._run_id: Optional[str] = None + self._uploads_per_step: Dict[int, int] = {} + self.steps_written = 0 + #: Records with no row in this table (eval work, bare traces, JSON). + self.skipped = 0 + #: Training episodes left out by the step cadence. Not a loss. + self.sampled_out = 0 + _fork.register(self) + + def reset_after_fork(self) -> None: + """Drop (not close) the inherited storage client, and stop sharing the + parent's upload counters: from here on ranges are drawn at random.""" + if self._owns_upload_client: + self._upload_client = None + else: + self._forked_with_injected_client = True + self._forked = True + + def start(self, run_id: str, context: Mapping[str, str]) -> None: + self._run_id = run_id + if self._encoder is None: + from ..projection import episodes_to_parquet_bytes, parquet_available + + if not parquet_available(): + logger.warning( + "Training samples sink off: pyarrow is not installed " + "(pip install 'prime-runs[train]'); episodes reach Prime Traces only" + ) + self.enabled = False + return + self._encoder = episodes_to_parquet_bytes + + def write(self, records: Sequence[Any]) -> None: + if not self.enabled or not records: + return + if self._run_id is None: + raise RuntimeError("RftSamplesSink.write called before start()") + send_each( + [ + (len(episodes), functools.partial(self._upload_step, step, episodes)) + for step, episodes in self._group(records) + ] + ) + + def _group(self, records: Sequence[Any]) -> List[Tuple[int, List[Any]]]: + """Training episodes by step, on the upload cadence, in step order.""" + by_step: Dict[int, List[Any]] = {} + skipped = 0 + for record in records: + step = training_step(record) + # The encoder projects attributes, so a JSON episode has no row either. + if step is None or isinstance(record, Mapping) or not is_episode(record): + skipped += 1 + continue + if step % self._step_interval: + self.sampled_out += 1 + continue + by_step.setdefault(step, []).append(record) + if skipped: + if not self.skipped: + logger.info( + "The training sample table holds training-work episodes only; %d " + "record(s) in this batch have no row there and reach Prime Traces only.", + skipped, + ) + self.skipped += skipped + return sorted(by_step.items()) + + def _upload_step(self, step: int, episodes: List[Any]) -> None: + assert self._encoder is not None and self._run_id is not None + payload = self._encoder( + episodes, self._run_id, step, sample_id_offset=self._next_range(step) * SAMPLE_ID_STRIDE + ) + if payload is None: # nothing with a trajectory + return + + # Replayable: a presign only mints a URL. + presign = self._client.post( + "/rft/samples/presign", + json_body={"run_id": self._run_id, "step": step}, + idempotent=True, + ) + data = presign.get("data") or presign + url = data.get("presignedUrl") or data.get("presigned_url") + key = data.get("s3Key") or data.get("s3_key") + if not url or not key: + raise APIError( + f"POST /rft/samples/presign returned no upload URL/key (keys: {sorted(data)})" + ) + self._put_object(url, payload) + # Replayable: confirm validates the key and refreshes the run's progress. + self._client.post( + "/rft/samples/confirm", + json_body={"run_id": self._run_id, "step": step, "s3_key": key}, + idempotent=True, + ) + self.steps_written += 1 + + def _next_range(self, step: int) -> int: + """The ``sample_id`` range for the next upload of ``step``. Counted per + attempt: an upload whose confirm was lost may still have landed.""" + if self._forked: + half = SAMPLE_ID_RANGES // 2 + return half + secrets.randbelow(half) + upload_index = self._uploads_per_step.get(step, 0) + self._uploads_per_step[step] = upload_index + 1 + return upload_index + + def _put_object(self, url: str, payload: bytes) -> None: + """PUT to storage with a bare client: the presigned URL carries its own + credentials and rejects the platform's auth headers.""" + client = self._upload() + for attempt in range(UPLOAD_ATTEMPTS): + error: APIError + try: + response = client.put( + url, + content=payload, + headers={"Content-Type": "application/parquet"}, + timeout=UPLOAD_TIMEOUT, + ) + except httpx.RequestError as exc: + error = TransportError(f"uploading samples failed: {type(exc).__name__}: {exc}") + else: + if response.is_success: + return + message = f"uploading samples failed: HTTP {response.status_code}" + if response.status_code >= 500 or response.status_code == 429: + error = RetryableAPIError(message, status_code=response.status_code) + else: + raise APIError(message, status_code=response.status_code) + if attempt == UPLOAD_ATTEMPTS - 1: + raise error + time.sleep(retry_delay(error, attempt)) + + def _upload(self) -> httpx.Client: + if self._forked_with_injected_client: + raise RuntimeError("an injected upload client cannot be reused after a fork") + if self._upload_client is None: + self._upload_client = httpx.Client(follow_redirects=False) + return self._upload_client + + def flush(self) -> None: + """Uploads are synchronous.""" + + def close(self) -> None: + client = self._upload_client + self._upload_client = None + if client is not None and self._owns_upload_client: + try: + client.close() + except Exception as exc: # noqa: BLE001 - teardown must not raise + logger.debug("Error closing the upload client: %s", exc) diff --git a/packages/prime-runs/src/prime_runs/worker.py b/packages/prime-runs/src/prime_runs/worker.py new file mode 100644 index 000000000..6e33c96d3 --- /dev/null +++ b/packages/prime-runs/src/prime_runs/worker.py @@ -0,0 +1,312 @@ +"""Background uploader: one daemon thread draining a bounded queue into sinks. + +Backpressure: a producer that outruns the uploader blocks briefly, then drops +(counted). Fork safety: a child starts over with an empty queue; the queued +records belong to the parent. Containment: record-specific failures drop only +their batch, transient failures get a few consecutive strikes before the sink +is retired (for good, or until ``retire_cooldown`` has passed), and any other +failure retires the sink at once. The thread never dies on one bad batch. +""" + +import logging +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_record_rejection, is_transient +from .sinks.base import Sink, SinkWriteError, is_episode + +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) + + +#: "Nothing pulled ahead"; ``None`` is the stop sentinel. +_NOTHING: Any = object() + + +def deadline_after(timeout: Optional[float]) -> Optional[float]: + return None if timeout is None else time.monotonic() + max(0.0, timeout) + + +def time_left(deadline: Optional[float]) -> Optional[float]: + return None if deadline is None else 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, + retire_cooldown: Optional[float] = None, + ) -> None: + self.sinks = sinks + self.max_queue_size = max_queue_size + self.put_timeout = put_timeout + self._on_error = on_error + #: Seconds after which a sink retired on transient failures is tried + #: again; ``None`` retires it for the rest of the run. + self._retire_cooldown = retire_cooldown + self._queue: "queue.Queue[Any]" = queue.Queue(maxsize=max_queue_size) + self._thread: Optional[threading.Thread] = None + self._stopping = threading.Event() + self._lock = threading.Lock() + #: Records never handed to any sink because the queue was full. + self.dropped = 0 + #: Records a sink could not store, by sink name. Kept apart from + #: ``dropped``: another sink may well have stored them. + self.failed_records: dict = {} + self._transient_failures: dict = {} + #: Sinks retired by this worker. Records that skip one are counted as + #: lost to it; a sink that switched itself off quietly is not in here. + self._retired: set = set() + #: Sink name -> monotonic time at which a paused sink is tried again. + self._cooldowns: dict = {} + _fork.register(self) + + def start(self) -> None: + with self._lock: + self._start_locked() + + def _start_locked(self) -> None: + if self._stopping.is_set(): # a closed worker is one-shot + return + if self._thread is not None and self._thread.is_alive(): + return + self._thread = threading.Thread(target=self._run, name="prime-runs-uploader", daemon=True) + self._thread.start() + + def _run(self) -> None: + # An item pulled off the queue while coalescing that could not join the + # batch; it is processed next, ahead of anything queued after it. + pending: Any = _NOTHING + while True: + item = pending if pending is not _NOTHING else self._queue.get() + pending = _NOTHING + try: + if item is None: + return + if isinstance(item, _Flush): + self._flush_sinks() + item.event.set() + else: + batch, pending = self._coalesce(item) + self._dispatch(batch) + except Exception as exc: # noqa: BLE001 - the thread must outlive one bad batch + logger.debug("Uploader iteration failed: %s", exc) + finally: + self._queue.task_done() + + def _coalesce(self, first: Sequence[Any]) -> "tuple[list, Any]": + """Merge everything already queued behind ``first`` into one batch, so + the request rate tracks upload latency rather than rollout throughput. + Stops at a control item (a flush barrier must not be reordered past the + records before it) and at a batch of the other record kind (the traces + sink infers the line format from the first record).""" + batch = list(first) + kind = is_episode(batch[0]) if batch else None + while True: + try: + item = self._queue.get_nowait() + except queue.Empty: + return batch, _NOTHING + if item is None or isinstance(item, _Flush): + return batch, item + if item and (kind is None or is_episode(item[0]) != kind): + return batch, item + batch.extend(item) + self._queue.task_done() + + def _dispatch(self, records: Sequence[Any]) -> None: + for sink in self.sinks: + self._maybe_revive(sink) + 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 SinkWriteError as exc: + self._fail_sink(sink, exc.cause, dropped=exc.failed_records) + except Exception as exc: # noqa: BLE001 - one sink failing must not stop the others + self._fail_sink(sink, exc, dropped=len(records)) + else: + self._transient_failures.pop(sink.name, None) + + def _maybe_revive(self, sink: Sink) -> None: + until = self._cooldowns.get(sink.name) + if until is None or time.monotonic() < until: + return + del self._cooldowns[sink.name] + self._retired.discard(sink.name) + self._transient_failures.pop(sink.name, None) + sink.enabled = True + logger.info("Sink %s re-enabled after its cooldown", sink.name) + + def _flush_sinks(self) -> None: + for sink in self.sinks: + if not sink.enabled: + 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: + name = sink.name + if dropped: + self.failed_records[name] = self.failed_records.get(name, 0) + dropped + + if dropped and is_record_rejection(exc): + logger.warning( + "Sink %s rejected a batch of %d record(s) (%s: %s); " + "still enabled for later records.", + name, + dropped, + type(exc).__name__, + exc, + ) + self._notify(name, exc) + return + + if is_transient(exc): + strikes = self._transient_failures.get(name, 0) + 1 + self._transient_failures[name] = strikes + 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) + if self._retire_cooldown is None: + logger.warning( + "Sink %s disabled after %d consecutive transient failures: %s: %s", + name, + strikes, + type(exc).__name__, + exc, + ) + else: + self._cooldowns[name] = time.monotonic() + self._retire_cooldown + logger.warning( + "Sink %s paused for %.0fs after %d consecutive transient failures: %s: %s", + name, + self._retire_cooldown, + strikes, + type(exc).__name__, + exc, + ) + else: + sink.enabled = False + self._retired.add(name) + 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) + + def submit(self, records: Sequence[Any]) -> bool: + """Queue a batch. ``False`` means it was dropped: the queue stayed full + for the whole put timeout.""" + if not self.sinks: # a disabled run: nowhere for the records to go + return True + # Serialize with close(): an accepted batch is queued before the stop sentinel. + with self._lock: + if self._stopping.is_set(): + return False + self._start_locked() + try: + self._queue.put(records, timeout=self.put_timeout) + return True + except queue.Full: + count = len(records) + self.dropped += count + logger.warning( + "Upload queue full after %.1fs; dropped %d item(s) (%d total). " + "The producer is outrunning the uploader.", + self.put_timeout, + count, + self.dropped, + ) + return False + + def flush(self, timeout: Optional[float] = None) -> bool: + """Block until everything queued so far has been written.""" + if self._thread is None or not self._thread.is_alive(): + self._flush_sinks() + return True + deadline = deadline_after(timeout) + barrier = _Flush() + try: + # The barrier gets the caller's drain budget, not the producer put timeout. + self._queue.put(barrier, timeout=time_left(deadline)) + except queue.Full: + logger.warning("Could not enqueue a flush barrier before the drain timeout") + return False + return barrier.event.wait(time_left(deadline)) + + def close(self, timeout: Optional[float] = 30.0) -> None: + """Drain, stop the thread, and close every sink.""" + with self._lock: + self._stopping.set() + thread = self._thread + if thread is not None and thread.is_alive(): + deadline = deadline_after(timeout) + try: + self._queue.put(None, timeout=time_left(deadline)) + except queue.Full: + logger.warning("Upload queue remained saturated through the close timeout") + thread.join(time_left(deadline)) + if thread.is_alive(): + # Closing the sinks now would pull a client out from under a request. + logger.warning( + "Uploader still running after %ss; leaving it and its sinks open. " + "Records still in flight may not finish before the process exits.", + 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) + + def reset_after_fork(self) -> None: + 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..1a419798a --- /dev/null +++ b/packages/prime-runs/tests/_fakes.py @@ -0,0 +1,163 @@ +"""Stand-ins for verifiers ``Trace``/``Episode``.""" + +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 EnvInfo: + id: str = "gsm8k" + name: Optional[str] = None + + +@dataclass +class TrainWorkInfo: + step: int + type: str = "train" + + +@dataclass +class EvalWorkInfo: + step: int + type: str = "eval" + + +@dataclass +class TrainRunInfo: + """verifiers stamps this on every episode a training run dispatches.""" + + id: str + work: Any + type: str = "train" + name: Optional[str] = None + + +@dataclass +class Episode: + id: str = "episode-1" + traces: List[Trace] = field(default_factory=list) + ok: bool = True + env: EnvInfo = field(default_factory=EnvInfo) + run: Optional[Any] = None + + def to_record(self) -> Dict[str, Any]: + return {"id": self.id, "traces": [trace.to_record() for trace in self.traces]} + + +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) + + +def make_train_episode( + episode_id: str = "episode-1", + *, + step: int = 10, + work: str = "train", + idx: int = 0, + reward: float = 1.0, + advantage: Optional[float] = None, + env_id: str = "gsm8k", + run_id: str = "run-1", +) -> Episode: + """An episode the way prime-rl hands it over: stamped with the run and the + training step it was dispatched at, ``advantage`` in the trace's info.""" + trace = make_trace(trace_id=f"{episode_id}-t", idx=idx, reward=reward) + if advantage is not None: + trace.info["advantage"] = advantage + work_info = TrainWorkInfo(step=step) if work == "train" else EvalWorkInfo(step=step) + return Episode( + id=episode_id, + traces=[trace], + env=EnvInfo(id=env_id), + run=TrainRunInfo(id=run_id, work=work_info), + ) diff --git a/packages/prime-runs/tests/conftest.py b/packages/prime-runs/tests/conftest.py new file mode 100644 index 000000000..a24a62587 --- /dev/null +++ b/packages/prime-runs/tests/conftest.py @@ -0,0 +1,195 @@ +"""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 + + +@pytest.fixture +def rft_routes() -> Dict[str, Any]: + """The happy path for a training run: register, metrics, samples, finalize.""" + + def presign(request: httpx.Request) -> httpx.Response: + import json + + step = json.loads(request.content)["step"] + return httpx.Response( + 200, + json={ + "data": { + "presignedUrl": f"http://storage.test/run-abc/step-{step}.parquet", + "s3Key": f"runs/run-abc/step-{step}.parquet", + "expiresIn": 900, + } + }, + ) + + return { + "POST /api/v1/rft/external-runs": lambda request: httpx.Response( + 201, + json={"run": {"id": "run-abc", "name": "test-run", "status": "RUNNING"}}, + ), + "POST /api/v1/rft/metrics": {"data": {"status": "success"}}, + "POST /api/v1/rft/samples/presign": presign, + "POST /api/v1/rft/samples/confirm": {"data": {"status": "success"}}, + "POST /api/v1/rft/finalize": { + "data": {"status": "success", "message": "Run finalized successfully"} + }, + "PUT /api/v1/rft/external-runs/run-abc/status": { + "run": {"id": "run-abc", "status": "FAILED"} + }, + } + + +class StorageHandler: + """Object storage for presigned PUTs: records every upload, answers 200.""" + + def __init__(self, status_codes: Sequence[int] = ()) -> None: + self.uploads: List[httpx.Request] = [] + self._status_codes = list(status_codes) + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.uploads.append(request) + code = self._status_codes.pop(0) if self._status_codes else 200 + return httpx.Response(code) + + def client(self) -> httpx.Client: + return httpx.Client(transport=httpx.MockTransport(self)) diff --git a/packages/prime-runs/tests/test_config_source.py b/packages/prime-runs/tests/test_config_source.py new file mode 100644 index 000000000..9daf1f2d8 --- /dev/null +++ b/packages/prime-runs/tests/test_config_source.py @@ -0,0 +1,193 @@ +"""The config a run is actually configured with: the file someone wrote, kept""" + +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, +) +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): + 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): + 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): + 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 + + +# ------------------------------------------------------- config normalization + + +def test_a_mapping_is_taken_exactly_as_given(): + assert _normalize_config({"a": 1, "b": None}) == {"a": 1, "b": None} + + +def test_a_path_becomes_a_stored_source(tmp_path): + 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(): + 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): + 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): + 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" 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..b8ccea974 --- /dev/null +++ b/packages/prime-runs/tests/test_evals_backend.py @@ -0,0 +1,192 @@ +"""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 +): + 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): + 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): + 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): + 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_a_pinned_environment_version_reaches_the_api(make_platform_client, eval_routes): + 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): + 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..7105a5978 --- /dev/null +++ b/packages/prime-runs/tests/test_http.py @@ -0,0 +1,208 @@ +"""Transport behaviour: error mapping and retry.""" + +import httpx +import pytest + +from prime_runs._http import PlatformClient +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): + 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(): + 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_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_retry_after_is_honoured(no_sleep): + 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_an_ambiguous_failure_does_not_replay_a_create(no_sleep): + 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): + 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): + 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): + 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): + 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): + 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..de301ce7d --- /dev/null +++ b/packages/prime-runs/tests/test_init.py @@ -0,0 +1,437 @@ +"""``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 RunSpec, RunStatus + +# -------------------------------------------------------------------- modes + + +def test_no_api_key_disables_the_run_and_says_so(caplog): + 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): + run = pr.init(mode="disabled") + + run.log_traces([{"id": "t1"}]) + # No sinks, so nothing was queued and no uploader thread was started for it. + assert run._worker._thread is None + run.finish(summary={"avg_reward": 1.0}) + + assert run.id.startswith("disabled-") + assert run.url is None + assert run.status is RunStatus.COMPLETED + assert run.dropped_records == 0 + assert not list(tmp_path.iterdir()) # tmp_path is $HOME here + + +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_eval_uploaders_keep_the_permanent_retirement(online): + run, _ = online() + + assert run._worker._retire_cooldown is None + run.finish() + + +def test_a_failed_create_closes_the_platform_client(monkeypatch): + class FailingClient: + def __init__(self): + self.closed = False + + def post(self, *args, **kwargs): + raise pr.APIError("create failed") + + def close(self): + self.closed = True + + client = FailingClient() + monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: client) + + with pytest.raises(pr.APIError, match="create failed"): + pr.init( + name="test-run", + environments=[{"id": "env-123"}], + api_key="test-key", + mode="online", + ) + + assert client.closed is True + + +def test_an_online_run_has_both_transports_by_default( + monkeypatch, make_platform_client, eval_routes +): + 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_the_end_to_end_shape_a_producer_writes(online): + 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)) + + # Episodes queued while a request is in flight go out together, so the + # request count depends on timing; every row landing exactly once does not. + posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples") + assert 1 <= len(posted) <= 3 + assert sorted(s["sample_id"] for body in posted for s in body["samples"]) == [ + "ep-0", + "ep-1", + "ep-2", + ] + finalize = handler.bodies_for("/api/v1/evaluations/eval-abc/finalize")[0] + assert finalize["metrics"]["avg_reward"] == 1.0 + assert run.status is RunStatus.COMPLETED + assert run.errors == [] + + +# --------------------------------------------------------------------- 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): + 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() + + +# ----------------------------------------------------------------- training + + +@pytest.fixture +def online_train(monkeypatch, make_platform_client, rft_routes): + """``init(kind="train", mode="online")`` wired to a MockTransport: traces + sink off, sample uploads to an in-memory store, a recording encoder.""" + from conftest import StorageHandler + + from prime_runs.sinks import RftSamplesSink + + storage = StorageHandler() + encoder_calls = [] + + def encoder(episodes, run_id, step, sample_id_offset=0): + encoder_calls.append((len(episodes), run_id, step)) + return b"parquet" + + def samples_sink(client, **kwargs): + return RftSamplesSink(client, encoder=encoder, upload_client=storage.client(), **kwargs) + + def _init(routes=None, **kwargs): + handler = RecordingHandler(routes or rft_routes) + monkeypatch.setattr( + "prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler) + ) + monkeypatch.setattr("prime_runs.run.TracesSink", lambda **_: _NullSink()) + monkeypatch.setattr("prime_runs.run.RftSamplesSink", samples_sink) + # The RFT create response carries no viewer URL; the SDK builds one. + monkeypatch.setenv("PRIME_FRONTEND_URL", "https://app.example") + kwargs.setdefault("name", "test-run") + kwargs.setdefault("model", "Qwen/Qwen3-8B") + kwargs.setdefault("environments", ["primeintellect/vf-math"]) + kwargs.setdefault("team_id", "team-1") + run = pr.init(kind="train", api_key="test-key", **kwargs) + return run, handler, storage, encoder_calls + + return _init + + +def test_a_training_run_registers_with_the_rft_api(online_train): + run, handler, _, _ = online_train( + training=pr.TrainingSpec(max_steps=50, batch_size=32), + config={"trainer": {"lr": 1e-6}}, + ) + + assert run.kind == "train" + assert run.id == "run-abc" + assert run.url == "https://app.example/dashboard/training/run-abc" + assert run.attached is False + body = handler.bodies_for("/api/v1/rft/external-runs")[0] + assert body["base_model"] == "Qwen/Qwen3-8B" + assert body["max_steps"] == 50 + assert body["batch_size"] == 32 + assert body["environments"] == [{"id": "primeintellect/vf-math"}] + assert body["run_config"] == {"trainer": {"lr": 1e-6}} + assert body["team_id"] == "team-1" + assert [sink.name for sink in run._worker.sinks] == ["traces", "rft_samples"] + assert [sink.name for sink in run._metrics_worker.sinks] == ["rft_metrics"] + run.finish() + + +def test_a_training_run_streams_metrics_and_step_samples(online_train): + from _fakes import make_train_episode + + run, handler, storage, encoder_calls = online_train() + + run.log_metrics({"loss": 0.5}, step=10) + run.log_episodes([make_train_episode("e1", step=10), make_train_episode("e2", step=10)]) + run.finish() + + metrics = handler.bodies_for("/api/v1/rft/metrics") + assert len(metrics) == 1 + assert metrics[0]["run_id"] == "run-abc" + assert metrics[0]["metrics"]["loss"] == 0.5 + assert metrics[0]["metrics"]["step"] == 10 + assert "_timestamp" in metrics[0]["metrics"] + assert encoder_calls == [(2, "run-abc", 10)] + assert len(storage.uploads) == 1 + assert handler.bodies_for("/api/v1/rft/samples/confirm")[0]["step"] == 10 + # Closed out through the idempotent finalize, not the status PUT. + assert handler.bodies_for("/api/v1/rft/finalize") == [{"run_id": "run-abc", "exit_code": 0}] + assert "PUT /api/v1/rft/external-runs/run-abc/status" not in handler.paths() + assert run.status is RunStatus.COMPLETED + assert run.errors == [] + + +def test_a_failed_training_run_is_marked_failed_with_the_reason(online_train): + run, handler, _, _ = online_train() + + with pytest.raises(ValueError): + with run: + raise ValueError("loss is NaN") + + body = handler.bodies_for("/api/v1/rft/external-runs/run-abc/status")[0] + assert body == {"status": "failed", "error_message": "ValueError: loss is NaN"} + assert "POST /api/v1/rft/finalize" not in handler.paths() + + +def test_training_uploaders_pause_rather_than_retire_after_an_outage(online_train): + from prime_runs.run import TRAIN_RETIRE_COOLDOWN + + run, _, _, _ = online_train() + + assert run._worker._retire_cooldown == TRAIN_RETIRE_COOLDOWN == 300.0 + assert run._metrics_worker._retire_cooldown == TRAIN_RETIRE_COOLDOWN + run.finish() + + +def test_attaching_to_a_managed_run_registers_nothing(online_train): + run, handler, _, _ = online_train(id="run-managed") + + assert run.id == "run-managed" + assert run.attached is True + assert run.url == "https://app.example/dashboard/training/run-managed" + assert "POST /api/v1/rft/external-runs" not in handler.paths() + run.finish() + # A clean exit still completes it. + assert handler.bodies_for("/api/v1/rft/finalize") == [{"run_id": "run-managed", "exit_code": 0}] + + +def test_an_attached_run_leaves_failure_marking_to_the_launcher(online_train, caplog): + run, handler, _, _ = online_train(id="run-managed") + + with caplog.at_level("INFO"): + run.fail("boom") + + assert not any("status" in path for path in handler.paths()) + assert "POST /api/v1/rft/finalize" not in handler.paths() + assert "leaving it to the launcher" in caplog.text + assert run.status is RunStatus.FAILED + + +def test_a_training_run_needs_a_team(monkeypatch, make_platform_client, rft_routes): + handler = RecordingHandler(rft_routes) + monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler)) + + with pytest.raises(ConfigurationError, match="team"): + pr.init(kind="train", model="Qwen/Qwen3-8B", api_key="test-key", mode="online") + assert handler.paths() == [] + + +def test_a_team_outside_the_allowlist_surfaces_the_platforms_reason( + monkeypatch, make_platform_client, rft_routes +): + import httpx + + routes = dict(rft_routes) + routes["POST /api/v1/rft/external-runs"] = lambda request: httpx.Response( + 403, json={"detail": "External training runs are not enabled for this team"} + ) + handler = RecordingHandler(routes) + monkeypatch.setattr("prime_runs.run.PlatformClient", lambda **_: make_platform_client(handler)) + + with pytest.raises(pr.ForbiddenError, match="not enabled for this team"): + pr.init(kind="train", model="m", team_id="team-1", api_key="test-key", mode="online") + + +def test_a_disabled_training_run_keeps_an_attached_id(tmp_path): + run = pr.init(kind="train", mode="disabled", id="run-managed", model="m") + + assert run.id == "run-managed" + assert run.kind == "train" + run.log_metrics({"loss": 0.1}, step=1) + assert run._metrics_worker._thread is None + run.finish() + assert run.status is RunStatus.COMPLETED + + +@pytest.mark.parametrize( + "kwargs", + [ + {"kind": "eval", "id": "eval-1"}, + {"kind": "eval", "training": pr.TrainingSpec()}, + {"kind": "serve"}, + ], + ids=["eval-with-id", "eval-with-training", "unknown-kind"], +) +def test_kind_arguments_are_checked_before_anything_else(kwargs): + with pytest.raises(ConfigurationError): + pr.init(mode="disabled", **kwargs) + + +def test_training_records_are_stamped_with_the_train_run_type(): + from prime_runs.run import _sink_context + from prime_runs.sinks.base import stamp_run + + spec = RunSpec(kind="train", model="m") + context = _sink_context(spec) + + assert context["run_type"] == "train" + assert stamp_run({"id": "t1"}, "run-abc", context["run_type"])["run"] == { + "id": "run-abc", + "type": "train", + } diff --git a/packages/prime-runs/tests/test_metrics.py b/packages/prime-runs/tests/test_metrics.py new file mode 100644 index 000000000..ca2888248 --- /dev/null +++ b/packages/prime-runs/tests/test_metrics.py @@ -0,0 +1,63 @@ +"""Run-level aggregates, moved here from verifiers.""" + +from _fakes import Reward, make_episode, make_trace + +from prime_runs import metrics + + +def test_only_trainable_traces_are_scored(): + 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(): + 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(): + 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(): + 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_metrics_sink.py b/packages/prime-runs/tests/test_metrics_sink.py new file mode 100644 index 000000000..905714b60 --- /dev/null +++ b/packages/prime-runs/tests/test_metrics_sink.py @@ -0,0 +1,93 @@ +"""The training metrics sink: one request per logged step.""" + +import httpx +import pytest +from conftest import RecordingHandler + +from prime_runs.exceptions import RetryableAPIError, UnauthorizedError +from prime_runs.sinks import RftMetricsSink +from prime_runs.sinks.base import SinkWriteError + + +def make_sink(make_platform_client, routes): + handler = RecordingHandler(routes) + sink = RftMetricsSink(make_platform_client(handler)) + sink.start("run-abc", {}) + return sink, handler + + +def test_each_metrics_dict_is_one_request(make_platform_client, rft_routes): + sink, handler = make_sink(make_platform_client, rft_routes) + + sink.write([{"step": 1, "loss": 0.9}, {"step": 2, "loss": 0.8}]) + + assert handler.bodies_for("/api/v1/rft/metrics") == [ + {"run_id": "run-abc", "metrics": {"step": 1, "loss": 0.9}}, + {"run_id": "run-abc", "metrics": {"step": 2, "loss": 0.8}}, + ] + assert sink.steps_written == 2 + + +def test_an_ambiguous_failure_is_replayed(make_platform_client, rft_routes, no_sleep): + calls = [] + + def flaky(request: httpx.Request) -> httpx.Response: + calls.append(request) + if len(calls) == 1: + return httpx.Response(502) + return httpx.Response(200, json={"data": {"status": "success"}}) + + routes = dict(rft_routes) + routes["POST /api/v1/rft/metrics"] = flaky + sink, handler = make_sink(make_platform_client, routes) + + sink.write([{"step": 1}, {"step": 2}]) + + assert len(calls) == 3 + assert sink.steps_written == 2 + + +def test_a_persistent_transient_failure_loses_only_its_step( + make_platform_client, rft_routes, no_sleep +): + import json + + def down_for_step_one(request: httpx.Request) -> httpx.Response: + if json.loads(request.content)["metrics"]["step"] == 1: + return httpx.Response(502) + return httpx.Response(200, json={"data": {"status": "success"}}) + + routes = dict(rft_routes) + routes["POST /api/v1/rft/metrics"] = down_for_step_one + sink, handler = make_sink(make_platform_client, routes) + + with pytest.raises(SinkWriteError) as info: + sink.write([{"step": 1}, {"step": 2}]) + + assert info.value.failed_records == 1 + assert isinstance(info.value.cause, RetryableAPIError) + assert sink.steps_written == 1 + + +def test_a_sink_wide_failure_stops_the_batch_and_counts_the_rest(make_platform_client, rft_routes): + routes = dict(rft_routes) + routes["POST /api/v1/rft/metrics"] = lambda request: httpx.Response(401) + sink, handler = make_sink(make_platform_client, routes) + + with pytest.raises(SinkWriteError) as info: + sink.write([{"step": 1}, {"step": 2}, {"step": 3}]) + + assert info.value.failed_records == 3 + assert isinstance(info.value.cause, UnauthorizedError) + assert len(handler.paths()) == 1 + + +def test_a_record_that_is_not_a_mapping_is_counted_not_sent(make_platform_client, rft_routes): + sink, handler = make_sink(make_platform_client, rft_routes) + + with pytest.raises(SinkWriteError) as info: + sink.write(["loss=0.5", {"step": 1}]) + + assert info.value.failed_records == 1 + assert isinstance(info.value.cause, TypeError) + assert len(handler.bodies_for("/api/v1/rft/metrics")) == 1 diff --git a/packages/prime-runs/tests/test_projection.py b/packages/prime-runs/tests/test_projection.py new file mode 100644 index 000000000..06073f7eb --- /dev/null +++ b/packages/prime-runs/tests/test_projection.py @@ -0,0 +1,135 @@ +"""The v0 eval-sample projection, moved here from verifiers.""" + +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(): + trace = make_trace(reward=0.25, rewards={"reward": Reward(9.0)}) + + assert trace_to_sample(trace)["reward"] == 0.25 + + +def test_non_finite_numbers_become_null_rather_than_losing_the_row(): + trace = make_trace(reward=float("nan"), metrics={"kl": float("inf")}) + trace.rewards = {"format": Reward(float("-inf"))} + + sample = trace_to_sample(trace) + (row,) = build_samples([make_episode("ep-1", [trace])]) + + assert sample["reward"] is None + assert sample["metrics"] == {"kl": None} + assert sample["format"] is None + assert row["info"]["native_wrapper"]["traces"][0]["reward"] is None + assert batch_samples([row]) # encodes: nothing non-finite is left + + +def test_trajectory_keeps_one_entry_per_branch(): + trace = make_trace(branches=3) + + 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(): + 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(): + 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(): + oversized = {"sample_id": "x", "blob": "a" * (MAX_SAMPLES_PAYLOAD_BYTES + 1)} + + with pytest.raises(ValueError, match="too large"): + batch_samples([oversized]) + + +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_rft_backend.py b/packages/prime-runs/tests/test_rft_backend.py new file mode 100644 index 000000000..8238e7171 --- /dev/null +++ b/packages/prime-runs/tests/test_rft_backend.py @@ -0,0 +1,214 @@ +"""Training run lifecycle against ``/api/v1/rft/*``.""" + +import httpx +import pytest +from conftest import RecordingHandler + +from prime_runs.backend import RftBackend +from prime_runs.exceptions import ConfigurationError, ForbiddenError, RetryableAPIError +from prime_runs.models import EnvironmentRef, RunSpec, RunStatus, TrainingSpec + + +def make_backend(make_platform_client, routes, **kwargs): + handler = RecordingHandler(routes) + client = make_platform_client(handler) + kwargs.setdefault("team_id", "team-1") + backend = RftBackend(client, frontend_url="https://app.example", **kwargs) + return backend, handler + + +def train_spec(**overrides) -> RunSpec: + values = dict( + name="test-run", + kind="train", + model="Qwen/Qwen3-8B", + environments=[EnvironmentRef.coerce("primeintellect/vf-math")], + training=TrainingSpec(max_steps=100, batch_size=64, rollouts_per_example=8, seq_len=4096), + config={"trainer": {"lr": 1e-6}}, + ) + values.update(overrides) + return RunSpec(**values) + + +def test_create_registers_an_external_run_for_the_team(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes) + + handle = backend.create(train_spec()) + + assert handler.paths() == ["POST /api/v1/rft/external-runs"] + body = handler.bodies_for("/api/v1/rft/external-runs")[0] + assert body["base_model"] == "Qwen/Qwen3-8B" + assert body["max_steps"] == 100 + assert body["batch_size"] == 64 + assert body["rollouts_per_example"] == 8 + assert body["seq_len"] == 4096 + assert body["run_config"] == {"trainer": {"lr": 1e-6}} + assert body["team_id"] == "team-1" + assert body["name"] == "test-run" + assert handle.id == "run-abc" + assert handle.url == "https://app.example/dashboard/training/run-abc" + + +@pytest.mark.parametrize( + ("environment", "expected"), + [ + ("primeintellect/vf-math", {"id": "primeintellect/vf-math"}), + ("agentic-judge+gsm8k", {"id": "agentic-judge+gsm8k"}), + ({"id": "env-123", "version_id": "v7"}, {"id": "env-123", "version_id": "v7"}), + ], + ids=["slug", "name", "id"], +) +def test_training_environments_are_passed_through_not_resolved( + make_platform_client, rft_routes, environment, expected +): + backend, handler = make_backend(make_platform_client, rft_routes) + + backend.create(train_spec(environments=[EnvironmentRef.coerce(environment)])) + + assert handler.bodies_for("/api/v1/rft/external-runs")[0]["environments"] == [expected] + assert "POST /api/v1/environmentshub/resolve" not in handler.paths() + + +def test_the_spec_team_wins_over_the_backend_default(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes, team_id="team-default") + + backend.create(train_spec(team_id="team-explicit")) + + assert handler.bodies_for("/api/v1/rft/external-runs")[0]["team_id"] == "team-explicit" + + +def test_create_needs_a_base_model(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes) + + with pytest.raises(ConfigurationError, match="model="): + backend.create(train_spec(model=None)) + assert handler.paths() == [] + + +def test_create_needs_a_team(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes, team_id=None) + + with pytest.raises(ConfigurationError, match="team"): + backend.create(train_spec()) + assert handler.paths() == [] + + +def test_a_team_outside_the_allowlist_is_a_forbidden_error(make_platform_client, rft_routes): + routes = dict(rft_routes) + routes["POST /api/v1/rft/external-runs"] = lambda request: httpx.Response( + 403, json={"detail": "External training runs are not enabled for this team"} + ) + backend, handler = make_backend(make_platform_client, routes) + + with pytest.raises(ForbiddenError, match="not enabled for this team"): + backend.create(train_spec()) + assert len(handler.paths()) == 1 + + +def test_create_is_not_replayed_after_an_ambiguous_failure(make_platform_client, rft_routes): + routes = dict(rft_routes) + routes["POST /api/v1/rft/external-runs"] = lambda request: httpx.Response(502) + backend, handler = make_backend(make_platform_client, routes) + + with pytest.raises(RetryableAPIError): + backend.create(train_spec()) + assert handler.paths() == ["POST /api/v1/rft/external-runs"] + + +def test_attach_makes_no_request(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes) + + handle = backend.attach("run-managed") + + assert handle.id == "run-managed" + assert handle.url == "https://app.example/dashboard/training/run-managed" + assert handler.paths() == [] + + +def test_update_is_a_no_op(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes) + + backend.update("run-abc", config={"a": 1}, summary={"loss": 0.1}) + + assert handler.paths() == [] + + +def test_a_completed_run_is_finalized_with_exit_code_zero(make_platform_client, rft_routes): + backend, handler = make_backend(make_platform_client, rft_routes) + + backend.finalize("run-abc", status=RunStatus.COMPLETED, summary={"loss": 0.1}) + + assert handler.paths() == ["POST /api/v1/rft/finalize"] + assert handler.bodies_for("/api/v1/rft/finalize")[0] == {"run_id": "run-abc", "exit_code": 0} + + +def test_finalize_is_replayed_after_an_ambiguous_failure(make_platform_client, rft_routes): + attempts = [] + + def flaky(request: httpx.Request) -> httpx.Response: + attempts.append(request) + if len(attempts) == 1: + return httpx.Response(504) + return httpx.Response(200, json={"data": {"status": "success"}}) + + routes = dict(rft_routes) + routes["POST /api/v1/rft/finalize"] = flaky + backend, handler = make_backend(make_platform_client, routes) + + backend.finalize("run-abc", status=RunStatus.COMPLETED) + + assert len(attempts) == 2 + + +def test_finalize_falls_back_to_a_completed_status_update(make_platform_client, rft_routes): + routes = dict(rft_routes) + routes["POST /api/v1/rft/finalize"] = lambda request: httpx.Response( + 400, json={"detail": "finalize unavailable"} + ) + backend, handler = make_backend(make_platform_client, routes) + + backend.finalize("run-abc", status=RunStatus.COMPLETED) + + assert handler.paths() == [ + "POST /api/v1/rft/finalize", + "PUT /api/v1/rft/external-runs/run-abc/status", + ] + assert handler.bodies_for("/api/v1/rft/external-runs/run-abc/status") == [ + {"status": "completed"} + ] + + +@pytest.mark.parametrize( + ("status", "error", "expected_message"), + [ + (RunStatus.FAILED, "ValueError: boom", "ValueError: boom"), + (RunStatus.CANCELLED, "interrupted", "cancelled: interrupted"), + (RunStatus.CRASHED, "process exited without finishing the run", "crashed: process exited"), + (RunStatus.FAILED, None, "failed"), + ], + ids=["failed", "cancelled", "crashed", "failed-no-reason"], +) +def test_other_terminal_states_are_reported_as_failed_with_the_reason( + make_platform_client, rft_routes, status, error, expected_message +): + backend, handler = make_backend(make_platform_client, rft_routes) + + backend.finalize("run-abc", status=status, error=error) + + assert handler.paths() == ["PUT /api/v1/rft/external-runs/run-abc/status"] + body = handler.bodies_for("/api/v1/rft/external-runs/run-abc/status")[0] + assert body["status"] == "failed" + assert body["error_message"].startswith(expected_message) + assert "POST /api/v1/rft/finalize" not in handler.paths() + + +def test_marking_an_already_closed_run_failed_is_not_an_error(make_platform_client, rft_routes): + routes = dict(rft_routes) + routes["PUT /api/v1/rft/external-runs/run-abc/status"] = lambda request: httpx.Response( + 409, json={"detail": "Cannot update status of run in 'COMPLETED' state."} + ) + backend, handler = make_backend(make_platform_client, routes) + + backend.finalize("run-abc", status=RunStatus.FAILED, error="boom") + + assert handler.paths() == ["PUT /api/v1/rft/external-runs/run-abc/status"] diff --git a/packages/prime-runs/tests/test_run.py b/packages/prime-runs/tests/test_run.py new file mode 100644 index 000000000..f6ba0a051 --- /dev/null +++ b/packages/prime-runs/tests/test_run.py @@ -0,0 +1,688 @@ +"""The run handle: lifecycle, containment, terminal status.""" + +import asyncio +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_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["run_type"] == "eval" + 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_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_non_finite_summary_values_are_dropped_rather_than_failing_the_request(): + 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(): + 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 + + +@pytest.mark.parametrize( + "operation", + [ + lambda run: run.log_traces([{"id": "t1"}]), + lambda run: run.log_episodes([{"id": "ep1", "traces": []}]), + lambda run: run.log_metrics({"loss": 0.5}), + lambda run: run.update_summary({"late": 1.0}), + ], + ids=["traces", "episodes", "metrics", "summary"], +) +def test_writes_after_finish_are_a_producer_bug(operation): + run = make_run(metrics_sinks=[FakeSink("metrics")]) + run.finish() + + with pytest.raises(RunFinishedError): + operation(run) + + +def test_logging_that_started_first_is_queued_before_concurrent_finish(): + class SignalingBackend(FakeBackend): + def __init__(self) -> None: + super().__init__() + self.finalize_started = threading.Event() + + def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> None: + self.finalize_started.set() + super().finalize( + run_id, + status=status, + summary=summary, + error=error, + config=config, + ) + + backend = SignalingBackend() + sink = FakeSink() + run = make_run(backend, sinks=[sink]) + submit_started = threading.Event() + release_submit = threading.Event() + finish_called = threading.Event() + original_submit = run._worker.submit + + def blocking_submit(records): + submit_started.set() + assert release_submit.wait(2.0) + return original_submit(records) + + run._worker.submit = blocking_submit + + def finish_run() -> None: + finish_called.set() + run.finish() + + with ThreadPoolExecutor(max_workers=2) as executor: + logging = executor.submit(run.log_traces, [{"id": "t1"}]) + assert submit_started.wait(2.0) + finishing = executor.submit(finish_run) + assert finish_called.wait(2.0) + try: + assert not backend.finalize_started.wait(0.1) + finally: + release_submit.set() + logging.result(timeout=2.0) + finishing.result(timeout=2.0) + + assert sink.batches == [[{"id": "t1"}]] + assert sink.closed is True + + +@pytest.mark.parametrize( + "operation", + [ + lambda run: run.log_traces([{"id": "t1"}]), + lambda run: run.log_episodes([{"id": "ep1", "traces": []}]), + lambda run: run.update_summary({"reward": 1.0}), + ], + ids=["traces", "episodes", "summary"], +) +def test_writes_during_finish_fail_without_waiting_for_teardown(operation): + class BlockingBackend(FakeBackend): + def __init__(self) -> None: + super().__init__() + self.finalize_started = threading.Event() + self.release_finalize = threading.Event() + + def finalize(self, run_id, *, status, summary=None, error=None, config=None) -> None: + self.finalize_started.set() + assert self.release_finalize.wait(2.0) + super().finalize( + run_id, + status=status, + summary=summary, + error=error, + config=config, + ) + + backend = BlockingBackend() + run = make_run(backend) + + with ThreadPoolExecutor(max_workers=2) as executor: + finishing = executor.submit(run.finish) + assert backend.finalize_started.wait(2.0) + writing = executor.submit(operation, run) + try: + with pytest.raises(RunFinishedError): + writing.result(timeout=0.2) + finally: + backend.release_finalize.set() + finishing.result(timeout=2.0) + + +def test_the_context_manager_completes_a_clean_run(): + backend = FakeBackend() + with make_run(backend) as run: + 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() + + +@pytest.mark.parametrize("interrupt", [KeyboardInterrupt, asyncio.CancelledError]) +def test_an_interrupt_is_recorded_as_a_decision_not_a_fault(interrupt): + backend = FakeBackend() + + with pytest.raises(interrupt): + with make_run(backend): + raise interrupt + + assert backend.finalized[0]["status"] is RunStatus.CANCELLED + assert backend.finalized[0]["error"] == "interrupted" + + +def test_cancelled_is_terminal_and_finalizes_like_failed(): + backend = FakeBackend() + run = make_run(backend) + + assert RunStatus.CANCELLED.is_terminal() + run.finish(status="cancelled", error="operator stopped it") + + assert run.status is RunStatus.CANCELLED + assert backend.finalized[0]["status"] is RunStatus.CANCELLED + assert backend.finalized[0]["error"] == "operator stopped it" + + +def test_finish_timeout_is_set_per_run_and_overridden_per_call(): + run = make_run(finish_timeout=1.5) + observed = {} + + def record_flush(timeout=None): + observed["flush"] = timeout + return True + + run._worker.flush = record_flush + run.finish(timeout=0.25) + + assert run._finish_timeout == 1.5 + assert 0.0 < observed["flush"] <= 0.25 + + +def test_finish_hands_the_full_config_to_finalize(): + 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): + 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(): + 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(): + 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(): + 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_an_upload_failure_reaches_the_caller_in_raise_mode(): + 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(): + 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} + + +# ------------------------------------------------------------------ metrics + + +def test_log_metrics_goes_to_the_metrics_sinks_only(): + records, metrics = FakeSink("records"), FakeSink("metrics") + run = make_run(sinks=[records], metrics_sinks=[metrics]) + + run.log_metrics({"loss": 0.5, "nan": float("nan")}, step=3) + run.log_traces([{"id": "t1"}]) + run.flush() + + assert records.batches == [[{"id": "t1"}]] + (batch,) = metrics.batches + (row,) = batch + assert row["loss"] == 0.5 + assert row["step"] == 3 + assert "nan" not in row + assert isinstance(row["_timestamp"], float) + run.finish() + + +def test_log_metrics_drops_non_finite_values_inside_sequences(): + metrics = FakeSink("metrics") + run = make_run(metrics_sinks=[metrics]) + + run.log_metrics( + { + "quantiles": [0.1, float("nan"), 0.9], + "nested": [{"loss": float("inf"), "reward": 1.0}, [float("-inf"), 2.0]], + } + ) + run.flush() + + row = metrics.batches[0][0] + assert row["quantiles"] == [0.1, 0.9] + assert row["nested"] == [{"reward": 1.0}, [2.0]] + run.finish() + + +def test_log_metrics_leaves_array_like_values_alone(): + class Arrayish: # numpy-style: comparisons are element-wise, truthiness ambiguous + def __eq__(self, other): + raise ValueError("ambiguous") + + __ne__ = __eq__ + __bool__ = __eq__ + + metrics = FakeSink("metrics") + run = make_run(metrics_sinks=[metrics]) + histogram = Arrayish() + + run.log_metrics({"hist": histogram, "nested": {"hist": histogram, "nan": float("nan")}}) + run.flush() + + row = metrics.batches[0][0] + assert row["hist"] is histogram and row["nested"] == {"hist": histogram} + run.finish() + + +def test_log_metrics_keeps_a_step_the_producer_already_set(): + metrics = FakeSink("metrics") + run = make_run(metrics_sinks=[metrics]) + + run.log_metrics({"step": 7, "loss": 0.5}, step=8) + run.flush() + + assert metrics.batches[0][0]["step"] == 7 + run.finish() + + +def test_log_metrics_without_a_metrics_sink_is_a_no_op(): + run = make_run() + + run.log_metrics({"loss": 0.5}, step=1) + + assert run._metrics_worker._thread is None + assert run.dropped_records == 0 + run.finish() + + +def test_finish_drains_and_closes_the_metrics_uploader_too(): + metrics = FakeSink("metrics") + run = make_run(metrics_sinks=[metrics]) + + run.log_metrics({"loss": 0.5}, step=1) + run.finish() + + assert len(metrics.batches) == 1 + assert metrics.closed is True + + +def test_losses_are_reported_across_both_uploaders(): + broken = FakeSink("metrics", fail_on_write=True) + run = make_run(metrics_sinks=[broken]) + + run.log_metrics({"loss": 0.5}, step=1) + run.flush() + + assert run.failed_records == {"metrics": 1} + run.finish() diff --git a/packages/prime-runs/tests/test_samples_sink.py b/packages/prime-runs/tests/test_samples_sink.py new file mode 100644 index 000000000..8349ba0ee --- /dev/null +++ b/packages/prime-runs/tests/test_samples_sink.py @@ -0,0 +1,156 @@ +"""The legacy sample sink that keeps today's viewer working.""" + +import httpx +import pytest +from _fakes import make_episode, make_trace +from conftest import RecordingHandler + +from prime_runs.sinks import EvalSamplesSink +from prime_runs.sinks.base import SinkWriteError +from prime_runs.worker import UploadWorker + + +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): + sink, handler = make_sink(make_platform_client, eval_routes) + + sink.write([make_episode("ep-1", [make_trace(idx=0)])]) + sink.write([make_episode("ep-2", [make_trace(idx=0)])]) + + posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples") + assert [body["samples"][0]["rollout_number"] for body in posted] == [1, 2] + + +def test_a_producer_that_already_speaks_v0_is_passed_through(make_platform_client, eval_routes): + sink, handler = make_sink(make_platform_client, eval_routes) + + sink.write([{"sample_id": "s1", "reward": 1.0}]) + + body = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")[0] + assert body["samples"] == [{"sample_id": "s1", "reward": 1.0}] + + +def test_records_this_sink_cannot_project_are_skipped_not_fatal( + make_platform_client, eval_routes, caplog +): + sink, handler = make_sink(make_platform_client, eval_routes) + + with caplog.at_level("WARNING"): + sink.write( + [ + make_episode("ep-1", [make_trace()]), + {"id": "ep-json", "traces": [{"id": "t"}]}, + ] + ) + sink.write([make_trace()]) + + body = handler.bodies_for("/api/v1/evaluations/eval-abc/samples")[0] + assert [row["sample_id"] for row in body["samples"]] == ["ep-1"] + assert len(handler.bodies_for("/api/v1/evaluations/eval-abc/samples")) == 1 + assert sink.skipped == 2 + assert sink.enabled is True + warnings = [ + record.getMessage() for record in caplog.records if "no projection" in record.getMessage() + ] + assert len(warnings) == 1 + assert "(dict)" in warnings[0] + + +def test_a_record_split_across_batches_is_counted_once( + make_platform_client, eval_routes, monkeypatch +): + eval_routes["POST /api/v1/evaluations/eval-abc/samples"] = lambda request: httpx.Response(502) + # Every episode falls back to one row per trace, and every row is its own batch. + monkeypatch.setattr("prime_runs.projection.MAX_SAMPLES_PAYLOAD_BYTES", 1) + monkeypatch.setattr( + "prime_runs.sinks.samples.batch_samples", lambda samples: [[sample] for sample in samples] + ) + sink, handler = make_sink(make_platform_client, eval_routes) + + with pytest.raises(SinkWriteError) as info: + sink.write([make_episode("ep-1", [make_trace(trace_id="a"), make_trace(trace_id="b")])]) + + assert len(handler.bodies_for("/api/v1/evaluations/eval-abc/samples")) == 2 + assert info.value.failed_records == 1 + + +def test_an_empty_batch_makes_no_request(make_platform_client, eval_routes): + sink, handler = make_sink(make_platform_client, eval_routes) + + sink.write([]) + + assert handler.requests == [] + + +def test_a_rejected_sample_batch_does_not_abort_later_batches( + make_platform_client, eval_routes, monkeypatch +): + statuses = iter([200, 422, 200, 200]) + + def respond(request): + status = next(statuses) + return httpx.Response(status, json={"detail": "invalid sample"} if status == 422 else {}) + + eval_routes["POST /api/v1/evaluations/eval-abc/samples"] = respond + monkeypatch.setattr( + "prime_runs.sinks.samples.batch_samples", + lambda samples: [[sample] for sample in samples], + ) + sink, handler = make_sink(make_platform_client, eval_routes) + worker = UploadWorker([sink]) + + worker.submit([{"sample_id": sample_id} for sample_id in ("a", "b", "c")]) + assert worker.flush(timeout=5.0) + worker.submit([{"sample_id": "d"}]) + assert worker.flush(timeout=5.0) + + posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples") + assert [body["samples"][0]["sample_id"] for body in posted] == ["a", "b", "c", "d"] + assert sink.samples_written == 3 + assert sink.enabled is True + assert worker.failed_records == {"eval_samples": 1} + worker.close() + + +def test_a_sink_wide_failure_counts_only_current_and_unattempted_records( + make_platform_client, eval_routes, monkeypatch +): + statuses = iter([200, 401]) + + def respond(request): + status = next(statuses) + return httpx.Response(status, json={"detail": "unauthorized"}) + + eval_routes["POST /api/v1/evaluations/eval-abc/samples"] = respond + monkeypatch.setattr( + "prime_runs.sinks.samples.batch_samples", + lambda samples: [[sample] for sample in samples], + ) + sink, handler = make_sink(make_platform_client, eval_routes) + worker = UploadWorker([sink]) + + worker.submit([{"sample_id": sample_id} for sample_id in ("a", "b", "c")]) + assert worker.flush(timeout=5.0) + + posted = handler.bodies_for("/api/v1/evaluations/eval-abc/samples") + assert [body["samples"][0]["sample_id"] for body in posted] == ["a", "b"] + assert sink.samples_written == 1 + assert sink.enabled is False + assert worker.failed_records == {"eval_samples": 2} + worker.close() diff --git a/packages/prime-runs/tests/test_traces_sink.py b/packages/prime-runs/tests/test_traces_sink.py new file mode 100644 index 000000000..9d59c8f10 --- /dev/null +++ b/packages/prime-runs/tests/test_traces_sink.py @@ -0,0 +1,192 @@ +"""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_type": "eval", "framework": "verifiers"}) + return sink + + +def test_records_go_out_with_provenance_but_not_the_join_key(): + client = FakeTracesClient() + sink = make_sink(client) + + sink.write([{"id": "t1", "run": {"id": "run-1"}}]) + + _, kwargs = client.calls[0] + assert kwargs["context"] == { + "source": "prime-runs", + "run_type": "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(): + 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_own_run_is_rekeyed_to_this_one(): + client = FakeTracesClient() + sink = make_sink(client) + original = {"id": "t1", "run": {"id": "theirs", "type": "train", "name": "local"}} + + sink.write([original]) + + assert client.calls[0][0][0]["run"] == {"id": "run-1", "type": "eval", "name": "local"} + assert original["run"] == {"id": "theirs", "type": "train", "name": "local"}, "not mutated" + + +def test_a_training_episode_keeps_its_dispatch_step_while_being_rekeyed(): + """prime-rl stamps ``TrainRunInfo(id=$PRL_RUN_ID, work=...)`` on every + dispatched episode; the platform's id replaces the launcher's, the work + (the dispatch step) stays.""" + client = FakeTracesClient() + sink = TracesSink(client=client) + sink.start("run-1", {"run_type": "train"}) + work = {"type": "train", "step": 10} + episode = {"id": "ep-1", "run": {"id": "prl-run", "type": "train", "work": work}, "traces": []} + + sink.write([episode]) + + assert client.calls[0][0][0]["run"] == {"id": "run-1", "type": "train", "work": work} + + +def test_an_episode_s_run_reaches_every_member_trace(): + 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_every_member_trace_is_rekeyed_too(): + client = FakeTracesClient() + sink = make_sink(client) + theirs = {"id": "other-run", "type": "eval", "name": "other"} + episode = { + "id": "ep-1", + "run": {"id": "env-run", "type": "eval", "name": "local"}, + "traces": [{"id": "a"}, {"id": "b", "run": theirs}], + } + + sink.write([episode]) + + sent = client.calls[0][0][0] + assert sent["run"] == {"id": "run-1", "type": "eval", "name": "local"} + assert sent["traces"][0]["run"] == sent["run"] + assert sent["traces"][1]["run"] == {"id": "run-1", "type": "eval", "name": "other"} + assert episode["traces"][0] == {"id": "a"}, "the caller's members were not mutated" + assert theirs == {"id": "other-run", "type": "eval", "name": "other"} + + +def test_an_account_outside_the_beta_retires_the_sink_without_a_failure(caplog): + 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): + 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(): + 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_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_a_missing_traces_client_disables_the_sink_and_reports_the_failure(monkeypatch, caplog): + 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_type": "eval"}) + + assert sink.enabled is False diff --git a/packages/prime-runs/tests/test_train_samples_sink.py b/packages/prime-runs/tests/test_train_samples_sink.py new file mode 100644 index 000000000..23401e54b --- /dev/null +++ b/packages/prime-runs/tests/test_train_samples_sink.py @@ -0,0 +1,423 @@ +"""The training samples sink: step-keyed Parquet through presign -> PUT -> confirm.""" + +import httpx +import pytest +from _fakes import make_episode, make_trace, make_train_episode +from conftest import RecordingHandler, StorageHandler + +from prime_runs.exceptions import APIError, RetryableAPIError +from prime_runs.sinks import RftSamplesSink, training_step +from prime_runs.sinks.base import SinkWriteError + + +class Encoder: + """Records what it was asked to encode; returns a recognizable payload.""" + + def __init__(self, payload=b"parquet-bytes"): + self.calls = [] + self.offsets = [] + self.payload = payload + + def __call__(self, episodes, run_id, step, sample_id_offset=0): + self.calls.append((list(episodes), run_id, step)) + self.offsets.append(sample_id_offset) + return self.payload + + +def make_sink(make_platform_client, routes, *, encoder=None, storage=None, **kwargs): + handler = RecordingHandler(routes) + storage = storage or StorageHandler() + encoder = encoder or Encoder() + sink = RftSamplesSink( + make_platform_client(handler), encoder=encoder, upload_client=storage.client(), **kwargs + ) + sink.start("run-abc", {}) + return sink, handler, storage, encoder + + +def test_each_step_on_the_cadence_is_one_upload(make_platform_client, rft_routes): + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes) + episodes = [ + make_train_episode("e1", step=10), + make_train_episode("e2", step=20), + make_train_episode("e3", step=10), + make_train_episode("e4", step=11), # off-cadence: not a loss + ] + + sink.write(episodes) + + assert [(step, [e.id for e in eps]) for eps, _, step in encoder.calls] == [ + (10, ["e1", "e3"]), + (20, ["e2"]), + ] + assert [body["step"] for body in handler.bodies_for("/api/v1/rft/samples/presign")] == [10, 20] + assert [str(u.url) for u in storage.uploads] == [ + "http://storage.test/run-abc/step-10.parquet", + "http://storage.test/run-abc/step-20.parquet", + ] + assert storage.uploads[0].content == b"parquet-bytes" + assert storage.uploads[0].headers["content-type"] == "application/parquet" + confirms = handler.bodies_for("/api/v1/rft/samples/confirm") + assert confirms[0] == { + "run_id": "run-abc", + "step": 10, + "s3_key": "runs/run-abc/step-10.parquet", + } + assert sink.steps_written == 2 + assert sink.sampled_out == 1 + assert sink.skipped == 0 + + +def test_the_cadence_is_configurable(make_platform_client, rft_routes): + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes, step_interval=1) + + sink.write([make_train_episode("e1", step=7)]) + + assert [step for _, _, step in encoder.calls] == [7] + with pytest.raises(ValueError): + RftSamplesSink(sink._client, step_interval=0) + + +def test_records_without_training_work_have_no_row_here(make_platform_client, rft_routes): + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes) + + sink.write( + [ + make_train_episode("eval-work", step=10, work="eval"), + make_episode("unstamped", [make_trace()]), + make_trace(), + { + "id": "json-episode", + "traces": [], + "run": {"type": "train", "work": {"type": "train", "step": 10}}, + }, + ] + ) + + assert encoder.calls == [] + assert handler.paths() == [] + assert sink.skipped == 4 + + +def test_nothing_to_encode_makes_no_request(make_platform_client, rft_routes): + sink, handler, storage, encoder = make_sink( + make_platform_client, rft_routes, encoder=Encoder(payload=None) + ) + + sink.write([make_train_episode("e1", step=10)]) + + assert encoder.calls and handler.paths() == [] + assert sink.steps_written == 0 + + +def test_a_failed_presign_loses_that_step_only(make_platform_client, rft_routes): + import json + + def flaky(request: httpx.Request) -> httpx.Response: + if json.loads(request.content)["step"] == 10: + return httpx.Response(502) + return rft_routes["POST /api/v1/rft/samples/presign"](request) + + routes = dict(rft_routes) + routes["POST /api/v1/rft/samples/presign"] = flaky + sink, handler, storage, encoder = make_sink(make_platform_client, routes) + + with pytest.raises(SinkWriteError) as info: + sink.write( + [ + make_train_episode("e1", step=10), + make_train_episode("e2", step=10), + make_train_episode("e3", step=20), + ] + ) + + # Step 10 (two episodes) was lost; step 20 went through. + assert info.value.failed_records == 2 + assert isinstance(info.value.cause, RetryableAPIError) + assert sink.steps_written == 1 + assert [str(u.url) for u in storage.uploads] == ["http://storage.test/run-abc/step-20.parquet"] + + +def test_the_storage_put_is_retried_then_reported_as_transient( + make_platform_client, rft_routes, no_sleep +): + storage = StorageHandler(status_codes=[503, 503, 503]) + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes, storage=storage) + + with pytest.raises(SinkWriteError) as info: + sink.write([make_train_episode("e1", step=10)]) + + assert len(storage.uploads) == 3 + assert isinstance(info.value.cause, RetryableAPIError) + assert handler.paths() == ["POST /api/v1/rft/samples/presign"] # never confirmed + + +def test_a_storage_put_that_recovers_is_confirmed(make_platform_client, rft_routes, no_sleep): + storage = StorageHandler(status_codes=[500, 200]) + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes, storage=storage) + + sink.write([make_train_episode("e1", step=10)]) + + assert len(storage.uploads) == 2 + assert "POST /api/v1/rft/samples/confirm" in handler.paths() + + +def test_a_storage_rejection_is_not_retried(make_platform_client, rft_routes): + storage = StorageHandler(status_codes=[403]) + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes, storage=storage) + + with pytest.raises(SinkWriteError) as info: + sink.write([make_train_episode("e1", step=10)]) + + assert len(storage.uploads) == 1 + assert isinstance(info.value.cause, APIError) + assert not isinstance(info.value.cause, RetryableAPIError) + + +def test_an_encoder_failure_is_a_record_rejection(make_platform_client, rft_routes): + def broken(episodes, run_id, step, sample_id_offset=0): + raise ValueError("cannot serialize") + + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes, encoder=broken) + + with pytest.raises(SinkWriteError) as info: + sink.write([make_train_episode("e1", step=10)]) + + assert isinstance(info.value.cause, ValueError) + assert handler.paths() == [] + + +def test_a_presign_without_a_url_is_an_api_error(make_platform_client, rft_routes): + routes = dict(rft_routes) + routes["POST /api/v1/rft/samples/presign"] = {"data": {"expiresIn": 900}} + sink, handler, storage, encoder = make_sink(make_platform_client, routes) + + with pytest.raises(SinkWriteError) as info: + sink.write([make_train_episode("e1", step=10)]) + + assert isinstance(info.value.cause, APIError) + assert "presign" in str(info.value.cause) + + +def test_a_step_logged_again_gets_its_own_sample_id_range(make_platform_client, rft_routes): + from prime_runs.sinks.train_samples import SAMPLE_ID_STRIDE + + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes) + + sink.write([make_train_episode("e1", step=10), make_train_episode("e2", step=20)]) + sink.write([make_train_episode("e3", step=10)]) # dispatched at 10, shipped later + + assert [step for _, _, step in encoder.calls] == [10, 20, 10] + assert encoder.offsets == [0, 0, SAMPLE_ID_STRIDE] + presigned = [body["step"] for body in handler.bodies_for("/api/v1/rft/samples/presign")] + assert presigned == [10, 20, 10] + assert sink.steps_written == 3 + + +def test_an_upload_that_failed_still_advances_the_sample_id_range(make_platform_client, rft_routes): + from prime_runs.sinks.train_samples import SAMPLE_ID_STRIDE + + storage = StorageHandler(status_codes=[403]) + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes, storage=storage) + + with pytest.raises(SinkWriteError): + sink.write([make_train_episode("e1", step=10)]) + sink.write([make_train_episode("e2", step=10)]) + + assert encoder.offsets == [0, SAMPLE_ID_STRIDE] + + +def test_a_forked_child_numbers_its_rows_apart_from_the_parent(make_platform_client, rft_routes): + from prime_runs.sinks.train_samples import SAMPLE_ID_RANGES, SAMPLE_ID_STRIDE + + assert SAMPLE_ID_RANGES * SAMPLE_ID_STRIDE <= 2**53 + sink, handler, storage, encoder = make_sink(make_platform_client, rft_routes) + sink.write([make_train_episode("e1", step=10)]) + + sink.reset_after_fork() # the child's copy of the sink... + # ...with its own storage client, as an owned one is rebuilt on the next upload. + sink._upload_client, sink._forked_with_injected_client = storage.client(), False + sink.write([make_train_episode("e2", step=10)]) + sink.write([make_train_episode("e3", step=10)]) + + parent, first, second = encoder.offsets + assert parent == 0 + band = range(SAMPLE_ID_RANGES // 2 * SAMPLE_ID_STRIDE, SAMPLE_ID_RANGES * SAMPLE_ID_STRIDE) + assert first in band and second in band and first != second + assert first % SAMPLE_ID_STRIDE == 0 + assert sink._uploads_per_step == {10: 1}, "the inherited counters are left to the parent" + + +def test_a_confirm_whose_response_was_lost_is_replayed(make_platform_client, rft_routes, no_sleep): + attempts = [] + + def flaky(request: httpx.Request) -> httpx.Response: + attempts.append(request) + if len(attempts) == 1: + return httpx.Response(504) + return httpx.Response(200, json={"data": {"status": "success"}}) + + routes = dict(rft_routes) + routes["POST /api/v1/rft/samples/confirm"] = flaky + sink, handler, storage, encoder = make_sink(make_platform_client, routes) + + sink.write([make_train_episode("e1", step=10)]) + + assert len(attempts) == 2 + assert len(storage.uploads) == 1 + assert sink.steps_written == 1 + + +def test_without_pyarrow_the_sink_turns_itself_off_quietly( + make_platform_client, rft_routes, monkeypatch, caplog +): + monkeypatch.setattr("prime_runs.projection.parquet_available", lambda: False) + handler = RecordingHandler(rft_routes) + sink = RftSamplesSink(make_platform_client(handler)) + + with caplog.at_level("WARNING"): + sink.start("run-abc", {}) + + assert sink.enabled is False + assert "prime-runs[train]" in caplog.text + sink.write([make_train_episode("e1", step=10)]) + assert handler.paths() == [] + + +def test_close_releases_only_an_owned_upload_client(make_platform_client, rft_routes): + handler = RecordingHandler(rft_routes) + injected = StorageHandler().client() + sink = RftSamplesSink(make_platform_client(handler), encoder=Encoder(), upload_client=injected) + sink.close() + assert not injected.is_closed + + owned = RftSamplesSink(make_platform_client(handler), encoder=Encoder()) + owned.start("run-abc", {}) + client = owned._upload() + owned.close() + assert client.is_closed + + +@pytest.mark.parametrize( + ("record", "expected"), + [ + (make_train_episode("e", step=12), 12), + (make_train_episode("e", step=12, work="eval"), None), + (make_episode("e"), None), + ({"run": {"type": "train", "work": {"type": "train", "step": 3}}}, 3), + ({"run": {"type": "eval"}}, None), + ({"run": {"type": "train", "work": {"type": "train", "step": True}}}, None), + ("not a record", None), + ], + ids=["train", "eval-work", "unstamped", "json", "json-eval", "bool-step", "string"], +) +def test_training_step_reads_the_episode_provenance(record, expected): + assert training_step(record) == expected + + +def test_the_default_encoder_writes_the_viewer_table(): + pa = pytest.importorskip("pyarrow") + pq = pytest.importorskip("pyarrow.parquet") + import io + import json + + from prime_runs.projection import episodes_to_parquet_bytes, train_sample_schema + + episodes = [ + make_train_episode("e1", step=10, idx=3, reward=1.0, advantage=0.25, env_id="gsm8k"), + make_train_episode("e2", step=10, idx=4, reward=0.0, advantage=-0.25, env_id="gsm8k"), + ] + + payload = episodes_to_parquet_bytes(episodes, "run-abc", 10) + + assert payload is not None + table = pq.read_table(io.BytesIO(payload)) + assert table.schema.equals(train_sample_schema()) + rows = table.to_pylist() + assert [row["run_id"] for row in rows] == ["run-abc", "run-abc"] + assert [row["step"] for row in rows] == [10, 10] + assert [row["problem_id"] for row in rows] == [3, 4] + assert [row["sample_id"] for row in rows] == [0, 1] + assert [row["reward"] for row in rows] == [1.0, 0.0] + assert [row["advantage"] for row in rows] == [0.25, -0.25] + assert [row["env_name"] for row in rows] == ["gsm8k", "gsm8k"] + trajectory = json.loads(rows[0]["trajectory"]) + assert trajectory[-1]["advantage"] == 0.25 + assert rows[0]["num_output_tokens"] == 5 + info = json.loads(rows[0]["info"]) + assert info["native_wrapper"]["id"] == "e1" + assert isinstance(pa.Table, type) + + +def test_the_default_encoder_writes_non_finite_numbers_as_null(): + pytest.importorskip("pyarrow") + import io + import json + import math + + import pyarrow.parquet as pq + + from prime_runs.projection import episodes_to_parquet_bytes + + episode = make_train_episode("e1", step=10, advantage=float("inf")) + episode.traces[0].metrics["kl"] = float("nan") + + payload = episodes_to_parquet_bytes([episode], "run-abc", 10) + + (row,) = pq.read_table(io.BytesIO(payload)).to_pylist() + strict = dict(parse_constant=lambda constant: pytest.fail(f"bare {constant} in a column")) + assert json.loads(row["metrics"], **strict) == {"kl": None} + assert json.loads(row["trajectory"], **strict)[-1]["advantage"] is None + assert math.isinf(row["advantage"]) # the float column keeps the value + + +def test_the_default_encoder_numbers_rows_from_the_offset(): + pytest.importorskip("pyarrow") + import io + + import pyarrow.parquet as pq + + from prime_runs.projection import episodes_to_parquet_bytes + + episodes = [make_train_episode("e1", step=10, idx=3), make_train_episode("e2", step=10, idx=4)] + + payload = episodes_to_parquet_bytes(episodes, "run-abc", 10, sample_id_offset=7) + + rows = pq.read_table(io.BytesIO(payload)).to_pylist() + assert [row["sample_id"] for row in rows] == [7, 8] + assert [row["problem_id"] for row in rows] == [3, 4] + + +def test_the_default_encoder_skips_episodes_without_a_trajectory(): + pytest.importorskip("pyarrow") + from _fakes import Episode, Trace + + from prime_runs.projection import episodes_to_parquet_bytes + + bare = Episode(id="no-branches", traces=[Trace(id="t", branches=[])]) + assert episodes_to_parquet_bytes([bare], "run-abc", 10) is None + assert episodes_to_parquet_bytes([], "run-abc", 10) is None + + +def test_a_forked_child_drops_the_inherited_storage_client(make_platform_client, rft_routes): + handler = RecordingHandler(rft_routes) + owned = RftSamplesSink(make_platform_client(handler), encoder=Encoder()) + owned.start("run-abc", {}) + inherited = owned._upload() + + owned.reset_after_fork() + + assert owned._upload_client is None + assert not inherited.is_closed # dropped, not closed + assert owned._upload() is not inherited + owned.close() + + injected = RftSamplesSink( + make_platform_client(handler), encoder=Encoder(), upload_client=StorageHandler().client() + ) + injected.start("run-abc", {}) + injected.reset_after_fork() + with pytest.raises(SinkWriteError) as info: + injected.write([make_train_episode("e1", step=10)]) + assert isinstance(info.value.cause, RuntimeError) + assert "after a fork" in str(info.value.cause) diff --git a/packages/prime-runs/tests/test_worker.py b/packages/prime-runs/tests/test_worker.py new file mode 100644 index 000000000..fbf6bf620 --- /dev/null +++ b/packages/prime-runs/tests/test_worker.py @@ -0,0 +1,489 @@ +"""The background uploader: backpressure, containment, fork safety.""" + +import queue +import threading + +import pytest +from conftest import FakeSink + +from prime_runs.exceptions import APIError, ValidationRejectedError +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_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_records_that_skip_a_sink_retired_by_error_are_counted_as_lost_to_it(): + 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(): + 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(): + 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(): + sink = FakeSink() + worker = UploadWorker([sink]) + worker.close() + worker.start() + + assert worker.submit([{"id": 1}]) is False + assert worker._thread is None + assert sink.batches == [] + + +def test_a_forked_child_starts_over_instead_of_re_uploading_the_parents_queue(): + 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): + 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(): + 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_a_transient_failure_drops_the_batch_but_keeps_the_sink(): + 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() + + +@pytest.mark.parametrize( + "error", + [ + ValueError("Out of range float values are not JSON compliant"), + ValidationRejectedError("invalid record", status_code=400), + APIError("unprocessable record", status_code=422), + ], + ids=["local-encoding", "traces-validation", "samples-validation"], +) +def test_a_record_rejection_drops_only_its_batch(error): + class RejectOnceSink(FakeSink): + def __init__(self) -> None: + super().__init__("reject-once") + self.calls = 0 + + def write(self, records) -> None: + self.calls += 1 + if self.calls == 1: + raise error + super().write(records) + + sink = RejectOnceSink() + reported = [] + worker = UploadWorker([sink], on_error=lambda name, exc: reported.append((name, exc))) + + # Drained between the two: queued batches coalesce into one write, and a + # rejection is per write. + worker.submit([{"id": "bad"}]) + drain(worker) + worker.submit([{"id": "good"}]) + drain(worker) + + assert sink.enabled is True + assert sink.calls == 2 + assert sink.batches == [[{"id": "good"}]] + assert worker.failed_records == {"reject-once": 1} + assert reported == [("reject-once", error)] + worker.close() + + +def test_a_sustained_outage_eventually_retires_the_sink(): + 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]) + + # One write per strike: drained between submits so they do not coalesce. + for _ in range(TRANSIENT_FAILURE_LIMIT): + worker.submit([{"id": 1}]) + drain(worker) + + assert sink.enabled is False + worker.close() + + +def test_a_cooldown_gives_a_retired_sink_another_chance(): + from prime_runs.exceptions import TransportError + from prime_runs.worker import TRANSIENT_FAILURE_LIMIT + + class OutageSink(FakeSink): + def __init__(self) -> None: + super().__init__("outage") + self.down = True + + def write(self, records) -> None: + if self.down: + raise TransportError("connection refused") + super().write(records) + + sink = OutageSink() + worker = UploadWorker([sink], retire_cooldown=0.05) + + for _ in range(TRANSIENT_FAILURE_LIMIT): + worker.submit([{"id": "lost"}]) + drain(worker) + assert sink.enabled is False + worker.submit([{"id": "during"}]) + drain(worker) + assert sink.batches == [] + + # A real wait: the autouse no_sleep fixture stubs out time.sleep. + threading.Event().wait(0.06) + sink.down = False + worker.submit([{"id": "after"}]) + drain(worker) + + assert sink.enabled is True + assert sink.batches == [[{"id": "after"}]] + assert worker.failed_records == {"outage": TRANSIENT_FAILURE_LIMIT + 1} + worker.close() + + +def test_a_permanent_failure_is_not_revived_by_the_cooldown(): + from prime_runs.exceptions import UnauthorizedError + + class DeniedSink(FakeSink): + def write(self, records) -> None: + raise UnauthorizedError("nope", status_code=401) + + sink = DeniedSink("denied") + worker = UploadWorker([sink], retire_cooldown=0.01) + + worker.submit([{"id": 1}]) + drain(worker) + threading.Event().wait(0.02) + worker.submit([{"id": 2}]) + drain(worker) + + assert sink.enabled is False + assert worker.failed_records == {"denied": 2} + worker.close() + + +def test_a_success_forgives_earlier_blips(): + 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. Drained one at + # a time: queued batches coalesce, and this is about consecutive *writes*. + for _ in range(TRANSIENT_FAILURE_LIMIT * 4): + worker.submit([{"id": 1}]) + drain(worker) + + 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(): + 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(): + 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() + + +def test_records_queued_during_an_upload_go_out_as_one_batch(): + sink = BlockingSink() + worker = UploadWorker([sink]) + + assert worker.submit([{"id": 0}]) + assert sink.entered.wait(1.0), "the uploader never reached the sink" + for n in (1, 2, 3): + assert worker.submit([{"id": n}]) + sink.released.set() + drain(worker) + + assert [[r["id"] for r in batch] for batch in sink.batches] == [[0], [1, 2, 3]] + worker.close() + + +def test_coalescing_keeps_traces_and_episodes_apart(): + sink = BlockingSink() + worker = UploadWorker([sink]) + + assert worker.submit([{"id": "t0"}]) + assert sink.entered.wait(1.0) + worker.submit([{"id": "t1"}]) + worker.submit([{"id": "e1", "traces": []}]) + worker.submit([{"id": "e2", "traces": []}]) + worker.submit([{"id": "t2"}]) + sink.released.set() + drain(worker) + + assert [[r["id"] for r in batch] for batch in sink.batches] == [ + ["t0"], + ["t1"], + ["e1", "e2"], + ["t2"], + ] + worker.close() + + +def test_a_flush_barrier_is_not_reordered_past_the_records_before_it(): + sink = BlockingSink() + worker = UploadWorker([sink]) + + assert worker.submit([{"id": 0}]) + assert sink.entered.wait(1.0) + worker.submit([{"id": 1}]) + release = threading.Timer(0.05, sink.released.set) + release.start() + try: + assert worker.flush(timeout=2.0) is True + finally: + release.join() + worker.submit([{"id": 2}]) + + assert [[r["id"] for r in batch] for batch in sink.batches] == [[0], [1]] + drain(worker) + assert [[r["id"] for r in batch] for batch in sink.batches] == [[0], [1], [2]] + worker.close() + + +def test_a_worker_with_no_sinks_never_starts_a_thread(): + worker = UploadWorker([]) + + assert worker.submit([{"id": 1}]) is True + assert worker.submit([{"id": 2}]) is True + + assert worker._thread is None + assert worker.dropped == 0 + assert worker.flush(timeout=1.0) is True + worker.close() 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..991446b92 100644 --- a/uv.lock +++ b/uv.lock @@ -8,13 +8,14 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-12T18:00:50.512194Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P7D" [options.exclude-newer-package] verifiers = false tasksets = false prime-traces = false +prime-runs = false prime-tunnel = false prime-sandboxes = false harnesses = false @@ -25,6 +26,7 @@ prime = false members = [ "prime", "prime-evals", + "prime-runs", "prime-sandboxes", "prime-traces", "prime-tunnel", @@ -1445,6 +1447,35 @@ toml = [ { name = "tomli" }, ] +[[package]] +name = "prime-runs" +source = { editable = "packages/prime-runs" } +dependencies = [ + { name = "httpx" }, + { name = "prime-traces" }, +] + +[package.optional-dependencies] +dev = [ + { name = "pytest" }, + { name = "pytest-asyncio" }, + { name = "ruff" }, +] +train = [ + { name = "pyarrow" }, +] + +[package.metadata] +requires-dist = [ + { name = "httpx", specifier = ">=0.25.0" }, + { name = "prime-traces", editable = "packages/prime-traces" }, + { name = "pyarrow", marker = "extra == 'train'", specifier = ">=15.0.0" }, + { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, + { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.23.0" }, + { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.13.1" }, +] +provides-extras = ["dev", "train"] + [[package]] name = "prime-sandboxes" source = { editable = "packages/prime-sandboxes" }