#!/usr/bin/env bash
# Purpose: start llama-server with a fixed number of parallel slots and continuous batching,
#          every setting that changes a load test stated on the command line, the load log
#          copied to a file you can search, and the process id written where you can find it
# Platform: all (spark, strix, mac, nvidia; CUDA, Vulkan and Metal builds); the only served
#           path on Track M and the certain one on Track X
# Minimum memory: 16 GB, because the context is allocated for every slot at load
# Assumes: llama.cpp built as in Part 6 (v0.4.0 or a build with the same options), a GGUF
#          file in $MODEL, nothing else answering on $PORT; the server runs in the foreground
#          so that Ctrl-C stops it
#
# Usage: MODEL=$HOME/models/unsloth/Qwen3-8B-GGUF/Qwen3-8B-Q4_K_M.gguf bash serve-llama-parallel.sh
#        SLOTS=8 CACHE_TYPE=q8_0 MODEL=... bash serve-llama-parallel.sh
#
# Environment:
#   MODEL         path to the GGUF file                   (default: the Qwen3-8B Q4_K_M file)
#   LLAMA_BIN     directory holding llama-server, or the binary itself
#                                                         (default: $HOME/llama.cpp/build/bin)
#   ALIAS         model name the API reports              (default: local-chat)
#   HOST          address to bind                         (default: 127.0.0.1)
#   PORT          port to listen on                       (default: 8080)
#   SLOTS         parallel slots, the most sequences decoded at once        (default: 20)
#   CTX_PER_SLOT  context per slot, in tokens              (default: 4096)
#   CACHE_TYPE    KV cache type for keys and values: f16, q8_0, ...        (default: f16)
#   CACHE_RAM     host-RAM prompt cache limit in MiB; 0 disables it         (default: 0)
#   NGL           layers to offload                        (default: 999, meaning all)
#   PIDFILE       where to write the server's pid          (default: ./llama-server.pid)
#   LOGFILE       copy of the server log                   (default: ./llama-server.log)
#
# What the options do, as read from llama-server --help at v0.4.0:
#   --ctx-size is the total across slots. With --parallel given explicitly the KV cache is not
#   unified ("default: enabled if number of slots is auto"), so each slot gets
#   ctx-size / parallel tokens and the load log says so: "n_slots = 20, n_ctx_slot = 4096,
#   kv_unified = 'false'". The whole cache is allocated when the model loads.
#   --cache-type-k/-v choose the bytes per cached token; a quantised value cache needs flash
#   attention, hence --flash-attn on.
#   --cache-prompt is on by default: each slot keeps its last prompt's KV and a new request
#   reuses the longest common prefix of the slot it lands on. That is the prefix caching the
#   lab measures. --cache-ram (default 8192 MiB) adds a second cache in host RAM that saves
#   idle slots and restores them; this script turns it off so that the baseline measures
#   batching rather than copies to and from host memory. Set CACHE_RAM=8192 to measure it.
#   --verbosity 4 is the level at which the load log prints the KV cache and buffer sizes.

set -euo pipefail

MODEL="${MODEL:-$HOME/models/unsloth/Qwen3-8B-GGUF/Qwen3-8B-Q4_K_M.gguf}"
LLAMA_BIN="${LLAMA_BIN:-$HOME/llama.cpp/build/bin}"
ALIAS="${ALIAS:-local-chat}"
HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-8080}"
SLOTS="${SLOTS:-20}"
CTX_PER_SLOT="${CTX_PER_SLOT:-4096}"
CACHE_TYPE="${CACHE_TYPE:-f16}"
CACHE_RAM="${CACHE_RAM:-0}"
NGL="${NGL:-999}"
PIDFILE="${PIDFILE:-./llama-server.pid}"
LOGFILE="${LOGFILE:-./llama-server.log}"

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

if [ -f "$LLAMA_BIN" ] && [ -x "$LLAMA_BIN" ]; then
    SERVER="$LLAMA_BIN"
elif [ -x "$LLAMA_BIN/llama-server" ]; then
    SERVER="$LLAMA_BIN/llama-server"
else
    die "no llama-server at $LLAMA_BIN; set LLAMA_BIN, or build llama.cpp as in Part 6"
fi
[ -f "$MODEL" ] || die "model file not found: $MODEL"
case "$SLOTS$CTX_PER_SLOT" in
    *[!0-9]*) die "SLOTS and CTX_PER_SLOT must be whole numbers" ;;
esac
[ "$SLOTS" -ge 1 ] || die "SLOTS must be at least 1"
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

TOTAL_CTX=$(( SLOTS * CTX_PER_SLOT ))

echo "==> llama-server, $SLOTS parallel slot(s)"
"$SERVER" --version 2>&1 | head -n 2 || true
cat <<INFO
    model             $MODEL
    served as         $ALIAS
    listening on      http://$HOST:$PORT/v1
    slots             $SLOTS
    context per slot  $CTX_PER_SLOT
    total context     $TOTAL_CTX   (--ctx-size is the total across slots)
    KV cache type     $CACHE_TYPE (keys and values)
    host prompt cache $CACHE_RAM MiB
    pid written to $PIDFILE, log copied to $LOGFILE
INFO

rm -f "$LOGFILE"
echo $$ > "$PIDFILE"

# exec keeps this shell's pid, so $PIDFILE names the server process itself.
exec "$SERVER" \
    --model "$MODEL" \
    --alias "$ALIAS" \
    --host "$HOST" \
    --port "$PORT" \
    --ctx-size "$TOTAL_CTX" \
    --parallel "$SLOTS" \
    --cont-batching \
    --n-gpu-layers "$NGL" \
    --flash-attn on \
    --cache-type-k "$CACHE_TYPE" \
    --cache-type-v "$CACHE_TYPE" \
    --cache-ram "$CACHE_RAM" \
    --metrics \
    --jinja \
    --verbosity 4 \
    --log-file "$LOGFILE"