chore: import upstream snapshot with attribution

2026-07-13 12:01:57 +08:00
commit 9dda3e2451
399 changed files with 118131 additions and 0 deletions
@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import asyncio
+import os
+import subprocess
+import tempfile
+from pathlib import Path
+
+from .config import TaskConfig
+from .manifest import (
+    RunManifest,
+    SessionEntry,
+    TbEntry,
+    TbTaskEntry,
+    write_manifest,
+)
+from .session import run_all_sessions, setup_codebase, setup_rtk
+from .terminal_bench import run_terminal_bench
+from .vm import create_vm_pool, destroy_vm_pool
+
+ROOT_DIR = Path(__file__).resolve().parent.parent
+
+
+def _create_tarball(source_dir: Path) -> str:
+    fd, tarball = tempfile.mkstemp(suffix=".tar.gz")
+    os.close(fd)
+    try:
+        subprocess.run(
+            ["tar", "czf", tarball, "-C", str(source_dir), "."],
+            check=True,
+        )
+    except Exception:
+        Path(tarball).unlink(missing_ok=True)
+        raise
+    return tarball
+
+
+def _print_step(step: int, total: int, msg: str):
+    print(f"\n[{step}/{total}] {msg}")
+
+
+def _session_to_entry(r) -> SessionEntry:
+    return SessionEntry(
+        vm_name=r.vm_name,
+        group=r.group,
+        stdout_json=f"{r.vm_name}-stdout.json",
+        otel_log=f"{r.vm_name}-otel.log",
+        rtk_db=f"{r.vm_name}-tracking.db" if r.rtk_db_path else None,
+        exit_code=r.exit_code,
+        error=r.error or None,
+    )
+
+
+def _tb_to_entry(r) -> TbEntry:
+    return TbEntry(
+        vm_name=r.vm_name,
+        group=r.group,
+        total=r.total,
+        passed=r.passed,
+        failed=r.failed,
+        tasks=[TbTaskEntry(name=t.name, passed=t.passed, duration_s=t.duration_s) for t in r.tasks],
+        error=r.error,
+    )
+
+
+async def run_benchmark(
+    task: TaskConfig,
+    vms: int,
+    api_key: str,
+    output_dir: Path,
+    cloud_init: Path | None = None,
+    terminal_bench: bool = False,
+    keep_vms: bool = False,
+) -> RunManifest:
+    if cloud_init is None:
+        cloud_init = ROOT_DIR / "cloud-init-base.yaml"
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    total_steps = 5 if terminal_bench else 4
+    vm_names: list[str] = []
+    local_tarball: str | None = None
+
+    manifest = RunManifest(
+        task_name=task.name,
+        model=task.model,
+        vm_count=vms,
+    )
+
+    try:
+        _print_step(1, total_steps, f"Creating {vms * 2} VMs ({vms} RTK ON + {vms} RTK OFF)")
+        vm_names = await create_vm_pool(vms, cloud_init)
+        print(f"  VMs ready: {', '.join(vm_names)}")
+
+        _print_step(2, total_steps, "Setting up codebases")
+        if not task.codebase.is_github:
+            local_tarball = _create_tarball(task.codebase.local_path())
+
+        await asyncio.gather(*(
+            setup_codebase(name, task.codebase, local_tarball)
+            for name in vm_names
+        ))
+        print("  Codebases deployed")
+
+        _print_step(3, total_steps, "Configuring RTK on ON VMs")
+        setup_script = ROOT_DIR / "setup-rtk.sh"
+        on_vms = [n for n in vm_names if "-on-" in n]
+        off_vms = [n for n in vm_names if "-off-" in n]
+        await asyncio.gather(*(setup_rtk(vm, setup_script) for vm in on_vms))
+        print(f"  RTK configured on {len(on_vms)} VMs")
+
+        _print_step(4, total_steps, f"Running Claude sessions (timeout: {task.timeout_minutes}min)")
+        results = await run_all_sessions(vm_names, task, api_key, output_dir)
+
+        on_ok = [r for r in results if r.group == "on" and not r.error]
+        off_ok = [r for r in results if r.group == "off" and not r.error]
+        errors = [r for r in results if r.error]
+        print(f"  Completed: {len(on_ok)} ON, {len(off_ok)} OFF, {len(errors)} errors")
+        for r in errors:
+            print(f"    {r.vm_name}: {r.error}")
+
+        manifest.sessions = [_session_to_entry(r) for r in results]
+
+        if terminal_bench:
+            _print_step(5, total_steps, "Running terminal-bench precision tests")
+            tb_on = await asyncio.gather(*(
+                run_terminal_bench(vm, "on", task.model, api_key)
+                for vm in on_vms
+            ))
+            tb_off = await asyncio.gather(*(
+                run_terminal_bench(vm, "off", task.model, api_key)
+                for vm in off_vms
+            ))
+
+            manifest.terminal_bench = [_tb_to_entry(r) for r in list(tb_on) + list(tb_off)]
+
+            ok_on = [r for r in tb_on if not r.error]
+            ok_off = [r for r in tb_off if not r.error]
+            if ok_on and ok_off:
+                on_total = sum(r.total for r in ok_on)
+                on_passed = sum(r.passed for r in ok_on)
+                off_total = sum(r.total for r in ok_off)
+                off_passed = sum(r.passed for r in ok_off)
+                on_rate = on_passed / on_total if on_total else 0
+                off_rate = off_passed / off_total if off_total else 0
+                print(f"  terminal-bench: ON pass rate={on_rate:.0%}, OFF pass rate={off_rate:.0%}, delta={on_rate - off_rate:+.0%}")
+
+            tb_errors = [r for r in list(tb_on) + list(tb_off) if r.error]
+            for r in tb_errors:
+                print(f"    {r.vm_name}: {r.error}")
+
+        write_manifest(manifest, output_dir)
+        print(f"\n  Manifest written to {output_dir / 'manifest.json'}")
+
+    finally:
+        if local_tarball:
+            Path(local_tarball).unlink(missing_ok=True)
+        if not keep_vms and vm_names:
+            print("\nCleaning up VMs...")
+            await destroy_vm_pool(vm_names)
+            print("  VMs destroyed")
+
+    return manifest