#!/bin/bash

# Parameter check
if [[ $# -ne 5 ]]; then
    echo "Usage: $0 <run_config_module> <partition> <gpu_type> <num_gpus> <num_cpus_per_gpu>"
    exit 1
fi

# Assign parameters
run_config=$1
partition=$2
gpu_type=$3
num_gpus=$4
num_cpus_per_gpu=$5

# Derive run name
run_name=$(basename "$run_config")
run_name="${run_name%.*}"

# Confirm inputs
echo "Starting run with configuration: $run_config"
echo "Partition: $partition"
echo "GPU type: $gpu_type"
echo "Number of GPUs: $num_gpus"
echo "Number of CPUs per GPU: $num_cpus_per_gpu"

# Set paths
current_dir_path=$(dirname "$(realpath "$0")")
venv_path=$(realpath "$current_dir_path/../../venv/bin/activate")

# Create log directory
log_dir="/work/rr41qemu-MA/logs/${run_name}_$(date +%Y-%m-%d_%H-%M-%S)"
mkdir -p "$log_dir"
echo "Logging to $log_dir"

# Submit job
sbatch <<EOF
#!/bin/bash
#SBATCH --job-name=$run_name
#SBATCH --output=$log_dir/log.out
#SBATCH --error=$log_dir/log.err
#SBATCH --time=48:00:00
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=$(($num_gpus * $num_cpus_per_gpu))
#SBATCH --mem=32G
#SBATCH --partition=$partition
#SBATCH --gpus=$gpu_type:$num_gpus

echo "Loading python virtual environment..."
source $venv_path

echo "Loading python 3.10..."
module load Python/3.10.4-GCCcore-11.3.0

cd $current_dir_path

torchrun --nproc_per_node=$num_gpus training_wrapper.py $run_config --item_limit -1
EOF