#!/bin/bash
#SBATCH --partition=scavenger-gpu
#SBATCH --array=1-100  # 50 configurations to test
#SBATCH --nodelist=dcc-allenlab-gpu-[01-12],dcc-biostat-gpu-[01-14],dcc-coganlab-gpu-[01-08],dcc-dhvi-gpu-[01-09],dcc-majoroslab-gpu-[01-08],dcc-yaolab-gpu-[01-08],dcc-youlab-gpu-[01-57]
#SBATCH --ntasks=1
#SBATCH --nodes=1
#SBATCH --cpus-per-task=1
#SBATCH --mem-per-cpu=24G
#SBATCH --gres=gpu:1
#SBATCH --time=72:00:00
#SBATCH --output=/work/jf381/code/code/ICL_LOG/data_generation/%A_%a_%j.out
#SBATCH --error=/work/jf381/code/code/ICL_LOG/data_generation/%A_%a_%j.err
#SBATCH --requeue

# SLURM_ARRAY_TASK_ID=1

# Optional arguments [MODEL_NAME] [NUM_GPUS] [NUM_NODES]
NUM_GPUS=${1:-4}                # default 4 GPUs per node
NUM_NODES=${2:-1}               # default 1 node

# Clear first 2 args to avoid hydra issues
shift 2

# CPUs and memory automatically determined 
CPUS_PER_TASK=8                             # 8 CPUs per GPU
TOTAL_CPUS=$((NUM_GPUS * CPUS_PER_TASK))
TOTAL_MEM=$((NUM_GPUS * 32))                # 32gb ram per GPU

# Extract model name from hydra overrides
MODEL_NAME="transformer"
for arg in "$@"; do
  if [[ "$arg" == model=* ]]; then
    MODEL_NAME="${arg#model=}"
    break
  fi
done

JOB_NAME="${MODEL_NAME}_train_medium"

DATE=$(date +%F)

echo "Launching training job: $JOB_NAME"
echo "Nodes: ${NUM_NODES}, GPUs per node: ${NUM_GPUS}"
echo "Resources per node: ${TOTAL_CPUS} CPUs, ${TOTAL_MEM} GB memory"

SBATCH_ARGS="--job-name=${JOB_NAME} \
             --nodes=${NUM_NODES} \
             --gres=gpu:h200:${NUM_GPUS} \
             --ntasks-per-node=${NUM_GPUS} \
             --cpus-per-task=${CPUS_PER_TASK} \
             --mem=${TOTAL_MEM}G \
             --output=../../logs/${MODEL_NAME}/${DATE}/%x-%j.out \
             --error=../../logs/${MODEL_NAME}/${DATE}/%x-%j.err"

# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_1_3b.sh \
    # size=1_3b \
    # training=1_3b \
    # dataset=tinygsm_resume \
    # "$@"

# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_353M.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=transformer_1_3B

# cp /work/jf381/code/lm-research/main_1_2.py /work/jf381/code/lm-research/main.py
bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_fst_353M_bert_prediction.sh \
    size=medium \
    training=medium \
    dataset=slimpajama_60b \
    +benchmark.sbatch_number=$SLURM_ARRAY_TASK_ID \
    "$@" \
    model=fst_353M_bert_prediction


# ork/jf381/code/lm-research/main_1_2.py /work/jf381/code/lm-research/main.py
# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_353M_bert.sh \
#     size=medium \
#     training=medium \
#     dataset=slimpajama_6b \
#     "$@" \
#     model=transformer_353M_bert

# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_1_3B.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=transformer_1_3B

# cp /work/jf381/code/lm-research/main_1_7.py /work/jf381/code/lm-research/main.py
# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_1_3B.sh \
#     size=1_3b \
#     training=1_3b \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=fst_1_3B


# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_353M.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=fst_353M

# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_353M.sh \
#     size=1_3b \
#     training=1_3b \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=fst_1_3B

# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_353M.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     "$@" \
#     model=transformer_353M
# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_1_3b.sh \
#     size=1_3b \
#     training=1_3b \
#     dataset=tinygsm_resume \
#     "$@"

# bash  /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_generate.sh \
#     size=medium \
#     training=medium \
#     dataset=tinygsm_resume \
#     "$@"
  ## left_small 02-
  ## 04 big-
  ## 03 fst big

# # cat << 'EOF' > 1_3b_transformer.txt
# bash /work/jf381/code/lm-research/scripts/jobs/run_job_h200_bash_resume_transformer_1_3b_generate.sh \
#     size=1_3b \
#     training=1_3b \
#     dataset=tinygsm_resume \
#     "$@"
# EOF