#!/usr/bin/env bash
# Purpose: run Part 9's load generator against whichever endpoint this lab has just
#          started, with the base URL, model name and notebook path filled in from .env, so
#          that the pair, the single machine and the two-card desktop are all measured by
#          the same tool with the same settings and produce comparable notebook lines
# Platform: all (it drives an HTTP endpoint; it can run on a third machine on the network)
# Minimum memory: none on the machine running this; the server needs the lab's floor
# Assumes: Part 9's load-test.py is on this machine and LOAD_TEST points at it; a server is
#          already answering on SERVE_HOST:SERVE_PORT with the name in SERVED_NAME; python3
#          on PATH; a .env filled in from env-example.txt. Pass a label as the first
#          argument: it is the tag written into the notebook line.
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

LABEL="${1:-}"
LOAD_TEST="${LOAD_TEST:-${HERE}/../part-09-vllm-and-sglang/load-test.py}"
SERVED_NAME="${SERVED_NAME:-local-cluster}"
SERVE_HOST="${SERVE_HOST:-127.0.0.1}"
SERVE_PORT="${SERVE_PORT:-8000}"
CONCURRENCY="${CONCURRENCY:-1,2,4}"
REQUESTS="${REQUESTS:-16}"
MAX_TOKENS="${MAX_TOKENS:-128}"
PROMPT_SET="${PROMPT_SET:-mixed}"
ENGINE="${ENGINE:-vllm}"
ENGINE_VERSION="${ENGINE_VERSION:-unknown}"
LABBOOK="${LABBOOK:-labbook.md}"
LOADTEST_API_KEY_ENV="${LOADTEST_API_KEY_ENV:-LOADTEST_API_KEY}"

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

[ -n "$LABEL" ] || fail "Usage: $(basename "$0") <label>   e.g. 'pair-tp2' or 'single-spark-120b'"
command -v python3 >/dev/null 2>&1 || fail "python3 is not on PATH."
[ -f "$LOAD_TEST" ] || fail "Part 9's load generator was not found at ${LOAD_TEST}. Copy load-test.py from the Part 9 lab files, or set LOAD_TEST to its path."

BASE_URL="http://${SERVE_HOST}:${SERVE_PORT}/v1"

if ! curl -sf --max-time 10 "${BASE_URL}/models" >/dev/null 2>&1; then
    fail "Nothing is answering at ${BASE_URL}/models. Start a server first, and wait for 'Application startup complete.'"
fi

cat <<INFO
==> Load test
    label             ${LABEL}
    endpoint          ${BASE_URL}
    model name        ${SERVED_NAME}
    concurrency       ${CONCURRENCY}
    requests / level  ${REQUESTS}
    max tokens        ${MAX_TOKENS}
    prompt set        ${PROMPT_SET}
    notebook          ${LABBOOK}

    On a very large model with almost no cache headroom, start at concurrency 1 and raise
    it only if the server survives. A refused or preempted request is a result too: record
    the concurrency at which it started happening.

INFO

exec python3 "$LOAD_TEST" \
    --base-url "$BASE_URL" \
    --model "$SERVED_NAME" \
    --concurrency "$CONCURRENCY" \
    --requests "$REQUESTS" \
    --max-tokens "$MAX_TOKENS" \
    --prompt-set "$PROMPT_SET" \
    --label "$LABEL" \
    --engine "$ENGINE" \
    --engine-version "$ENGINE_VERSION" \
    --api-key-env "$LOADTEST_API_KEY_ENV" \
    --labbook "$LABBOOK"