| """Tiny GPU-slot scheduler: runs a list of training configs across the local T4s, |
| N concurrent processes per GPU, skipping configs whose result file already exists.""" |
|
|
| import argparse |
| import json |
| import os |
| import subprocess |
| import sys |
| import threading |
| import time |
| from queue import Queue |
|
|
| PY = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), |
| ".venvr", "bin", "python") |
| TRAIN = os.path.join(os.path.dirname(os.path.abspath(__file__)), "train.py") |
|
|
|
|
| def worker(slot, gpu, q, logdir): |
| while True: |
| try: |
| name, args = q.get_nowait() |
| except Exception: |
| return |
| out = os.path.join(logdir, f"{name}.json") |
| if os.path.exists(out) and os.path.getsize(out) > 0: |
| print(f"[slot{slot}] skip {name}", flush=True) |
| q.task_done() |
| continue |
| env = dict(os.environ, CUDA_VISIBLE_DEVICES=str(gpu), OMP_NUM_THREADS="4") |
| cmd = [PY, TRAIN] + args + ["--out", out] |
| log = os.path.join(logdir, f"{name}.log") |
| t0 = time.time() |
| print(f"[slot{slot}/gpu{gpu}] START {name}", flush=True) |
| with open(log, "w") as fh: |
| rc = subprocess.call(cmd, stdout=fh, stderr=subprocess.STDOUT, env=env) |
| print(f"[slot{slot}/gpu{gpu}] DONE {name} rc={rc} " |
| f"({time.time()-t0:.0f}s)", flush=True) |
| q.task_done() |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--jobs", required=True, help="json file: {name: [args...]}") |
| ap.add_argument("--logdir", required=True) |
| ap.add_argument("--gpus", default="0,1") |
| ap.add_argument("--per-gpu", type=int, default=2) |
| a = ap.parse_args() |
| os.makedirs(a.logdir, exist_ok=True) |
| jobs = json.load(open(a.jobs)) |
| q = Queue() |
| for name, args in jobs.items(): |
| q.put((name, args)) |
| threads = [] |
| gpus = [g for g in a.gpus.split(",") if g] |
| for gi, gpu in enumerate(gpus): |
| for s in range(a.per_gpu): |
| t = threading.Thread(target=worker, args=(f"{gi}.{s}", gpu, q, a.logdir)) |
| t.start() |
| threads.append(t) |
| time.sleep(3) |
| for t in threads: |
| t.join() |
| print("ALL DONE", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
