#!/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_sharegpt4v_vlm
#SBATCH --gpus-per-node 8
#SBATCH --cpus-per-task 16
#SBATCH --mem=1024G
#SBATCH -e workdir/train_sharegpt4v_vlm_mdm_fix_eos_error.log
#SBATCH -o workdir/train_sharegpt4v_vlm_mdm_fix_eos_out.log
#SBATCH --requeue
#SBATCH --signal=B:USR1@120

set -euxo pipefail

# Signal handler for requeue
on_signal() {
    echo "Caught signal; requeue current job $SLURM_JOB_ID"
    REQUEUED=1
    scontrol requeue $SLURM_JOB_ID || true
}
trap on_signal SIGTERM SIGUSR1

# 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_Qwen2VL VLM Training (ShareGPT4V, single-node)"
echo "=========================================="
echo "Node: $(hostname)"
echo "Date: $(date)"
echo "Working Directory: $(pwd)"
echo ""

# Activate environment
# source ~/.bashrc
# conda activate shiyi-lmflow || true

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 (merge from interactive_gpu.slurm and run_deepspeed_sharegpt4v.sh)
# export PYTHONPATH="${PROJECT_ROOT}:${PROJECT_ROOT}/chengyue_script:${PROJECT_ROOT}/src:${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

# Qwen2.5-VL processor dynamic resolution controls (read by examples/finetune_qwenvl.py)
# Reduced to ~448x448 fixed to limit vision tokens (avoid OOM and extreme prefill lengths)
export MIN_PIXELS="${MIN_PIXELS:-200704}"
export MAX_PIXELS="${MAX_PIXELS:-200704}"
# Optional fixed resize override (must be multiples of 28 if set)
export RESIZED_HEIGHT="${RESIZED_HEIGHT:-}"
export RESIZED_WIDTH="${RESIZED_WIDTH:-}"

# Configurable parameters (override via env before submitting)
MODEL_PATH="${MODEL_PATH:-/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow/src/models/Qwen2.5-VL-3B-Instruct-mdm}"
TOKENIZER_NAME="${TOKENIZER_NAME:-Qwen/Qwen2.5-VL-3B-Instruct}"

DATASET_JSON="${DATASET_JSON:-/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/dataset/vila-sft/ShareGPT4V/jason-filter-sharegpt4v_mix665k_cap23k_coco-ap9k_lcs3k_sam9k_div2k.json}"
IMAGE_FOLDER="${IMAGE_FOLDER:-/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/dataset/vila-sft/ShareGPT4V/data}"

EPOCHS="${EPOCHS:-1}"
BLOCK_SIZE="${BLOCK_SIZE:-2048}"
BATCH="${BATCH:-1}"
GRAD_ACC="${GRAD_ACC:-8}"
LR="${LR:-2e-5}"
WARMUP_RATIO="${WARMUP_RATIO:-0.03}"
BF16="${BF16:-1}"
GRADIENT_CHECKPOINTING="${GRADIENT_CHECKPOINTING:-1}"
DISABLE_GROUP_TEXTS="${DISABLE_GROUP_TEXTS:-1}"
SAVE_STEPS="${SAVE_STEPS:-50}"
SAVE_TOTAL_LIMIT="${SAVE_TOTAL_LIMIT:-20}"
LOGGING_STEPS="${LOGGING_STEPS:-50}"
WORKERS="${WORKERS:-8}"
PREP_WORKERS="${PREP_WORKERS:-32}"
MDM="${MDM:-1}"
BD_SIZE="${BD_SIZE:-8}"
LEARN_PADDING="${LEARN_PADDING:-1}"
CONPLE_MASK="${CONPLE_MASK:-1}"
TEMPLATE="${TEMPLATE:-qwen2_5_no_reasoning}"
RUN_ID="${RUN_ID:-sharegpt4v_vlm_mdm_fix_eos}"
OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/output_models/${RUN_ID}}"
LOG_DIR="${LOG_DIR:-${PROJECT_ROOT}/log/${RUN_ID}}"
DS_CONFIG="${DS_CONFIG:-configs/ds_config_zero2_no_offload.json}"
NGPUS="${NGPUS:-8}"
SKIP_ADAPT="${SKIP_ADAPT:-0}"

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

echo "=========================================="
echo "Configuration:"
echo "  Model: ${MODEL_PATH}"
echo "  Tokenizer: ${TOKENIZER_NAME}"
echo "  Dataset JSON: ${DATASET_JSON}"
echo "  Image Folder: ${IMAGE_FOLDER}"
echo "  Output: ${OUTPUT_DIR}"
echo "  Log Dir: ${LOG_DIR}"
echo "  Epochs: ${EPOCHS}"
echo "  Block Size: ${BLOCK_SIZE}"
echo "  Batch: ${BATCH}"
echo "  Grad Acc: ${GRAD_ACC}"
echo "  LR: ${LR}"
echo "  Warmup Ratio: ${WARMUP_RATIO}"
echo "  BF16: ${BF16}"
echo "  MDM: ${MDM}  BD_SIZE: ${BD_SIZE}  LearnPadding: ${LEARN_PADDING}  ComplementaryMask: ${CONPLE_MASK}"
echo "  DS Config: ${DS_CONFIG}"
echo "  Processor MinPixels/MaxPixels: ${MIN_PIXELS}/${MAX_PIXELS}"
echo "  Processor Fixed WxH (opt): ${RESIZED_WIDTH:-auto}x${RESIZED_HEIGHT:-auto}"
echo "  Num GPUs: ${NGPUS}"
echo "  Auto Resume: ${AUTO_RESUME}"
echo "=========================================="
echo ""

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

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

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

# Dataset schema preflight & adaptation (optional)
DATASET_TO_USE="${DATASET_JSON}"
ADAPTED_JSON="${OUTPUT_DIR}/sharegpt4v_adapted.json"
if [ "${SKIP_ADAPT}" != "1" ]; then
    echo "Checking dataset schema compatibility..."
    ADAPTED_JSON="${ADAPTED_JSON}" DATASET_JSON="${DATASET_JSON}" IMAGE_FOLDER="${IMAGE_FOLDER}" python3 - <<'PY' || { echo "Schema preflight failed"; exit 1; }
import json, os, sys
dataset_json = os.environ.get('DATASET_JSON')
image_folder = os.environ.get('IMAGE_FOLDER')
adapted_json = os.environ.get('ADAPTED_JSON')

with open(dataset_json, 'r') as f:
    data = json.load(f)

if not isinstance(data, list) or len(data) == 0:
    print("❌ JSON must be a list with at least one item", flush=True)
    sys.exit(2)

sample = data[0]
has_conversations = 'conversations' in sample or 'messages' in sample
if not has_conversations:
    print("❌ Missing conversations/messages field in dataset items", flush=True)
    sys.exit(3)

need_adapt = False
field_conversations = 'conversations' if 'conversations' in sample else 'messages'
msg0 = sample[field_conversations][0] if sample[field_conversations] else {}
if 'from' not in msg0 and 'role' in msg0:
    need_adapt = True
if 'image' not in sample and any(k in sample for k in ['image_path', 'image_paths', 'images']):
    need_adapt = True

if not need_adapt:
    print("✅ Schema looks compatible (conversations/from + image)", flush=True)
    sys.exit(0)

print("ℹ️ Adapting dataset schema to LLaVA format...", flush=True)
out = []
def normalize_role_to_from(role):
    if role in ('user', 'human'): return 'human'
    if role in ('assistant', 'gpt', 'bot'): return 'gpt'
    return role or 'human'

for item in data:
    new_item = dict(item)
    conv_key = 'conversations' if 'conversations' in item else ('messages' if 'messages' in item else None)
    if conv_key is None:
        continue
    conv = []
    for m in item[conv_key]:
        if 'from' in m and 'value' in m:
            conv.append(m)
        else:
            role = m.get('role')
            text = m.get('value') or m.get('text') or m.get('content') or ''
            conv.append({'from': normalize_role_to_from(role), 'value': text})
    new_item['conversations'] = conv
    if 'image' in item:
        pass
    elif 'image_path' in item and isinstance(item['image_path'], str):
        new_item['image'] = item['image_path']
    elif 'image_paths' in item and isinstance(item['image_paths'], list) and len(item['image_paths'])>0:
        new_item['image'] = item['image_paths'][0]
    elif 'images' in item and isinstance(item['images'], list) and len(item['images'])>0:
        new_item['image'] = item['images'][0]
    else:
        continue
    out.append(new_item)

with open(adapted_json, 'w') as f:
    json.dump(out, f)

print("✅ Adapted to:", adapted_json)
PY

    if [ -s "${ADAPTED_JSON}" ]; then
        DATASET_TO_USE="${ADAPTED_JSON}"
    fi
fi

echo "Using dataset file: ${DATASET_TO_USE}"
echo ""

# AutoResume: detect latest checkpoint
if [ "${AUTO_RESUME}" = "1" ] && [ -z "${RESUME_FROM_CHECKPOINT}" ]; then
    if [ -d "${OUTPUT_DIR}" ]; then
        # Find the latest checkpoint directory
        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
elif [ -n "${RESUME_FROM_CHECKPOINT}" ]; then
    echo "📍 AutoResume: Using explicitly specified checkpoint: ${RESUME_FROM_CHECKPOINT}"
fi

# Build deepspeed command
CMD=(
    deepspeed --num_gpus "${NGPUS}" examples/finetune_qwenvl.py
    --model_name_or_path "${MODEL_PATH}"
    --tokenizer_name "${TOKENIZER_NAME}"
    --conversation_template "${TEMPLATE}"
    --trust_remote_code 1
    --dataset_path "${DATASET_TO_USE}"
    --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}"
    --disable_group_texts "${DISABLE_GROUP_TEXTS}"
    --block_size "${BLOCK_SIZE}"
    --per_device_train_batch_size "${BATCH}"
    --gradient_accumulation_steps "${GRAD_ACC}"
    --run_name "${RUN_ID}"
    --logging_steps "${LOGGING_STEPS}"
    --do_train
    --validation_split_percentage 0
    --ddp_timeout 72000
    --save_steps "${SAVE_STEPS}"
    --save_total_limit "${SAVE_TOTAL_LIMIT}"
    --dataloader_num_workers "${WORKERS}"
    --preprocessing_num_workers "${PREP_WORKERS}"
    --bd_size "${BD_SIZE}"
    --mdm "${MDM}"
    --learn_padding "${LEARN_PADDING}"
    --conplemenrary_mask "${CONPLE_MASK}"
    --deepspeed "${DS_CONFIG}"
)

# 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
    # Only add overwrite_output_dir if NOT resuming
    CMD+=( --overwrite_output_dir )
    echo "🆕 Starting fresh training (will overwrite output dir)"
fi

if [ "${BF16}" = "1" ]; then
    CMD+=( --bf16 )
fi
if [ "${GRADIENT_CHECKPOINTING}" = "1" ]; then
    CMD+=( --gradient_checkpointing )
fi

echo ""
echo "=========================================="
echo "Starting Deepspeed training..."
echo "=========================================="

# Run training with proper error handling
set +e
"${CMD[@]}" 2>&1 | tee "${LOG_DIR}/train.log"
EXIT_CODE=${PIPESTATUS[0]}
set -e

echo ""
echo "=========================================="
if [ $EXIT_CODE -eq 0 ]; then
    echo "✅ ShareGPT4V VLM 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 "❌ ShareGPT4V VLM Training Failed (exit code: ${EXIT_CODE})"
    echo -e "Please check logs:\n  Output: ${LOG_DIR}/train.log\n  Error: workdir/train_sharegpt4v_vlm_error.log"
fi
echo "=========================================="

# Handle requeue logic (optional: add auto-resubmit on failure if desired)
SCRIPT_PATH="${PROJECT_ROOT}/chengyue_script/dvlm/train_sharegpt4v_vlm_autoresume.sh"
if [ "${REQUEUED:-0}" = "1" ]; then
    echo "Job was requeued via signal. Exit code: ${EXIT_CODE}"
elif [ "$EXIT_CODE" -ne 0 ]; then
    echo "Training failed with exit code: ${EXIT_CODE}"
    # Optionally add auto-resubmit logic here if needed
    # echo "Resubmitting due to non-zero exit..."
    # sbatch --dependency=afterany:$SLURM_JOB_ID --export=ALL ${SCRIPT_PATH}
fi

exit $EXIT_CODE


