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))")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 cmds = " ".join(f"(cd ~/wp28/repo/src && {PY} -m pamle.worker --connect {coord_ip}:{port} --id {wid + i} > /dev/null 2>&1 &)" for i in range(n))28 wid += n29 procs.append(subprocess.Popen(["ssh", "-o", "ConnectTimeout=10", alias, cmds + "; wait"], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True))30 return procs313233def main():34 ap = argparse.ArgumentParser(); ap.add_argument("--config", required=True); ap.add_argument("--out", required=True); ap.add_argument("--only", default=None)35 a = ap.parse_args(); cfg = yaml.safe_load(open(a.config)); out = pathlib.Path(a.out); out.parent.mkdir(parents=True, exist_ok=True)36 coord_alias, coord_ip = cfg["coordinator"]["alias"], cfg["coordinator"]["lan_ip"]37 port = int(cfg.get("port", 50600))38 for i, d in enumerate(cfg["designs"]):39 if a.only and str(i) not in a.only.split(","):40 continue41 nodes = d["nodes"]; wpn = d["workers_per_node"]; K = d["K"]42 assign = []; left = K43 for alias in nodes:44 n = min(wpn, left); assign.append((alias, n)); left -= n45 if left == 0:46 break47 if left:48 print(f"design {i}: not enough worker slots for K={K}", file=sys.stderr); continue49 for rep in range(d.get("repeats", 3)):50 port += 151 seed = int(hashlib.sha256(f"{i}-{rep}".encode()).hexdigest(), 16) % 10 ** 652 co = launch_coordinator(coord_alias, port, K, d["model"], d["p"], d["N"], d["splits"], d["q"], d["hessian_mode"], seed, d.get("bits"))53 time.sleep(1.5)54 t0 = time.time(); wk = launch_workers(assign, coord_ip, port)55 try:56 sout, serr = co.communicate(timeout=1800)57 except subprocess.TimeoutExpired:58 co.kill(); print(f"design {i} rep {rep}: TIMEOUT", file=sys.stderr); continue59 for w in wk:60 w.wait(timeout=60)61 if co.returncode != 0 or not sout.strip():62 print(f"design {i} rep {rep}: coordinator failed: {serr[-800:]}", file=sys.stderr); continue63 res = json.loads(sout.strip().splitlines()[-1])64 res.update({"design_index": i, "repeat": rep, "launch_wall_s": time.time() - t0, "assignment": assign, "label": d.get("label", "")})65 with open(out, "a") as f:66 f.write(json.dumps(res) + "\n")67 st = res["stats"]68 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 "69 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)707172if __name__ == "__main__":73 main()74