| |
| """Run a plan of generation jobs on a pool of GPUs, one vLLM worker per GPU. |
| |
| python scripts/run_jobs.py plans/trainer.json --gpus 0-7 |
| |
| A plan lists runs: {"name", "model", "data", "n", "shards", "est" (millions of tokens, for ordering), |
| optional "priority" (higher starts first), "preset", "set" and "args"}. "model" is "base", a folder name under |
| models/ ("qwen35-p50c") or a path; "data" is a file name under data/ without .jsonl. Each run is split into |
| shards; the biggest shards start first. Finished shards (a .done.json) are skipped, so the script can be restarted |
| any time. When every shard of a run is done, the run is merged into runs/<name>/merged.jsonl and summarized in |
| summary.json. |
| """ |
| import argparse |
| import fnmatch |
| import json |
| import os |
| import subprocess |
| import sys |
| import time |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.insert(0, ROOT) |
| from wsp.summary import summarize |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("plan") |
| ap.add_argument("--gpus", default="0-7") |
| ap.add_argument("--lab", default="/nvme/wsp-lab") |
| ap.add_argument("--only", nargs="*", help="run names or globs") |
| ap.add_argument("--retries", type=int, default=1) |
| ap.add_argument("--dry", action="store_true") |
| a = ap.parse_args() |
|
|
| if "-" in a.gpus: |
| lo, hi = map(int, a.gpus.split("-")) |
| GPUS = list(range(lo, hi + 1)) |
| else: |
| GPUS = [int(g) for g in a.gpus.split(",")] |
|
|
| plan = json.load(open(a.plan)) |
| runs = [r for r in plan["runs"] if not a.only or any(fnmatch.fnmatch(r["name"], p) for p in a.only)] |
|
|
|
|
| def model_path(m): |
| if m == "base": |
| return os.path.join(a.lab, "models", "Qwen3.5-0.8B") |
| return m if m.startswith("/") else os.path.join(a.lab, "models", m) |
|
|
|
|
| def run_dir(r): |
| return os.path.join(a.lab, "runs", r["name"]) |
|
|
|
|
| def shard_out(r, i): |
| return os.path.join(run_dir(r), f"shard{i}of{r['shards']}") |
|
|
|
|
| def merge(r): |
| rows = [json.loads(line) for line in open(os.path.join(a.lab, "data", r["data"] + ".jsonl"))] |
| recs, done = {}, [] |
| for i in range(r["shards"]): |
| for line in open(shard_out(r, i) + ".jsonl"): |
| x = json.loads(line) |
| recs[x["id"]] = x |
| done.append(json.load(open(shard_out(r, i) + ".done.json"))) |
| missing = [x["id"] for x in rows if x["id"] not in recs] |
| if missing and "--limit" not in r.get("args", []): |
| raise KeyError(f"{len(missing)} problems have no answers, the first is {missing[0]}") |
| merged = [recs[x["id"]] for x in rows if x["id"] in recs] |
| with open(os.path.join(run_dir(r), "merged.jsonl"), "w") as f: |
| for x in merged: |
| f.write(json.dumps(x) + "\n") |
| s = summarize(merged) |
| s["run"] = r |
| s["sampling"] = done[0]["sampling"] |
| s["edit_note"] = done[0]["edit_note"] |
| s["tokens_total"] = sum(d["tokens"] for d in done) |
| s["gpu_seconds"] = round(sum(d["gen_s"] + d["load_s"] + d["grade_s"] for d in done)) |
| s["tok_per_s_per_gpu"] = round(sum(d["tokens"] for d in done) / max(sum(d["gen_s"] for d in done), 1e-9)) |
| json.dump(s, open(os.path.join(run_dir(r), "summary.json"), "w"), indent=1) |
| print(f"[{time.strftime('%H:%M:%S')}] MERGED {r['name']}: acc={s['acc']:.3f} cut={s['cut_share']:.3f} " |
| f"mean_tokens={s['mean_tokens']:.0f} tok/s/gpu={s['tok_per_s_per_gpu']}", flush=True) |
|
|
|
|
| units = [] |
| for r in runs: |
| for i in range(r["shards"]): |
| if not os.path.exists(shard_out(r, i) + ".done.json"): |
| units.append((r, i)) |
| units.sort(key=lambda u: (-u[0].get("priority", 0), -u[0].get("est", 1) / u[0]["shards"])) |
| print(f"{len(units)} shards to run on GPUs {GPUS}:", flush=True) |
| for r, i in units: |
| print(f" {r['name']} {i}/{r['shards']} ~{r.get('est', 0) / r['shards']:.1f}M tokens", flush=True) |
| if a.dry: |
| sys.exit(0) |
|
|
| L = a.lab |
| env = dict(os.environ, NLTK_DATA=f"{L}/nltk_data", HF_HOME=f"{L}/hf", VLLM_CACHE_ROOT=f"{L}/vllm-cache", |
| FLASHINFER_WORKSPACE_BASE=L, TRITON_CACHE_DIR=f"{L}/triton-cache", |
| TORCHINDUCTOR_CACHE_DIR=f"{L}/inductor-cache", XDG_CACHE_HOME=f"{L}/xdg-cache", |
| MPLCONFIGDIR=f"{L}/mpl", TOKENIZERS_PARALLELISM="false", PYTHONUNBUFFERED="1") |
| free, running, tries = list(GPUS), {}, {} |
| failed = [] |
|
|
|
|
| def launch(r, i, g): |
| os.makedirs(run_dir(r), exist_ok=True) |
| cmd = [sys.executable, os.path.join(ROOT, "scripts", "gen_worker.py"), "--model", model_path(r["model"]), |
| "--data", os.path.join(a.lab, "data", r["data"] + ".jsonl"), "--out", shard_out(r, i), |
| "--n", str(r["n"]), "--shard", f"{i}/{r['shards']}", "--preset", r.get("preset", "think35")] |
| for kv in r.get("set", []): |
| cmd += ["--set", kv] |
| cmd += r.get("args", []) |
| log = open(shard_out(r, i) + ".log", "a") |
| p = subprocess.Popen(cmd, env=dict(env, CUDA_VISIBLE_DEVICES=str(g)), stdout=log, stderr=subprocess.STDOUT) |
| running[p] = (r, i, g, time.time()) |
| print(f"[{time.strftime('%H:%M:%S')}] START {r['name']} {i}/{r['shards']} on GPU {g}", flush=True) |
|
|
|
|
| while units or running: |
| while units and free: |
| r, i = units.pop(0) |
| launch(r, i, free.pop(0)) |
| time.sleep(10) |
| for p in list(running): |
| if p.poll() is None: |
| continue |
| r, i, g, t0 = running.pop(p) |
| free.append(g) |
| ok = p.returncode == 0 and os.path.exists(shard_out(r, i) + ".done.json") |
| mins = (time.time() - t0) / 60 |
| print(f"[{time.strftime('%H:%M:%S')}] {'DONE' if ok else 'FAIL'} {r['name']} {i}/{r['shards']} " |
| f"on GPU {g} after {mins:.1f} min (exit {p.returncode})", flush=True) |
| if not ok: |
| k = (r["name"], i) |
| tries[k] = tries.get(k, 0) + 1 |
| if tries[k] <= a.retries: |
| units.insert(0, (r, i)) |
| else: |
| failed.append(k) |
| continue |
| if all(os.path.exists(shard_out(r, j) + ".done.json") for j in range(r["shards"])): |
| try: |
| merge(r) |
| except Exception as e: |
| print("MERGE FAILED", r["name"], repr(e), flush=True) |
|
|
| for r in runs: |
| if all(os.path.exists(shard_out(r, j) + ".done.json") for j in range(r["shards"])) and \ |
| not os.path.exists(os.path.join(run_dir(r), "summary.json")): |
| merge(r) |
| print("ALL DONE" if not failed else f"FINISHED WITH FAILURES: {failed}", flush=True) |
|
|
