#!/usr/bin/env bash
# Purpose: serve one model across two GPUs in a single desktop, in either the tensor-split
#          or the pipeline-split arrangement, printing the PCIe generation and link width
#          each card negotiated so the measurement can be read against the link it ran on
# Platform: nvidia (one machine, two or more cards, no NVLink on recent GeForce parts)
# Minimum memory: 24 GB per GPU for the 70B-class AWQ model; less for a 32B-class one
# Assumes: vllm on PATH in an activated environment; both cards visible to nvidia-smi; the
#          model already downloaded; nothing else listening on SERVE_PORT. Set SPLIT=tp
#          for a tensor split or SPLIT=pp for a pipeline split, and run it both ways.
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="${TWO_CARD_MODEL:-Qwen/Qwen3-32B-AWQ}"
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}"
SPLIT="${SPLIT:-tp}"
DEVICES="${DEVICES:-2}"

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

command -v vllm >/dev/null 2>&1 || fail "vllm is not on PATH. Activate the environment from Part 9's install lesson."
command -v nvidia-smi >/dev/null 2>&1 || fail "nvidia-smi is not on PATH; the driver is the first thing to check."

printf '==> The link these cards actually negotiated\n'
nvidia-smi --query-gpu=index,name,memory.total,pcie.link.gen.current,pcie.link.width.current \
    --format=csv | sed 's/^/    /'
printf '    Read this under load as well as at idle: the current values drop when a GPU is\n'
printf '    idle, so an idle reading understates the link.\n\n'

case "$SPLIT" in
    tp) TP_SIZE="$DEVICES"; PP_SIZE=1 ;;
    pp) TP_SIZE=1; PP_SIZE="$DEVICES" ;;
    *)  fail "SPLIT must be 'tp' (tensor parallel) or 'pp' (pipeline parallel); got '${SPLIT}'." ;;
esac

cat <<INFO
==> vLLM across ${DEVICES} GPUs in one machine
    model             ${MODEL}
    served as         ${SERVED_NAME}
    listening on      http://${SERVE_HOST}:${SERVE_PORT}
    split             ${SPLIT}  (tensor parallel ${TP_SIZE}, pipeline parallel ${PP_SIZE})
    max model length  ${CTX}
    memory fraction   ${MEM_FRACTION}
    max sequences     ${MAX_SEQS}

    Run this twice, once with SPLIT=tp and once with SPLIT=pp, and load-test each. Without
    NVLink the tensor split runs a collective on every layer over PCIe, and vLLM's own
    documentation suggests pipeline parallel on machines without that interconnect. Which
    wins on your slot is a measurement, and both numbers belong in the notebook next to
    the link width printed above.

INFO

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