#!/usr/bin/env bash
# Provisional quality-only launcher: sealed r17 + exact G64 Q-only FP8 reader.
# This is not the full-CUDA-graph NVFP4 speed-qualified release profile.
set -Eeuo pipefail

ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
RUNTIME="$ROOT/runtime/r17-g64-q-only"

IMAGE="${IMAGE:-voipmonitor/vllm@sha256:c5e96c5bcc5a073f7ce6b56173d88538de3a416900cff97c88b4bf7967fe1dc0}"
NAME="${NAME:-glm52-r17-g64-q-only}"
MODEL_DIR="${MODEL_DIR:-$ROOT}"
CACHE_DIR="${CACHE_DIR:-$HOME/.cache/glm52-r17-g64-q-only}"
PORT="${PORT:-8000}"
GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.97}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-262144}"
MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-1280}"
MAX_NUM_SEQS="${MAX_NUM_SEQS:-4}"
GRAPH="${GRAPH:-16}"
MTP="${MTP:-3}"
EAGER="${EAGER:-1}"
KV_CACHE_MEMORY_BYTES="${KV_CACHE_MEMORY_BYTES:-4294967296}"

required=(
  "$MODEL_DIR/config.json"
  "$MODEL_DIR/model.safetensors.index.json"
  "$RUNTIME/exl3_overlay.py"
  "$RUNTIME/gate_linear.py"
  "$RUNTIME/decode_math.py"
)

require_files() {
  local path
  for path in "${required[@]}"; do
    [[ -f "$path" ]] || { printf 'required file missing: %s\n' "$path" >&2; exit 2; }
  done
  mkdir -p "$CACHE_DIR" "$CACHE_DIR/tmp"
}

start() {
  require_files
  docker image inspect "$IMAGE" >/dev/null 2>&1 || docker pull "$IMAGE"
  docker rm -f "$NAME" >/dev/null 2>&1 || true
  local extra_args=(--enable-bf16x3-router-gemm --kv-cache-memory-bytes "$KV_CACHE_MEMORY_BYTES")
  [[ "$EAGER" == "1" ]] && extra_args+=(--enforce-eager)
  docker run -d --name "$NAME" \
    --entrypoint /usr/local/bin/serve-gilded-gnosis.sh \
    --gpus all --runtime nvidia --network host --ipc host --privileged --init \
    --shm-size 64g --ulimit memlock=-1:-1 --ulimit nofile=1048576:1048576 \
    --ulimit stack=67108864 --restart unless-stopped \
    -v "$MODEL_DIR:/model:ro" \
    -v "$CACHE_DIR:/cache:rw" \
    -v "$CACHE_DIR/tmp:/container-tmp:rw" \
    -v "$RUNTIME/exl3_overlay.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/quantization/exl3.py:ro" \
    -v "$RUNTIME/exl3_overlay.py:/opt/infernal-invocation/vllm/vllm/model_executor/layers/quantization/exl3.py:ro" \
    -v "$RUNTIME/exl3_overlay.py:/opt/venv/lib/python3.12/site-packages/vllm/model_executor/layers/quantization/exl3.py:ro" \
    -v "$RUNTIME/gate_linear.py:/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/router/gate_linear.py:ro" \
    -v "$RUNTIME/gate_linear.py:/opt/infernal-invocation/vllm/vllm/model_executor/layers/fused_moe/router/gate_linear.py:ro" \
    -v "$RUNTIME/gate_linear.py:/opt/venv/lib/python3.12/site-packages/vllm/model_executor/layers/fused_moe/router/gate_linear.py:ro" \
    -v "$RUNTIME/decode_math.py:/opt/infernal-invocation/b12x/b12x/attention/_shared/mla/decode_math.py:ro" \
    -v "$RUNTIME/decode_math.py:/opt/venv/lib/python3.12/site-packages/b12x/attention/_shared/mla/decode_math.py:ro" \
    -e MODEL_FAMILY=glm52-exl3 -e MODEL=/model \
    -e SERVED_MODEL_NAME=GLM-5.2-EXL3-TR3v4-3.5bpw \
    -e GPUS="${GPUS:-3,1,2,0}" -e PORT="$PORT" \
    -e TP=4 -e DCP=4 -e DCP_KV_CACHE_INTERLEAVE_SIZE=64 -e MTP="$MTP" \
    -e MAX_MODEL_LEN="$MAX_MODEL_LEN" -e MAX_BATCHED_TOKENS="$MAX_NUM_BATCHED_TOKENS" \
    -e MAX_NUM_SEQS="$MAX_NUM_SEQS" -e GRAPH="$GRAPH" \
    -e GPU_MEMORY_UTILIZATION="$GPU_MEMORY_UTILIZATION" \
    -e QUANTIZATION=exl3 -e ONLINE_QUANT=none -e LOAD_FORMAT=safetensors \
    -e KV_CACHE_DTYPE=fp8 -e MOE_MODE=a16 -e MOE_BACKEND=b12x \
    -e MTP_MOE_BACKEND=triton -e MTP_DRAFT_SAMPLE_METHOD=greedy \
    -e ASYNC_SCHEDULING=0 -e PCIE_CALIBRATION=off \
    -e VLLM_DCP_Q_REPLICATE=0 -e VLLM_DCP_INDEXER_SHARDS=4 \
    -e VLLM_DCP_QUERY_SPLIT=0 -e VLLM_DCP_GLOBAL_TOPK=1 \
    -e VLLM_DCP_SHARD_DRAFT=1 -e VLLM_USE_FUSED_MOE_GROUPED_TOPK=1 \
    -e VLLM_EXL3_R7_FUSED=1 -e VLLM_EXL3_R7_FUSED_LAYERS=48 \
    -e VLLM_EXL3_PREFILL_CAPACITY=1024 -e VLLM_EXL3_TRELLIS_MIN_M=1 \
    -e VLLM_ROUTER_WEIGHT_FP32=1 -e VLLM_EXL3_ROUTE_WEIGHT_FP32=0 \
    -e VLLM_FP8_MLA_SCALE_BANK4=0 -e VLLM_FP8_MLA_SCALE_BANK16=0 \
    -e VLLM_FP8_MLA_BF16_G64=0 -e B12X_FP8_MLA_BF16_QK=0 \
    -e B12X_FP8_MLA_BF16_PV=0 -e B12X_FP8_MLA_ARBITRARY_Q_SCALE=1 \
    -e B12X_FP8_MLA_ARBITRARY_Q_G64=1 -e KV_FP8_ROPE=0 \
    -e B12X_CUTE_COMPILE_CACHE_DIR=/cache/b12x-cute/r17-g64-q-only \
    "$IMAGE" "${extra_args[@]}" >/dev/null
  printf 'started %s on http://127.0.0.1:%s\n' "$NAME" "$PORT"
  printf 'PROVISIONAL quality-only profile; not the speed-qualified release default\n'
  printf 'profile: r17, TP4/DCP4, MTP-%s, FP8 MLA KV, exact G64 Q, BF16 RoPE, eager=%s\n' "$MTP" "$EAGER"
  printf 'follow startup: %s logs\n' "$0"
}

case "${1:-start}" in
  start) start ;;
  restart) docker rm -f "$NAME" >/dev/null 2>&1 || true; start ;;
  stop) docker rm -f "$NAME" >/dev/null ;;
  logs) docker logs -f --tail 150 "$NAME" ;;
  status)
    docker ps -a --filter "name=^/${NAME}$" --format 'table {{.Names}}\t{{.Status}}\t{{.Image}}'
    curl -fsS "http://127.0.0.1:$PORT/v1/models" || true
    printf '\n'
    ;;
  pull) docker pull "$IMAGE" ;;
  *) printf 'usage: %s [start|restart|stop|logs|status|pull]\n' "$0" >&2; exit 2 ;;
esac