#!/usr/bin/env bash
# Serve google/gemma-4-12B-it-qat-w4a16-ct as frank-small on Dankclaw's RTX
# 3090 via vLLM (since 2026-09-18), with the google bf16 MTP assistant
# for speculative decoding.
#
# This script is the CONTAINER ENTRYPOINT of the frank-small quadlet: the
# vllm/vllm-openai image's own entrypoint is `vllm serve`, and the quadlet
# replaces it with this file (mounted read-only out of the repo checkout) so
# the flag set stays in git. The image pins the gemma4-unified vLLM build,
# which is what loads gemma4_unified natively and runs the Gemma 4 MTP path.
#
# MTP: the drafter is a separate 4-layer assistant checkpoint in bf16
# (Google's original gemma-4-12B-it-assistant). It goes in
# --speculative-config as method "mtp" — vLLM's Gemma 4 MTP path, not a
# generic draft model. The w4a16-ct re-quant of the assistant was tried
# first and drafts at acceptance 1.00 (the drafter path does not apply the
# Marlin kernel to its linears); the bf16 assistant accepts at 2.80.
#
# Sampling: greedy by default (fleet contract) via
# --override-generation-config; clients that send their own temperature get
# their own, and the harness overrides per model from models.yml.
#
# Tunables come from frank-small.env (mounted beside this file by the
# quadlet) unless FRANK_SMALL_ENV_FILE overrides the path.
set -euo pipefail

ENV_FILE="${FRANK_SMALL_ENV_FILE:-/scripts/frank-small.env}"
if [[ -f "$ENV_FILE" ]]; then
    set -a
    # shellcheck disable=SC1090
    source "$ENV_FILE"
    set +a
fi

# --- defaults (used when the env file does not set them) -------------------
: "${MODEL_DIR:=/models/main}"
: "${DRAFT_DIR:=/models/assistant}"
: "${ALIAS:=frank-small}"
: "${BIND_ADDR:=0.0.0.0}"
: "${LISTEN_PORT:=8000}"
: "${MAX_MODEL_LEN:=262144}"
: "${GPU_MEM_UTIL:=0.95}"
: "${KV_CACHE_DTYPE:=auto}"
: "${MAX_NUM_SEQS:=16}"
: "${MAX_NUM_BATCHED_TOKENS:=16384}"
: "${SPEC_NUM_TOKENS:=4}"
: "${ENFORCE_EAGER:=0}"

# Fail loudly rather than half-loading a model into VRAM: a missing file here
# is a download that has not finished, and the unit's restart policy would
# otherwise turn that into a crash loop that reads like a runtime fault.
for path in "$MODEL_DIR/model.safetensors" "$MODEL_DIR/config.json" \
            "$DRAFT_DIR/model.safetensors" "$DRAFT_DIR/config.json"; do
    [[ -e "$path" ]] || { echo "frank-small: missing $path" >&2; exit 1; }
done

SPEC_CONFIG="{\"method\":\"mtp\",\"model\":\"${DRAFT_DIR}\",\"num_speculative_tokens\":${SPEC_NUM_TOKENS}}"

VLLM_FLAGS=(
    serve "$MODEL_DIR"
--served-model-name "$ALIAS"
--host "$BIND_ADDR"
--port "$LISTEN_PORT"
--max-model-len "$MAX_MODEL_LEN"
--gpu-memory-utilization "$GPU_MEM_UTIL"
--kv-cache-dtype "$KV_CACHE_DTYPE"
--max-num-seqs "$MAX_NUM_SEQS"
--max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS"
--speculative-config "$SPEC_CONFIG"
--reasoning-parser gemma4
--enable-auto-tool-choice
--tool-call-parser gemma4
--override-generation-config '{"temperature": 0}'
)
# ENFORCE_EAGER=1 skips Inductor/Triton compilation and CUDA graph capture
# (slower steady-state decode, faster and more robust boot). The 3090 is
# sm_86 (Ampere) — no native fp8 — so KV stays bf16 and the default
# graphed path is the right one; the knob is a boot-debugging lever.
if [[ "$ENFORCE_EAGER" == "1" ]]; then
    VLLM_FLAGS+=(--enforce-eager)
fi

exec vllm "${VLLM_FLAGS[@]}"