File size: 2,336 Bytes
e7d8c8b | 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 | #!/usr/bin/env python
"""Persistent GPU orchestrator. Keeps a job queue; whenever a GPU frees, launches the
next job on it. Survives session teardown (setsid + cron anchor). Never idles a GPU."""
import os,sys,json,time,subprocess,shlex
AD="/workspace/aditi"; QF=f"{AD}/jobqueue.json"; LOG=f"{AD}/logs/orchestrator.log"
def log(m):
with open(LOG,"a") as f: f.write(f"{time.strftime('%FT%T')} {m}\n")
def gpu_busy():
try:
out=subprocess.run(["nvidia-smi","--query-compute-apps=gpu_uuid","--format=csv,noheader"],
capture_output=True,text=True,timeout=20).stdout
uu=subprocess.run(["nvidia-smi","--query-gpu=index,uuid","--format=csv,noheader"],
capture_output=True,text=True,timeout=20).stdout
idx={}
for l in uu.strip().splitlines():
i,u=[x.strip() for x in l.split(",")]; idx[u]=int(i)
busy=set()
for l in out.strip().splitlines():
u=l.strip()
if u in idx: busy.add(idx[u])
return busy
except Exception: return {0,1}
def load():
try: return json.load(open(QF))
except Exception: return {"jobs":[]}
def save(q): json.dump(q,open(QF,"w"),indent=2)
def main():
while True:
q=load(); jobs=q.get("jobs",[])
pend=[j for j in jobs if j.get("state")=="pending"]
if not pend:
time.sleep(60); continue
busy=gpu_busy(); free=[g for g in (0,1) if g not in busy]
for g in free:
pend=[j for j in load().get("jobs",[]) if j.get("state")=="pending"]
if not pend: break
j=pend[0]
cmd=j["cmd"]; name=j["name"]
env=dict(os.environ); env["CUDA_VISIBLE_DEVICES"]=str(g)
lf=f"{AD}/logs/{name}.log"
log(f"LAUNCH gpu{g} {name}")
subprocess.Popen(f"setsid {cmd} </dev/null > {lf} 2>&1 &",
shell=True,env=env,cwd=AD,start_new_session=True)
q=load()
for jj in q["jobs"]:
if jj["name"]==name and jj["state"]=="pending":
jj["state"]="launched"; jj["gpu"]=g; jj["started"]=time.strftime('%FT%T'); break
save(q)
time.sleep(90) # let it claim the GPU before checking again
time.sleep(45)
if __name__=="__main__": main()
|