from __future__ import annotations
import os
import shlex
import signal
import time
from datetime import UTC, datetime
from pathlib import Path
from prime_sandboxes import APIClient, CreateSandboxRequest, SandboxClient
from prime_sandboxes.models import BackgroundJob, Sandbox
from github import GitHub
from report import render
from schema import (
ROOT,
UV_VERSION,
Report,
Request,
Result,
SandboxUsage,
report_complete,
trial_errors,
write_json,
)
REMOTE = "/opt/prime-benchmark"
OWNER_LABEL = "prime-agent-benchmarks-v1"
FILES = ("schema.py", "terminal.py", "kernel.py", "worker.py", "ui.py", "pyproject.toml", "uv.lock")
def elapsed_seconds(start: datetime) -> float:
return max(0, (datetime.now(UTC) - start.replace(tzinfo=start.tzinfo or UTC)).total_seconds())
class Canceled(Exception):
pass
def labels(repository: str, run_id: int, attempt: int) -> list[str]:
return [OWNER_LABEL, f"repository:{repository}", f"run:{run_id}", f"attempt:{attempt}"]
def cleanup(client: SandboxClient, repository: str, run_id: int, attempt: int) -> list[str]:
expected = labels(repository, run_id, attempt)
owned = []
page = 1
while True:
response = client.list(labels=expected, page=page, per_page=100, exclude_terminated=True)
owned.extend(sandbox.id for sandbox in response.sandboxes if set(expected) <= set(sandbox.labels))
if not response.has_next:
break
page += 1
failed = []
for sandbox_id in owned:
try:
client.delete(sandbox_id)
except Exception:
failed.append(sandbox_id)
if failed:
raise RuntimeError(f"Cleanup failed for {len(failed)} sandbox(es); TTL remains active")
return owned
class Controller:
def __init__(self, report: Report, results: Path, *, live_github: bool = True):
self.report = report
self.results = results
self.results.mkdir(parents=True, exist_ok=True)
self.harness = self.results / "harness"
self.harness.mkdir(exist_ok=True)
for name in FILES:
(self.harness / name).write_bytes((ROOT / name).read_bytes())
self.client = SandboxClient(APIClient(api_key=os.environ["PRIME_SANDBOX_API_KEY"]))
self.github = GitHub(report.repository) if live_github else None
self.sandboxes: dict[str, Sandbox] = {}
self.requests: dict[str, Request] = {}
self.started = time.monotonic()
self.deadline = self.started + report.config.timeout_seconds
self.last_fresh_check = 0.0
def checkpoint(self) -> None:
now = time.monotonic()
if now > self.deadline:
raise TimeoutError("Benchmark exceeded its runtime limit")
if self.github and now - self.last_fresh_check > 60:
self.last_fresh_check = now
if not self.github.fresh(self.report):
raise Canceled("A newer PR head or benchmark run superseded this run")
def save(self) -> None:
self.report.errors = self.report.errors[:30]
self.report.warnings = self.report.warnings[:30]
if self.report.status == "completed" and not report_complete(self.report):
self.report.status = "partial"
self.report = Report.model_validate(self.report.model_dump())
write_json(self.results / "report.json", self.report)
(self.results / "comment.md").write_text(render(self.report))
def command(self, sandbox: Sandbox, args: list[str], timeout: int = 60) -> None:
self.checkpoint()
result = self.client.execute_command(sandbox.id, shlex.join(args), timeout=timeout)
if result.exit_code:
raise RuntimeError(f"Sandbox setup command failed with exit code {result.exit_code}")
def start(self, role: str) -> Sandbox:
started = time.monotonic()
side = self.report.main if role == "main" else self.report.pr_head
config = self.report.config
sha = self.report.base_sha if role == "main" else self.report.head_sha
request = Request(
repository=self.report.repository,
source_repository=self.report.repository if role == "main" else self.report.head_repository,
sha=sha,
harness_sha=self.report.harness_sha,
pr=self.report.pr,
run_id=self.report.run_id,
attempt=self.report.attempt,
role=role,
config=config,
)
self.requests[role] = request
self.checkpoint()
sandbox = self.client.create(
CreateSandboxRequest(
name=f"agent-bench-{self.report.run_id}-{self.report.attempt}-{role}",
docker_image=config.image,
cpu_cores=config.cpu_cores,
memory_gb=config.memory_gb,
disk_size_gb=config.disk_gb,
vm=False,
region=config.region,
timeout_minutes=config.ttl_minutes,
labels=labels(self.report.repository, self.report.run_id, self.report.attempt)
+ [f"role:{role}"],
idempotency_key=f"agent-bench-{self.report.repository}-{self.report.run_id}-{self.report.attempt}-{role}",
)
)
self.sandboxes[role] = sandbox
print(f"Created {role} sandbox {sandbox.id}", flush=True)
while self.client.get(sandbox.id).status != "RUNNING":
self.checkpoint()
state = self.client.get(sandbox.id).status
if state in ("ERROR", "FAILED", "TERMINATED", "STOPPED", "TIMEOUT"):
raise RuntimeError(f"{role} sandbox failed during provisioning")
time.sleep(2)
side.timings["provision"] = time.monotonic() - started
started = time.monotonic()
self.command(sandbox, ["mkdir", "-p", REMOTE])
for name in FILES:
self.client.upload_file(sandbox.id, f"{REMOTE}/{name}", str(self.harness / name))
request_path = self.results / f"request-{role}.json"
write_json(request_path, request)
self.client.upload_file(sandbox.id, f"{REMOTE}/request.json", str(request_path))
setup = (
"set -eu\n"
"apt-get update -qq\n"
"DEBIAN_FRONTEND=noninteractive apt-get install -y -qq python3-venv "
"libcairo2-dev libpango1.0-dev libjpeg-dev libgif-dev librsvg2-dev fd-find ripgrep\n"
"ln -sf /usr/bin/fdfind /usr/local/bin/fd\n"
"useradd --create-home --uid 1500 --shell /bin/bash builder\n"
"python3 -m venv /opt/benchmark-bootstrap\n"
f"/opt/benchmark-bootstrap/bin/pip -q install uv=={UV_VERSION}\n"
f"/opt/benchmark-bootstrap/bin/uv sync --locked --no-dev --project {REMOTE}\n"
)
command = f"{shlex.join(['sh', '-c', setup])} > {REMOTE}/setup.log 2>&1"
job = self.client.start_background_job(sandbox.id, command)
self.wait(sandbox, job, "setup")
side.timings["setup"] = time.monotonic() - started
return sandbox
def wait(self, sandbox: Sandbox, job: BackgroundJob, label: str) -> None:
while True:
self.checkpoint()
status = self.client.get_background_job(sandbox.id, job, timeout=30)
if status.completed:
if status.exit_code:
raise RuntimeError(f"{label} failed with exit code {status.exit_code}")
return
time.sleep(2)
def phase(self, role: str, phase: str, trial: int = 0) -> None:
started = time.monotonic()
sandbox = self.sandboxes[role]
args = [f"{REMOTE}/.venv/bin/python", f"{REMOTE}/worker.py", phase, "--trial", str(trial)]
command = f"{shlex.join(args)} > {REMOTE}/{phase}-{trial}.log 2>&1"
job = self.client.start_background_job(sandbox.id, command)
try:
self.wait(sandbox, job, f"{role} {phase} {trial}")
except (Canceled, TimeoutError):
raise
except Exception:
try:
self.collect_result(role)
except Exception:
pass
raise
else:
self.collect_result(role)
finally:
if phase == "prepare":
side = self.report.main if role == "main" else self.report.pr_head
side.timings["build"] = time.monotonic() - started
self.save()
def collect_result(self, role: str) -> None:
path = self.results / f"result-{role}.json"
self.client.download_file(self.sandboxes[role].id, f"{REMOTE}/results/result.json", str(path))
if path.stat().st_size > 2_000_000:
raise ValueError("Sandbox result exceeds size limit")
result = Result.model_validate_json(path.read_text())
if result.request != self.requests[role] or result.side.sha != self.requests[role].sha:
raise ValueError("Sandbox result does not match this run's request")
previous = self.report.main if role == "main" else self.report.pr_head
result.side.timings = previous.timings.copy()
if role == "main":
self.report.main = result.side
else:
self.report.pr_head = result.side
self.save()
def logs(self, role: str) -> None:
sandbox = self.sandboxes[role]
# Never unpack PR-produced archives on the GitHub runner.
args = [
"python3",
"-c",
"import pathlib,tarfile; "
f"root=pathlib.Path({REMOTE!r}); "
"paths=[p for p in root.glob('*.log') if p.is_file() and p.stat().st_size < 8000000]; "
"paths += [p for p in (root/'results').glob('*') "
"if p.is_file() and p.stat().st_size < 8000000]; "
"archive=tarfile.open(root/'logs.tar.gz','w:gz'); "
"[archive.add(p,arcname=str(p.relative_to(root)),recursive=False) for p in paths[:1000]]; "
"archive.close()",
]
result = self.client.execute_command(sandbox.id, shlex.join(args), timeout=30)
if result.exit_code != 0:
raise RuntimeError(f"{role} log archive failed with exit code {result.exit_code}")
self.client.download_file(
sandbox.id, f"{REMOTE}/logs.tar.gz", str(self.results / f"{role}-logs.tar.gz")
)
def run(self) -> None:
old_handler = signal.signal(signal.SIGTERM, cancel)
try:
self.save()
while time.monotonic() - self.started < self.report.config.debounce_seconds:
self.checkpoint()
time.sleep(1)
ready = []
for role in ("main", "pr"):
try:
self.start(role)
self.phase(role, "prepare")
server = shlex.join(
[
"runuser",
"-u",
"builder",
"--",
"python3",
"-m",
"http.server",
"18741",
"--bind",
"127.0.0.1",
"--directory",
f"{REMOTE}/www",
]
)
self.client.start_background_job(self.sandboxes[role].id, server)
ready.append(role)
except (Canceled, TimeoutError):
raise
except Exception as error:
self.report.errors.append(f"{role} setup: {error_message(error)}"[:500])
for phase, count in (
("install", self.report.config.install_trials),
("measure", self.report.config.trials),
("runtime", self.report.config.trials),
("ui", self.report.config.ui_trials),
):
blocked: set[str] = set()
failures: dict[str, tuple[str, int]] = {}
if phase in ("measure", "runtime", "ui"):
for role in ready:
side = self.report.main if role == "main" else self.report.pr_head
if not any(
s.trial == 0 and s.value is not None for s in side.metrics.get("install", [])
):
blocked.add(role)
self.report.errors.append(
f"{role} {phase}: skipped; first installation did not succeed"
)
for trial in range(count):
for role in ready if trial % 2 == 0 else list(reversed(ready)):
if role in blocked:
continue
print(f"{role}: {phase} trial {trial + 1}/{count}", flush=True)
failure = ""
try:
self.phase(role, phase, trial)
except (Canceled, TimeoutError):
raise
except Exception as error:
failure = error_message(error)
side = self.report.main if role == "main" else self.report.pr_head
causes = trial_errors(side, phase, trial)
if causes:
failure = (
"; ".join(causes)
if not failure or any("measurement missing" not in cause for cause in causes)
else failure
)
if failure:
previous, repeated = failures.get(role, ("", 0))
repeated = repeated + 1 if failure == previous else 1
failures[role] = (failure, repeated)
if len(self.report.errors) < 25:
self.report.errors.append(f"{role} {phase} trial {trial}: {failure}"[:500])
if repeated >= self.report.config.failure_limit and trial + 1 < count:
blocked.add(role)
self.report.errors.append(
(
f"{role} {phase}: skipped {count - trial - 1} remaining trials after "
f"{repeated} identical consecutive failures; last cause: {failure}"
)[:500]
)
else:
failures.pop(role, None)
self.save()
compute = sum(
elapsed_seconds(sandbox.created_at) * self.report.config.hourly_cost() / 3600
for sandbox in self.sandboxes.values()
)
if compute >= self.report.config.budget_usd:
raise TimeoutError("Reached the estimated run budget")
self.report.status = "completed" if report_complete(self.report) else "partial"
if not ready:
self.report.status = "failed"
except Canceled as error:
self.report.status = "canceled"
self.report.errors.append(error_message(error))
except Exception as error:
self.report.status = "failed"
self.report.errors.append(error_message(error))
finally:
signal.signal(signal.SIGTERM, signal.SIG_IGN)
for role, sandbox in self.sandboxes.items():
deleted = False
try:
self.logs(role)
except Exception as error:
self.report.warnings.append(
f"{role} logs could not be collected: {error_message(error)}"[:500]
)
try:
self.client.delete(sandbox.id)
deleted = True
except Exception as error:
self.report.warnings.append(
(
f"{role} sandbox cleanup deferred to the cleanup workflow or TTL: "
f"{error_message(error)}"
)[:500]
)
seconds = elapsed_seconds(sandbox.created_at)
if not deleted:
seconds = max(seconds, self.report.config.ttl_minutes * 60)
self.report.sandboxes.append(
SandboxUsage(
id=sandbox.id,
role=role,
seconds=seconds,
deleted=deleted,
estimated_usd=seconds * self.report.config.hourly_cost() / 3600,
)
)
self.report.finished_at = datetime.now(UTC)
self.save()
signal.signal(signal.SIGTERM, old_handler)
def error_message(error: Exception) -> str:
message = f"{type(error).__name__}: {error}"
for key in ("PRIME_SANDBOX_API_KEY", "GITHUB_TOKEN"):
if os.environ.get(key):
message = message.replace(os.environ[key], "[REDACTED]")
return message[:500]
def cancel(_signum: int, _frame: object) -> None:
raise Canceled("Workflow was canceled")