1#!/usr/bin/env python32"""Distribute Monte Carlo cells over MacLustr nodes (greedy by estimated cost, weighted by worker slots).34Writes experiments/cluster_assignment.json, syncs the repo to each node, and launches `xargs -P` worker pools.5Usage: python3 experiments/scatter_cluster.py --configs a.yaml b.yaml --launch6Collect later with experiments/collect_cluster.sh.7"""8import argparse, json, os, pathlib, subprocess, sys9sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1] / "experiments"))10import yaml11from run_monte_carlo import expand_cells1213ROOT = pathlib.Path(__file__).resolve().parents[1]14NODES = {"M3U96a": 6, "M3U96b": 6, "M2U64": 4, "M4M64a": 4, "M4M64b": 4, "M4M64c": 4, "M5M36": 4, "M4BP48": 3, "M4BP36": 3,15 "M2M32c": 3, "m4mm": 3, "m4mn": 3, "m6ma": 3, "m6mb": 3, "m5ma": 3, "m4md": 2, "m4ml": 2, "M3BA24": 2}161718def cost(c):19 extras = 1.6 if c["extras"] else 1.020 return c["N"] * (c["p"] ** 1.5) * c["reps"] * (1 + 0.5 * c["splits_max"]) * extras212223def main():24 ap = argparse.ArgumentParser(); ap.add_argument("--configs", nargs="+", required=True); ap.add_argument("--launch", action="store_true")25 a = ap.parse_args()26 jobs = []27 for cfg in a.configs:28 cells = expand_cells(yaml.safe_load(open(cfg)))29 for i, c in enumerate(cells):30 jobs.append({"config": os.path.relpath(cfg, ROOT), "cell": i, "cost": cost(c), "name": f'{c["label"]}_p{c["p"]}_N{c["N"]}_K{c["K"]}'})31 jobs.sort(key=lambda j: -j["cost"])32 load = {n: 0.0 for n in NODES}; assign = {n: [] for n in NODES}33 for j in jobs:34 n = min(NODES, key=lambda k: load[k] / NODES[k])35 assign[n].append(j); load[n] += j["cost"]36 total = sum(j["cost"] for j in jobs)37 print(f"{len(jobs)} jobs; est. share per node (%):", {n: round(100 * load[n] / total, 1) for n in NODES})38 json.dump({"nodes": NODES, "assignment": assign}, open(ROOT / "experiments" / "cluster_assignment.json", "w"), indent=1)39 if not a.launch:40 return41 for n, js in assign.items():42 if not js:43 continue44 script = "\n".join(f'{j["config"]} {j["cell"]}' for j in js) + "\n"45 (ROOT / "logs" / f"jobs_{n}.txt").write_text(script)46 subprocess.run(["rsync", "-az", "--exclude", "results", "--exclude", "logs", "--exclude", "references", "--exclude", "paper",47 "--exclude", ".git", f"{ROOT}/", f"{n}:~/wp28/repo/"], check=True)48 subprocess.run(["scp", "-q", str(ROOT / "logs" / f"jobs_{n}.txt"), f"{n}:~/wp28/jobs.txt"], check=True)49 W = NODES[n]50 cmd = (f"cd ~/wp28/repo && mkdir -p logs results/raw && nohup sh -c 'cat ~/wp28/jobs.txt | xargs -P {W} -L 1 sh -c "51 f"\"OMP_NUM_THREADS=2 ~/wp28/venv/bin/python experiments/run_monte_carlo.py --config \\$0 --cells \\$1 "52 f">> logs/cell_\\$(basename \\$0 .yaml)_\\$1.log 2>&1\"' > logs/pool.log 2>&1 &")53 subprocess.run(["ssh", n, cmd], check=True)54 print(f"launched {len(js)} jobs on {n} with {W} workers")555657if __name__ == "__main__":58 main()59