#!/usr/bin/env bash
# Purpose: start trtllm-serve on the same model the portable engine is serving, inside the NGC
#          release container with the flags NVIDIA's documentation requires, with the KV cache
#          capped at the lab's context length, the Qwen3 tool and reasoning parsers selected,
#          and guided decoding enabled so that JSON-schema requests can be honoured
# Platform: spark (also valid on nvidia when the card has the memory for bf16 weights)
# Minimum memory: 12 GB
# Assumes: an NVIDIA GPU with a working driver; Docker with the NVIDIA Container Toolkit when
#          trtllm-serve is not on PATH; the checkpoint already downloaded under $MODELS_DIR;
#          network access to nvcr.io the first time the image is pulled; the port free
#
# Usage: TRTLLM_IMAGE=nvcr.io/nvidia/tensorrt-llm/release:<tag> bash serve-trtllm.sh
#        MODEL=$HOME/models/Qwen/Qwen3-8B PORT=8082 TRTLLM_IMAGE=... bash serve-trtllm.sh
#
# Environment:
#   MODEL         checkpoint directory on the host      (default: $HOME/models/Qwen/Qwen3-8B)
#   MODELS_DIR    directory mounted read-only as /models (default: $HOME/models)
#   TRTLLM_IMAGE  NGC image, when trtllm-serve is not on PATH (no default; record the tag)
#   HOST          address to publish on                 (default: 127.0.0.1)
#   PORT          port to listen on                     (default: 8082)
#   MAX_SEQ_LEN   longest request, prompt plus output   (default: 8192)
#   KV_TOKENS     tokens the KV cache may hold          (default: 8192)
#   CONFIG        YAML options file this script writes  (default: ./trtllm-lab.yml)
#   CONTAINER     name given to the container           (default: trtllm-lab)
#
# Why the YAML file: TensorRT-LLM sizes its KV cache as a fraction of free GPU memory (0.9 by
# default, per its KV cache documentation). On a 128 GB unified-memory machine that is tens of
# gigabytes, and the memory column would measure the fraction rather than the engine.
# kv_cache_config.max_tokens caps it: the documentation says the lesser of max_tokens and the
# fraction is allocated. guided_decoding_backend: xgrammar is how the guided-decoding page says
# to enable response_format constraints for trtllm-serve.
# No pid file is written in the container case: the process this shell would record is the
# docker client, not the server, so memory for this engine comes from the idle-snapshot delta.

set -euo pipefail

MODEL="${MODEL:-$HOME/models/Qwen/Qwen3-8B}"
MODELS_DIR="${MODELS_DIR:-$HOME/models}"
HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-8082}"
MAX_SEQ_LEN="${MAX_SEQ_LEN:-8192}"
KV_TOKENS="${KV_TOKENS:-8192}"
CONFIG="${CONFIG:-./trtllm-lab.yml}"
CONTAINER="${CONTAINER:-trtllm-lab}"

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

command -v nvidia-smi >/dev/null || die "nvidia-smi not found; this script needs an NVIDIA GPU"
[ -f "$MODEL/config.json" ] || die "$MODEL/config.json not found; download the checkpoint first"
case "$MODEL" in
  "$MODELS_DIR"/*) ;;
  *) die "MODEL ($MODEL) must be inside MODELS_DIR ($MODELS_DIR), which is what the container sees" ;;
esac
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

cat > "$CONFIG" <<EOF
guided_decoding_backend: xgrammar
kv_cache_config:
  max_tokens: $KV_TOKENS
EOF
CONFIG_ABS="$(cd "$(dirname "$CONFIG")" && pwd)/$(basename "$CONFIG")"

echo "==> trtllm-serve"
echo "    model       $MODEL"
echo "    listen      http://$HOST:$PORT/v1"
echo "    max_seq_len $MAX_SEQ_LEN, max_batch_size 1, KV cache capped at $KV_TOKENS tokens"
echo "    parsers     tool qwen3, reasoning qwen3; guided decoding xgrammar"
echo "    options     $CONFIG_ABS"

if command -v trtllm-serve >/dev/null; then
  echo "    running trtllm-serve from PATH"
  echo $$ > ./trtllm-serve.pid
  # exec keeps this shell's pid, so the pid file names the server process itself.
  exec trtllm-serve "$MODEL" \
      --host "$HOST" \
      --port "$PORT" \
      --max_batch_size 1 \
      --max_seq_len "$MAX_SEQ_LEN" \
      --tool_parser qwen3 \
      --reasoning_parser qwen3 \
      --extra_llm_api_options "$CONFIG_ABS"
fi

[ -n "${TRTLLM_IMAGE:-}" ] || die "trtllm-serve is not on PATH; set TRTLLM_IMAGE to an nvcr.io/nvidia/tensorrt-llm/release tag"
command -v docker >/dev/null || die "docker is not installed, and TensorRT-LLM is distributed as a container"
if docker ps --all --quiet --filter "name=^${CONTAINER}$" | grep -q .; then
  die "a container named $CONTAINER already exists; docker rm -f $CONTAINER first"
fi

IN_CONTAINER_MODEL="/models/${MODEL#"$MODELS_DIR"/}"
echo "    image       $TRTLLM_IMAGE (container name $CONTAINER)"
echo "    in container $IN_CONTAINER_MODEL"
echo "    stop it with: docker stop $CONTAINER"

# --ipc=host and the two ulimits are the documented docker run flags for these images; the
# server binds 0.0.0.0 inside the container, and Docker publishes it on $HOST only.
exec docker run --rm \
    --name "$CONTAINER" \
    --gpus=all \
    --ipc=host \
    --ulimit memlock=-1 \
    --ulimit stack=67108864 \
    --publish "$HOST:$PORT:$PORT" \
    --volume "$MODELS_DIR:/models:ro" \
    --volume "$CONFIG_ABS:/config/trtllm-lab.yml:ro" \
    "$TRTLLM_IMAGE" \
    trtllm-serve "$IN_CONTAINER_MODEL" \
        --host 0.0.0.0 \
        --port "$PORT" \
        --max_batch_size 1 \
        --max_seq_len "$MAX_SEQ_LEN" \
        --tool_parser qwen3 \
        --reasoning_parser qwen3 \
        --extra_llm_api_options /config/trtllm-lab.yml