#!/bin/bash
# =============================================================================
# AnomSeer / TimerPO on Time-RA RATs-Uni — 2-GPU train + eval launcher.
#
# Task : 15-class anomaly classification + reasoning (no localization).
# Tuning: LoRA fine-tuning (GRPO) by default.
#
# Modes (STAGE):
#   train       : train only (with periodic in-training validation)   [default]
#   eval        : evaluate MODEL_PATH only (val_only)
#   train_eval  : train -> auto-merge latest checkpoint -> evaluate it
#
# Examples:
#   bash example/timerpo_trainer/run_rats_2gpu.sh                      # train
#   STAGE=eval MODEL_PATH=Qwen/Qwen2.5-VL-3B-Instruct \
#       bash example/timerpo_trainer/run_rats_2gpu.sh                  # zero-shot eval
#   STAGE=train_eval bash example/timerpo_trainer/run_rats_2gpu.sh     # train then eval
# =============================================================================
set -euo pipefail

# ---- scratch on /dev/shm (NOT /tmp) ----------------------------------------
# The CIFS share (/mnt/share01) cannot host Ray's unix sockets, and the local root
# disk (/, where /tmp and ~/.cache live) is ~100% full. So Ray's runtime/spill and
# the torch/triton/vLLM/HF caches go to /dev/shm: local tmpfs, supports sockets,
# ~120 GB free. Persistent outputs (checkpoints, logs) still go under the project.
SCRATCH_ROOT="${SCRATCH_ROOT:-/dev/shm/anomseer}"
export RAY_TMPDIR="${RAY_TMPDIR:-${SCRATCH_ROOT}/ray}"
export TMPDIR="${TMPDIR:-${SCRATCH_ROOT}/tmp}"
export XDG_CACHE_HOME="${XDG_CACHE_HOME:-${SCRATCH_ROOT}/cache}"
export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-${SCRATCH_ROOT}/cache/torchinductor}"
export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${SCRATCH_ROOT}/cache/triton}"
export VLLM_CACHE_ROOT="${VLLM_CACHE_ROOT:-${SCRATCH_ROOT}/cache/vllm}"
export HF_HOME="${HF_HOME:-${SCRATCH_ROOT}/cache/hf}"
mkdir -p "$RAY_TMPDIR" "$TMPDIR" "$XDG_CACHE_HOME" "$TORCHINDUCTOR_CACHE_DIR" \
         "$TRITON_CACHE_DIR" "$VLLM_CACHE_ROOT" "$HF_HOME"

# ---- knobs (override via env) ----------------------------------------------
PYTHON_BIN="${PYTHON_BIN:-/home/suiqk/anaconda3/envs/scalerag-ts-v4/bin/python}"
MODEL_PATH=${MODEL_PATH:-/mnt/share01/sqk/models/Qwen2.5-VL-3B-Instruct}   # local 3B; set to the 7B path / HF id to switch
STAGE=${STAGE:-train}                                   # train | eval | train_eval
N_GPUS=${N_GPUS:-2}                                     # number of GPUs (this script targets 2)
# TP = rollout tensor-parallel size; must divide N_GPUS. Default 2 splits vLLM's model
# across both cards (safer on memory). Set TP=1 for data-parallel rollout (full replica
# per card, no cross-card TP comm — faster on these no-NVLink cards; fine for 3B on 48 GB).
TP=${TP:-2}
GPU_MEM_UTIL=${GPU_MEM_UTIL:-0.4}                       # vllm KV-cache fraction (lower if OOM)
TRAIN_BATCH=${TRAIN_BATCH:-16}                         # prompts fetched per training step (rollout batch)
MICRO_BSZ=${MICRO_BSZ:-2}                               # actor ppo micro batch per GPU (lower if OOM)
LOGP_MICRO_BSZ=${LOGP_MICRO_BSZ:-8}                     # log-prob micro batch per GPU
PARAM_OFFLOAD=${PARAM_OFFLOAD:-False}                   # set True to offload actor params (saves VRAM)
OPTIM_OFFLOAD=${OPTIM_OFFLOAD:-False}                   # set True to offload optimizer (saves VRAM)
MAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-384}         # covers the expert explanations without runaway output
KL_COEF=${KL_COEF:-0.01}                                # resist reward-driven mode collapse
# remove-padding (rmpad) is a throughput optimization that monkey-patches Qwen2VL
# FlashAttention2 internals. Those classes were removed in transformers>=4.52, so it
# must stay False on this env (transformers 4.54.1). Set True only on transformers<=4.51.1.
USE_RMPAD=${USE_RMPAD:-False}
EPOCHS=${EPOCHS:-1}

# ---- LoRA (optional) -------------------------------------------------------
# LORA_RANK=16 (default) -> LoRA fine-tuning.
# LORA_RANK=0            -> full-parameter fine-tuning.
LORA_RANK=${LORA_RANK:-16}
LORA_ALPHA=${LORA_ALPHA:-16}
LORA_DROPOUT=${LORA_DROPOUT:-0.0}
# NOTE: PEFT 'all-linear' tries to wrap whole Qwen2_5_VLVisionBlock modules and errors on
# this VL model, so default to the explicit LLM (+vision MLP) projection names instead.
LORA_TARGET_MODULES=${LORA_TARGET_MODULES:-q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj}

# LoRA usually wants a larger LR than full fine-tuning; default accordingly (override with LR=).
if [ "$LORA_RANK" -gt 0 ]; then
  LR=${LR:-1e-5}
else
  LR=${LR:-1e-6}
fi
LOGGER=${LOGGER:-console}                               # 'console' or "console','wandb" (needs wandb login)

PROJECT=${PROJECT:-anomseer}
EXP=${EXP:-anomseer_rats_uni_2gpu}
CKPT_ROOT=${CKPT_ROOT:-checkpoints/${PROJECT}/${EXP}}

# Default trains on a random 1/4 subset (7566 samples) for faster runs; override with
# TRAIN_FILE=./data/rats_uni_processed/train_full.parquet for the full 30266.
TRAIN_FILE=${TRAIN_FILE:-./data/rats_uni_processed/train_quarter.parquet}
# verl sends the WHOLE val set to vLLM in one batch, so the full 6034-sample test set
# OOMs system RAM. Use a stratified 474-sample subset for in-training validation; for a
# final full-test eval pass VAL_FILE=./data/rats_uni_processed/test_full.parquet.
VAL_FILE=${VAL_FILE:-./data/rats_uni_processed/test_small.parquet}

# ---- core launcher ---------------------------------------------------------
# args: $1 = val_only (True/False)   $2 = model path
run_verl () {
  local VAL_ONLY="$1"; local MPATH="$2"
  local ACTIVE_LORA_RANK="$LORA_RANK"
  local RESUME_MODE="auto"
  if [ "$VAL_ONLY" = "True" ]; then
    # Evaluation consumes a complete HF model, including any already-merged adapter.
    ACTIVE_LORA_RANK=0
    RESUME_MODE="disable"
  fi
  "$PYTHON_BIN" -m verl.trainer.main_ppo \
      algorithm.adv_estimator=grpo \
      data.train_files="$TRAIN_FILE" \
      data.val_files="$VAL_FILE" \
      data.train_batch_size="$TRAIN_BATCH" \
      data.max_prompt_length=1024 \
      data.max_response_length="$MAX_RESPONSE_LENGTH" \
      data.filter_overlong_prompts=True \
      data.truncation='error' \
      data.image_key=images \
      actor_rollout_ref.model.path="$MPATH" \
      actor_rollout_ref.model.lora_rank="$ACTIVE_LORA_RANK" \
      actor_rollout_ref.model.lora_alpha="$LORA_ALPHA" \
      actor_rollout_ref.model.lora_dropout="$LORA_DROPOUT" \
      actor_rollout_ref.model.lora_target_modules="'$LORA_TARGET_MODULES'" \
      actor_rollout_ref.actor.optim.lr="$LR" \
      actor_rollout_ref.model.use_remove_padding="$USE_RMPAD" \
      actor_rollout_ref.actor.ppo_mini_batch_size="$TRAIN_BATCH" \
      actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu="$MICRO_BSZ" \
      actor_rollout_ref.actor.use_kl_loss=True \
      actor_rollout_ref.actor.kl_loss_coef="$KL_COEF" \
      actor_rollout_ref.actor.kl_loss_type=low_var_kl \
      actor_rollout_ref.model.enable_gradient_checkpointing=True \
      actor_rollout_ref.actor.fsdp_config.param_offload="$PARAM_OFFLOAD" \
      actor_rollout_ref.actor.fsdp_config.optimizer_offload="$OPTIM_OFFLOAD" \
      actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu="$LOGP_MICRO_BSZ" \
      actor_rollout_ref.rollout.tensor_model_parallel_size="$TP" \
      actor_rollout_ref.rollout.name=vllm \
      actor_rollout_ref.rollout.gpu_memory_utilization="$GPU_MEM_UTIL" \
      actor_rollout_ref.rollout.stop='["</class>"]' \
      actor_rollout_ref.rollout.include_stop_str_in_output=True \
      actor_rollout_ref.rollout.enable_chunked_prefill=False \
      actor_rollout_ref.rollout.enforce_eager=False \
      actor_rollout_ref.rollout.free_cache_engine=False \
      actor_rollout_ref.rollout.n=5 \
      actor_rollout_ref.rollout.val_kwargs.do_sample=False \
      actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu="$LOGP_MICRO_BSZ" \
      actor_rollout_ref.ref.fsdp_config.param_offload=True \
      algorithm.kl_ctrl.kl_coef="$KL_COEF" \
      trainer.critic_warmup=0 \
      trainer.logger="['${LOGGER}']" \
      trainer.project_name="$PROJECT" \
      trainer.experiment_name="$EXP" \
      trainer.default_local_dir="$CKPT_ROOT" \
      trainer.n_gpus_per_node="$N_GPUS" \
      trainer.nnodes=1 \
      trainer.save_freq=500 \
      trainer.test_freq=10 \
      trainer.val_only="$VAL_ONLY" \
      trainer.val_before_train=True \
      trainer.resume_mode="$RESUME_MODE" \
      trainer.total_epochs="$EPOCHS" \
      ts.use_sem_orth=True \
      ts.adv_mix=0.3 \
      ts.similarity_method=ot \
      ts.ot_eps=0.08 \
      ts.ot_n_iter=50 "${@:3}"
}

merge_latest_ckpt () {
  local latest
  latest=$(ls -d "${CKPT_ROOT}"/global_step_* 2>/dev/null | sort -t_ -k3 -n | tail -1 || true)
  if [ -z "$latest" ]; then
    echo "[ERROR] no checkpoint found under ${CKPT_ROOT}" >&2; exit 1
  fi
  if has_hf_weights "${latest}/actor/huggingface"; then
    echo "[merge] reusing existing HF weights: ${latest}/actor/huggingface" >&2
    echo "${latest}/actor/huggingface"
    return
  fi
  echo "[merge] latest checkpoint: ${latest}/actor" >&2
  "$PYTHON_BIN" scripts/model_merger.py \
      --local_dir "${latest}/actor" \
      --lora-alpha "$LORA_ALPHA" >&2
  echo "${latest}/actor/huggingface"   # only the path goes to stdout (captured by caller)
}

has_hf_weights () {
  local path="$1"
  [ -f "${path}/model.safetensors" ] ||
    [ -f "${path}/model.safetensors.index.json" ] ||
    [ -f "${path}/pytorch_model.bin" ] ||
    [ -f "${path}/pytorch_model.bin.index.json" ]
}

prepare_eval_model () {
  local model_path="$1"
  if has_hf_weights "$model_path"; then
    echo "$model_path"
    return
  fi

  local actor_dir=""
  if [ "$(basename "$model_path")" = "huggingface" ]; then
    actor_dir="$(dirname "$model_path")"
  elif [ -d "${model_path}/huggingface" ]; then
    actor_dir="$model_path"
  fi

  if [ -z "$actor_dir" ] ||
     ! find "$actor_dir" -maxdepth 1 -name 'model_world_size_*_rank_0.pt' -print -quit | grep -q .; then
    echo "[ERROR] '${model_path}' has no HF model weights and is not a mergeable FSDP checkpoint." >&2
    return 1
  fi

  echo "[merge] HF weights missing; merging checkpoint ${actor_dir}" >&2
  "$PYTHON_BIN" scripts/model_merger.py \
      --local_dir "$actor_dir" \
      --lora-alpha "$LORA_ALPHA" >&2

  model_path="${actor_dir}/huggingface"
  if ! has_hf_weights "$model_path"; then
    echo "[ERROR] merge completed without producing HF model weights under '${model_path}'." >&2
    return 1
  fi
  echo "$model_path"
}

# ---- log file --------------------------------------------------------------
# Tee all output (terminal + file). Override path with LOG_FILE=, or LOG_DIR=.
LOG_DIR=${LOG_DIR:-/mnt/share01/sqk/AnomSeer/logs}
mkdir -p "$LOG_DIR"
LOG_FILE=${LOG_FILE:-${LOG_DIR}/rats_${STAGE}$([ "$LORA_RANK" -gt 0 ] && echo _lora)_$(date +%Y%m%d_%H%M%S).log}
exec > >(tee -a "$LOG_FILE") 2>&1
echo "[log] saving full output to: $LOG_FILE"

# ---- dispatch --------------------------------------------------------------
if [ "$LORA_RANK" -gt 0 ]; then
  FT_MODE="LoRA (r=${LORA_RANK}, alpha=${LORA_ALPHA}, target=${LORA_TARGET_MODULES})"
else
  FT_MODE="full-parameter fine-tuning"
fi
if [ "$TP" -le 1 ]; then PARALLEL="data-parallel (DP=${N_GPUS}, full replica per card)"; else PARALLEL="tensor-parallel (TP=${TP})"; fi
echo "[config] tuning=${FT_MODE} | lr=${LR} | gpus=${N_GPUS} | ${PARALLEL}"
echo "[config] model=${MODEL_PATH}"
echo "[config] python=${PYTHON_BIN} | stage=${STAGE}"
if [ "$LORA_RANK" -eq 0 ] && [ "$PARAM_OFFLOAD" != "True" ] && [[ "$MODEL_PATH" == *7[bB]* ]]; then
  echo "[hint] full-parameter 7B on 2x48GB will likely OOM. Either:"
  echo "       (a) LoRA:    LORA_RANK=16 bash $0"
  echo "       (b) offload: PARAM_OFFLOAD=True OPTIM_OFFLOAD=True bash $0"
fi

case "$STAGE" in
  train)
    echo "[stage] TRAIN  (model=${MODEL_PATH}, gpus=${N_GPUS}, tp=${TP})"
    run_verl False "$MODEL_PATH" "$@"
    ;;
  eval)
    eval_model=$(prepare_eval_model "$MODEL_PATH")
    echo "[stage] EVAL   (model=${eval_model}, gpus=${N_GPUS}, tp=${TP}, lora_rank=0)"
    run_verl True "$eval_model" "$@"
    ;;
  train_eval)
    echo "[stage] TRAIN  (model=${MODEL_PATH}, gpus=${N_GPUS}, tp=${TP})"
    run_verl False "$MODEL_PATH"
    merged=$(merge_latest_ckpt)
    echo "[stage] EVAL   (merged checkpoint=${merged})"
    run_verl True "$merged" "$@"
    ;;
  *)
    echo "[ERROR] unknown STAGE='$STAGE' (use train | eval | train_eval)" >&2; exit 1
    ;;
esac