#!/usr/bin/env bash
# Purpose: start the vLLM instance that receives key-value blocks from the prefill
#          instance and generates the answer, with the same model and context length so
#          that the blocks it receives are blocks it can use
# Platform: spark, nvidia (vLLM's GPU path; Track X only where your ROCm build works)
# Minimum memory: 24 GB for Qwen3-8B at bf16 on the two-machine path; use SMALL_MODEL and
#          SPLIT_MEM_FRACTION for two processes on one 16 GB device
# Assumes: vLLM installed as in Part 9 and on PATH; a .env copied from env-example.txt and
#          filled in with the SAME MODEL and CTX the prefill instance uses; for
#          CONNECTOR=shared, KV_SHARED_PATH readable from this machine; nothing else
#          listening on DECODE_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="${MODEL:-Qwen/Qwen3-8B}"
SERVED_NAME="${SERVED_NAME:-local-chat}"
SERVE_HOST="${SERVE_HOST:-127.0.0.1}"
DECODE_PORT="${DECODE_PORT:-8200}"
CTX="${CTX:-8192}"
MAX_SEQS="${MAX_SEQS:-8}"
MEM_FRACTION="${MEM_FRACTION:-0.85}"
CONNECTOR="${CONNECTOR:-shared}"
KV_SHARED_PATH="${KV_SHARED_PATH:-}"
SIDE_CHANNEL_ADDR="${SIDE_CHANNEL_ADDR:-}"
DECODE_SIDE_CHANNEL_PORT="${DECODE_SIDE_CHANNEL_PORT:-5601}"
CLUSTER_IFACE="${CLUSTER_IFACE:-}"
SPLIT="${SPLIT_MEM_FRACTION:-0.40}"
SINGLE_MACHINE="${SINGLE_MACHINE:-0}"
MONOLITHIC="${MONOLITHIC:-0}"

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

command -v vllm >/dev/null 2>&1 || fail "vllm is not on PATH. Install it as Part 9 describes."

if [ "$SINGLE_MACHINE" = "1" ]; then
    MODEL="${SMALL_MODEL:-Qwen/Qwen3-1.7B}"
    MEM_FRACTION="$SPLIT"
fi

# The baseline run: one ordinary server, no connector, everything on one machine. This is
# the number the disaggregated pair is compared against, and it must use the same model,
# the same context length and the same served name so that only one thing differs.
if [ "$MONOLITHIC" = "1" ]; then
    cat <<INFO
==> vLLM monolithic baseline (no connector, no split)
    model             ${MODEL}
    served as         ${SERVED_NAME}
    listening on      http://${SERVE_HOST}:${DECODE_PORT}
    max model length  ${CTX}
    memory fraction   ${MEM_FRACTION}
    max sequences     ${MAX_SEQS}

    Point the load generator at THIS port for the baseline, and at the proxy port for the
    disaggregated run. Nothing else about the two runs should differ.

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

case "$CONNECTOR" in
    shared)
        [ -n "$KV_SHARED_PATH" ] || fail "CONNECTOR=shared needs KV_SHARED_PATH set to the SAME directory the prefill instance writes to."
        [ -d "$KV_SHARED_PATH" ] || fail "KV_SHARED_PATH ($KV_SHARED_PATH) is not a directory on this machine. Is the shared mount up?"
        [ -r "$KV_SHARED_PATH" ] || fail "KV_SHARED_PATH ($KV_SHARED_PATH) is not readable by this account."
        KV_CONFIG="{\"kv_connector\":\"ExampleConnector\",\"kv_role\":\"kv_consumer\",\"kv_connector_extra_config\":{\"shared_storage_path\":\"${KV_SHARED_PATH}\"}}"
        ;;
    nixl)
        [ -n "$SIDE_CHANNEL_ADDR" ] || fail "CONNECTOR=nixl needs SIDE_CHANNEL_ADDR: the address the OTHER machine can reach this one on."
        KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_consumer"}'
        export VLLM_NIXL_SIDE_CHANNEL_HOST="$SIDE_CHANNEL_ADDR"
        export VLLM_NIXL_SIDE_CHANNEL_PORT="$DECODE_SIDE_CHANNEL_PORT"
        [ -n "$CLUSTER_IFACE" ] && export UCX_NET_DEVICES="$CLUSTER_IFACE"
        ;;
    mooncake)
        [ -n "$SIDE_CHANNEL_ADDR" ] || fail "CONNECTOR=mooncake needs SIDE_CHANNEL_ADDR."
        KV_CONFIG='{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer"}'
        ;;
    *)
        fail "CONNECTOR must be one of: shared, nixl, mooncake. Got '${CONNECTOR}'."
        ;;
esac

cat <<INFO
==> vLLM decode instance (the consumer)
    model             ${MODEL}
    served as         ${SERVED_NAME}
    listening on      http://${SERVE_HOST}:${DECODE_PORT}
    connector         ${CONNECTOR}
    kv role           kv_consumer
    max model length  ${CTX}
    memory fraction   ${MEM_FRACTION}
    max sequences     ${MAX_SEQS}

    This instance is where the measurement lives. Its /metrics endpoint carries
    vllm:time_to_first_token_seconds, vllm:prefix_cache_hits and
    vllm:prefix_cache_queries, and a hit rate near zero here means the blocks are not
    arriving: the instance is quietly prefilling every prompt itself and returning
    correct answers, which is the failure this lab is designed to catch.

INFO

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