#!/bin/bash
#SBATCH --job-name=hoot_align
#SBATCH --output=/app/suno/slurm/logs/run_hoot_align_%j.txt
#SBATCH --error=/app/suno/slurm/logs/run_hoot_align_%j.err
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --gres=gpu:8

# Create logs directory if it doesn't exist
export CUDA_LAUNCH_BLOCKING=0
export NCCL_DEBUG=WARN
export TORCH_DISTRIBUTED_DEBUG=OFF
export TORCH_CPP_LOG_LEVEL=WARNING

export OMP_NUM_THREADS=1
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=12835
export TRITON_CACHE_DIR=/mnt/localdisk/.triton_cache_$USER
export SLURM_NTASKS_PER_NODE=8


TRAIN_PATH=/home/tony/Work/tony/hoot
echo working from $TRAIN_PATH
cd $TRAIN_PATH
pkill -f 'spawn_main'
# Calculate chunk size based on total dataset size
# We'll pass this as an argument to our script
START_IDX="${2:-0}"    # Default to 0 if not specified

# Launch 8 processes, one for each GPU
for GPU_ID in {0..7}; do
    # Calculate the start index for this GPU
    # Each GPU will handle its own portion of the data
    PROCESS_START_IDX=$((START_IDX + GPU_ID * 10000))
    echo "Processing GPU $GPU_ID with start index $PROCESS_START_IDX"
    # Launch the process in the background
    CUDA_VISIBLE_DEVICES=$GPU_ID /home/tony/anaconda3/envs/suno_env_dev/bin/python p_1_hoot_alignment.py \
        --dataset discogs \
        --start_index $PROCESS_START_IDX &
done

# Wait for all background processes to complete
wait