Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
18 kB · 394 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395from __future__ import annotations
import osimport shleximport signalimport timefrom datetime import UTC, datetimefrom pathlib import Path
from prime_sandboxes import APIClient, CreateSandboxRequest, SandboxClientfrom prime_sandboxes.models import BackgroundJob, Sandbox
from github import GitHubfrom report import renderfrom 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")