#!/usr/bin/env bash
# Purpose: serve one model across both DGX Sparks with a single vllm serve command inside
#          the Ray head container, printing the configuration first so the run is
#          reproducible from the terminal log, and reminding you which two startup lines
#          to copy into the notebook
# Platform: spark (run on the head node only, after both Ray shells are up)
# Minimum memory: 128 GB per node; the pair's 256 GB is the ceiling for weights plus cache
# Assumes: start-ray-head.sh and start-ray-worker.sh are both running and "ray status"
#          reports two nodes; the model is already downloaded into HF_HOME on both nodes,
#          or the cache is shared; a .env filled in from env-example.txt. Nothing else is
#          listening on SERVE_PORT.
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="${PAIR_MODEL:-}"
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}"
TP_SIZE="${TP_SIZE:-2}"
PP_SIZE="${PP_SIZE:-1}"

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

command -v docker >/dev/null 2>&1 || fail "docker is not on PATH."
[ -n "$MODEL" ] || fail "PAIR_MODEL is not set in .env."

CONTAINER="${VLLM_CONTAINER:-$(docker ps --format '{{.Names}}' | grep -E '^node-[0-9]+$' | head -n 1)}"
[ -n "$CONTAINER" ] || fail "No Ray container found. Is start-ray-head.sh still running in its own shell?"

printf '==> Ray sees:\n'
docker exec "$CONTAINER" ray status | sed 's/^/    /'
printf '\n'

if [ "$TP_SIZE" -ne 1 ] && [ "$PP_SIZE" -ne 1 ]; then
    printf '    NOTE: both parallel sizes are above one. On a pair of one-GPU machines that\n'
    printf '          asks for more devices than exist. Set one of TP_SIZE or PP_SIZE to 1.\n\n'
fi

cat <<INFO
==> vLLM across the pair
    container         ${CONTAINER}
    model             ${MODEL}
    served as         ${SERVED_NAME}
    listening on      http://${SERVE_HOST}:${SERVE_PORT}
    tensor parallel   ${TP_SIZE}
    pipeline parallel ${PP_SIZE}
    max model length  ${CTX}
    memory fraction   ${MEM_FRACTION}
    max sequences     ${MAX_SEQS}

    Two lines in the startup log are the ones to record. One reports the GPU KV cache size
    in tokens, summed across the cluster. The other reports the maximum concurrency that
    implies at your context length. If the second is below the concurrency you intend to
    serve, stop now and change something: it will not improve when requests arrive.

    Startup on a very large model is slow, and most of it is loading weights. Watch for
    "Application startup complete."

INFO

EXEC_ENV=()
if [ -n "${NCCL_DEBUG:-}" ]; then
    EXEC_ENV+=(-e "NCCL_DEBUG=${NCCL_DEBUG}")
fi

exec docker exec -it "${EXEC_ENV[@]+"${EXEC_ENV[@]}"}" "$CONTAINER" \
    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"