#!/bin/bash
#SBATCH -A nvr_elm_llm  
#SBATCH -p batch_block1,batch_block3,batch_block4
#SBATCH -t 01:00:00
#SBATCH -N 1
#SBATCH -J dvlm_debug_1node
#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)"

# Single job (not array), so no idx/total
echo "[driver] SLURM_JOB_ID=${SLURM_JOB_ID:-}" || 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=single

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

# deepspeed args (match reference style)
deepspeed_args="--master_port=23456 --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_debug_1node
dataset_path=/home/chengyuew/chengyuew/workspace/LMFlow/data/nemotron_no_thinking
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=
trust_remote_code=1

# dataset
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

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_zero2_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 all nodes
sleep 60;

# run cluster ssh setup
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

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 (fallback to python3 -m deepspeed)
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 \
    --trust_remote_code '"${trust_remote_code}"' \
    --dataset_path '"${dataset_path}"' \
    --output_dir '"${output_dir}"' \
    '"${resume_arg}"' \
    --max_steps 50 \
    --learning_rate 2e-5 \
    --disable_group_texts 0 \
    --block_size 2048 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 1 \
    --deepspeed '"${ds_config}"' \
    --bf16 \
    --run_name '"${exp_id}"' \
    --validation_split_percentage 0 \
    --logging_steps 5 \
    --do_train \
    --ddp_timeout 72000 \
    --save_steps 500 \
    --dataloader_num_workers 4 \
    --preprocessing_num_workers 16 \
    --bd_size 128 \
    --mdm 1 \
    --mdm_loss 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_${slurm_name}.log -e ${log_dir}/slurm_log/train/error_%A_${slurm_name}.log \
  bash -c "$cmd"

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