#!/usr/bin/env bash
# Purpose: start vLLM from NVIDIA's container image with the settings this lab measures,
#          publishing the port on the loopback address only, mounting the model library
#          read-only, and copying the log to a file on the host
# Platform: spark (NVIDIA's NGC vLLM image, the documented path on DGX Spark); usable on
#           nvidia with IMAGE overridden. Track X uses serve-vllm-native.sh; Track M has no
#           vLLM path
# Minimum memory: 16 GB of accelerator-visible memory for the four-bit 8B checkpoint
# Assumes: docker able to reach the GPU (--gpus all), the image already pulled, the
#          checkpoint downloaded under $MODELS_DIR on the host, nothing else answering on
#          $PORT; runs in the foreground, and Ctrl-C or "docker rm -f vllm-lab" stops it
#
# Usage: KV_BYTES=12079595520 bash serve-vllm-container.sh
#        MAX_SEQS=8 KV_BYTES=12079595520 bash serve-vllm-container.sh
#        KV_BYTES= MEM_FRACTION=0.30 bash serve-vllm-container.sh
#
# Environment:
#   IMAGE           container image                         (default: nvcr.io/nvidia/vllm:26.08-py3)
#   MODELS_DIR      host model library, mounted at /models  (default: $HOME/models)
#   MODEL           checkpoint path inside the container    (default: /models/Qwen/Qwen3-8B-AWQ)
#   SERVED_NAME     model name clients send                 (default: local-chat)
#   PORT            host port, bound to 127.0.0.1           (default: 8000)
#   CTX, MAX_SEQS, KV_BYTES, MEM_FRACTION, KV_DTYPE, PREFIX_CACHING
#                   as in serve-vllm-native.sh              (defaults: 4096, 20, empty, 0.90,
#                                                            auto, on)
#   CONTAINER_NAME  container name                          (default: vllm-lab)
#   LOGFILE         copy of the server log on the host      (default: ./vllm-server.log)
#
# On a unified-memory machine --gpu-memory-utilization is a share of memory the operating
# system also uses, so this lab sets KV_BYTES instead wherever it can. --ipc=host: vLLM's
# Docker documentation says PyTorch needs the host's shared memory to pass data between
# processes.

set -euo pipefail

IMAGE="${IMAGE:-nvcr.io/nvidia/vllm:26.08-py3}"
MODELS_DIR="${MODELS_DIR:-$HOME/models}"
MODEL="${MODEL:-/models/Qwen/Qwen3-8B-AWQ}"
SERVED_NAME="${SERVED_NAME:-local-chat}"
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}"
CONTAINER_NAME="${CONTAINER_NAME:-vllm-lab}"
LOGFILE="${LOGFILE:-./vllm-server.log}"

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

command -v docker >/dev/null 2>&1 || die "docker is not on PATH"
docker image inspect "$IMAGE" >/dev/null 2>&1 || die "image $IMAGE not present; run: docker pull $IMAGE"
case "$MODEL" in
    /models/*) [ -f "$MODELS_DIR/${MODEL#/models/}/config.json" ] \
        || die "$MODELS_DIR/${MODEL#/models/}/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://127.0.0.1:$PORT/"; then
    die "something is already answering on 127.0.0.1:$PORT; stop it first (one server at a time)"
fi
if docker container inspect "$CONTAINER_NAME" >/dev/null 2>&1; then
    die "a container named $CONTAINER_NAME exists; remove it with: docker rm -f $CONTAINER_NAME"
fi

cat <<INFO
==> vllm serve, in a container
    image             $IMAGE
    model             $MODEL  (host: $MODELS_DIR, read-only)
    served as         $SERVED_NAME
    published on      http://127.0.0.1:$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
    log copied to     $LOGFILE
    Stop it with Ctrl-C, or from another terminal: docker rm -f $CONTAINER_NAME
INFO

docker run --rm \
    --name "$CONTAINER_NAME" \
    --gpus all \
    --ipc=host \
    -v "$MODELS_DIR:/models:ro" \
    -p "127.0.0.1:$PORT:8000" \
    "$IMAGE" \
    vllm serve "$MODEL" \
    --host 0.0.0.0 \
    --port 8000 \
    --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" \
    --enable-prompt-tokens-details 2>&1 | tee "$LOGFILE"