#!/usr/bin/env bash
# Purpose: start vLLM as an OpenAI-compatible server with the settings this lab measures -
#          context per sequence, sequence cap, KV cache size or memory fraction, KV cache
#          dtype and prefix caching - printing the configuration first and copying the log
#          to a file, so the run is reproducible from the terminal and searchable afterwards
# Platform: nvidia and strix (a native vLLM install in an activated virtual environment);
#           Track S uses serve-vllm-container.sh, Track M has no vLLM path
# Minimum memory: 16 GB of accelerator-visible memory for the four-bit 8B checkpoint
# Assumes: vllm on PATH (the install lesson's environment, activated), the checkpoint in
#          $MODEL (a local directory with config.json, or a Hugging Face repository id that
#          vLLM downloads), nothing else answering on $PORT; runs in the foreground
#
# Usage: MODEL=$HOME/models/Qwen/Qwen3-8B-AWQ KV_BYTES=12079595520 bash serve-vllm-native.sh
#        MAX_SEQS=8 MODEL=... bash serve-vllm-native.sh
#        KV_BYTES= MEM_FRACTION=0.80 MODEL=... bash serve-vllm-native.sh
#
# Environment:
#   MODEL           checkpoint directory or repository id  (default: $HOME/models/Qwen/Qwen3-8B-AWQ)
#   SERVED_NAME     model name clients send                (default: local-chat)
#   HOST            address to bind                        (default: 127.0.0.1)
#   PORT            port to listen on                      (default: 8000)
#   CTX             --max-model-len, tokens per sequence   (default: 4096)
#   MAX_SEQS        --max-num-seqs, sequences in flight    (default: 20)
#   KV_BYTES        --kv-cache-memory-bytes; empty to size the cache from MEM_FRACTION
#                                                          (default: empty)
#   MEM_FRACTION    --gpu-memory-utilization, used when KV_BYTES is empty (default: 0.90)
#   KV_DTYPE        --kv-cache-dtype: auto (the model's dtype) or fp8     (default: auto)
#   PREFIX_CACHING  on or off                               (default: on, vLLM 0.28.0's default)
#   TP_SIZE         --tensor-parallel-size                  (default: 1)
#   LOGFILE         copy of the server log                  (default: ./vllm-server.log)
#
# Read from vLLM 0.28.0's source: kv_cache_memory_bytes, when set, "ignores
# gpu_memory_utilization" (config/cache.py), and the worker then logs "reserved ... GiB memory
# for KV Cache as specified by kv_cache_memory_bytes config and skipped memory profiling";
# without it the worker logs "Available KV cache memory: N GiB". Either way the engine logs
# "GPU KV cache size: N tokens, Maximum concurrency for N tokens per request: N.NNx".
# enable_prefix_caching defaults to True. --enable-prompt-tokens-details adds
# usage.prompt_tokens_details.cached_tokens to responses, which load-test.py records.

set -euo pipefail

MODEL="${MODEL:-$HOME/models/Qwen/Qwen3-8B-AWQ}"
SERVED_NAME="${SERVED_NAME:-local-chat}"
HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-8000}"
CTX="${CTX:-4096}"
MAX_SEQS="${MAX_SEQS:-20}"
KV_BYTES="${KV_BYTES:-}"
MEM_FRACTION="${MEM_FRACTION:-0.90}"
KV_DTYPE="${KV_DTYPE:-auto}"
PREFIX_CACHING="${PREFIX_CACHING:-on}"
TP_SIZE="${TP_SIZE:-1}"
LOGFILE="${LOGFILE:-./vllm-server.log}"

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

command -v vllm >/dev/null 2>&1 \
    || die "vllm is not on PATH; activate the environment from the install lesson first"
case "$MODEL" in
    /*|./*|../*) [ -f "$MODEL/config.json" ] || die "$MODEL/config.json not found; download the checkpoint first" ;;
esac
case "$PREFIX_CACHING" in
    on) PREFIX_FLAG="--enable-prefix-caching" ;;
    off) PREFIX_FLAG="--no-enable-prefix-caching" ;;
    *) die "PREFIX_CACHING must be on or off" ;;
esac
if [ -n "$KV_BYTES" ]; then
    MEMORY_FLAG="--kv-cache-memory-bytes"
    MEMORY_VALUE="$KV_BYTES"
else
    MEMORY_FLAG="--gpu-memory-utilization"
    MEMORY_VALUE="$MEM_FRACTION"
fi
if command -v curl >/dev/null 2>&1 \
    && curl --silent --max-time 2 --output /dev/null "http://$HOST:$PORT/"; then
    die "something is already answering on $HOST:$PORT; stop it first (one server at a time)"
fi
if [ "$HOST" != "127.0.0.1" ] && [ "$HOST" != "localhost" ]; then
    echo "NOTE: binding to $HOST exposes an unauthenticated generation endpoint." >&2
    echo "      Set an API key and put it behind TLS before doing this on a real network." >&2
fi

echo "==> vllm serve, $(vllm --version 2>/dev/null | tail -n 1)"
cat <<INFO
    model             $MODEL
    served as         $SERVED_NAME
    listening on      http://$HOST:$PORT/v1
    max model length  $CTX tokens per sequence
    max sequences     $MAX_SEQS
    KV cache memory   $MEMORY_FLAG $MEMORY_VALUE
    KV cache dtype    $KV_DTYPE
    prefix caching    $PREFIX_CACHING
    tensor parallel   $TP_SIZE
    log copied to     $LOGFILE
    The first start profiles memory and captures CUDA graphs; wait for
    "Application startup complete." before sending requests.
INFO

# The log goes to the terminal and to $LOGFILE; exec keeps this shell's pid for the server.
exec > >(tee "$LOGFILE") 2>&1
exec vllm serve "$MODEL" \
    --host "$HOST" \
    --port "$PORT" \
    --served-model-name "$SERVED_NAME" \
    --max-model-len "$CTX" \
    --max-num-seqs "$MAX_SEQS" \
    "$MEMORY_FLAG" "$MEMORY_VALUE" \
    --kv-cache-dtype "$KV_DTYPE" \
    "$PREFIX_FLAG" \
    --tensor-parallel-size "$TP_SIZE" \
    --enable-prompt-tokens-details