#!/bin/bash
#SBATCH --job-name="tony_diff"
#SBATCH --nodes=16
#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)
#SBATCH --output=/app/suno/slurm/logs/run_%x_%j.txt
#SBATCH --error=/app/suno/slurm/logs/run_%x_%j_err.txt

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

# this is fairly arbity
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=12885

# cluster specific defaults
export OMP_NUM_THREADS=1
export NCCL_CROSS_NIC=2

# explicit cache dirs to not get user conflicts
export TRITON_CACHE_DIR=/mnt/localdisk/.triton_cache_$USER

pkill -f 'spawn_main'
cd /home/tony/Work/neon/sunoDiff
# using bare python rather than torchrun prevents hanging on nightly
# -K1 means slurm job will crash if run crashes rather than just hang
# -u means python will stream stoud as it comes rather than buffer
# reference run: https://wandb.ai/suno/harmonai_train_vt/runs/s2bcr31c
# starting run: "/app/suno/checkpoints/2024-12-05_03-42-43_s4538"
echo "Starting train script..."
srun -K1 /home/tony/anaconda3/envs/suno_env_dev/bin/python -u train_dpo.py \
    --config_path="/home/tony/Work/tony/slurm/diffusion/diffv2_real_v3.json" \
    --master_addr=$MASTER_ADDR \
    --master_port=$MASTER_PORT