#!/bin/bash

set -e

echo "=========================================="
echo "Deepspeed Runner (Interactive Session)"
echo "=========================================="
echo "Node: $(hostname)"
echo "Date: $(date)"
echo ""

# Resolve project root based on this script location
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
cd "${PROJECT_ROOT}"

# Environment (assumes interactive_gpu.slurm already set these; re-apply for safety)
source "${PROJECT_ROOT}/scripts/cluster_env.sh"
lmflow_activate_conda
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:-}"

# Defaults (override by env)
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:-1000}"
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:-128}"
LEARN_PADDING="${LEARN_PADDING:-1}"
CONPLE_MASK="${CONPLE_MASK:-1}"
TEMPLATE="${TEMPLATE:-qwen2_5_no_reasoning}"
RUN_ID="${RUN_ID:-sharegpt4v_interactive_$(date +%Y%m%d_%H%M%S)}"
OUTPUT_DIR="${OUTPUT_DIR:-output_models/${RUN_ID}}"
LOG_DIR="${LOG_DIR:-log/${RUN_ID}}"
DS_CONFIG="${DS_CONFIG:-configs/ds_config_zero2_no_offload.json}"
NGPUS="${NGPUS:-}"
SKIP_ADAPT="${SKIP_ADAPT:-0}"

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

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

# Determine GPU count if not provided
if [ -z "${NGPUS}" ]; then
    if command -v nvidia-smi >/dev/null 2>&1; then
        NGPUS=$(nvidia-smi -L | wc -l | awk '{print $1}')
    else
        NGPUS=1
    fi
fi
echo "  Num GPUs: ${NGPUS}"

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

export MDM_VL_KEEP_VISION=1

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}"
    --overwrite_output_dir
    --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}"
)

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

echo ""
echo "Starting Deepspeed training..."
"${CMD[@]}" 2>&1 | tee "${LOG_DIR}/train.log"

exit_code=${PIPESTATUS[0]}

echo ""
echo "=========================================="
if [ $exit_code -eq 0 ]; then
    echo "✅ Deepspeed Training Success!"
else
    echo "❌ Deepspeed Training Failed (exit code: ${exit_code})"
    echo -e "Please check logs:\n  Output: ${LOG_DIR}/train.log"
fi
echo "=========================================="

exit $exit_code


