1 files changed, 131 insertions, 0 deletions
diff --git a/scripts/run_sweep.py b/scripts/run_sweep.py
new file mode 100644
index 0000000..9be5592
--- /dev/null
+++ b/scripts/run_sweep.py
@@ -0,0 +1,131 @@
+#!/usr/bin/env python
+"""Parallel experiment sweep runner.
+
+Usage:
+  .venv/bin/python -m scripts.run_sweep --exp seeds --jobs jobs.csv
+  .venv/bin/python -m scripts.run_sweep --exp seeds --jobs jobs.csv --post-only   # eval/plot only
+
+jobs.csv columns: job,model,seed,flags   (flags = extra train.py args, may be empty)
+
+For each job: trains runs/<exp>/<job>/<model>/seed<seed>, then runs eval + plot.
+Skips jobs whose results.json already exists (idempotent resume).
+Concurrency default 2 (2 cores, one thread per child). Logs to runs/<exp>/logs/.
+Appends a compact row per job to runs/<exp>/summary.csv.
+"""
+import argparse
+import csv
+import json
+import os
+import subprocess
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+PY = os.path.join(ROOT, ".venv", "bin", "python")
+
+
+def parse_jobs(path: str) -> list[dict]:
+    jobs = []
+    with open(path) as fh:
+        for row in csv.DictReader(fh):
+            jobs.append({
+                "job": row["job"].strip(),
+                "model": row["model"].strip(),
+                "seed": row["seed"].strip(),
+                "flags": (row.get("flags") or "").strip().split(),
+            })
+    return jobs
+
+
+def out_dir(exp: str, job: dict) -> str:
+    return os.path.join(ROOT, "runs", exp, job["job"], job["model"], f"seed{job['seed']}")
+
+
+def run_one(exp: str, job: dict, post_only: bool) -> dict:
+    od = out_dir(exp, job)
+    logdir = os.path.join(ROOT, "runs", exp, "logs")
+    os.makedirs(logdir, exist_ok=True)
+    tag = f"{job['job']}__{job['model']}_s{job['seed']}"
+    logpath = os.path.join(logdir, tag + ".log")
+    results_path = os.path.join(od, "results.json")
+    res = {"job": job["job"], "model": job["model"], "seed": job["seed"]}
+
+    if os.path.exists(results_path):
+        res["skipped"] = True
+        return res
+
+    t0 = time.time()
+    if not post_only:
+        cmd = [PY, "-m", "src.train", job["model"], job["seed"],
+               "--out_dir", f"runs/{exp}/{job['job']}"] + job["flags"]
+        env = dict(os.environ, OMP_NUM_THREADS="1", MKL_NUM_THREADS="1")
+        with open(logpath, "w") as lf:
+            p = subprocess.run(cmd, cwd=ROOT, env=env, stdout=lf, stderr=subprocess.STDOUT)
+        if p.returncode != 0:
+            res["error"] = f"train rc={p.returncode} (see {logpath})"
+            return res
+
+    for cmd in ([PY, "-m", "src.eval", job["model"], job["seed"], "--runs-dir", f"runs/{exp}/{job['job']}"],
+                [PY, "-m", "scripts.plot", job["model"], job["seed"], "--runs-dir", f"runs/{exp}/{job['job']}"]):
+        with open(logpath, "a") as lf:
+            subprocess.run(cmd, cwd=ROOT, stdout=lf, stderr=subprocess.STDOUT)
+
+    if not os.path.exists(results_path):
+        res["error"] = f"eval produced no results.json (see {logpath})"
+        return res
+
+    with open(results_path) as fh:
+        r = json.load(fh)
+    sig = r.get("signature", {})
+    final = r.get("final", {})
+    probe = final.get("probe", {})
+    hlt = final.get("halting") or {}
+    res.update({
+        "minutes": round((time.time() - t0) / 60, 1),
+        "signature": sig.get("code"),
+        "probe": probe.get("code"),
+        "halting": hlt.get("code"),
+        "val_em_last": final.get("val_exact_match"),
+        "val_em_best": r.get("val_selected", {}).get("val_exact_match"),
+        "probe_acc": probe.get("acc"),
+        "halt_mean": hlt.get("mean_steps"),
+    })
+    return res
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--exp", required=True)
+    ap.add_argument("--jobs", required=True)
+    ap.add_argument("--post-only", action="store_true")
+    ap.add_argument("--concurrency", type=int, default=2)
+    a = ap.parse_args()
+    jobs = parse_jobs(os.path.join(ROOT, a.jobs))
+    print(f"{len(jobs)} jobs, concurrency {a.concurrency}, post_only={a.post_only}")
+
+    results = []
+    with ThreadPoolExecutor(max_workers=a.concurrency) as pool:
+        futs = [pool.submit(run_one, a.exp, j, a.post_only) for j in jobs]
+        for f in as_completed(futs):
+            r = f.result()
+            results.append(r)
+            print(f"[{r['job']}/{r['model']}/s{r['seed']}] "
+                  + (r.get("error") or r.get("skipped") and "skipped (exists)" or
+                     f"{r.get('signature')}/{r.get('probe')}/{r.get('halting')} "
+                     f"valEM {r.get('val_em_last')} probe {r.get('probe_acc')} "
+                     f"({r.get('minutes')}min)"))
+
+    summary_path = os.path.join(ROOT, "runs", a.exp, "summary.csv")
+    with open(summary_path, "w", newline="") as fh:
+        keys = ["job", "model", "seed", "minutes", "signature", "probe", "halting",
+                "val_em_last", "val_em_best", "probe_acc", "halt_mean", "error", "skipped"]
+        w = csv.DictWriter(fh, fieldnames=keys, extrasaction="ignore")
+        w.writeheader()
+        for r in sorted(results, key=lambda x: (x["job"], x["model"])):
+            w.writerow(r)
+    print(f"summary: {summary_path}")
+
+
+if __name__ == "__main__":
+    main()