diff --git a/pyproject.toml b/pyproject.toml index e3b3d969ec..7b98bc9bd1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ "openai-agents>=0.8.2", "prime-tunnel>=0.1.8", "prime-sandboxes>=0.2.37", + "prime-runs>=0.1.0", "pydantic>=2.12.3", "requests", "rich>=11.0.0", @@ -163,6 +164,8 @@ url = "https://pypi.org/simple" default = true [tool.uv.sources] +# TEMPORARY: prime-runs is not on PyPI yet +prime-runs = { git = "https://github.com/PrimeIntellect-ai/prime.git", subdirectory = "packages/prime-runs", branch = "feature/prime-runs-sdk" } compact = { path = "environments/compact", editable = true } glossary = { path = "environments/glossary", editable = true } deepwiki = { path = "environments/deepwiki", editable = true } @@ -192,6 +195,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/uv.lock b/uv.lock index 46fa9dd710..a54a06acf9 100644 --- a/uv.lock +++ b/uv.lock @@ -18,13 +18,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 [manifest] @@ -3331,6 +3332,15 @@ toml = [ { name = "tomli" }, ] +[[package]] +name = "prime-runs" +version = "0.1.0" +source = { git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-runs&branch=feature%2Fprime-runs-sdk#34faa4aaf4e45ed4b0fe387089585ddb698423fd" } +dependencies = [ + { name = "httpx" }, + { name = "prime-traces" }, +] + [[package]] name = "prime-sandboxes" version = "0.2.37" @@ -3350,6 +3360,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c3/4c/346abc72f6891267f490abfbe9de867c3ba2ed64ae8a4f964ffc4065492a/prime_sandboxes-0.2.37-py3-none-any.whl", hash = "sha256:9687e1b698c183798138b5e2ca116554773e3e300307a599beb7fbe8ed65e783", size = 49587, upload-time = "2026-08-17T17:07:17.403Z" }, ] +[[package]] +name = "prime-traces" +version = "0.0.2" +source = { git = "https://github.com/PrimeIntellect-ai/prime.git?subdirectory=packages%2Fprime-traces&branch=feature%2Fprime-runs-sdk#34faa4aaf4e45ed4b0fe387089585ddb698423fd" } +dependencies = [ + { name = "httpx" }, + { name = "pydantic" }, +] + [[package]] name = "prime-tunnel" version = "0.1.10" @@ -5100,6 +5119,7 @@ dependencies = [ { name = "openai" }, { name = "openai-agents" }, { 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 497eacb1c5..7c6c8ed947 100644 --- a/verifiers/v1/cli/dashboard/eval.py +++ b/verifiers/v1/cli/dashboard/eval.py @@ -279,17 +279,27 @@ 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` status line under the rollouts. The run opens before the first rollout and + its traces stream up as they land, so the line carries the run's URL for the whole eval: + dim `Pushing traces ()` while rollouts are still going, then white `Traces pushed + ()` — with anything that degraded along the way appended in yellow — or red `Trace + push failed ()` when there is no run to show. `None` (no line) 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 everything that 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 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") + if push.warning: # pushed, but not all of it - say what went wrong + line.append(f" {push.warning}", style="yellow") return Group(Rule(style="dim"), line) diff --git a/verifiers/v1/cli/eval/main.py b/verifiers/v1/cli/eval/main.py index 68d4bdba36..a8d8d33ab0 100644 --- a/verifiers/v1/cli/eval/main.py +++ b/verifiers/v1/cli/eval/main.py @@ -18,6 +18,7 @@ write_config, ) from verifiers.v1.cli.resolve import ( + config_file_ref, extract_id, narrow_config, plugin_errors, @@ -68,6 +69,8 @@ def main(argv: list[str] | None = None) -> None: *argv, ] # let prime-pydantic-config render help/errors config = cli(config_type) + # The `@ eval.toml` this run was launched from. + config.run.record_source(config_file_ref(argv)) # A named run directory is re-entered only by `--resume` or wiped by `--clean`: any # other write into it — the dry-run config.toml included, which would clobber the # config a resume typically re-runs — would overwrite the previous run. @@ -146,10 +149,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 not config.rich: - from verifiers.v1.utils.platform import push_traces - - push_traces(episodes, config) if not config.rich: # --rich is the whole output; otherwise dump each trace as JSON for episode in episodes: for trace in episode.traces: diff --git a/verifiers/v1/cli/eval/runner.py b/verifiers/v1/cli/eval/runner.py index 5444ed1817..39ac4fc3f0 100644 --- a/verifiers/v1/cli/eval/runner.py +++ b/verifiers/v1/cli/eval/runner.py @@ -4,7 +4,8 @@ import contextlib import logging import time -from typing import cast +from collections.abc import Awaitable, Iterable +from typing import TypeVar, cast from verifiers.v1.cli.dashboard import dashboard from verifiers.v1.cli.eval import resume @@ -18,9 +19,36 @@ from verifiers.v1.configs.cli.eval import EvalConfig from verifiers.v1.env import Env, RunSlot from verifiers.v1.episode import Episode, EvalRunInfo +from verifiers.v1.utils.platform import PushState, abort_run, finish_run, 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 stops the others too. + + Plain `gather` raises the first error and leaves the rest running. They then + keep going while the caller is already handling that error — still uploading + to a run it has just closed, still using an env it is tearing down. + Cancelling them here, and waiting for each one to finish unwinding, keeps + those two things from overlapping. + + The error is re-raised exactly as it arrived, which is why this is not an + `asyncio.TaskGroup`: a TaskGroup wraps everything in an `ExceptionGroup`, and + `main` would stop recognizing 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 so this waits for all of them; without it the first + # cancellation would raise and the rest would be left running again. + await asyncio.gather(*tasks, return_exceptions=True) + raise + async def run_eval(env: Env, config: EvalConfig) -> list[Episode]: logger.info("eval config:\n%s", config.model_dump_json(indent=2)) @@ -76,38 +104,45 @@ async def run_eval(env: Env, config: EvalConfig) -> list[Episode]: logger.info("results: %s", out) write_lock = asyncio.Lock() + push_state = PushState() if config.push and config.rich else None + + # Opened before the first rollout so the platform's id is the run's id + run = open_run(config, push_state) + config.run.adopt_id(run.id) + # A resume's kept rollouts are part of this run too, so they carry its id and + # go up with the rest + for episode in finished: + episode.record_run(EvalRunInfo(id=config.run.id, name=config.run.name)) + run.log_episodes(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(run.log_episodes, [episode]) # Serving resources (shared tool servers, interception) come up once for the - # run; plan slots inside so the env's agents borrow them. - async with env.serving(): - planned = [slot for task, n in plan for slot in env.slots(task, n=n)] - slots = [RunSlot.finished(episode) for episode in finished] + planned - push_state = None - if config.push and config.rich: - from verifiers.v1.utils.platform import PushState - - push_state = PushState() - display = ( - dashboard(slots, config, start, push=push_state) - if config.rich - else contextlib.nullcontext() - ) - async with display: - results = await asyncio.gather( - *(env.run_slot(slot, ctx, semaphore, on_complete) for slot in planned) + # run; plan slots inside so the env's agents borrow them. Everything from + # bringing those up to tearing them down is inside the try: a run that was + # opened is closed out whatever breaks, so none of them sits at running. + try: + async with env.serving(): + planned = [slot for task, n in plan for slot in env.slots(task, n=n)] + slots = [RunSlot.finished(episode) for episode in finished] + planned + display = ( + dashboard(slots, config, start, push=push_state) + if config.rich + else contextlib.nullcontext() ) - 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( + env.run_slot(slot, ctx, semaphore, on_complete) for slot in planned + ) + episodes = finished + list(results) + # Drain and close out off the event loop so the view keeps refreshing. + await 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 @@ -203,6 +238,12 @@ async def run_eval_server(config: EvalConfig) -> list[Episode]: ) write_lock = asyncio.Lock() + run = open_run(config) + config.run.adopt_id(run.id) + for episode in finished: + episode.record_run(EvalRunInfo(id=config.run.id, name=config.run.name)) + run.log_episodes(finished) + async def run_unit(payload: dict) -> list[Episode]: async with semaphore or contextlib.nullcontext(): episode = await client.run( @@ -213,13 +254,20 @@ async def run_unit(payload: dict) -> list[Episode]: ) episode.record_run(EvalRunInfo(id=config.run.id, name=config.run.name)) await append_episode(out, episode, write_lock) + await asyncio.to_thread(run.log_episodes, [episode]) return [cast(Episode, episode)] # Each rollout is its own `run` request, dispatched least-busy across workers. units = [run_unit(payload) for payload, n in plan for _ in range(n)] - results = await asyncio.gather(*units) - await client.close() - return finished + [record for unit in results for record in unit] + try: + results = await gather_rollouts(units) + await client.close() + episodes = finished + [record for unit in results for record in unit] + await asyncio.to_thread(finish_run, run, episodes) + except BaseException as e: + await asyncio.to_thread(abort_run, run, e) + raise + return episodes finally: proc.terminate() with contextlib.suppress(Exception): diff --git a/verifiers/v1/cli/resolve.py b/verifiers/v1/cli/resolve.py index dc2b671943..5441c2995f 100644 --- a/verifiers/v1/cli/resolve.py +++ b/verifiers/v1/cli/resolve.py @@ -46,6 +46,16 @@ def references_config_file(argv: list[str]) -> bool: return any(arg.startswith("@") for arg in argv) +def config_file_ref(argv: list[str]) -> str | None: + """The file a run was launched from — the root-level `@ ` or None.""" + paths = [ + argv[i + 1] + for i, arg in enumerate(argv) + if arg == "@" and i + 1 < len(argv) and not (i and argv[i - 1].startswith("--")) + ] + return paths[0] if len(paths) == 1 else None + + def extract_id(argv: list[str], field: str, default: str = "") -> str: """The chosen `.id` from `--.id ` (or `=`) on the CLI, before the typed parse (the positional taskset shorthand is applied upstream). Two diff --git a/verifiers/v1/configs/cli/eval.py b/verifiers/v1/configs/cli/eval.py index 95017fa9fa..6ff697eedb 100644 --- a/verifiers/v1/configs/cli/eval.py +++ b/verifiers/v1/configs/cli/eval.py @@ -40,13 +40,33 @@ 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())) + """The platform's run id once `prime_runs.init()` has opened the run (see + `adopt_id`); the local uuid until then, and for a run that stays local.""" + + _source: str | None = PrivateAttr(default=None) + """The `@ file.toml` this run was launched from, recorded by the CLI.""" @property def id(self) -> str: return self._id + @property + def source(self) -> str | None: + return self._source + + def adopt_id(self, run_id: str) -> None: + """Take the platform's run id as this run's id. + + Called once, before the first rollout, so that every trace is stamped + with the id the platform knows the run by.""" + self._id = run_id + + def record_source(self, path: str | None) -> None: + """Remember the config file the run was launched from, so it can be + uploaded verbatim with the run.""" + self._source = path + class EvalConfig(BaseConfig): env: SerializeAsAny[EnvConfig] = SingleAgentEnvConfig() diff --git a/verifiers/v1/utils/platform.py b/verifiers/v1/utils/platform.py index d81ff7fd92..5427b95009 100644 --- a/verifiers/v1/utils/platform.py +++ b/verifiers/v1/utils/platform.py @@ -1,317 +1,174 @@ -"""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 dataclasses import dataclass from typing import Any -import httpx +import prime_runs as pr +from prime_runs.projection import build_samples 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 - +FRAMEWORK = "verifiers" -def json_bytes(value: Any) -> int: - return len( - json.dumps( - value, - ensure_ascii=False, - separators=(",", ":"), - allow_nan=False, - ).encode("utf-8") - ) +__all__ = [ + "PushState", + "abort_run", + "build_samples", + "finish_run", + "open_run", +] @dataclass class PushState: - """Mutable upload status shared with the dashboard.""" - - started: bool = False - done: bool = False - url: str | None = None - error: str | None = None - + """The dashboard's view of the upload. Owns no I/O — the run does. -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 + Reads through to the live run, so the footer can show the run's URL from the + moment it opens rather than only once everything has been uploaded.""" + run: "pr.Run | None" = None + error: str | None = None -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, - } - - -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, + @property + def url(self) -> str | None: + """Where to watch the run. `None` for a run that stays local.""" + return self.run.url if self.run is not None else None + + @property + def finished(self) -> bool: + """Whether the run has been closed out.""" + return self.run is not None and self.run.finished + + @property + def started(self) -> bool: + """Whether there is anything to report: a live run, or why there isn't one.""" + return self.error is not None or self.url is not None + + @property + def warning(self) -> str | None: + """What went wrong without sinking the upload, or `None` if nothing did. + + Records that were lost first, and how: dropped records reached no sink at + all (the rollouts outran the uploader), while a sink failing says nothing + about the others — with traces and samples both on, those records are + usually still safe in the one that worked. Failing that, the first thing + the SDK contained (`on_error="warn"`), which would otherwise be visible + only in the run's log file.""" + if self.run is None: + return None + parts = ( + [f"{self.run.dropped_records} dropped"] if self.run.dropped_records else [] ) - 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 + parts += [ + f"{count} failed via {sink}" + for sink, count in sorted(self.run.failed_records.items()) + if count + ] + if parts: + return ", ".join(parts) + return self.run.errors[0] if self.run.errors else None + + +def open_run(config: EvalConfig, state: PushState | None = None) -> "pr.Run": + """Open the run this eval streams into, before the first rollout.""" + identity: dict[str, Any] = { + "name": config.run.name, + # The environment is resolved by name through the hub's get-or-create, so + # a local env uploads without a prior `prime env push`. A run with no + # taskset has nothing to attach to and can only be a local run — say so + # by passing none, rather than asking the hub to resolve an empty name. + "environments": [config.env.taskset.id] if config.env.taskset.id else [], + "model": config.model, + "framework": FRAMEWORK, + "config": run_config(config), + } + if config.push: + try: + run = pr.init(mode="online", **identity) + if state is not None: + state.run = run + return run + 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, + ) + if state is not None: + state.error = f"{type(e).__name__}: {e}" + run = pr.init(mode="disabled", **identity) + if state is not None: + state.run = run + return run + + +def run_config(config: EvalConfig) -> dict[str, Any]: + """What the run was configured with — the fields somebody actually set, plus + the file it was launched from, kept byte for byte.""" + values: dict[str, Any] = config.model_dump(mode="json", exclude_unset=True) + source = config.run.source + if source is not None: + try: + values[pr.CONFIG_SOURCE_KEY] = pr.ConfigSource.from_file(source).to_dict() + except pr.ConfigurationError as e: + logger.warning("--push: not recording the run's config file (%s)", e) + return values + + +def finish_run( + run: "pr.Run", episodes: list[Episode], state: PushState | None = None +) -> None: + """Drain the queued episodes, write the run's aggregates and close it out. + + Blocking — call it off the event loop (`asyncio.to_thread`) so the dashboard + keeps refreshing while the last uploads land.""" + try: + summary = pr.metrics.from_episodes(episodes) + except Exception as e: # noqa: BLE001 - close the run even without its headline logger.warning( - "Episode %s exceeds the Platform sample limit; uploading projected traces", - episode.id, - ) - samples.extend( - trace_to_sample(candidate, number, episode.id) - for candidate in episode.traces + "--push: could not aggregate the run's metrics (%s: %s)", + type(e).__name__, + e, ) - 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: + summary = None + _close(run, state, summary=summary) + + +def abort_run( + run: "pr.Run", error: BaseException, state: PushState | None = None +) -> None: + """Close the run out after the eval broke, so it doesn't sit at running.""" + if run.finished: + return + if isinstance(error, (KeyboardInterrupt, asyncio.CancelledError)): + status, message = pr.RunStatus.CRASHED, "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 | None, + summary: dict[str, Any] | None = None, + status: "pr.RunStatus" = pr.RunStatus.COMPLETED, + error: str | None = None, +) -> None: + """`run.finish()` with the same best-effort contract as the rest of this + module: the eval's results are already on disk, so nothing here may raise.""" + try: + run.finish(summary, status=status, error=error) + except Exception as e: # noqa: BLE001 - the run is over; report, don't raise logger.warning( - "--push: no PRIME_API_KEY (set it or run `prime login`); skipping upload" + "--push: could not close out the run (%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. - 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) + if state is not None and state.error is None: + state.error = f"{type(e).__name__}: {e}" + else: + if run.url: + logger.info("--push: %s -> %s", status.value, run.url)