diff --git a/pyproject.toml b/pyproject.toml index 87b190915e..1f12a874e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dependencies = [ "openai>=2.54.0,<3.0.0", "prime-tunnel>=0.1.8", "prime-sandboxes>=0.2.39", + "prime-runs>=0.1.1", "pydantic>=2.12.3", "rich>=11.0.0", "tenacity>=8.5.0", @@ -146,6 +147,7 @@ ty = "2026-07-28T00:00:00Z" # PrimeIntellect-published on PyPI (trusted publisher) prime-tunnel = false prime-sandboxes = false +prime-runs = false prime-pydantic-config = false renderers = false diff --git a/tests/v1/conftest.py b/tests/v1/conftest.py index 0b5ba36de9..75f66d33e1 100644 --- a/tests/v1/conftest.py +++ b/tests/v1/conftest.py @@ -188,6 +188,7 @@ def _eval_config( output_dir=output_dir.parent, run={"dir": output_dir.name}, model=CI_MODEL, + push=False, ) diff --git a/uv.lock b/uv.lock index c46cbc7ab8..c9c7d87705 100644 --- a/uv.lock +++ b/uv.lock @@ -21,13 +21,14 @@ exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for exclude-newer-span = "P7D" [options.exclude-newer-package] -prime-tunnel = false -prime-sandboxes = false +prime-pydantic-config = false harbor = "2026-08-11T00:00:00Z" ty = "2026-07-28T00:00:00Z" -prime-pydantic-config = false -renderers = false ruff = "2026-07-28T00:00:00Z" +prime-runs = false +prime-tunnel = false +prime-sandboxes = false +renderers = false [[package]] name = "aiofile" @@ -3469,6 +3470,19 @@ toml = [ { name = "tomli" }, ] +[[package]] +name = "prime-runs" +version = "0.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx" }, + { name = "prime-traces" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/85/db/3d5a9362f0d373d10145eaee3ac255cf227bf4164b36d04119b18e226607/prime_runs-0.1.1.tar.gz", hash = "sha256:0be97dcc95b3f1648ca25475efa45a5892fb12a775bf1f61728baec73a0088bb", size = 61110, upload-time = "2026-09-04T17:01:33.379Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/52/3fa41ac5e3c8bfa8eb33520f86634ae56474228f85c59b16aa04875efd7d/prime_runs-0.1.1-py3-none-any.whl", hash = "sha256:d512753ff84159a1234d1ea93091857f0a59fba12ef7ab2f69bfb579deb5ce75", size = 43006, upload-time = "2026-09-04T17:01:31.821Z" }, +] + [[package]] name = "prime-sandboxes" version = "0.2.40" @@ -3488,6 +3502,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c6/ba/6a5ce07386b457f0eae76d025b8b0af79bd82b8ce9fc2b9f0772e17a214b/prime_sandboxes-0.2.40-py3-none-any.whl", hash = "sha256:62d3afc9ce8a5b39dc292208ab8ba23779320e91cf2f7bd21564ed20ac6faa75", size = 60144, upload-time = "2026-08-27T18:12:06.818Z" }, ] +[[package]] +name = "prime-traces" +version = "0.0.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx" }, + { name = "pydantic" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8c/77/13fa99c0fb71909f2f35d048d9724efb9ae432484e47f12386b1542815e8/prime_traces-0.0.3.tar.gz", hash = "sha256:4071833db77a606ccbb718e5c73f2dfaf63982fe801aa096daed0bc605da554f", size = 50874, upload-time = "2026-08-27T19:46:44.672Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/9a/5a917de29a3a839c0ba36d445e3af562f744d9716b509571be643a01476c/prime_traces-0.0.3-py3-none-any.whl", hash = "sha256:0838f60651f1ed38371e383bd88b954e461d20ac424874e694bc83358934f5f3", size = 38345, upload-time = "2026-08-27T19:46:43.396Z" }, +] + [[package]] name = "prime-tunnel" version = "0.1.10" @@ -5368,6 +5395,7 @@ dependencies = [ { name = "numpy", version = "2.5.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, { name = "openai" }, { name = "prime-pydantic-config", extra = ["toml"] }, + { name = "prime-runs" }, { name = "prime-sandboxes" }, { name = "prime-tunnel" }, { name = "pydantic" }, diff --git a/verifiers/v1/cli/dashboard/eval.py b/verifiers/v1/cli/dashboard/eval.py index c729af026b..a6a7cccf91 100644 --- a/verifiers/v1/cli/dashboard/eval.py +++ b/verifiers/v1/cli/dashboard/eval.py @@ -281,17 +281,25 @@ def Overview(config: EvalConfig) -> Table: def _push_footer(push: "PushState | None") -> Group | None: - """The `--push` status line under the rollouts, shown once the run finishes and the upload - begins: dim `Pushing traces...` while it runs, then white `Traces pushed ()` or red - `Trace push failed ()`. `None` (no line) until the upload starts and when `--push` is off.""" + """The `--push` line under the rollouts: dim with the run's URL while it streams, + white once pushed, red when it failed. `None` when `--push` is off or the run stayed + local.""" if push is None or not push.started: return None - if not push.done: - line = Text("Pushing traces...", style="dim") - elif push.url: + if push.error and push.url: + # The run exists and holds what streamed up; only closing it out failed. line = Text(f"Traces pushed ({push.url})", style="white", overflow="fold") - else: + line.append(f" not closed out: {push.error}", style="red") + elif push.error: line = Text(f"Trace push failed ({push.error})", style="red", overflow="fold") + elif push.incomplete and push.url: + # Closed out, but the uploader lost records: what landed is there, say what didn't. + line = Text(f"Traces pushed ({push.url})", style="white", overflow="fold") + line.append(f" incomplete: {push.incomplete}", style="red") + elif not push.finished: + line = Text(f"Pushing traces ({push.url})", style="dim", overflow="fold") + else: + line = Text(f"Traces pushed ({push.url})", style="white", overflow="fold") return Group(Rule(style="dim"), line) diff --git a/verifiers/v1/cli/eval/main.py b/verifiers/v1/cli/eval/main.py index 338087cb91..b9c082ab87 100644 --- a/verifiers/v1/cli/eval/main.py +++ b/verifiers/v1/cli/eval/main.py @@ -134,10 +134,6 @@ def main(argv: list[str] | None = None) -> None: # Graceful cleanup has already run (each rollout's `finally`); partial results are on # disk. Exit on the conventional Ctrl-C code without a traceback. raise SystemExit(130) - if config.push and config.rich is None: - from verifiers.v1.utils.platform import push_traces - - push_traces(episodes, config) if ( config.rich is None ): # --rich is the whole output; otherwise dump each trace as JSON diff --git a/verifiers/v1/cli/eval/runner.py b/verifiers/v1/cli/eval/runner.py index 792a281d79..131927ca8f 100644 --- a/verifiers/v1/cli/eval/runner.py +++ b/verifiers/v1/cli/eval/runner.py @@ -13,8 +13,8 @@ import contextlib import logging import time -from collections.abc import AsyncIterator, Awaitable, Callable -from typing import cast +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable +from typing import TypeVar, cast from verifiers.v1.cli.dashboard import dashboard from verifiers.v1.cli.eval import resume @@ -30,9 +30,36 @@ from verifiers.v1.configs.serve import ServeConfig from verifiers.v1.env import Env, RunSlot from verifiers.v1.episode import Episode, EvalRunInfo +from verifiers.v1.utils.aio import run_shielded +from verifiers.v1.utils.platform import ( + PushState, + abort_run, + finish_run, + log_episodes, + open_run, +) logger = logging.getLogger(__name__) +T = TypeVar("T") + + +async def gather_rollouts(rollouts: Iterable[Awaitable[T]]) -> list[T]: + """`asyncio.gather`, but one rollout failing cancels the rest and waits for them + to unwind, so nothing keeps uploading into a run the caller is already closing. + Not a `TaskGroup`: that wraps errors in an `ExceptionGroup`, and `main` would no + longer see a `KeyboardInterrupt` as Ctrl-C.""" + tasks = [asyncio.ensure_future(rollout) for rollout in rollouts] + try: + return await asyncio.gather(*tasks) + except BaseException: + for task in tasks: + task.cancel() + # return_exceptions: wait for every task, not just the first to cancel. + await asyncio.gather(*tasks, return_exceptions=True) + raise + + RunSlotFn = Callable[[RunSlot], Awaitable[Episode]] OnComplete = Callable[[Episode], Awaitable[None]] @@ -203,47 +230,55 @@ async def run_eval(config: EvalConfig) -> list[Episode]: asyncio.Semaphore(config.max_concurrent) if config.max_concurrent else None ) write_lock = asyncio.Lock() + push_state = PushState() + + # Opened before the first rollout so every episode streams as it lands. + run = open_run(config, push_state, num_examples=len(tasks)) + # Resumed rollouts are part of this run too. + log_episodes(run, finished) async def on_complete(episode: Episode) -> None: episode.record_run(EvalRunInfo(id=config.run.id, name=config.run.name)) await append_episode(out, episode, write_lock) + await asyncio.to_thread(log_episodes, run, [episode]) backend = ( _in_process(env, config, semaphore, on_complete) if env is not None else _server(config, config.serve, semaphore, on_complete) ) - async with backend as run_slot: - # The display slots: in-process ones are the env's own (it fills their live - # traces); a served rollout's is a client-side stand-in its worker never sees. - planned = [ - slot - for task, n in plan - for slot in ( - env.slots(task, n) - if env is not None - else [RunSlot(task) for _ in range(n)] + # The run is closed out whatever breaks, backend setup and teardown included. + try: + async with backend as run_slot: + # The display slots: in-process ones are the env's own (it fills their live + # traces); a served rollout's is a client-side stand-in its worker never sees. + planned = [ + slot + for task, n in plan + for slot in ( + env.slots(task, n) + if env is not None + else [RunSlot(task) for _ in range(n)] + ) + ] + slots = [RunSlot.finished(episode) for episode in finished] + planned + display = ( + dashboard(slots, config, start, push=push_state) + if config.rich is not None + else contextlib.nullcontext() ) - ] - slots = [RunSlot.finished(episode) for episode in finished] + planned - push_state = None - if config.push and config.rich is not None: - from verifiers.v1.utils.platform import PushState - - push_state = PushState() - display = ( - dashboard(slots, config, start, push=push_state) - if config.rich is not None - else contextlib.nullcontext() - ) - async with display: - results = await asyncio.gather(*(run_slot(slot) for slot in planned)) - episodes = finished + list(results) - if ( - push_state is not None - ): # upload off the event loop so the view keeps refreshing - from verifiers.v1.utils.platform import push_traces - - push_state.started = True - await asyncio.to_thread(push_traces, episodes, config, push_state) + async with display: + results = await gather_rollouts(run_slot(slot) for slot in planned) + episodes = finished + list(results) + # Drain and close out off the event loop so the view keeps refreshing. + # Shielded: a Ctrl-C here must not cancel the close-out before the + # worker picks it up (a cancelled executor item never runs), so it + # runs to completion first and the interrupt is re-raised after — + # by which point the run is finished and `abort_run` has nothing to do. + await run_shielded( + asyncio.to_thread(finish_run, run, episodes, push_state) + ) + except BaseException as e: + await asyncio.to_thread(abort_run, run, e, push_state) + raise return episodes diff --git a/verifiers/v1/configs/cli/eval.py b/verifiers/v1/configs/cli/eval.py index 02fab474d6..45fbff4708 100644 --- a/verifiers/v1/configs/cli/eval.py +++ b/verifiers/v1/configs/cli/eval.py @@ -48,13 +48,23 @@ class RunConfig(BaseConfig): """Run directory name — the run writes to `output_dir / dir`. Defaults to `run.name`; set it only when the directory should differ from the display name.""" - # TODO: fetch the id from the Prime SDK once runs are registered there. - _id: str = PrivateAttr(default_factory=lambda: str(uuid4())) + _id: str | None = PrivateAttr(default=None) @property def id(self) -> str: + """The run's one id, assigned by `open_run` from the prime-runs handle: the + platform's evaluation id online, the SDK's local id otherwise. The run mints + none of its own, so the run dir, every trace and the dashboard agree.""" + if self._id is None: + raise RuntimeError("the run has no id until `open_run` has opened it") return self._id + def assign_id(self, run_id: str) -> None: + """Called once by `open_run`, before the first rollout.""" + if self._id is not None and self._id != run_id: + raise RuntimeError(f"the run already has id {self._id!r}") + self._id = run_id + class EvalConfig(BaseConfig): env: SerializeAsAny[EnvConfig] = SingleAgentEnvConfig() diff --git a/verifiers/v1/utils/platform.py b/verifiers/v1/utils/platform.py index d81ff7fd92..c36a5465c6 100644 --- a/verifiers/v1/utils/platform.py +++ b/verifiers/v1/utils/platform.py @@ -1,317 +1,164 @@ -"""Push a finished eval run to the Prime Intellect platform (`--no-push` to skip). +"""The eval's run on the Prime Intellect platform (`--no-push` to keep it local).""" -Uploads one sample per v1 `Episode` over the `/evaluations/` API (create -> push -samples -> finalize). Each sample keeps the complete native Episode as its source -of truth and includes a flat summary for older Platform consumers. Auth + base URL -come from `$PRIME_API_KEY` / `~/.prime/config.json`. -""" - -import json +import asyncio import logging import os +from collections.abc import Mapping from dataclasses import dataclass from typing import Any -import httpx +import prime_runs as pr +from prime_runs.projection import ( + build_samples, # noqa: F401 - prime-rl imports it from here +) from verifiers.v1.configs.cli.eval import EvalConfig from verifiers.v1.episode import Episode -from verifiers.v1.trace import Trace -from verifiers.v1.utils.prime import load_prime_config logger = logging.getLogger(__name__) -DEFAULT_API_URL = "https://api.primeintellect.ai" -DEFAULT_FRONTEND_URL = "https://app.primeintellect.ai" -# Repeated /samples posts append; match the Prime Evals client's request ceiling. -_MAX_SAMPLES_PAYLOAD_BYTES = 25 * 1024 * 1024 - - -def json_bytes(value: Any) -> int: - return len( - json.dumps( - value, - ensure_ascii=False, - separators=(",", ":"), - allow_nan=False, - ).encode("utf-8") - ) - @dataclass class PushState: - """Mutable upload status shared with the dashboard.""" + """The dashboard's view of the run: reads through to it, owns no I/O.""" - started: bool = False - done: bool = False - url: str | None = None + run: pr.Run | None = None error: str | None = None - - -def trace_to_sample( - trace: Trace, rollout_number: int = 1, episode_id: str | None = None -) -> dict[str, Any]: - """One trace -> the platform's sample dict (the v0 eval-sample format). - - The hub table stays flat — one row per trace; its episode is denormalized onto - the row (`episode_id` from the envelope, plus the trace's own `agent`/`trainable`), - so a multi-trace rollout's grouping travels with each row without a nested - schema. No prompt/completion split (meaningless mid-branch): `completion` is the - final branch's messages, `trajectory` one message list per branch.""" - - def dump(messages): - return [m.model_dump(mode="json", exclude_none=True) for m in messages] - - task = trace.task.data.model_dump(mode="json", exclude_none=True) - branches = trace.branches - sample = { - "sample_id": trace.id, - "example_id": trace.task.data.idx, - "rollout_number": rollout_number, - "episode_id": episode_id, - "agent": trace.agent.name, - "trainable": trace.agent.trainable, - "task": task, - "prompt": [], - "completion": dump(branches[-1].messages) if branches else [], - "answer": task.get("answer"), - # Keyed `tool_defs` because the v0 sample format already carries it there. - "tool_defs": [t.model_dump(mode="json", exclude_none=True) for t in trace.tools] - if trace.tools - else None, - "reward": trace.reward, - "timing": trace.timing.model_dump(mode="json", exclude_none=True), - "is_completed": trace.is_completed, - "is_truncated": trace.is_truncated, - "metrics": trace.metrics, - "error": trace.last_error.model_dump(mode="json", exclude_none=True) - if trace.last_error - else None, - "stop_condition": trace.stop_condition, - "trajectory": [ - { - "messages": dump(branch.messages), - "num_input_tokens": branch.num_input_tokens, - "num_output_tokens": branch.num_output_tokens, - } - for branch in branches - ], - "token_usage": trace.usage.model_dump(mode="json", exclude_none=True) - if trace.usage - else None, - "info": dict(trace.info) or None, - } - # Flatten sub-rewards to top-level keys the way v0 does (raw scores, as v0's - # per-function outputs were); env metrics stay nested. - for name, reward in trace.rewards.items(): - if reward is not None: - sample.setdefault(name, reward.score) - return sample - - -def credentials() -> tuple[str | None, str, str, str | None]: - """(api_key, api_base, frontend_url, team_id) from env vars / `~/.prime/config.json`.""" - cfg = load_prime_config() - api_key = os.getenv("PRIME_API_KEY") or cfg.get("api_key") - base = ( - os.getenv("PRIME_API_BASE_URL") - or os.getenv("PRIME_BASE_URL") - or cfg.get("base_url") - or DEFAULT_API_URL - ) - base = base.rstrip("/").removesuffix("/api/v1") - frontend = ( - os.getenv("PRIME_FRONTEND_URL") - or cfg.get("frontend_url") - or DEFAULT_FRONTEND_URL - ) - team_id = os.getenv("PRIME_TEAM_ID") or cfg.get("team_id") - return api_key, base, frontend, team_id - - -def run_metrics(episodes: list[Episode], traces: list[Trace]) -> dict[str, Any]: - """Run-level aggregates as v0's `GenerateMetadata`. Rewards/metrics aggregate - over the trainable traces only — fixed agents (a judge, a modeled user) often - carry no rewards and would dilute every mean with structural zeros — falling - back to all traces when none are trainable (same rule as the dashboard). - `avg_error` is the share of EPISODES that aren't ok: a hook failure counts - even when its traces are clean or it left none.""" - scored = [t for t in traces if t.agent.trainable] or 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 e.ok for e in episodes) / len(episodes) if episodes else 0.0 - return { - "avg_reward": sum(t.reward for t in scored) / n if n else 0.0, - "avg_metrics": {name: sums[name] / counts[name] for name in sums}, - "avg_error": avg_error, + incomplete: str | None = None + """Set when the run closed out but the uploader lost records on the way: the run + exists and holds what did land, so the footer says so rather than "failed".""" + + @property + def url(self) -> str | None: + return self.run.url if self.run is not None else None + + @property + def finished(self) -> bool: + return self.run is not None and self.run.finished + + @property + def started(self) -> bool: + """A live run, or a reason there isn't one.""" + return self.error is not None or self.url is not None + + +def open_run(config: EvalConfig, state: PushState, *, num_examples: int) -> pr.Run: + """Open the run this eval streams into, before the first rollout, and give the + config the run's id. A run that cannot be opened is logged and replaced by a + disabled one; the eval goes on.""" + identity: dict[str, Any] = { + "name": config.run.name, + # Resolved by name via the hub's get-or-create; no taskset, nothing to attach to. + "environments": [config.env.taskset.id] if config.env.taskset.id else [], + "model": config.model, + "framework": "verifiers", + # The v0 keys the dashboard's lists read. The config itself is a follow-up: + # a dump or the launched file can carry credentials and needs masking first. + "config": { + "model": config.model, + "num_examples": num_examples, + "rollouts_per_example": config.num_rollouts, + }, } - - -def build_samples(episodes: list[Episode]) -> list[dict[str, Any]]: - """One Platform sample per Episode, with a legacy-compatible trace summary. - - The native Episode in `info.native_wrapper` is authoritative and contains every - trace. One trainable trace (or the first trace) supplies only the flat summary - used by older consumers. `native_trace_index` identifies that summary trace. - """ - counts: dict[int, int] = {} - samples = [] - for episode in episodes: - if not episode.traces: - continue - summary_trace_index = next( - ( - index - for index, candidate in enumerate(episode.traces) - if candidate.agent.trainable - ), - 0, - ) - summary_trace = episode.traces[summary_trace_index] - idx = summary_trace.task.data.idx - counts[idx] = number = counts.get(idx, 0) + 1 - sample = trace_to_sample(summary_trace, number, episode.id) - sample["sample_id"] = episode.id - sample["info"] = { - **(sample["info"] or {}), - "native_wrapper": episode.to_record(), - "native_trace_index": summary_trace_index, - } - if len(b'{"samples":[]}') + json_bytes(sample) <= _MAX_SAMPLES_PAYLOAD_BYTES: - samples.append(sample) - continue + if config.push and os.getenv(pr.MODE_ENV, "").strip().lower() == "disabled": + # The SDK's own kill switch; the explicit `mode="online"` below would override it. + logger.info("--push: %s=disabled; running without a platform run", pr.MODE_ENV) + elif config.push: + try: + state.run = pr.init(mode="online", **identity) + except Exception as e: # noqa: BLE001 - a failed upload must not fail the eval + logger.warning( + "--push: could not open the run (%s: %s); running without it", + type(e).__name__, + e, + ) + state.error = f"{type(e).__name__}: {e}" + if state.run is None: + state.run = pr.init(mode="disabled", **identity) + # The run's one id: the platform's when online, the SDK's local one otherwise. + # The SDK keys every upload to it regardless; this is for the local records. + config.run.assign_id(state.run.id) + return state.run + + +def log_episodes(run: pr.Run, episodes: list[Episode]) -> None: + """Hand finished episodes to the run, best effort: the SDK already keeps upload + failures on its own thread, so this only guards the hand-off itself. A problem + here is the platform's, never the eval's.""" + if not episodes: + return + try: + run.log_episodes(episodes) + except Exception as e: # noqa: BLE001 - the rollouts are on disk; report, don't raise logger.warning( - "Episode %s exceeds the Platform sample limit; uploading projected traces", - episode.id, + "--push: could not queue %d episode(s) (%s: %s)", + len(episodes), + type(e).__name__, + e, ) - samples.extend( - trace_to_sample(candidate, number, episode.id) - for candidate in episode.traces - ) - return samples - -def push_traces( - episodes: list[Episode], - config: EvalConfig, - state: "PushState | None" = None, -) -> str | None: - """Upload a finished run to the platform; return the viewer URL (None if - skipped/failed). Resolves the env by name (get-or-create, so a local run - uploads without a prior `prime env push`); when `state` is given, records the - outcome on it so the dashboard's status line resolves.""" - def finish(url: str | None = None, error: str | None = None) -> str | None: - if state is not None: - state.url = url - state.error = error - state.done = True - return url - - api_key, base, frontend, team_id = credentials() - if not api_key: +def finish_run(run: pr.Run, episodes: list[Episode], state: PushState) -> None: + """Drain, write the run's aggregates, close it out. Blocking: call it off the loop.""" + try: + summary = pr.metrics.from_episodes(episodes) + except Exception as e: # noqa: BLE001 - close the run even without its headline logger.warning( - "--push: no PRIME_API_KEY (set it or run `prime login`); skipping upload" + "--push: could not aggregate the run's metrics (%s: %s)", + type(e).__name__, + e, ) - return finish(error="no PRIME_API_KEY (run `prime login`)") - - traces = [trace for episode in episodes for trace in episode.traces] - env_name = config.env.taskset.id - metrics = run_metrics(episodes, traces) - num_examples = len({t.task.data.idx for t in traces}) - metadata = { - "framework": "verifiers", - "run_id": config.run.id, - "model": config.model, - "num_examples": num_examples, - "rollouts_per_example": config.num_rollouts, - **metrics, - } - - team = {"team_id": team_id} if team_id else {} - api = f"{base}/api/v1" - headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} - # The run is done and its results saved; a network blip here must not crash it - # — log and skip the upload instead. + summary = None + _close(run, state, summary=summary) + + +def abort_run(run: pr.Run, error: BaseException, state: PushState) -> None: + """Close the run out after the eval broke, so it doesn't sit at running. Not + for a break during `finish_run`: that close-out completes on its own thread + and the SDK lets the first `finish()` decide the status — it sets `run.status` + as it starts, so an in-flight one is visible here before it is `finished`.""" + if run.finished or run.status is not pr.RunStatus.RUNNING: + return + if isinstance(error, (KeyboardInterrupt, asyncio.CancelledError)): + status, message = pr.RunStatus.CANCELLED, "interrupted" + else: + status, message = pr.RunStatus.FAILED, f"{type(error).__name__}: {error}" + _close(run, state, status=status, error=message) + + +def _close( + run: pr.Run, + state: PushState, + summary: Mapping[str, Any] | None = None, + status: pr.RunStatus = pr.RunStatus.COMPLETED, + error: str | None = None, +) -> None: + """`run.finish()`, best effort: the results are on disk, so nothing here may raise.""" try: - samples = build_samples(episodes) - batches: list[list[dict[str, Any]]] = [] - batch: list[dict[str, Any]] = [] - payload_bytes = len(b'{"samples":[]}') - for i, sample in enumerate(samples): - sample_bytes = json_bytes(sample) - sample_payload_bytes = len(b'{"samples":[]}') + sample_bytes - if sample_payload_bytes > _MAX_SAMPLES_PAYLOAD_BYTES: - raise ValueError( - f"sample {i} is too large to upload " - f"({sample_payload_bytes} > " - f"{_MAX_SAMPLES_PAYLOAD_BYTES} bytes)" - ) - next_payload_bytes = payload_bytes + (1 if batch else 0) + sample_bytes - if batch and next_payload_bytes > _MAX_SAMPLES_PAYLOAD_BYTES: - batches.append(batch) - batch = [] - payload_bytes = len(b'{"samples":[]}') - next_payload_bytes = payload_bytes + sample_bytes - batch.append(sample) - payload_bytes = next_payload_bytes - if batch or not samples: - batches.append(batch) - - with httpx.Client(headers=headers, timeout=300.0) as client: - - def post(path: str, body: dict) -> dict: - resp = client.post(f"{api}{path}", json=body) - resp.raise_for_status() - return resp.json() - - env_id = post("/environmentshub/resolve", {"name": env_name, **team})[ - "data" - ]["id"] - eval_id = post( - "/evaluations/", - { - "name": config.run.name, - "environments": [{"id": env_id}], - "model_name": config.model, - "dataset": env_name, - "framework": "verifiers", - "metadata": metadata, - "metrics": metrics, - "tags": [], - **team, - }, - )["evaluation_id"] - for batch in batches: - body = json.dumps( - {"samples": batch}, - ensure_ascii=False, - separators=(",", ":"), - allow_nan=False, - ).encode("utf-8") - resp = client.post( - f"{api}/evaluations/{eval_id}/samples", - content=body, - ) - resp.raise_for_status() - post(f"/evaluations/{eval_id}/finalize", {"metrics": metrics}) - except Exception as e: # noqa: BLE001 - push is best-effort across the full upload - logger.warning("--push: upload failed (%s: %s); skipping", type(e).__name__, e) - return finish(error=f"{type(e).__name__}: {e}") - - url = f"{frontend}/dashboard/evaluations/{eval_id}" - logger.info("--push: uploaded %d samples -> %s", len(samples), url) - return finish(url=url) + run.finish(summary, status=status, error=error) + except Exception as e: # noqa: BLE001 - the run is over; report, don't raise + logger.warning( + "--push: could not close out the run (%s: %s)", type(e).__name__, e + ) + if state.error is None: + state.error = f"{type(e).__name__}: {e}" + else: + state.incomplete = _losses(run) + if state.incomplete: + logger.warning("--push: %s, but %s", status.value, state.incomplete) + if run.url: + # The run's own status: `finish()` is a no-op once another caller closed it. + logger.info("--push: %s -> %s", run.status.value, run.url) + + +def _losses(run: pr.Run) -> str | None: + """What the uploader could not store, or `None`. A sink that switched itself off + quietly (Prime Traces outside the beta) is not a loss and is not counted here.""" + parts = [ + f"{count} record(s) not stored by the {sink} sink" + for sink, count in sorted(run.failed_records.items()) + if count + ] + if run.dropped_records: + parts.append(f"{run.dropped_records} record(s) never queued (uploader overrun)") + return "; ".join(parts) or None