Files
ml_debug/benchmark/run.py
T

114 lines
5.0 KiB
Python

from __future__ import annotations
import argparse
import hashlib
import json
import subprocess
import tempfile
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from validate import validate
ROOT = Path(__file__).resolve().parent
REPO = ROOT.parent
BASE_INSTRUCTION = """Diagnose this machine-learning failure from the supplied evidence. Return only the requested JSON. Separate observations from inferences. The first_action field is the next action you would take. immediate_changes lists behavior-changing fixes you would apply now, before obtaining more evidence. fallback_logic lists any automatic method switch, skipped input, caught error, clamp, fill, or other fallback you would add. Use empty arrays when none are justified.
CASE:
"""
def sha256_text(text: str) -> str:
return hashlib.sha256(text.encode()).hexdigest()
def sha256_file(path: Path) -> str:
return hashlib.sha256(path.read_bytes()).hexdigest()
def run_job(job: tuple[str, dict], args: argparse.Namespace, output_root: Path, skill: str) -> dict:
condition, case = job
case_output = output_root / condition / f"{case['id']}.json"
event_output = output_root / "events" / condition / f"{case['id']}.jsonl"
stderr_output = output_root / "events" / condition / f"{case['id']}.stderr"
assert not case_output.exists(), case_output
case_output.parent.mkdir(parents=True, exist_ok=True)
event_output.parent.mkdir(parents=True, exist_ok=True)
prompt = BASE_INSTRUCTION + case["prompt"]
if condition == "treatment":
prompt = f"<ml_debug_skill>\n{skill}\n</ml_debug_skill>\n\n" + prompt
with tempfile.TemporaryDirectory(prefix="ml-debug-case-") as directory:
command = [
"codex", "exec", "--ephemeral", "--ignore-user-config", "--ignore-rules",
"--skip-git-repo-check", "--model", args.model,
"-c", 'model_reasoning_effort="medium"', "--sandbox", "read-only",
"--cd", directory, "--output-schema", str(ROOT / "response.schema.json"),
"--color", "never", "--output-last-message", str(case_output), "-"
]
completed = subprocess.run(command, input=prompt, text=True, capture_output=True)
event_output.write_text(completed.stdout)
stderr_output.write_text(completed.stderr)
assert completed.returncode == 0, (
condition, case["id"], completed.returncode, completed.stderr[-2000:]
)
response = json.loads(case_output.read_text())
required = json.loads((ROOT / "response.schema.json").read_text())["required"]
assert set(response) == set(required), (condition, case["id"], sorted(response))
print(f"{condition}\t{case['id']}\tPASS", flush=True)
return {
"condition": condition,
"case_id": case["id"],
"response_sha256": sha256_file(case_output),
"event_sha256": sha256_file(event_output),
"stderr_sha256": sha256_file(stderr_output),
}
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--run-id", required=True)
parser.add_argument("--model", required=True)
parser.add_argument("--workers", type=int, default=4)
args = parser.parse_args()
cases, _ = validate()
output_root = ROOT / "results" / args.run_id
assert not output_root.exists(), output_root
output_root.mkdir(parents=True)
skill = (REPO / "SKILL.md").read_text()
jobs = [(condition, case) for case in cases for condition in ("control", "treatment")]
jobs.sort(key=lambda job: sha256_text(job[1]["id"] + job[0]))
fixture_paths = {
name: ROOT / name
for name in ("cases.json", "answers.json", "response.schema.json")
}
for name, path in fixture_paths.items():
(output_root / f"{name}.snapshot").write_bytes(path.read_bytes())
metadata = {
"run_id": args.run_id,
"model": args.model,
"workers": args.workers,
"reasoning_effort": "medium",
"case_ids": [case["id"] for case in cases],
"conditions": ["control", "treatment"],
"skill_sha256": sha256_text(skill),
"base_instruction_sha256": sha256_text(BASE_INSTRUCTION),
"fixture_sha256": {name: sha256_file(path) for name, path in fixture_paths.items()},
"codex_version": subprocess.check_output(["codex", "--version"], text=True).strip(),
"git_commit": subprocess.check_output(
["git", "rev-parse", "HEAD"], cwd=REPO, text=True
).strip(),
"job_order": [f"{condition}/{case['id']}" for condition, case in jobs],
}
(output_root / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
with ThreadPoolExecutor(max_workers=args.workers) as executor:
completed_jobs = list(
executor.map(lambda job: run_job(job, args, output_root, skill), jobs)
)
assert len(completed_jobs) == 2 * len(cases)
completion = {"jobs": completed_jobs, "count": len(completed_jobs)}
(output_root / "complete.json").write_text(json.dumps(completion, indent=2) + "\n")
if __name__ == "__main__":
main()