import os
import sys
GPU_PER_NODE = 6
TASK_PER_GPU = 8

if __name__ == "__main__":
    split = sys.argv[1]
    start_idx = sys.argv[2]
    end_idx = sys.argv[3]
    job_id = os.environ["SLURM_JOBID"]
    gpu_id = 0
    job_bash = f"temp/job/{job_id}.sh"
    with open(job_bash, "w") as f:
        f.write("export TRANSFORMERS_OFFLINE=1\n")
        for i, idx in enumerate(range(int(start_idx), int(end_idx))):
            zfill_idx = str(idx).zfill(6)
            f.write(f"CUDA_VISIBLE_DEVICES={gpu_id} python3 offline_grounding_dino.py {split} {zfill_idx} &> temp/log/{split}_{zfill_idx}_{job_id}_{gpu_id}.txt &\n")
            gpu_id = (gpu_id + 1) % GPU_PER_NODE
        f.write("sleep 7200\n")
    print("run!")
    os.system(f"bash {job_bash}")
    print("end!")