#!/bin/bash
#SBATCH --job-name=ear-v2
#SBATCH --nodes=1
#SBATCH --gres=gpu:8
#SBATCH --ntasks-per-node=8  # match this to n_gpu if python
#SBATCH --cpus-per-task=4  # n_cpu*n_task has to be <= cpus per node (~64)

# other useful slurm commands
# --exclude=h100-ord01-03-[]
# --nodelist=h100-ord01-03-[]

# Get the node list and format it for torchrun
export NODELIST=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | tr '\n' ',' | sed 's/,$/\n/')

# Get the master node (first node in the list)
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
export MASTER_PORT=29500

# Calculate total number of processes
export WORLD_SIZE=$((SLURM_NNODES * SLURM_NTASKS_PER_NODE))

# Launch the training script using torchrun
srun /home/$USER/miniconda3/envs/suno_env/bin/torchrun \
    --nnodes=$SLURM_NNODES \
    --nproc_per_node=$SLURM_NTASKS_PER_NODE \
    --rdzv_id=$SLURM_JOB_ID \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
    train_ear.py  # Your training script