#!/usr/bin/env bash
# Purpose: start vLLM's OpenAI-compatible server on a four-bit checkpoint of the lab's model,
#          as the third engine on Track N, with one sequence, the lab's context length, a KV
#          cache sized for that context rather than for a fraction of the card, and the Qwen3
#          tool-call and reasoning parsers enabled
# Platform: nvidia (vLLM's CUDA wheels; on Track S use the NGC container from Part 9 instead)
# Minimum memory: 12 GB
# Assumes: vLLM installed into the active Python environment as this lab's Track N tab shows,
#          an NVIDIA GPU with a working driver, the checkpoint downloaded to $MODEL, the port
#          free and nothing else on the GPU; the server runs in the foreground
#
# Usage: MODEL=$HOME/models/Qwen/Qwen3-8B-AWQ bash serve-vllm.sh
#
# Environment:
#   MODEL       checkpoint directory                     (default: $HOME/models/Qwen/Qwen3-8B-AWQ)
#   SERVED_AS   name the API reports for the model       (default: the directory's base name)
#   HOST        address to bind                          (default: 127.0.0.1)
#   PORT        port to listen on                        (default: 8083)
#   MAX_LEN     maximum model length, in tokens          (default: 8192)
#   KV_BYTES    KV cache size in bytes                   (default: 1342177280, 1.25 GiB)
#   PIDFILE     where to write the process id            (default: ./vllm-server.pid)
#
# Sizing: Qwen3-8B's KV cache is 2 x 36 layers x 8 KV heads x 128 x 2 bytes = 147,456 bytes per
# token, so 8,192 tokens need 1,207,959,552 bytes; 1.25 GiB leaves room for whole blocks.
# vLLM's cache configuration documents that --kv-cache-memory-bytes, when set, ignores
# --gpu-memory-utilization, which otherwise reserves a fraction of the whole card.
# Parsers: vLLM's tool-calling page gives --enable-auto-tool-choice with --tool-call-parser
# hermes for Qwen models, and its reasoning page lists the qwen3 reasoning parser. Part 9
# teaches both properly. --enable-prompt-tokens-details adds usage.prompt_tokens_details, the
# cached-token count the measuring script uses to confirm that prefix caching was not hit.

set -euo pipefail

MODEL="${MODEL:-$HOME/models/Qwen/Qwen3-8B-AWQ}"
SERVED_AS="${SERVED_AS:-$(basename "$MODEL")}"
HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-8083}"
MAX_LEN="${MAX_LEN:-8192}"
KV_BYTES="${KV_BYTES:-1342177280}"
PIDFILE="${PIDFILE:-./vllm-server.pid}"

die() { echo "serve-vllm: $*" >&2; exit 1; }

command -v vllm >/dev/null || die "vllm not found; activate the environment you installed it into"
command -v nvidia-smi >/dev/null || die "nvidia-smi not found; vLLM here expects an NVIDIA GPU"
[ -f "$MODEL/config.json" ] || die "$MODEL/config.json not found; download the checkpoint first"
if command -v curl >/dev/null && curl --silent --max-time 2 "http://$HOST:$PORT/health" >/dev/null 2>&1; then
  die "something is already answering on $HOST:$PORT; stop it first (one server at a time)"
fi

echo "==> vllm serve"
vllm --version || true
echo "    model    $MODEL"
echo "    served   $SERVED_AS"
echo "    listen   http://$HOST:$PORT/v1"
echo "    context  $MAX_LEN tokens, 1 sequence, KV cache $KV_BYTES bytes"
echo "    pid written to $PIDFILE"
echo "    The first start profiles and captures CUDA graphs; wait for the startup-complete line."

echo $$ > "$PIDFILE"

# exec keeps this shell's pid, so $PIDFILE names the server process itself.
exec vllm serve "$MODEL" \
    --served-model-name "$SERVED_AS" \
    --host "$HOST" \
    --port "$PORT" \
    --max-model-len "$MAX_LEN" \
    --max-num-seqs 1 \
    --kv-cache-memory-bytes "$KV_BYTES" \
    --enable-auto-tool-choice \
    --tool-call-parser hermes \
    --reasoning-parser qwen3 \
    --enable-prompt-tokens-details