File size: 2,257 Bytes
2188a91 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | """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()
|