wehub-resource-sync
164 lines
5.3 KiB
Python
164 lines
5.3 KiB
Python
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
|