#!/usr/bin/env bash
# Purpose: serve the 120B-class mixture-of-experts comparison model on ONE DGX Spark, with
#          the same served name and port as the pair, so the same load-test command line
#          measures both and the two notebook lines are directly comparable
# Platform: spark (one machine; this is also the fallback path for readers with one Spark)
# Minimum memory: 128 GB of unified memory
# Assumes: vllm on PATH in an activated environment, or set USE_CONTAINER=1 and VLLM_IMAGE
#          to run it inside the NGC image instead; the model already downloaded into
#          HF_HOME; nothing else listening on SERVE_PORT. No Ray cluster is involved.
set -euo pipefail

HERE="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
if [ -f "${HERE}/.env" ]; then
    set -a
    # shellcheck disable=SC1091  # written by the reader from env-example.txt
    . "${HERE}/.env"
    set +a
fi

MODEL="${SINGLE_MODEL:-openai/gpt-oss-120b}"
SERVED_NAME="${SERVED_NAME:-local-cluster}"
SERVE_HOST="${SERVE_HOST:-127.0.0.1}"
SERVE_PORT="${SERVE_PORT:-8000}"
CTX="${CTX:-4096}"
MEM_FRACTION="${MEM_FRACTION:-0.90}"
MAX_SEQS="${MAX_SEQS:-4}"
USE_CONTAINER="${USE_CONTAINER:-0}"
HF_HOME="${HF_HOME:-$HOME/.cache/huggingface}"

fail() { printf '%s\n' "$*" >&2; exit 1; }

cat <<INFO
==> vLLM on one Spark, the single-machine comparison
    model             ${MODEL}
    served as         ${SERVED_NAME}
    listening on      http://${SERVE_HOST}:${SERVE_PORT}
    max model length  ${CTX}
    memory fraction   ${MEM_FRACTION}
    max sequences     ${MAX_SEQS}
    container         ${USE_CONTAINER}

    Record the same two startup lines you recorded for the pair: the KV cache size in
    tokens and the maximum concurrency at this context length. On one machine the cache is
    smaller and the model is smaller, and which way those two facts push the comparison is
    the question this lab exists to answer.

INFO

if [ "$USE_CONTAINER" = "1" ]; then
    command -v docker >/dev/null 2>&1 || fail "docker is not on PATH."
    [ -n "${VLLM_IMAGE:-}" ] || fail "USE_CONTAINER=1 needs VLLM_IMAGE set in .env."
    exec docker run --rm -it \
        --gpus all \
        --ipc=host \
        --network host \
        -v "${HF_HOME}:/root/.cache/huggingface" \
        "$VLLM_IMAGE" \
        vllm serve "$MODEL" \
            --host "$SERVE_HOST" \
            --port "$SERVE_PORT" \
            --served-model-name "$SERVED_NAME" \
            --max-model-len "$CTX" \
            --gpu-memory-utilization "$MEM_FRACTION" \
            --max-num-seqs "$MAX_SEQS"
fi

command -v vllm >/dev/null 2>&1 || fail "vllm is not on PATH. Activate the environment from Part 9's install lesson, or set USE_CONTAINER=1."

exec vllm serve "$MODEL" \
    --host "$SERVE_HOST" \
    --port "$SERVE_PORT" \
    --served-model-name "$SERVED_NAME" \
    --max-model-len "$CTX" \
    --gpu-memory-utilization "$MEM_FRACTION" \
    --max-num-seqs "$MAX_SEQS"