

from itertools import product
import os 
import subprocess 
import shutil
import itertools
DIRECTORY = "/data/user_data/gghosal/grokkreason/results/sweep_ratio/"
args = []
args = []
for dir in os.listdir(DIRECTORY):
    args.append(f"--dir {DIRECTORY}{dir} --final_outfile sweepratio_{dir}.out")
num_args = len(args)
args_string = "ARGS=("
for arg in args[:-1]:
    args_string+=f""""{arg}" """
args_string+=f""""{args[-1]}")"""
print(args_string)


SLURMSCRIPT = "#!/bin/bash" \
                    + f"\n#SBATCH --job-name=grokeval" \
                    +"\n#SBATCH --output=./logs/grokeval.out" \
                    +"\n#SBATCH --error=./logs/grokeval.err" \
                    +"\n#SBATCH --partition=general" \
                    +f"\n#SBATCH --array=0-{num_args-1}%8"\
                    +"\n#SBATCH --time=7:00:00" \
                    +"\n#SBATCH --gres=gpu:1" \
                    +"\n#SBATCH --mem=64G" \
                    +"\n#SBATCH --nodes=1" \
                    +"\n#SBATCH --requeue"\
                    +"\nsource ../../miniconda3/etc/profile.d/conda.sh"\
                    +"\nconda activate grokkedtransformer"\
                    +f"\n{args_string}"\
                    +f"\nnvidia-smi" \
                    +"\npython eval_qa.py ${ARGS[$SLURM_ARRAY_TASK_ID]}"
f = open("jobarrayscript.sbatch", "w")
f.write(SLURMSCRIPT)
f.close()
subprocess.run(f'sbatch jobarrayscript.sbatch', shell=True)





