#!/usr/bin/env bash
# Purpose: start one server with speculative decoding enabled, on llama.cpp or on vLLM,
#          using the same target weights as the baseline server so that the pair can be
#          measured against each other by measure-speculative.py
# Platform: llamacpp: spark, strix, mac, nvidia. vllm: spark and nvidia (vLLM's GPU
#           installation page lists ROCm wheels for the Ryzen AI Max+ that the course has
#           not exercised, and no mainline macOS GPU path).
# Minimum memory: the target weights plus the draft plus a KV cache. For an 8B target at
#           Q4_K_M with a 0.6B draft at Q4_K_M that is roughly 6 GB of weights before the
#           cache, which is why this lab's floor is 24 GB rather than 8.
# Assumes: llama-server on PATH (built in Part 6) for ENGINE=llamacpp, or vllm on PATH
#          (Part 9) for ENGINE=vllm; the target and draft model paths exist. Nothing here
#          writes to a model directory or deletes anything.
#
# Usage: bash serve-with-draft.sh ENGINE TARGET DRAFT [PORT]
#   ENGINE   llamacpp | vllm
#   TARGET   GGUF file (llamacpp) or Hugging Face id / directory (vllm)
#   DRAFT    draft GGUF file (llamacpp) or draft model id / directory (vllm);
#            pass the literal word none to start the baseline server with no draft
#   PORT     default 8081 with a draft, 8080 without
#
# Environment: HOST (default 127.0.0.1), CTX (default 4096), NGL (default 999),
#              NGLD (draft layers on the accelerator, default 999), DRAFT_MAX (default 4),
#              DRAFT_MIN (default 0), DRAFT_P_MIN (default 0.75), ALIAS (default local-chat),
#              GPU_UTIL for vLLM (default 0.85).
set -euo pipefail

ENGINE="${1:-}"
TARGET="${2:-}"
DRAFT="${3:-}"

HOST="${HOST:-127.0.0.1}"
CTX="${CTX:-4096}"
NGL="${NGL:-999}"
NGLD="${NGLD:-999}"
DRAFT_MAX="${DRAFT_MAX:-4}"
DRAFT_MIN="${DRAFT_MIN:-0}"
DRAFT_P_MIN="${DRAFT_P_MIN:-0.75}"
ALIAS="${ALIAS:-local-chat}"
GPU_UTIL="${GPU_UTIL:-0.85}"

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

[[ -n "$ENGINE" && -n "$TARGET" && -n "$DRAFT" ]] ||
  die "usage: bash serve-with-draft.sh ENGINE TARGET DRAFT [PORT]"
[[ "$ENGINE" == "llamacpp" || "$ENGINE" == "vllm" ]] ||
  die "ENGINE must be llamacpp or vllm (got '$ENGINE')"

if [[ "$DRAFT" == "none" ]]; then
  PORT="${4:-8080}"
  ROLE="baseline, no draft"
else
  PORT="${4:-8081}"
  ROLE="speculative, drafting up to $DRAFT_MAX token(s) per step"
fi

echo "==> $ENGINE on $HOST:$PORT - $ROLE"
echo "    target $TARGET"
[[ "$DRAFT" == "none" ]] || echo "    draft  $DRAFT"
echo "    Serve the baseline and the speculative server at the same time, on two ports,"
echo "    so measure-speculative.py can compare them without a reload in between."
echo

if [[ "$ENGINE" == "llamacpp" ]]; then
  command -v llama-server >/dev/null || die "llama-server is not on PATH; build it as in Part 6"
  [[ -f "$TARGET" ]] || die "$TARGET is not a file; llama.cpp wants a GGUF path"

  if [[ "$DRAFT" == "none" ]]; then
    llama-server \
      --model "$TARGET" \
      --alias "$ALIAS" \
      --host "$HOST" \
      --port "$PORT" \
      --ctx-size "$CTX" \
      --n-gpu-layers "$NGL" \
      --metrics \
      --slots
  else
    [[ -f "$DRAFT" ]] || die "$DRAFT is not a file; llama.cpp wants a GGUF path for the draft too"
    # The draft must share the target's tokeniser. A draft from another family will either
    # be refused at start-up or accepted and never agree with the target; both are wasted time.
    llama-server \
      --model "$TARGET" \
      --alias "$ALIAS" \
      --host "$HOST" \
      --port "$PORT" \
      --ctx-size "$CTX" \
      --n-gpu-layers "$NGL" \
      --metrics \
      --slots \
      --spec-draft-model "$DRAFT" \
      --spec-draft-n-max "$DRAFT_MAX" \
      --spec-draft-n-min "$DRAFT_MIN" \
      --spec-draft-p-min "$DRAFT_P_MIN" \
      -ngld "$NGLD"
  fi
else
  command -v vllm >/dev/null || die "vllm is not on PATH; install it as in Part 9"

  if [[ "$DRAFT" == "none" ]]; then
    vllm serve "$TARGET" \
      --served-model-name "$ALIAS" \
      --host "$HOST" \
      --port "$PORT" \
      --max-model-len "$CTX" \
      --gpu-memory-utilization "$GPU_UTIL" \
      --seed 0
  else
    # One JSON object carries the whole speculative configuration. "draft_model" is the
    # method for a separate small model; the lab page gives the ngram and eagle3 shapes.
    SPEC_CONFIG="{\"method\": \"draft_model\", \"model\": \"$DRAFT\", \"num_speculative_tokens\": $DRAFT_MAX}"
    echo "    speculative-config: $SPEC_CONFIG"
    vllm serve "$TARGET" \
      --served-model-name "$ALIAS" \
      --host "$HOST" \
      --port "$PORT" \
      --max-model-len "$CTX" \
      --gpu-memory-utilization "$GPU_UTIL" \
      --seed 0 \
      --speculative-config "$SPEC_CONFIG"
  fi
fi