1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
|
#!/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__)))
_venv_py = os.path.join(ROOT, ".venv", "bin", "python")
PY = _venv_py if os.path.exists(_venv_py) else sys.executable
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"]
# CPU-resident jobs get 2 threads (sequential tied cell benefits); GPU jobs
# stay at 1 so the feeder process doesn't contend with concurrent CPU jobs.
n_threads = "2" if "--device cpu" in job["flags"] else "1"
env = dict(os.environ, OMP_NUM_THREADS=n_threads, MKL_NUM_THREADS=n_threads)
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()
|
