#!/bin/bash
#SBATCH --job-name=cs_diff
#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)
# --nodelist=h100-ord01-03-[120,129,143,173,202,206,240,278,292,305-306,328,375,378,398,402,409,421,431,453,478,523,548,565,606,676,687,704,713,741,794,818,839,844,873,886,900,902,920,930,947,966,984,986-987,994,1012,1021]
# h100-ord01-03-[120,129,143,173,202,206,240,278,292,305-306,328,375,378,398,402,409,421,431,453,478,523,548,565,606,676,687,704,713,741,794,818,839,844,873,886,900,902,920,930,947,966,984,986-987,994,1012,1021]
# other useful slurm commands


# "/app/suno/data/diffusion_v5/v0",
# /app/suno/data/diffusion_mix/dac_vae_tuned_25hz

# /app2/suno/data/christian/models/mert_768d_centroids_4x4k_50.npy
# /app2/suno/data/christian/models/mert_768d_centroids_12x4k_50.npy

# -K1 /home/$USER/miniconda3/envs/suno_diff45/bin/python

# 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

# suno_clone_diff
# suno_env_fa2

# 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
echo "Starting train script..."
srun -K1 /home/$USER/miniconda3/envs/suno_clone_diff/bin/python -u train.py \
    --master_addr=$MASTER_ADDR \
    --master_port=$MASTER_PORT \
    --config /home/christian/code/neon/sunoDiff/config/25hz_v3_flow_shared_pretrain.json