#!/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()