1#!/usr/bin/env python32"""Dispatch a job list over the MacLustr nodes (core-proportional shards) and collect results.3usage: dispatch_nodes.py launch --jobs results/manifests/jobs_main.json --name main [--nodes M3U96a:26,M3U96b:26,...]4 dispatch_nodes.py status --name main5 dispatch_nodes.py collect --name main # rsync <node>:~/wp46-cache/forecasts_<name>/ -> results/forecasts/<name>/6Panels (daily returns + realized measures) must exist on every node under ~/wp46-cache/panels (synced here)."""7from __future__ import annotations89import argparse10import json11import subprocess12from pathlib import Path1314ROOT = Path(__file__).resolve().parents[1]15DEFAULT_NODES = "M3U96a:24,M3U96b:26,M2U64:22,M5M36:16,M4M64a:14,M4M64c:14"161718def sh(cmd: list[str], timeout=600) -> str:19 return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout).stdout202122def main() -> None:23 ap = argparse.ArgumentParser()24 ap.add_argument("cmd", choices=["launch", "status", "collect", "sync-panels"])25 ap.add_argument("--jobs", default="")26 ap.add_argument("--name", required=True)27 ap.add_argument("--nodes", default=DEFAULT_NODES)28 a = ap.parse_args()29 nodes = [(n.split(":")[0], int(n.split(":")[1])) for n in a.nodes.split(",")]30 if a.cmd == "sync-panels":31 for n, _ in nodes:32 print(n, sh(["rsync", "-aq", "--delete", str(Path.home() / "wp46-cache" / "panels") + "/", f"{n}:~/wp46-cache/panels/"]))33 return34 if a.cmd == "launch":35 jobs = json.loads(Path(a.jobs).read_text())36 tot = sum(c for _, c in nodes)37 # proportional contiguous shards over the (heavy-first) job order -> interleave: node k takes jobs with38 # index % N in its share; simpler and balanced: assign job i to node by weighted round robin39 assign = {n: [] for n, _ in nodes}40 weights = [c / tot for _, c in nodes]41 acc = [0.0] * len(nodes)42 for i, j in enumerate(jobs):43 k = max(range(len(nodes)), key=lambda q: weights[q] * (i + 1) - acc[q])44 acc[k] += 145 assign[nodes[k][0]].append(j)46 for n, cores in nodes:47 jf = ROOT / "results" / "manifests" / f"jobs_{a.name}_{n}.json"48 jf.write_text(json.dumps(assign[n]))49 subprocess.run(["scp", "-q", str(jf), f"{n}:~/wp46_uqo/results/manifests/"], check=True)50 print(n, cores, len(assign[n]), sh(["ssh", n, f"bash ~/wp46_uqo/scripts/node_run.sh results/manifests/jobs_{a.name}_{n}.json 0/1 {cores} {a.name}"]).strip())51 return52 if a.cmd == "status":53 for n, _ in nodes:54 out = sh(["ssh", n, f"cd ~/wp46_uqo && (wc -l < ~/wp46-cache/forecasts_{a.name}/_runs.jsonl 2>/dev/null || echo 0); "55 f"python3 -c \"import json;print(len(json.load(open('results/manifests/jobs_{a.name}_{n}.json'))))\" 2>/dev/null; "56 f"grep -c error ~/wp46-cache/forecasts_{a.name}/_runs.jsonl 2>/dev/null; uptime | sed 's/.*load/load/'; "57 f"pgrep -fc run_benchmark.py"])58 print(n, " | ".join(x.strip() for x in out.splitlines()))59 return60 if a.cmd == "collect":61 dest = ROOT / "results" / "forecasts" / a.name62 dest.mkdir(parents=True, exist_ok=True)63 for n, _ in nodes:64 print(n, sh(["rsync", "-aq", f"{n}:~/wp46-cache/forecasts_{a.name}/", str(dest) + "/"], timeout=3600))65 runs = [json.loads(l) for f in sorted(dest.glob("_runs*.jsonl")) for l in f.read_text().splitlines()]66 print("collected runs:", len(runs), "errors:", sum("error" in r for r in runs))676869if __name__ == "__main__":70 main()71