#!/bin/bash
#SBATCH --job-name=encode_semantic
#SBATCH --nodes=4                   # Number of nodes (adjust as needed)
#SBATCH --ntasks-per-node=8         # 8 GPUs per node
#SBATCH --gres=gpu:8                # Request 8 GPUs per node
#SBATCH --cpus-per-task=8           # CPUs per GPU task
#SBATCH --mem=256G                  # Memory per node
#SBATCH --time=24:00:00             # Max 24 hours
#SBATCH --output=logs/encode_semantic_%j.out
#SBATCH --error=logs/encode_semantic_%j.err
# #SBATCH --partition=gpu           # Uncomment if your cluster requires a partition

# Exit on error
set -e

echo "=========================================="
echo "SLURM Job: Semantic Encoding Pipeline"
echo "=========================================="
echo "Job ID: $SLURM_JOB_ID"
echo "Nodes: $SLURM_JOB_NUM_NODES"
echo "Tasks per node: $SLURM_NTASKS_PER_NODE"
echo "Total tasks: $SLURM_NTASKS"
echo "=========================================="
echo ""

# Create logs directory if it doesn't exist
mkdir -p logs

# Environment setup
echo "Setting up environment..."
source ~/.bashrc

# Activate conda environment (adjust environment name as needed)
# Uncomment and modify the line below based on your setup
# conda activate your_env_name

# Set environment variables for distributed training
export MASTER_ADDR=$(scontrol show hostname $SLURM_NODELIST | head -n 1)
export MASTER_PORT=29500
export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=0
export NCCL_SOCKET_IFNAME=ib0

echo "Master address: $MASTER_ADDR"
echo "Master port: $MASTER_PORT"
echo ""

# Configuration
METADATA_PATH="/home/tony/Data/Preference/RealGen/metas_v5_val_sampled_diverse.jsonl"
OUTPUT_DIR="/app2/suno/data/semantic_code/sft"
BATCH_SIZE=16

echo "Configuration:"
echo "  Metadata: $METADATA_PATH"
echo "  Output:   $OUTPUT_DIR"
echo "  Batch:    $BATCH_SIZE"
echo ""

# Create output directory
mkdir -p $OUTPUT_DIR

# Launch distributed encoding job
echo "Launching distributed encoding..."
echo ""

srun python /home/tony/Work/tony/RealGen/encode_semantic_codes.py \
    --metadata_path $METADATA_PATH \
    --output_dir $OUTPUT_DIR \
    --batch_size $BATCH_SIZE

echo ""
echo "=========================================="
echo "Job completed!"
echo "=========================================="