| |
| """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) |
| time.sleep(45) |
| if __name__=="__main__": main() |
|
|
