#!/bin/bash
#SBATCH -A nvr_elm_llm  
#SBATCH -p batch_block1,batch_block3,batch_block4
#SBATCH -t 04:00:00
#SBATCH -N 8
#SBATCH -J dvlm_qwen2.5_vl_3b_text_sft
#SBATCH --array=1-100%1
#SBATCH --gpus-per-node 8
#SBATCH -e workdir/%x_%j.err
#SBATCH -o workdir/%x_%j.out
#SBATCH --exclusive

set -euo pipefail

echo "[driver] start at $(date), host=$(hostname)"
echo "[driver] SLURM_JOB_ID=${SLURM_JOB_ID:-} SLURM_ARRAY_TASK_ID=${SLURM_ARRAY_TASK_ID:-} SLURM_NNODES=${SLURM_NNODES:-}"

idx=$SLURM_ARRAY_TASK_ID
total=$SLURM_ARRAY_TASK_COUNT

echo "[driver] SLURM_JOB_ID=${SLURM_JOB_ID:-} ARRAY=${idx:-}-of-${total:-}" || true

echo "[driver] resolve nodes..."
nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) )
nodes_array=(${nodes})
head_node=${nodes_array[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)
slurm_name=${idx}-of-${total}

echo "[driver] head_node=${head_node} head_node_ip=${head_node_ip}"

export NCCL_BLOCKING_WAIT=1
export NCCL_IB_DISABLE=1
export PDSH_RCMD_TYPE=ssh
export NCCL_SOCKET_IFNAME=eth0
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

# deepspeed args (match reference style)
# choose a per-job, per-array-task master port to avoid EADDRINUSE
# Use larger range and hash to minimize conflicts
MASTER_PORT=$(( 30000 + (${SLURM_JOB_ID:-0}%20000) + (${SLURM_ARRAY_TASK_ID:-0} * 100) ))
deepspeed_args="--master_port=${MASTER_PORT} --num_nodes=$SLURM_NNODES --node_rank=$SLURM_PROCID"

# Do NOT cd via $(dirname $0); rely on sbatch submission directory as CWD
echo "[driver] PWD=$(pwd) PATH=$PATH"

# job name and paths (relative)
exp_id=dvlm_qwen2_5_vl_3b_text_sft
# Allow cache-only UX: user can export DATASET_CACHE_DIR; dataset_path defaults internally
dataset_path=${DATASET_PATH:-/home/chengyuew/chengyuew/workspace/LMFlow/data/nemotron_no_thinking_subset}
dataset_cache_dir=${DATASET_CACHE_DIR:-$(pwd)/.cache/lmf_preprocessed/dvlm_qwen2_5_vl_3b_no_reasoning}
output_dir=output_models/${exp_id}
log_dir=log/${exp_id}
model_name_or_path=/home/chengyuew/chengyuew/workspace/shiyi/LMFlow/src/models/Qwen2_5VL_DFM
conversation_template=qwen2_5_no_reasoning
trust_remote_code=1

# dataset
if [ -n "${dataset_cache_dir:-}" ]; then
  echo "[driver] dataset_cache_dir set: ${dataset_cache_dir}; skip dataset_path check"
  dataset_path_arg=""
else
  if [ ! -d "${dataset_path}" ]; then
    echo "[driver] dataset path not found: ${dataset_path}" 1>&2
    exit 1
  else
    echo "[driver] dataset ready: ${dataset_path}"
  fi
  dataset_path_arg="--dataset_path ${dataset_path}"
fi

echo "[driver] ensure output/log dirs..."
mkdir -p ${output_dir} ${log_dir}
ls -ld "${output_dir}" "${log_dir}" || true

# auto resume
echo "[driver] checking latest checkpoint in ${output_dir}..."
latest_checkpoint=""
if [ -d "${output_dir}" ]; then
  latest_checkpoint=$(find "${output_dir}" -name "checkpoint-*" -type d | sort -V | tail -1 || true)
fi
if [ -n "${latest_checkpoint}" ]; then
  echo "[driver] Found latest checkpoint: ${latest_checkpoint}"
else
  echo "[driver] No checkpoint found under ${output_dir}, will start from scratch."
fi
resume_arg=""
if [ -n "${latest_checkpoint}" ]; then
  resume_arg="--resume_from_checkpoint ${latest_checkpoint}"
fi

echo "[driver] resume_arg='${resume_arg}'"

# ds config (relative)
ds_config=configs/ds_config_zero3_no_offload.json
[ -f "${ds_config}" ] && echo "[driver] ds_config ok: ${ds_config}" || echo "[driver] ds_config missing: ${ds_config}"

cmd='
echo "[cmd] node=$SLURMD_NODENAME host=$(hostname) date=$(date)"
echo "[cmd] PATH=$PATH"

# activate conda environment
source "${PROJECT_ROOT}/scripts/cluster_env.sh"
lmflow_activate_conda

which deepspeed && deepspeed --version || echo "[cmd] deepspeed not found"
python3 -V || true
nvidia-smi -L || true

# identify master host and prepare hostfile (head node only)
echo "[cmd] preparing hostfile at '"${log_dir}"'/temp-hostfile (head='"${head_node}"')"
if [ "$SLURMD_NODENAME" = '"${head_node}"' ]; then
  echo $SLURM_JOB_NODELIST > '"${log_dir}"'/temp-hostfile;
  read master_node slave_node <<< $(./expand_hostfile.sh '"${log_dir}"');
fi

# wait for all nodes to be ready (critical for SSH setup)
echo "[cmd] Waiting 60s for all nodes to initialize..."
sleep 60

# run cluster ssh setup (must be before cache/training to enable pdsh)
echo "[cmd] running setup_ssh.sh"
./setup_ssh.sh '"${log_dir}"'
echo "[cmd] setup_ssh.sh done; listing '"${log_dir}"'"
ls -l '"${log_dir}"' || true

# pre-generate HF cache on head node only (after SSH is ready)
if [ "$SLURMD_NODENAME" = '"${head_node}"' ]; then
  echo "[cmd] [HEAD NODE] Pre-generating HF transformers cache for Qwen2_5VL_DFM"
  export MODEL_PATH=${model_name_or_path}
  python3 -c "import os; from transformers import AutoConfig; AutoConfig.from_pretrained(os.environ['MODEL_PATH'], trust_remote_code=True)" 2>&1 | head -n 5 || echo "[WARN] Cache pre-generation had warnings (may be OK)"
  echo "[cmd] [HEAD NODE] Cache pre-generation completed"
  
  # create a marker file to signal cache is ready
  touch '"${log_dir}"'/cache_ready.marker
  echo "[cmd] [HEAD NODE] Cache marker created at '"${log_dir}"'/cache_ready.marker"
fi

# wait for head node to complete cache generation (all nodes)
echo "[cmd] Waiting for HF cache to be ready (checking for marker file)..."
for i in {1..60}; do
  if [ -f '"${log_dir}"'/cache_ready.marker ]; then
    echo "[cmd] Cache ready marker found after ${i} seconds"
    break
  fi
  sleep 1
done

# additional safety wait
sleep 5

echo "MASTER_NODE:= $master_node" 1>&2;

# ensure hostfile exists
if [ -f '"${log_dir}"'/deepspeed-hostfile ]; then
  echo "[cmd] deepspeed-hostfile present, head -n 50:"
  head -n 50 '"${log_dir}"'/deepspeed-hostfile || true
else
  echo "[cmd] deepspeed-hostfile missing"
fi

# select deepspeed launcher (hostfile-based, like reference script)
LAUNCH=deepspeed
which deepspeed >/dev/null 2>&1 || LAUNCH="python3 -m deepspeed"
echo "[cmd] using launcher: ${LAUNCH}"

# sanity checks before launch
[ -d '"${dataset_path}"' ] && echo "[cmd] dataset ok: '"${dataset_path}"'" || echo "[cmd] dataset missing: '"${dataset_path}"'"
[ -f '"${ds_config}"' ] && echo "[cmd] ds_config ok: '"${ds_config}"'" || echo "[cmd] ds_config missing: '"${ds_config}"'"
mkdir -p '"${output_dir}"' && ls -ld '"${output_dir}"' || true

echo "[cmd] resume_arg='"${resume_arg}"'"

if [ "$SLURMD_NODENAME" = '"${head_node}"' ]; then
  echo "[cmd] launching deepspeed..."
  set -x
  ${LAUNCH} --master_addr=${master_node} --hostfile '"${log_dir}"'/deepspeed-hostfile ${deepspeed_args} \
    examples/finetune.py \
    --model_name_or_path '"${model_name_or_path}"' \
    --tokenizer_name Qwen/Qwen2.5-VL-3B-Instruct \
    --conversation_template '"${conversation_template}"' \
    --trust_remote_code '"${trust_remote_code}"' \
    ${dataset_path_arg} \
    --output_dir '"${output_dir}"' \
    '"${resume_arg}"' \
    --num_train_epochs 1 \
    --learning_rate 1e-5 \
    --lr_scheduler_type constant_with_warmup \
    --warmup_ratio 0.03 \
    --disable_group_texts 1 \
    --block_size 2048 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 4 \
    --deepspeed '"${ds_config}"' \
    --bf16 \
    --run_name '"${exp_id}"' \
    --validation_split_percentage 0 \
    --logging_steps 20 \
    --do_train \
    --ddp_timeout 72000 \
    --save_steps 100 \
    --save_total_limit 30 \
    --dataloader_num_workers 8 \
    --preprocessing_num_workers 1 \
    --dataset_cache_dir '"${dataset_cache_dir}"' \
    --bd_size 128 \
    --gradient_checkpointing 1 \
    --mdm 1 \
    --learn_padding 1 \
    --conplemenrary_mask 1 \
    > >(tee '"${log_dir}"'/train.log) \
    2> >(tee '"${log_dir}"'/train.err >&2)
  set +x
else
  echo "[cmd] not master node: $SLURMD_NODENAME"
fi
'

echo "[driver] submit srun at $(date)"
mkdir -p ${log_dir}/slurm_log/train
srun -o ${log_dir}/slurm_log/train/slurm_%A_%a_${slurm_name}.log -e ${log_dir}/slurm_log/train/error_%A_%a_${slurm_name}.log \
  bash -c "$cmd"

echo "[driver] end at $(date)"
