#!/bin/bash
#SBATCH -A nvr_elm_llm
#SBATCH -p batch_block1,batch_block3,batch_block4
#SBATCH -t 4:00:00
#SBATCH -N 1
#SBATCH --ntasks-per-node=1
#SBATCH -J train_dvlmad_waymo
#SBATCH --gpus-per-node 8
#SBATCH --cpus-per-task 16
#SBATCH --mem=1024G
#SBATCH -e workdir/train_dvlmad_waymo_error.log
#SBATCH -o workdir/train_dvlmad_waymo_out.log
#SBATCH --open-mode=append
#SBATCH --requeue
#SBATCH --signal=B:USR1@120
#SBATCH --array=1-30%1

# ============================================
# 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 "=========================================="

    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 $TRAIN_PID 2>/dev/null || true
    fi

    echo "Requesting SLURM to requeue job..."
    scontrol requeue $SLURM_JOB_ID
    exit 0
}

trap handle_timeout USR1

# Resolve project root
PROJECT_ROOT="${SLURM_SUBMIT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
cd "${PROJECT_ROOT}"

echo "=========================================="
echo "dVLM-AD Waymo MDM Training (single-node)"
echo "=========================================="
echo "  Node: $(hostname)  Date: $(date)  Job: ${SLURM_JOB_ID:-N/A}"
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__)')"
# nvidia-smi || true
# echo ""

# ============================================
# Environment variables (aligned with reference)
# ============================================
export PYTHONPATH="${PROJECT_ROOT}/src:${PROJECT_ROOT}/third_party/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 TOKENIZERS_PARALLELISM=false

# Qwen2.5-VL processor dynamic resolution controls
export MIN_PIXELS="${MIN_PIXELS:-200704}"
export MAX_PIXELS="${MAX_PIXELS:-200704}"
export RESIZED_HEIGHT="${RESIZED_HEIGHT:-}"
export RESIZED_WIDTH="${RESIZED_WIDTH:-}"

# ============================================
# Training parameters (aligned with reference)
# ============================================
# MODEL_PATH="${MODEL_PATH:-models/Qwen2.5-VL-3B-Instruct}"
# MODEL_PATH=output_models_shiyi/sharegpt4v_8node_bsz256_anneal_bd16_block_causal_lr_1e-6_unfreeze_vision/checkpoint-2997
export WANDB_API_KEY=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
MODEL_PATH=output_models/dvlmad_waymo_dvlm_v2_32_lr1e-6_3361679/checkpoint-902
TOKENIZER_NAME="${TOKENIZER_NAME:-Qwen/Qwen2.5-VL-3B-Instruct}"
MAX_LENGTH="${MAX_LENGTH:-2048}"

# Waymo dataset configuration
IMAGE_FOLDER="${IMAGE_FOLDER:-dvlmad_data/waymo_train_joint}"
DATASET_JSON="${DATASET_JSON:-dvlmad_data/dvlm-ad_waymo_training_cot_joint_filtered.json}"
TEMPLATE="${TEMPLATE:-qwen2_5_no_reasoning}"

EPOCHS="${EPOCHS:-5}"
BATCH="${BATCH:-1}"
GRAD_ACC="${GRAD_ACC:-4}"                # 8 GPUs * 1 * 4 = 32 effective (single-node)
LR="${LR:-1e-6}"
WARMUP_RATIO="${WARMUP_RATIO:-0.03}"
SAVE_STEPS="${SAVE_STEPS:-50}"
SAVE_TOTAL_LIMIT="${SAVE_TOTAL_LIMIT:-20}"
LOGGING_STEPS="${LOGGING_STEPS:-10}"
WORKERS="${WORKERS:-8}"
BLOCK_SIZE="${BLOCK_SIZE:-2048}"

# MDM specific (aligned with reference)
MDM="${MDM:-1}"
BD_SIZE="${BD_SIZE:-32}"
LEARN_PADDING="${LEARN_PADDING:-1}"
COMPLEMENTARY_MASK="${COMPLEMENTARY_MASK:-1}"
ALWAYS_MASK_IM_END="${ALWAYS_MASK_IM_END:-1}"
USE_BLOCK_CAUSAL_MASK="${USE_BLOCK_CAUSAL_MASK:-1}"
ANNEAL_BLOCK_SIZE="${ANNEAL_BLOCK_SIZE:-0}"
FLEXIBLE_BD_SIZE="${FLEXIBLE_BD_SIZE:-0}"

# Activation checkpointing (saves memory)
ACTIVATION_CHECKPOINTING="${ACTIVATION_CHECKPOINTING:-1}"

NGPUS_PER_NODE="${NGPUS_PER_NODE:-8}"
NNODES=1
MASTER_ADDR="${MASTER_ADDR:-localhost}"
MASTER_PORT="${MASTER_PORT:-$(python3 -c 'import socket; s=socket.socket(); s.bind(("",0)); p=s.getsockname()[1]; s.close(); print(p)')}"
# MASTER_PORT=29501


RUN_ID="${RUN_ID:-dvlmad_waymo_dvlm_v2_${BD_SIZE}_lr${LR}_${SLURM_JOB_ID:-$(date +%Y%m%d_%H%M%S)}}"
OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/output_models/${RUN_ID}}"
OUTPUT_DIR=output_models/dvlmad_waymo_dvlm_v2_32_lr1e-6_3361679
LOG_DIR="${LOG_DIR:-${PROJECT_ROOT}/log/${RUN_ID}}"

mkdir -p "${OUTPUT_DIR}" "${LOG_DIR}"

# ============================================
# 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)
fi

resume_arg=""
WANDB_RUN_ID_FILE="${OUTPUT_DIR}/wandb_run_id.txt"
if [ -n "${latest_checkpoint}" ]; then
    resume_arg="--resume_from_checkpoint ${latest_checkpoint}"
    echo "Resuming from: ${latest_checkpoint}"
    # When resuming, use the latest checkpoint as MODEL_PATH if the original doesn't exist
    if [ -z "${MODEL_PATH:-}" ] || [ ! -d "${MODEL_PATH}" ]; then
        echo "MODEL_PATH not set or not found, using latest checkpoint: ${latest_checkpoint}"
        MODEL_PATH="${latest_checkpoint}"
    fi
    # Resume the same WandB run
    if [ -f "${WANDB_RUN_ID_FILE}" ]; then
        export WANDB_RUN_ID=$(cat "${WANDB_RUN_ID_FILE}")
        export WANDB_RESUME=must
        echo "Resuming WandB run: ${WANDB_RUN_ID}"
    fi
else
    resume_arg="--overwrite_output_dir"
    echo "Training from scratch"
    # Generate a new WandB run ID and save it
    export WANDB_RUN_ID=$(python3 -c "import hashlib,os; print(hashlib.md5(os.path.basename('${OUTPUT_DIR}').encode()).hexdigest()[:12])")
    export WANDB_RESUME=allow
    echo "${WANDB_RUN_ID}" > "${WANDB_RUN_ID_FILE}"
    echo "New WandB run ID: ${WANDB_RUN_ID}"
fi

# ============================================
# Check paths
# ============================================
if [ ! -d "${MODEL_PATH}" ]; then
    echo "ERROR: Model path does not exist: ${MODEL_PATH}"
    exit 1
fi

if [ ! -f "${DATASET_JSON}" ]; then
    echo "ERROR: Dataset json not found: ${DATASET_JSON}"
    exit 1
fi

if [ ! -d "${IMAGE_FOLDER}" ]; then
    echo "ERROR: Image folder not found: ${IMAGE_FOLDER}"
    exit 1
fi

ac_arg=""
if [ "${ACTIVATION_CHECKPOINTING}" = "1" ]; then
    ac_arg="--gradient_checkpointing"
fi

echo ""
echo "=========================================="
echo "Config:"
echo "  Model: ${MODEL_PATH}"
echo "  Dataset JSON: ${DATASET_JSON}"
echo "  Image Folder: ${IMAGE_FOLDER}"
echo "  Output: ${OUTPUT_DIR}"
echo "  Epochs: ${EPOCHS}  Batch: ${BATCH}  GradAcc: ${GRAD_ACC}"
echo "  LR: ${LR}  MDM: ${MDM}  BD_SIZE: ${BD_SIZE}"
echo "  BlockCausal: ${USE_BLOCK_CAUSAL_MASK}  Complementary: ${COMPLEMENTARY_MASK}"
echo "  AlwaysMaskImEnd: ${ALWAYS_MASK_IM_END}  AnnealBlockSize: ${ANNEAL_BLOCK_SIZE}"
echo "  TotalGPUs: $((NNODES * NGPUS_PER_NODE))  EffBatch: $((NNODES * NGPUS_PER_NODE * BATCH * GRAD_ACC))"
echo "  Workers: ${WORKERS}"
echo "  Processor MinPixels/MaxPixels: ${MIN_PIXELS}/${MAX_PIXELS}"
echo "=========================================="
echo ""

# ============================================
# Launch training (FSDP + torchrun, aligned with reference)
# ============================================
torchrun \
    --nnodes=${NNODES} \
    --nproc_per_node=${NGPUS_PER_NODE} \
    --rdzv_id="${SLURM_JOB_ID:-1}" \
    --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_JSON}" \
    --image_folder "${IMAGE_FOLDER}" \
    --output_dir "${OUTPUT_DIR}" \
    --num_train_epochs ${EPOCHS} \
    --lr_scheduler_type constant_with_warmup \
    --warmup_ratio ${WARMUP_RATIO} \
    --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 ${COMPLEMENTARY_MASK} \
    --bf16 \
    --fsdp "shard_grad_op" \
    --always_mask_im_end ${ALWAYS_MASK_IM_END} \
    --flexible_bd_size ${FLEXIBLE_BD_SIZE} \
    --use_block_causal_mask ${USE_BLOCK_CAUSAL_MASK} \
    --anneal_block_size ${ANNEAL_BLOCK_SIZE} \
    --run_name "${RUN_ID}" \
    --validation_split_percentage 0 \
    ${ac_arg} \
    ${resume_arg} \
    2>&1 | tee "${LOG_DIR}/train_${SLURM_ARRAY_TASK_ID:-0}-of-${SLURM_ARRAY_TASK_COUNT:-1}.log" &
TRAIN_PID=$!

wait $TRAIN_PID
exit_code=$?

echo ""
echo "=========================================="
if [ $exit_code -eq 0 ]; then
    echo "Training completed successfully!"
else
    echo "Training failed with exit code: ${exit_code}"
    echo "Check logs: ${LOG_DIR}/"
fi
echo "=========================================="

exit $exit_code
