Buckets:
| """Concurrent no-lock runner for the Stage-B/C queue (owner directive 2026-09-10 13:58 EDT: the exclusive GPU | |
| lock is no longer required; run 2-3 short jobs in parallel with launch checks). | |
| python aide/train/run_concurrent.py [--max-parallel 3] [--only a,b] [--dry-run] [--skip-post] | |
| Launch checks before every job (re-checked every 30 s until satisfied): | |
| nvidia-smi free VRAM >= job need + 4 GB (need: repro 6 GB, cwm3d 2 GB) | |
| host `free -g` available >= 14 GB + job RSS estimate (3 GB) | |
| OMP/MKL threads = 4 in the child environment (train_kp also calls torch.set_num_threads(4)) | |
| Each job is `train_kp.py <argv> --no-lock` (evaluation inside the job); gated configs use run_stage_bc.gate_status. | |
| Post: interventions for run_stage_bc.INTERVENTION_RUNS that lack an interventions json, then report_stage_bc.py. | |
| Failures are logged (FAILED <name> rc) and do not stop the queue. Log: results/runs/stage_bc_driver.log. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| for _v in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS"): | |
| os.environ.setdefault(_v, "4") | |
| import argparse # noqa: E402 | |
| import subprocess # noqa: E402 | |
| import sys # noqa: E402 | |
| import time # noqa: E402 | |
| from pathlib import Path # noqa: E402 | |
| ROOT = Path(__file__).resolve().parents[2] | |
| sys.path.insert(0, str(ROOT)) | |
| from aide.data import protocol as P # noqa: E402 | |
| from aide.train import run_stage_bc as Q # noqa: E402 | |
| from common.runmeta import write_run_json # noqa: E402 | |
| PY = sys.executable | |
| RUNS = P.RESULTS / "runs" | |
| NEED_GB = {"repro": 6.0, "cwm3d": 2.0, "cwm3d_vis": 4.0} # attention-head variants use ~3.8 GB | |
| RSS_GB = 3.0 | |
| def gpu_free_gb(): | |
| out = subprocess.run(["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"], capture_output=True, text=True) | |
| return float(out.stdout.strip().split("\n")[0]) / 1024.0 | |
| def ram_avail_gb(): | |
| with open("/proc/meminfo") as f: | |
| for line in f: | |
| if line.startswith("MemAvailable"): | |
| return int(line.split()[1]) / 1024 / 1024 | |
| return 0.0 | |
| def checks_ok(family): | |
| g, r = gpu_free_gb(), ram_avail_gb() | |
| ok = g >= NEED_GB.get(family, 2.0) + 4.0 and r >= 14.0 + RSS_GB | |
| return ok, f"vram_free {g:.1f} GB, ram_avail {r:.1f} GB" | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--max-parallel", type=int, default=3) | |
| ap.add_argument("--only", default="") | |
| ap.add_argument("--dry-run", action="store_true") | |
| ap.add_argument("--skip-post", action="store_true") | |
| ap.add_argument("--max-minutes", type=float, default=30.0) | |
| a = ap.parse_args() | |
| t0 = time.time() | |
| queue = [c for c in Q.QUEUE if not a.only or c["name"] in a.only.split(",")] | |
| todo = [c for c in queue if not Q.is_done(c["name"])] | |
| Q.log(f"concurrent runner start (no lock, max {a.max_parallel} parallel): {len(todo)} to run: {[c['name'] for c in todo]}") | |
| if a.dry_run: | |
| for c in todo: | |
| print(c["name"], " ".join(c["argv"])) | |
| return 0 | |
| env = {**os.environ, "OMP_NUM_THREADS": "4", "MKL_NUM_THREADS": "4", "OPENBLAS_NUM_THREADS": "4"} | |
| running, results, skipped, retried = {}, [], [], set() | |
| while todo or running: | |
| # reap | |
| for name, (proc, t_s, logf) in list(running.items()): | |
| rc = proc.poll() | |
| if rc is not None: | |
| logf.close() | |
| dur = (time.time() - t_s) / 60 | |
| results.append({"name": name, "rc": rc, "minutes": round(dur, 2)}) | |
| Q.log(f"{'end' if rc == 0 else 'FAILED'} {name}: rc {rc}, {dur:.2f} min") | |
| del running[name] | |
| if rc != 0 and "CUDA out of memory" in (RUNS / f"{name}_train.log").read_text()[-4000:] \ | |
| and name not in retried: | |
| retried.add(name) | |
| todo.append(next(c for c in queue if c["name"] == name)) | |
| Q.log(f"requeue {name} once (CUDA OOM caused by concurrent GPU use)") | |
| # launch | |
| while todo and len(running) < a.max_parallel: | |
| c = todo[0] | |
| if c["gated"]: | |
| g = Q.gate_status() | |
| if g is None and any(r["name"] == Q.GATE_RUN for r in results) is False and not Q.is_done(Q.GATE_RUN): | |
| if Q.GATE_RUN in running: | |
| break # wait for the gate run | |
| if g is None or not g[0]: | |
| Q.log(f"skip {c['name']}: gate {'not ready' if g is None else 'FAILED ' + str(g[1])}") | |
| skipped.append(c["name"]) | |
| todo.pop(0) | |
| continue | |
| ok, msg = checks_ok(c["family"]) | |
| if not ok: | |
| Q.log(f"launch check failed for {c['name']} ({msg}); retry in 30 s") | |
| break | |
| todo.pop(0) | |
| argv = [x for x in c["argv"] if x != "--defer-eval"] + ["--max-minutes", str(a.max_minutes)] | |
| logf = open(RUNS / f"{c['name']}_train.log", "a") | |
| proc = subprocess.Popen([PY, str(ROOT / "aide/train/train_kp.py")] + argv, cwd=str(ROOT), env=env, | |
| stdout=logf, stderr=subprocess.STDOUT) | |
| running[c["name"]] = (proc, time.time(), logf) | |
| Q.log(f"start {c['name']} (pid {proc.pid}; {msg}; {len(running)} running)") | |
| time.sleep(5 if running else 30) | |
| Q.log(f"queue finished: {len(results)} jobs, failures {[r['name'] for r in results if r['rc'] != 0]}, skipped {skipped}") | |
| if not a.skip_post: | |
| todo_int = [r for r in Q.INTERVENTION_RUNS if Q.is_done(r) and not Q.interventions_done(r) | |
| and (RUNS / r / f"run_{r}_latest.json").exists()] | |
| if todo_int: | |
| Q.log(f"interventions for {todo_int}") | |
| cmd = [PY, str(ROOT / "aide/eval/interventions.py")] + sum([["--run", r] for r in todo_int], []) | |
| rc = subprocess.call(cmd, cwd=str(ROOT), env=env) | |
| Q.log(f"interventions exit {rc}") | |
| rc = subprocess.call([PY, str(ROOT / "aide/eval/report_stage_bc.py"), "--out", "stage_bc"], cwd=str(ROOT), env=env, | |
| stdout=open(RUNS / "report_stage_bc.log", "a"), stderr=subprocess.STDOUT) | |
| Q.log(f"report_stage_bc exit {rc}") | |
| write_run_json(P.RESULTS / "tables", "stage_bc_concurrent", [__file__, Q.__file__, P.__file__], [P.CACHE_NPZ], seed=None, | |
| t_start=t0, input_listing=False, extra={"results": results, "skipped": skipped, "max_parallel": a.max_parallel}) | |
| Q.log("driver exit 0") | |
| return 0 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |
Xet Storage Details
- Size:
- 6.61 kB
- Xet hash:
- 28f01d1fc8d06cfc67bb668b035305de2386b299d4afe2198b4ae4bf3af33540
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.