#!/usr/bin/env bash
set -euo pipefail

# 配置区域（可根据实际情况修改）
MODEL_BASE_PATH="${MODEL_BASE_PATH:-/mnt/Qwen3-Omni-30B-A3B-Instruct-NVFP4}"
THINKER_PATH="${MODEL_BASE_PATH}/thinker"
CONTAINER_NAME="${CONTAINER_NAME:-vllm-thinker}"
HOST_PORT="${HOST_PORT:-8091}"
GPU_DEVICES="${GPU_DEVICES:-all}"          # 例如 "0,1" 或 "all"
TENSOR_PARALLEL_SIZE="${TENSOR_PARALLEL_SIZE:-1}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-32768}"
VLLM_IMAGE="${VLLM_IMAGE:-vllm/vllm-openai:latest}"

# 检查模型路径是否存在
if [[ ! -d "$THINKER_PATH" ]]; then
  echo "错误: 模型路径不存在: $THINKER_PATH"
  echo "请设置环境变量 MODEL_BASE_PATH 指向正确的模型根目录（包含 thinker 子目录）"
  exit 1
fi

# 检查 Docker 可用性
if command -v docker >/dev/null 2>&1 && docker info >/dev/null 2>&1; then
  DOCKER_CMD="docker"
else
  DOCKER_CMD="sudo docker"
fi

# 停止并移除可能已存在的同名容器
${DOCKER_CMD} rm -f "${CONTAINER_NAME}" >/dev/null 2>&1 || true

echo "启动 Thinker 服务..."
echo "模型路径: $THINKER_PATH"
echo "服务端口: $HOST_PORT"
echo "Tensor Parallel Size: $TENSOR_PARALLEL_SIZE"

# 启动容器
${DOCKER_CMD} run -d \
  --name "${CONTAINER_NAME}" \
  --gpus "${GPU_DEVICES}" \
  --ipc=host \
  --network host \
  -v "${THINKER_PATH}:/model:ro" \
  -e VLLM_USE_MODELSCOPE=false \
  "${VLLM_IMAGE}" \
  sh -c "pip install --no-cache-dir nvidia-modelopt && \
         vllm serve /model \
           --port ${HOST_PORT} \
           --quantization modelopt_fp4 \
           --dtype bfloat16 \
           --max-model-len ${MAX_MODEL_LEN} \
           --tensor-parallel-size ${TENSOR_PARALLEL_SIZE} \
           --trust-remote-code \
           --enforce-eager"

# 等待服务启动
sleep 3

# 健康检查
echo "健康检查 (等待几秒后):"
echo "  curl http://127.0.0.1:${HOST_PORT}/health"
echo "查看日志:"
echo "  ${DOCKER_CMD} logs -f ${CONTAINER_NAME}"