1#!/usr/bin/env python32"""Real distributed benchmarks on MacLustr (Phase 8). The coordinator runs on `--coordinator HOST` (LAN IP), workers are launched3by ssh from this laptop on the listed nodes (several worker processes per node allowed: K partitions need not equal K hosts).45Usage:6 python3 experiments/run_benchmarks.py --config experiments/configs/benchmarks.yaml --out results/raw/benchmarks.jsonl7Each design row: model, p, N, K, splits, q, hessian_mode, repeats, workers_per_node, nodes (list of alias:lan_ip).8Requires ~/wp28/repo synced and ~/wp28/venv on every node (experiments/scatter_cluster.py does the sync).9"""10import argparse, json, pathlib, subprocess, sys, time, socket, yaml, shlex, hashlib11ROOT = pathlib.Path(__file__).resolve().parents[1]12PY = "~/wp28/venv/bin/python"131415def launch_coordinator(host_alias, port, K, model, p, N, splits, q, hmode, seed, bits):16 code = (f"import sys,json; sys.path.insert(0,'{'~/wp28/repo/src'}'); import os; sys.path[0]=os.path.expanduser(sys.path[0]); "17 f"from pamle.protocol import coordinator_run; "18 f"r=coordinator_run({port},{K},'{model}',{p},{N},{tuple(splits)},{q},'{hmode}',{seed},{bits},log=lambda *a: None); print(json.dumps(r, default=lambda o: o.item() if hasattr(o, 'item') else str(o)))")19 cmd = f"{PY} -c {shlex.quote(code)}"20 return subprocess.Popen(["ssh", "-o", "ConnectTimeout=10", host_alias, cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)212223def launch_workers(assign, coord_ip, port):24 procs = []25 wid = 026 for alias, n in assign:27 # workers are children of the ssh shell and `wait` keeps the session (and the workers) alive until they exit28 cmds = "cd ~/wp28/repo/src && " + " ".join(f"{PY} -m pamle.worker --connect {coord_ip}:{port} --id {wid + i} > /tmp/wp28_worker_{wid + i}.log 2>&1 &" for i in range(n)) + " wait"29 wid += n30 procs.append(subprocess.Popen(["ssh", "-o", "ConnectTimeout=10", alias, cmds], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True))31 return procs323334def main():35 ap = argparse.ArgumentParser(); ap.add_argument("--config", required=True); ap.add_argument("--out", required=True); ap.add_argument("--only", default=None)36 a = ap.parse_args(); cfg = yaml.safe_load(open(a.config)); out = pathlib.Path(a.out); out.parent.mkdir(parents=True, exist_ok=True)37 coord_alias, coord_ip = cfg["coordinator"]["alias"], cfg["coordinator"]["lan_ip"]38 port = int(cfg.get("port", 50600))39 for i, d in enumerate(cfg["designs"]):40 if a.only and str(i) not in a.only.split(","):41 continue42 nodes = d["nodes"]; wpn = d["workers_per_node"]; K = d["K"]43 assign = []; left = K44 for alias in nodes:45 n = min(wpn, left); assign.append((alias, n)); left -= n46 if left == 0:47 break48 if left:49 print(f"design {i}: not enough worker slots for K={K}", file=sys.stderr); continue50 for rep in range(d.get("repeats", 3)):51 port += 152 seed = int(hashlib.sha256(f"{i}-{rep}".encode()).hexdigest(), 16) % 10 ** 653 co = launch_coordinator(coord_alias, port, K, d["model"], d["p"], d["N"], d["splits"], d["q"], d["hessian_mode"], seed, d.get("bits"))54 time.sleep(1.5)55 t0 = time.time(); wk = launch_workers(assign, coord_ip, port)56 try:57 sout, serr = co.communicate(timeout=1800)58 except subprocess.TimeoutExpired:59 co.kill(); print(f"design {i} rep {rep}: TIMEOUT", file=sys.stderr); continue60 for w in wk:61 w.wait(timeout=60)62 if co.returncode != 0 or not sout.strip():63 print(f"design {i} rep {rep}: coordinator failed: {serr[-800:]}", file=sys.stderr); continue64 res = json.loads(sout.strip().splitlines()[-1])65 res.update({"design_index": i, "repeat": rep, "launch_wall_s": time.time() - t0, "assignment": assign, "label": d.get("label", "")})66 with open(out, "a") as f:67 f.write(json.dumps(res) + "\n")68 st = res["stats"]69 print(f"design {i} rep {rep}: K={K} hosts={res['n_hosts']} p={d['p']} q={d['q']} {d['hessian_mode']} | fit+upload {st['stages']['A_B_local_fit_and_upload']['wall']:.3f}s "70 f"| rounds {[round(v['wall'],3) for k,v in st['stages'].items() if k.startswith('D_F')]} | up {st['bytes_up']} B down {st['bytes_down']} B | total {res['t_total_after_connect']:.3f}s", flush=True)717273if __name__ == "__main__":74 main()75