#!/usr/bin/env bash
#
# Gemma 4 家族统一构建 + 启动（Iluvatar BI-V150）
#
# 自动读取模型目录下的 config.json，判断是 dense 还是 MoE，然后选择：
#   Gemma-4-31B-it   : TP4 / 4 卡 / 8095 / gemma-4-31B-it / eager0 / 0.90 / 4 seqs
#   GEV-26B-Decide   : TP2 / 2 卡 / 8094 / GEV-26B-Decide / eager1 / 0.95 / 2 seqs
#
# 用法：
#   ./run.sh /mnt/disk0/models/google/gemma-4-31B-it
#   ./run.sh /mnt/disk0/models/autotrust/GEV-26B-Decide
#   MODEL_HOST_DIR=/path/to/model ./run.sh
#   ./run.sh                    # 依次尝试两个默认模型目录
#
# 常用覆盖项：
#   GPUS=12,13,14,15 ./run.sh /path/to/gemma4
#   EXTRA_DOCKER_ARGS="-v /mnt/disk1:/data:ro" ./run.sh /path/to/model
#   EXTRA_ARGS="--max-model-len 16384" ./run.sh /path/to/model
#   HOST_PORT=24397 ./run.sh /path/to/model             # 换宿主机端口（默认 24396）
#   TP=4 SERVED_NAME=llm ./run.sh /path/to/model
#   BUILD=0 IMAGE=<ci-image> ./run.sh /path/to/model   # 使用已有镜像
#   DRY_RUN=1 ./run.sh /path/to/model                  # 只打印 docker run 命令
#
set -euo pipefail

cd "$(dirname "$0")"

usage() {
  cat >&2 <<'EOF'
用法: ./run.sh [MODEL_DIR]

不传 MODEL_DIR 时，按顺序尝试：
  /mnt/disk0/models/google/gemma-4-31B-it
  /mnt/disk0/models/autotrust/GEV-26B-Decide

模型类型由 MODEL_DIR/config.json 自动判断，不需要 gemma4-31b / gev-26b-decide 参数。

环境变量覆盖：
  IMAGE NAME MODEL_HOST_DIR MODEL_MOUNT SERVED_NAME TP GPUS
  HOST_PORT（宿主机端口，默认 24396；容器内固定 80）
  MAX_MODEL_LEN MAX_NUM_SEQS GPU_MEM_UTIL MAX_IMAGES ENFORCE_EAGER
  AUTO_DETECT EXTRA_ARGS EXTRA_DOCKER_ARGS BUILD=0/1 DRY_RUN=0/1
EOF
}

if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
  usage
  exit 0
fi

MODEL_HOST_DIR="${1:-${MODEL_HOST_DIR:-}}"
if [[ -z "${MODEL_HOST_DIR}" ]]; then
  for candidate in \
    "/mnt/disk0/models/google/gemma-4-31B-it" \
    "/mnt/disk0/models/autotrust/GEV-26B-Decide"; do
    if [[ -f "${candidate}/config.json" ]]; then
      MODEL_HOST_DIR="${candidate}"
      break
    fi
  done
fi

if [[ -z "${MODEL_HOST_DIR}" ]]; then
  echo "ERROR: 未指定模型目录，也没有找到默认模型目录。" >&2
  usage
  exit 2
fi

if [[ ! -f "${MODEL_HOST_DIR}/config.json" ]]; then
  echo "ERROR: 模型目录不完整: ${MODEL_HOST_DIR}/config.json" >&2
  echo "请先下载模型，或把 MODEL_DIR 指向正确目录。" >&2
  exit 1
fi

DETECT_ARGS=()
if [[ -n "${TP:-}" ]]; then
  DETECT_ARGS=(--tp "${TP}")
fi
DETECT_OUTPUT="$(python3 detect_model.py "${MODEL_HOST_DIR}" "${DETECT_ARGS[@]}")" || exit 1
eval "${DETECT_OUTPUT}"

MODEL_KIND="${MODEL_KIND:-${DETECTED_MODEL_KIND}}"
IMAGE="${IMAGE:-gemma4-bi150:2026-10-27}"
NAME="${NAME:-${MODEL_KIND}}"
HOST="${HOST:-0.0.0.0}"
# 容器内端口固定 80；宿主机端口是 HOST_PORT，默认 24396（串行部署复用同一端口）。
# PORT 保留为兼容别名，若显式给 PORT 则视作 HOST_PORT。
CONTAINER_PORT=80
HOST_PORT="${HOST_PORT:-${PORT:-24396}}"
MODEL_MOUNT="${MODEL_MOUNT:-/models/gemma4}"
SERVED_NAME="${SERVED_NAME:-llm}"
TP="${TP:-${DETECTED_TP}}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-8192}"
MAX_NUM_SEQS="${MAX_NUM_SEQS:-${DETECTED_MAX_NUM_SEQS}}"
GPU_MEM_UTIL="${GPU_MEM_UTIL:-${DETECTED_GPU_MEM_UTIL}}"
MAX_IMAGES="${MAX_IMAGES:-4}"
ENFORCE_EAGER="${ENFORCE_EAGER:-${DETECTED_ENFORCE_EAGER}}"
AUTO_DETECT="${AUTO_DETECT:-1}"
EXTRA_ARGS="${EXTRA_ARGS:-}"
EXTRA_DOCKER_ARGS="${EXTRA_DOCKER_ARGS:-}"
BUILD="${BUILD:-1}"
DRY_RUN="${DRY_RUN:-0}"

GPU_COUNT="${DETECTED_GPU_COUNT}"

# 默认优先使用序号最大的卡（12-15），避免和 16 卡大模型抢 0-11。
if [[ -n "${GPUS:-}" ]]; then
  GPUS="${GPUS}"
elif [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
  GPUS="${CUDA_VISIBLE_DEVICES}"
elif (( GPU_COUNT == 4 )); then
  GPUS="15,14,13,12"
else
  GPUS="$(seq -s, $((16 - GPU_COUNT)) 15)"
fi
GPUS="${GPUS// /}"

IFS=',' read -r -a GPU_LIST <<< "${GPUS}"
if (( ${#GPU_LIST[@]} < GPU_COUNT )); then
  echo "ERROR: ${MODEL_KIND} 需要 ${GPU_COUNT} 张卡，但 GPUS 只给了 ${#GPU_LIST[@]} 张: ${GPUS}" >&2
  exit 1
fi
if (( ${#GPU_LIST[@]} > GPU_COUNT )); then
  echo "WARN: ${MODEL_KIND} 只需要 ${GPU_COUNT} 张卡，但 GPUS 给了 ${#GPU_LIST[@]} 张；vLLM 只会使用其中 ${GPU_COUNT} 张。" >&2
fi

DOCKER_EXTRA=()
if [[ -n "${EXTRA_DOCKER_ARGS}" ]]; then
  # 按空白拆分；路径中不要带空格。
  read -r -a DOCKER_EXTRA <<< "${EXTRA_DOCKER_ARGS}"
fi

DOCKER_RUN=(
  docker run -d --name "${NAME}"
  --entrypoint /usr/local/bin/gemma4-entrypoint.sh
  --workdir /opt/gemma4
  --expose "${CONTAINER_PORT}"
  --privileged
  --pid host
  --ipc host
  -p "${HOST_PORT}:${CONTAINER_PORT}"
  -v /dev:/dev
  -v "${MODEL_HOST_DIR}:${MODEL_MOUNT}:ro"
  --ulimit memlock=-1
  --ulimit stack=67108864
  -e CUDA_VISIBLE_DEVICES="${GPUS}"
  -e MODEL_PATH="${MODEL_MOUNT}"
  -e AUTO_DETECT="${AUTO_DETECT}"
  -e MODEL_KIND="${MODEL_KIND}"
  -e HOST="${HOST}"
  -e SERVED_NAME="${SERVED_NAME}"
  -e PORT="${CONTAINER_PORT}"
  -e TP="${TP}"
  -e MAX_MODEL_LEN="${MAX_MODEL_LEN}"
  -e MAX_NUM_SEQS="${MAX_NUM_SEQS}"
  -e GPU_MEM_UTIL="${GPU_MEM_UTIL}"
  -e MAX_IMAGES="${MAX_IMAGES}"
  -e ENFORCE_EAGER="${ENFORCE_EAGER}"
  -e PYTHONUNBUFFERED=1
  -e EXTRA_ARGS="${EXTRA_ARGS}"
  "${DOCKER_EXTRA[@]}"
  --restart unless-stopped
  "${IMAGE}"
)

echo "==> detected ${MODEL_KIND}: ${DETECTED_DESCRIPTION}"
echo "    model : ${MODEL_HOST_DIR}"
echo "    TP    : ${TP} (${GPU_COUNT} GPUs)"
echo "    port  : 127.0.0.1:${HOST_PORT} -> container:${CONTAINER_PORT}"
echo "    name  : ${SERVED_NAME}"
echo "    GPUs  : ${GPUS}"
echo "    eager : ${ENFORCE_EAGER}, util=${GPU_MEM_UTIL}, seqs=${MAX_NUM_SEQS}"

if [[ "${DRY_RUN}" == "1" ]]; then
  printf '%q ' "${DOCKER_RUN[@]}"
  printf '\n'
  exit 0
fi

if [[ "${BUILD}" == "1" ]]; then
  echo "==> build ${IMAGE}"
  docker build -t "${IMAGE}" .
else
  echo "==> skip build (BUILD=0), image=${IMAGE}"
fi

echo "==> recreate ${NAME}"
docker rm -f "${NAME}" 2>/dev/null || true

echo "==> run ${MODEL_KIND}"
"${DOCKER_RUN[@]}"

# 冷启动：Gemma-4-31B 约 59 GB、GEV 约 49 GB，首次读盘实测 5-8 分钟。
HEALTH_TIMEOUT="${HEALTH_TIMEOUT:-900}"
echo "==> wait for http://127.0.0.1:${HOST_PORT}/health (up to ${HEALTH_TIMEOUT}s)"
DEADLINE=$((SECONDS + HEALTH_TIMEOUT))
while (( SECONDS < DEADLINE )); do
  if curl -sf "http://127.0.0.1:${HOST_PORT}/health" >/dev/null 2>&1; then
    echo "ready: http://127.0.0.1:${HOST_PORT}/v1/models"
    exit 0
  fi
  sleep 5
done

echo "ERROR: 服务 ${HEALTH_TIMEOUT}s 未就绪，查看日志：" >&2
echo "  docker logs -f ${NAME}" >&2
exit 1