#!/usr/bin/env bash
#
# Gemma 4 统一容器入口：
#   1. 幂等应用 CoreX vLLM 的 Gemma4 / MoE 补丁
#   2. 读取模型 config.json，自动判断 Gemma-4-31B（dense）/ GEV-26B-Decide（MoE）
#   3. 选择对应的 TP / 端口 / served name / eager / 显存，再启动 vLLM
#
# AUTO_DETECT=1（默认）按 config.json 自动选择；显式设置的
# TP / PORT / SERVED_NAME / MAX_NUM_SEQS / GPU_MEM_UTIL / ENFORCE_EAGER
# 会覆盖检测结果。AUTO_DETECT=0 完全手动。
#
set -euo pipefail

PATCH_SCRIPT="${GEMMA4_PATCH_SCRIPT:-/opt/gemma4/apply_patch.py}"
DETECT_SCRIPT="${GEMMA4_DETECT_SCRIPT:-/opt/gemma4/detect_model.py}"
VLLM_BIN="${VLLM_BIN:-vllm}"

python3 "${PATCH_SCRIPT}"

MODEL_PATH="${MODEL_PATH:-/models/gemma4}"
AUTO_DETECT="${AUTO_DETECT:-1}"
MODEL_DESCRIPTION=""

if [[ ! -f "${MODEL_PATH}/config.json" ]]; then
  echo "ERROR: 模型目录不完整：${MODEL_PATH}/config.json 不存在" >&2
  echo "请确认已挂载模型，例如：" >&2
  echo "  -v /mnt/disk0/models/google/gemma-4-31B-it:/models/gemma4:ro" >&2
  echo "  -v /mnt/disk0/models/autotrust/GEV-26B-Decide:/models/gemma4:ro" >&2
  exit 1
fi

if [[ "${AUTO_DETECT}" == "1" ]]; then
  DETECT_ARGS=()
  if [[ -n "${TP:-}" ]]; then
    DETECT_ARGS=(--tp "${TP}")
  fi
  DETECT_OUTPUT="$(python3 "${DETECT_SCRIPT}" "${MODEL_PATH}" "${DETECT_ARGS[@]}")" || exit 1
  eval "${DETECT_OUTPUT}"
  MODEL_KIND="${MODEL_KIND:-${DETECTED_MODEL_KIND}}"
  TP="${TP:-${DETECTED_TP}}"
  PORT="${PORT:-${DETECTED_PORT}}"
  SERVED_NAME="${SERVED_NAME:-${DETECTED_SERVED_NAME}}"
  MAX_NUM_SEQS="${MAX_NUM_SEQS:-${DETECTED_MAX_NUM_SEQS}}"
  GPU_MEM_UTIL="${GPU_MEM_UTIL:-${DETECTED_GPU_MEM_UTIL}}"
  ENFORCE_EAGER="${ENFORCE_EAGER:-${DETECTED_ENFORCE_EAGER}}"
  MODEL_DESCRIPTION="${DETECTED_DESCRIPTION:-}"
else
  MODEL_KIND="${MODEL_KIND:-unknown}"
fi

HOST="${HOST:-0.0.0.0}"
TP="${TP:-4}"
PORT="${PORT:-8095}"
SERVED_NAME="${SERVED_NAME:-gemma4}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-8192}"
MAX_NUM_SEQS="${MAX_NUM_SEQS:-4}"
GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.90}"
MAX_IMAGES="${MAX_IMAGES:-4}"
ENFORCE_EAGER="${ENFORCE_EAGER:-0}"
EXTRA_ARGS="${EXTRA_ARGS:-}"

EAGER_ARGS=()
if [[ "${ENFORCE_EAGER}" == "1" ]]; then
  EAGER_ARGS=(--enforce-eager)
fi

echo "============================================================"
echo "Gemma 4 family (BI-V150 unified image)"
echo "  model kind : ${MODEL_KIND}"
echo "  model      : ${MODEL_PATH}"
echo "  served name: ${SERVED_NAME}"
echo "  listen     : ${HOST}:${PORT}"
echo "  TP         : ${TP}  (${TP} GPUs)"
echo "  max len    : ${MAX_MODEL_LEN}"
echo "  max seqs   : ${MAX_NUM_SEQS}"
echo "  gpu util   : ${GPU_MEM_UTIL}"
echo "  max images : ${MAX_IMAGES}"
echo "  eager      : ${ENFORCE_EAGER}"
echo "  visible    : ${CUDA_VISIBLE_DEVICES:-all}"
echo "  extra      : ${EXTRA_ARGS}"
if [[ -n "${MODEL_DESCRIPTION}" ]]; then
  echo "  detected   : ${MODEL_DESCRIPTION}"
fi
echo "============================================================"

# shellcheck disable=SC2086
exec "${VLLM_BIN}" serve "${MODEL_PATH}" \
  --host "${HOST}" \
  --port "${PORT}" \
  --served-model-name "${SERVED_NAME}" \
  --tensor-parallel-size "${TP}" \
  --max-model-len "${MAX_MODEL_LEN}" \
  --max-num-seqs "${MAX_NUM_SEQS}" \
  --gpu-memory-utilization "${GPU_MEM_UTIL}" \
  --limit-mm-per-prompt "{\"image\": ${MAX_IMAGES}}" \
  "${EAGER_ARGS[@]}" \
  ${EXTRA_ARGS}