#!/bin/bash
# Prompt: VinciCoder-1.6M -> subset VinciCoder 150k (5类, 每类30k)
#
#SBATCH -t 04:00:00
#SBATCH -N 8
#SBATCH --ntasks-per-node=1
#SBATCH -J vincicoder_150k_bd4_8node_bsz256
#SBATCH --gpus-per-node 8
#SBATCH --array=1-5%1
#SBATCH --cpus-per-task 16
#SBATCH --mem=1024G
#SBATCH -e workdir/vincicoder_150k_bd4_8node_bsz256_error.log
#SBATCH -o workdir/vincicoder_150k_bd4_8node_bsz256_out.log
#SBATCH --open-mode=append
#SBATCH --requeue
#SBATCH --signal=B:USR1@120

# ============================================
# SIGNAL HANDLER FOR SLURM TIMEOUT/PREEMPTION
# ============================================
TRAIN_PID=""

handle_timeout() {
    echo ""
    echo "=========================================="
    echo "Received USR1 signal (SLURM timeout warning)"
    echo "Gracefully stopping training for checkpoint save..."
    echo "=========================================="
    
    # Send SIGTERM to training process to trigger graceful shutdown
    if [ -n "$TRAIN_PID" ] && kill -0 $TRAIN_PID 2>/dev/null; then
        echo "Sending SIGTERM to training process (PID: $TRAIN_PID)..."
        kill -TERM $TRAIN_PID 2>/dev/null || true
        # Wait for process to finish saving checkpoint
        wait $TRAIN_PID 2>/dev/null || true
    fi
    
    echo "Requesting SLURM to requeue job..."
    scontrol requeue $SLURM_JOB_ID
    exit 0
}

# Set up signal handler
trap handle_timeout USR1

# Auto-detect project root and change to it immediately
LMFLOW_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../../../.." && pwd)"
cd "${LMFLOW_ROOT}"

# Auto-detect account & partitions, resubmit via sbatch if not in a SLURM job
source "${LMFLOW_ROOT}/chengyue_script/dvlm/utils/auto_submit.sh"

# ============================================
# MULTI-NODE CONFIGURATION (8 NODES)
# ============================================
# Get master node address from SLURM
MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
MASTER_PORT=${MASTER_PORT:-29500}
NNODES=${SLURM_NNODES:-8}
NGPUS_PER_NODE=${NGPUS_PER_NODE:-8}

export MASTER_ADDR
export MASTER_PORT

echo "=========================================="
echo "VinciCoder-150k MDM Training (5 categories x 30k)"
echo "=========================================="
echo "Multi-Node Configuration (8 Nodes):"
echo "  NNODES: ${NNODES}"
echo "  NGPUS_PER_NODE: ${NGPUS_PER_NODE}"
echo "  MASTER_ADDR: ${MASTER_ADDR}"
echo "  MASTER_PORT: ${MASTER_PORT}"
echo "  SLURM_JOB_NODELIST: ${SLURM_JOB_NODELIST}"
echo ""
echo "Node: $(hostname)"
echo "Date: $(date)"
echo "SLURM Job ID: ${SLURM_JOB_ID:-N/A}"
echo "Working Directory: $(pwd)"
echo ""

# Activate environment
source "${LMFLOW_ROOT}/scripts/cluster_env.sh"
lmflow_activate_conda

echo "Python: $(python3 --version)"
echo "PyTorch: $(python3 -c 'import torch; print(torch.__version__)')"
echo "CUDA available: $(python3 -c 'import torch; print(torch.cuda.is_available())')"
echo "nvidia-smi:"
nvidia-smi || true
echo ""

# Set environment variables
export PYTHONPATH="${LMFLOW_ROOT}/src:${LMFLOW_ROOT}/third_party/vila:${LMFLOW_ROOT}/examples:${PYTHONPATH:-}"
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-16}"

export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True,max_split_size_mb:64,garbage_collection_threshold:0.9"

export TORCHDYNAMO_DISABLE=1
export TORCH_COMPILE=0
export MDM_VL_KEEP_VISION=1
export TOKENIZERS_PARALLELISM=false

# VILA dataset registry
# VILA_DATASETS — loaded from cluster_env.sh
# DATASET_BACKEND — loaded from cluster_env.sh

# Configurable parameters (override via env before submitting)
# NOTE: Must use MDM version of the model!
MODEL_PATH="${MODEL_PATH:-${LMFLOW_ROOT}/src/models/Qwen2.5-VL-3B-Instruct-mdm}"
TOKENIZER_NAME="${TOKENIZER_NAME:-Qwen/Qwen2.5-VL-3B-Instruct}"
MAX_LENGTH="${MAX_LENGTH:-2048}"
# VinciCoder-150k balanced subset (chart2code, web2html, img2svg, img2latex, plot2code)
DATASET_MIXTURE="${DATASET_MIXTURE:-vincicoder_150k}"

EPOCHS="${EPOCHS:-1}"
BATCH="${BATCH:-1}"
# With 8 nodes (64 GPUs), reduce grad_acc to maintain similar effective batch size
# 8-node: 64 GPUs * 1 batch * 4 grad_acc = 256 effective batch
GRAD_ACC="${GRAD_ACC:-4}"
LR="${LR:-2e-5}"
SAVE_STEPS="${SAVE_STEPS:-50}"
SAVE_TOTAL_LIMIT="${SAVE_TOTAL_LIMIT:-20}"
LOGGING_STEPS="${LOGGING_STEPS:-10}"
WORKERS="${WORKERS:-1}"
BLOCK_SIZE="${BLOCK_SIZE:-2048}"

# MDM specific parameters
MDM="${MDM:-1}"
BD_SIZE="${BD_SIZE:-4}"
LEARN_PADDING="${LEARN_PADDING:-1}"
CONPLE_MASK="${CONPLE_MASK:-1}"
ALWAYS_MASK_IM_END="${ALWAYS_MASK_IM_END:-1}"

# Enable Activation Checkpointing (saves memory by recomputing activations)
ACTIVATION_CHECKPOINTING="${ACTIVATION_CHECKPOINTING:-1}"

TEMPLATE="${TEMPLATE:-qwen2_5_no_reasoning}"
# Use fixed RUN_ID for auto-resume across sbatch submissions
RUN_ID="${RUN_ID:-vincicoder_150k_bd4_8node_bsz256}"
OUTPUT_DIR="${OUTPUT_DIR:-${LMFLOW_ROOT}/output_models/${RUN_ID}}"
LOG_DIR="${LOG_DIR:-${LMFLOW_ROOT}/log/${RUN_ID}}"

echo "Debug: Current PWD=$(pwd)"
echo "Debug: OUTPUT_DIR=${OUTPUT_DIR}"
echo "Debug: LOG_DIR=${LOG_DIR}"
mkdir -p "${OUTPUT_DIR}" "${LOG_DIR}"

# Sanity-check pixel settings stay compatible with the model patch size (14x14).
if [[ "${MIN_PIXELS}" =~ ^[0-9]+$ ]] && (( MIN_PIXELS % 196 != 0 )); then
    echo "Warning: MIN_PIXELS=${MIN_PIXELS} is not divisible by 14*14; expect padding overhead."
fi
if [[ "${MAX_PIXELS}" =~ ^[0-9]+$ ]] && (( MAX_PIXELS % 196 != 0 )); then
    echo "Warning: MAX_PIXELS=${MAX_PIXELS} is not divisible by 14*14; expect padding overhead."
fi

# ============================================
# AUTO-RESUME: Detect latest checkpoint
# ============================================
latest_checkpoint=""
if [ -d "${OUTPUT_DIR}" ]; then
    latest_checkpoint=$(find "${OUTPUT_DIR}" -maxdepth 1 -name "checkpoint-*" -type d | sort -V | tail -1)
    if [ -n "${latest_checkpoint}" ]; then
        echo "Found latest checkpoint: ${latest_checkpoint}"
    else
        echo "No checkpoint found in ${OUTPUT_DIR}"
        latest_checkpoint=""
    fi
else
    echo "Output dir ${OUTPUT_DIR} does not exist, training from scratch"
    latest_checkpoint=""
fi

# Build resume_from_checkpoint argument
resume_arg=""
if [ -n "${latest_checkpoint}" ]; then
    resume_arg="--resume_from_checkpoint ${latest_checkpoint}"
    echo "Resuming from checkpoint"
else
    resume_arg="--overwrite_output_dir"
    echo "Training from scratch"
fi

# Activation Checkpointing configuration
ac_arg=""
fsdp_wrap_arg=""
if [ "${ACTIVATION_CHECKPOINTING}" = "1" ]; then
    ac_arg="--gradient_checkpointing"
    fsdp_wrap_arg="" 
    echo "Activation Checkpointing enabled"
fi

echo ""
echo "=========================================="
echo "Configuration (VinciCoder-150k MDM - 8 Nodes - Bsz 256):"
echo "  Model: ${MODEL_PATH}"
echo "  Tokenizer: ${TOKENIZER_NAME}"
echo "  Dataset Mixture: ${DATASET_MIXTURE}"
echo "  VILA Registry: ${VILA_DATASETS}"
echo "  Output: ${OUTPUT_DIR}"
echo "  Log Dir: ${LOG_DIR}"
echo "  Epochs: ${EPOCHS}"
echo "  Batch: ${BATCH}"
echo "  Grad Acc: ${GRAD_ACC}"
echo "  LR: ${LR}"
echo "  MDM: ${MDM}  BD_SIZE: ${BD_SIZE}  LearnPadding: ${LEARN_PADDING}  ComplementaryMask: ${CONPLE_MASK}"
echo "  Processor MinPixels/MaxPixels: ${MIN_PIXELS}/${MAX_PIXELS}"
echo "  FSDP: shard_grad_op (ZeRO-2)"
echo "  Num Nodes: ${NNODES}"
echo "  GPUs per Node: ${NGPUS_PER_NODE}"
echo "  Total GPUs: $((NNODES * NGPUS_PER_NODE))"
echo "  Effective Batch Size: $((NNODES * NGPUS_PER_NODE * BATCH * GRAD_ACC))"
echo "  Resume Arg: ${resume_arg}"
echo "  SLURM Requeue: enabled"
echo "=========================================="
echo ""

# Check model path
echo "Checking paths..."
if [ ! -d "${MODEL_PATH}" ]; then
    echo "Model path does not exist: ${MODEL_PATH}"
    echo "Make sure you're using the MDM version of the model!"
    exit 1
fi
echo "Model path: ${MODEL_PATH}"

echo ""
echo "=========================================="
echo "Starting VinciCoder-150k MDM training with FSDP (8 nodes, 64 GPUs)..."
echo "=========================================="

# ============================================
# Detect interactive vs sbatch mode
# ============================================
# SLURM_SCRIPT_CONTEXT is set by sbatch jobs (they have SLURM_JOB_NAME matching #SBATCH -J).
# Interactive sessions (salloc/srun --pty) set SLURM_JOB_ID but NOT SLURM_ARRAY_TASK_ID (unless array job).
# Heuristic: if SLURM_NNODES matches #SBATCH -N (multi-node), we're in sbatch → use srun.
#            Otherwise (interactive, usually 1 node) → run torchrun directly.
IS_INTERACTIVE=0
if [ "${NNODES}" -eq 1 ] && [ -z "${SBATCH_MODE:-}" ]; then
    # Single-node: likely interactive. Override with SBATCH_MODE=1 to force srun.
    IS_INTERACTIVE=1
fi

# Build the torchrun command (shared between interactive and sbatch paths)
TORCHRUN_CMD="torchrun \
    --nnodes=${NNODES} \
    --nproc_per_node=${NGPUS_PER_NODE} \
    --rdzv_id=${SLURM_JOB_ID:-0} \
    --rdzv_backend=c10d \
    --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
    examples/finetune_qwenvl.py \
    --model_name_or_path ${MODEL_PATH} \
    --tokenizer_name ${TOKENIZER_NAME} \
    --conversation_template ${TEMPLATE} \
    --trust_remote_code 1 \
    --dataset_path ${DATASET_MIXTURE} \
    --output_dir ${OUTPUT_DIR} \
    --num_train_epochs ${EPOCHS} \
    --learning_rate ${LR} \
    --per_device_train_batch_size ${BATCH} \
    --gradient_accumulation_steps ${GRAD_ACC} \
    --logging_steps ${LOGGING_STEPS} \
    --save_steps ${SAVE_STEPS} \
    --save_total_limit ${SAVE_TOTAL_LIMIT} \
    --do_train \
    --dataloader_num_workers ${WORKERS} \
    --ddp_timeout 7200 \
    --mdm ${MDM} \
    --bd_size ${BD_SIZE} \
    --block_size ${BLOCK_SIZE} \
    --learn_padding ${LEARN_PADDING} \
    --complementary_mask ${CONPLE_MASK} \
    --bf16 \
    --fsdp shard_grad_op \
    --max_length ${MAX_LENGTH} \
    --always_mask_im_end ${ALWAYS_MASK_IM_END} \
    --report_to none \
    ${ac_arg} \
    ${fsdp_wrap_arg} \
    ${resume_arg}"

LOG_SUFFIX="${SLURM_ARRAY_TASK_ID:-1}-of-${SLURM_ARRAY_TASK_COUNT:-1}"

if [ "${IS_INTERACTIVE}" -eq 1 ]; then
    # ============================================
    # Interactive mode: run torchrun directly (no srun)
    # ============================================
    echo "[Interactive mode] Running torchrun directly on $(hostname) ..."
    ${TORCHRUN_CMD} 2>&1 | tee "${LOG_DIR}/train_${LOG_SUFFIX}.log" &
    TRAIN_PID=$!
else
    # ============================================
    # Sbatch mode: use srun to launch on all nodes
    # ============================================
    echo "[Sbatch mode] Launching via srun on ${NNODES} nodes ..."
    srun --nodes=${NNODES} --ntasks-per-node=1 \
        --export=ALL,MASTER_ADDR=${MASTER_ADDR},MASTER_PORT=${MASTER_PORT} \
        bash -c '
            cd '"${LMFLOW_ROOT}"'
            source '"${LMFLOW_ROOT}"'/scripts/cluster_env.sh
            lmflow_activate_conda

            export PYTHONPATH="'"${LMFLOW_ROOT}"'/src:'"${LMFLOW_ROOT}"'/third_party/vila:'"${LMFLOW_ROOT}"'/examples"
            export MIN_PIXELS="'"${MIN_PIXELS}"'"
            export MAX_PIXELS="'"${MAX_PIXELS}"'"
            export MDM_VL_KEEP_VISION=1
            export OMP_NUM_THREADS=16
            export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True,max_split_size_mb:64,garbage_collection_threshold:0.9"
            export TORCHDYNAMO_DISABLE=1
            export TORCH_COMPILE=0
            export TOKENIZERS_PARALLELISM=false

            torchrun \
                --nnodes='"${NNODES}"' \
                --nproc_per_node='"${NGPUS_PER_NODE}"' \
                --rdzv_id='"${SLURM_JOB_ID}"' \
                --rdzv_backend=c10d \
                --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
                examples/finetune_qwenvl.py \
                --model_name_or_path "'"${MODEL_PATH}"'" \
                --tokenizer_name "'"${TOKENIZER_NAME}"'" \
                --conversation_template "'"${TEMPLATE}"'" \
                --trust_remote_code 1 \
                --dataset_path "'"${DATASET_MIXTURE}"'" \
                --output_dir "'"${OUTPUT_DIR}"'" \
                --num_train_epochs '"${EPOCHS}"' \
                --learning_rate '"${LR}"' \
                --per_device_train_batch_size '"${BATCH}"' \
                --gradient_accumulation_steps '"${GRAD_ACC}"' \
                --logging_steps '"${LOGGING_STEPS}"' \
                --save_steps '"${SAVE_STEPS}"' \
                --save_total_limit '"${SAVE_TOTAL_LIMIT}"' \
                --do_train \
                --dataloader_num_workers '"${WORKERS}"' \
                --ddp_timeout 7200 \
                --mdm '"${MDM}"' \
                --bd_size '"${BD_SIZE}"' \
                --block_size '"${BLOCK_SIZE}"' \
                --learn_padding '"${LEARN_PADDING}"' \
                --complementary_mask '"${CONPLE_MASK}"' \
                --bf16 \
                --fsdp "shard_grad_op" \
                --max_length '"${MAX_LENGTH}"' \
                --always_mask_im_end '"${ALWAYS_MASK_IM_END}"' \
                --report_to none \
                '"${ac_arg}"' \
                '"${fsdp_wrap_arg}"' \
                '"${resume_arg}"'
        ' 2>&1 | tee "${LOG_DIR}/train_${LOG_SUFFIX}.log" &
    TRAIN_PID=$!
fi

# Wait for training to complete (or be interrupted by signal)
wait $TRAIN_PID
exit_code=$?

echo ""
echo "=========================================="
if [ $exit_code -eq 0 ]; then
    echo "VinciCoder-150k MDM Training (8-node, Bsz 256) Success!"
    
    # ============================================
    # AUTO EVAL: submit eval after training
    # ============================================
    echo ""
    echo "=========================================="
    echo "Auto-submitting evaluation jobs..."
    echo "=========================================="
    
    if [ -f "${LMFLOW_ROOT}/chengyue_script/dvlm/evaluation/submit_batch_eval.sh" ]; then
        bash "${LMFLOW_ROOT}/chengyue_script/dvlm/evaluation/submit_batch_eval.sh" \
            "${OUTPUT_DIR}" \
            --skip-existing 1 \
            2>&1 | tee "${LOG_DIR}/auto_eval_submit.log"
        echo "Evaluation jobs submitted. Check: ${OUTPUT_DIR}/eval_vlmeval/"
    else
        echo "submit_batch_eval.sh not found, skipping auto eval"
    fi
else
    echo "VinciCoder-150k MDM Training Failed (exit code: ${exit_code})"
    echo -e "Please check logs:\n  Output: ${LOG_DIR}/train.log\n  Error: workdir/vincicoder_150k_bd4_8node_bsz256_error.log"
fi
echo "=========================================="

exit $exit_code
