#!/usr/bin/env bash
# Purpose: run vLLM's own example disaggregated-prefill proxy in front of the prefill and
#          decode instances, after checking that both are answering, and tell you exactly
#          where to obtain the proxy if you do not have it
# Platform: all (the proxy is plain Python; it may run on any machine that can reach both
#          instances, including your workstation)
# Minimum memory: 1 GB; the proxy holds no model
# Assumes: python3 and curl on PATH; serve-prefill.sh and serve-decode.sh already running
#          and answering /v1/models; a .env copied from env-example.txt. The proxy file is
#          examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py in the vLLM
#          repository; it is NOT installed by `pip install vllm`.
#
# Usage: bash run-proxy.sh            start the proxy (the file must already be here)
#        bash run-proxy.sh --fetch    download the proxy from the tag in VLLM_TAG first
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

SERVED_NAME="${SERVED_NAME:-local-chat}"
SERVE_HOST="${SERVE_HOST:-127.0.0.1}"
PREFILL_HOST="${PREFILL_HOST:-127.0.0.1}"
DECODE_HOST="${DECODE_HOST:-127.0.0.1}"
PREFILL_PORT="${PREFILL_PORT:-8100}"
DECODE_PORT="${DECODE_PORT:-8200}"
PROXY_PORT="${PROXY_PORT:-8000}"
VLLM_TAG="${VLLM_TAG:-main}"
PROXY_FILE="${PROXY_FILE:-${HERE}/disagg_proxy_demo.py}"

REPO_PATH="examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py"
RAW_BASE="https://raw.githubusercontent.com/vllm-project/vllm"

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

command -v python3 >/dev/null 2>&1 || fail "python3 is not on PATH."
command -v curl >/dev/null 2>&1 || fail "curl is not on PATH."

if [ "${1:-}" = "--fetch" ]; then
    printf '==> Fetching the proxy from the vLLM repository\n'
    printf '    tag   %s\n' "$VLLM_TAG"
    printf '    path  %s\n' "$REPO_PATH"
    printf '    into  %s\n\n' "$PROXY_FILE"
    printf '    Use the SAME tag as the vLLM you installed. A proxy from a different\n'
    printf '    version than the engines is a supported way to waste an evening.\n\n'
    curl -fsSL "${RAW_BASE}/${VLLM_TAG}/${REPO_PATH}" -o "$PROXY_FILE"
    printf '    Downloaded. Read it before running it; it is a demonstration file with no\n'
    printf '    authentication of any kind.\n\n'
fi

if [ ! -f "$PROXY_FILE" ]; then
    cat >&2 <<MISSING
The proxy is not here and it is not installed by the vLLM wheel. It lives in the
vLLM source tree at:

    ${REPO_PATH}

Get it in whichever way suits you, at the SAME tag as your installed vLLM:

    git clone --depth 1 --branch <your vllm tag> https://github.com/vllm-project/vllm
    cp vllm/${REPO_PATH} ${PROXY_FILE}

or let this script fetch just that file:

    VLLM_TAG=<your vllm tag> bash run-proxy.sh --fetch

Then run this script again.
MISSING
    exit 1
fi

# --- both instances must be answering before the proxy is worth starting ---------------
check_instance() {
    local role="$1" host="$2" port="$3"
    if curl -fsS --max-time 5 "http://${host}:${port}/v1/models" >/dev/null 2>&1; then
        printf '    OK: the %s instance answers on %s:%s\n' "$role" "$host" "$port"
    else
        fail "The ${role} instance is not answering on ${host}:${port}. Start it first."
    fi
}

printf '==> Checking both instances\n'
check_instance prefill "$PREFILL_HOST" "$PREFILL_PORT"
check_instance decode "$DECODE_HOST" "$DECODE_PORT"
printf '\n'

cat <<INFO
==> Proxy
    model     ${SERVED_NAME}
    prefill   ${PREFILL_HOST}:${PREFILL_PORT}
    decode    ${DECODE_HOST}:${DECODE_PORT}
    listening http://${SERVE_HOST}:${PROXY_PORT}

    Send every measured request to the proxy port and nothing to the two instances
    directly. The proxy calls prefill first and decode second; if you see requests
    arriving only at the decode instance, you have measured a single-machine deployment
    wearing a disaggregated hat.

    No authentication, on any of the three. Keep them on your cluster network and put
    Part 9's gateway in front if anything outside this machine will reach them.

INFO

exec python3 "$PROXY_FILE" \
    --model "$SERVED_NAME" \
    --prefill "${PREFILL_HOST}:${PREFILL_PORT}" \
    --decode "${DECODE_HOST}:${DECODE_PORT}" \
    --port "$PROXY_PORT"