#!/bin/bash
#SBATCH -A nvr_elm_llm
#SBATCH -p batch_block1,batch_block3,batch_block4
#SBATCH -t 4:00:00
#SBATCH -N 1
#SBATCH -J train_j63_mdm
#SBATCH --gpus-per-node 8
#SBATCH --cpus-per-task 16
#SBATCH --mem=1024G
#SBATCH -e workdir/train_j63_mdm_error.log
#SBATCH -o workdir/train_j63_mdm_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

# Resolve project root and change to it immediately
PROJECT_ROOT="/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow"
cd "${PROJECT_ROOT}"

echo "=========================================="
echo "MDM_Qwen2.5-VL Training (VILA j63 Mixture)"
echo "=========================================="
echo "Node: $(hostname)"
echo "Date: $(date)"
echo "SLURM Job ID: ${SLURM_JOB_ID:-N/A}"
echo "Working Directory: $(pwd)"
echo ""

# Activate environment
source "${PROJECT_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="${PROJECT_ROOT}/src:${PROJECT_ROOT}/VILA:${PROJECT_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 (cs-oci-ord has the correct paths for this cluster)
# VILA_DATASETS — loaded from cluster_env.sh
# DATASET_BACKEND — loaded from cluster_env.sh
# Qwen2.5-VL processor dynamic resolution controls
# Default to ~448x448 (200704 pixels) to stay aligned with 14x14 patching.
# Override via env vars only after profiling memory usage.
export MIN_PIXELS="${MIN_PIXELS:-200704}"
export MAX_PIXELS="${MAX_PIXELS:-200704}"

# Configurable parameters (override via env before submitting)
# NOTE: Must use MDM version of the model!
MODEL_PATH="${MODEL_PATH:-${PROJECT_ROOT}/src/models/Qwen2.5-VL-3B-Instruct-mdm}"
TOKENIZER_NAME="${TOKENIZER_NAME:-Qwen/Qwen2.5-VL-3B-Instruct}"

# VILA mixture name (j63_recipe, j63_no_sft, j63_small, j63_tiny, etc.)
DATASET_MIXTURE="${DATASET_MIXTURE:-j63_recipe}"

EPOCHS="${EPOCHS:-1}"
BATCH="${BATCH:-1}"
GRAD_ACC="${GRAD_ACC:-8}"
LR="${LR:-2e-5}"
SAVE_STEPS="${SAVE_STEPS:-500}"
SAVE_TOTAL_LIMIT="${SAVE_TOTAL_LIMIT:-20}"
LOGGING_STEPS="${LOGGING_STEPS:-10}"
WORKERS="${WORKERS:-4}"

# MDM specific parameters
MDM="${MDM:-1}"
BD_SIZE="${BD_SIZE:-8}"
LEARN_PADDING="${LEARN_PADDING:-1}"
CONPLE_MASK="${CONPLE_MASK:-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
# Override with RUN_ID env var if you want a different experiment name
RUN_ID="${RUN_ID:-j63_mdm}"
OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/output_models/${RUN_ID}}"
LOG_DIR="${LOG_DIR:-${PROJECT_ROOT}/log/${RUN_ID}}"
NGPUS="${NGPUS:-8}"

# AutoResume: check for existing checkpoints
AUTO_RESUME="${AUTO_RESUME:-1}"
RESUME_FROM_CHECKPOINT="${RESUME_FROM_CHECKPOINT:-}"

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

echo "=========================================="
echo "Configuration (MDM Mode):"
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: full_shard"
echo "  Num GPUs: ${NGPUS}"
echo "  Auto Resume: ${AUTO_RESUME}"
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 ""

# AutoResume: detect latest checkpoint
if [ "${AUTO_RESUME}" = "1" ] && [ -z "${RESUME_FROM_CHECKPOINT}" ]; then
    if [ -d "${OUTPUT_DIR}" ]; then
        LATEST_CHECKPOINT=$(find "${OUTPUT_DIR}" -maxdepth 1 -type d -name "checkpoint-*" | sort -V | tail -n 1)
        if [ -n "${LATEST_CHECKPOINT}" ]; then
            RESUME_FROM_CHECKPOINT="${LATEST_CHECKPOINT}"
            echo "🔄 AutoResume: Found existing checkpoint: ${RESUME_FROM_CHECKPOINT}"
        else
            echo "ℹ️ AutoResume: No checkpoint found, starting from scratch"
        fi
    else
        echo "ℹ️ AutoResume: Output directory does not exist, starting from scratch"
    fi
fi

# Activation Checkpointing configuration
ac_arg=""
fsdp_wrap_arg=""
if [ "${ACTIVATION_CHECKPOINTING}" = "1" ]; then
    # In HF Trainer, --gradient_checkpointing enables activation checkpointing
    ac_arg="--gradient_checkpointing"
    # NOTE: We rely on the model's _no_split_modules attribute for auto-wrapping
    # Explicitly specifying the class name via CLI fails because the class isn't in the global namespace
    fsdp_wrap_arg="" 
    echo "✅ Activation Checkpointing enabled (relying on model._no_split_modules for FSDP wrapping)"
fi

# Build torchrun command with FSDP (as per working command)
CMD=(
    torchrun --nproc_per_node="${NGPUS}" --master_port=29503
    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}"
    --learn_padding "${LEARN_PADDING}"
    --conplemenrary_mask "${CONPLE_MASK}"
    --bf16
    --fsdp "full_shard"
    ${ac_arg}
    ${fsdp_wrap_arg}
)

# Add resume_from_checkpoint if available
if [ -n "${RESUME_FROM_CHECKPOINT}" ]; then
    CMD+=( --resume_from_checkpoint "${RESUME_FROM_CHECKPOINT}" )
    echo "📍 Resuming from checkpoint: ${RESUME_FROM_CHECKPOINT}"
else
    CMD+=( --overwrite_output_dir )
    echo "🆕 Starting fresh training (will overwrite output dir)"
fi

echo ""
echo "=========================================="
echo "Starting MDM training with FSDP..."
echo "Command: ${CMD[*]}"
echo "=========================================="

# Run training in background to allow signal handling
"${CMD[@]}" 2>&1 | tee "${LOG_DIR}/train.log" &
TRAIN_PID=$!

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

echo ""
echo "=========================================="
if [ $exit_code -eq 0 ]; then
    echo "✅ J63 MDM Training Success!"
    # ============================================
    # AUTO EVAL: 训练成功后自动提交评测任务
    # ============================================
    echo ""
    echo "=========================================="
    echo "🚀 Auto-submitting evaluation jobs..."
    echo "=========================================="
    
    if [ -f "${PROJECT_ROOT}/chengyue_script/dvlm/evaluation/submit_batch_eval.sh" ]; then
        bash "${PROJECT_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 "❌ J63 MDM Training Failed (exit code: ${exit_code})"
    echo -e "Please check logs:\n  Output: ${LOG_DIR}/train.log\n  Error: workdir/train_j63_mdm_error.log"
fi
echo "=========================================="

exit $exit_code
