#!/usr/bin/env bash
# Purpose: start one of two vLLM instances that share a key-value store through a
#          directory both can read and write, so that a prefix one instance computed can
#          be reused by the other; run it twice, once as instance a and once as b
# Platform: spark, nvidia (vLLM's GPU path; Track X only where your ROCm build works)
# Minimum memory: 16 GB of device memory for two instances of a small model at
#          SPLIT_MEM_FRACTION each, or one instance per machine at MEM_FRACTION
# Assumes: vLLM installed as in Part 9 and on PATH; a .env copied from env-example.txt
#          with KV_SHARED_PATH set to a directory BOTH instances can read and write; on
#          two machines that is the shared mount from Part 18, on one machine any
#          directory will do. Nothing else listening on the port chosen.
#
# Usage: bash serve-shared.sh a      first instance,  port PREFILL_PORT
#        bash serve-shared.sh b      second instance, port 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

INSTANCE="${1:-}"

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."

SERVE_HOST="${SERVE_HOST:-127.0.0.1}"
SERVED_NAME="${SERVED_NAME:-local-chat}"
CTX="${CTX:-8192}"
MAX_SEQS="${MAX_SEQS:-8}"
KV_SHARED_PATH="${KV_SHARED_PATH:-}"
SINGLE_MACHINE="${SINGLE_MACHINE:-0}"

case "$INSTANCE" in
    a) PORT="${PREFILL_PORT:-8100}" ;;
    b) PORT="${DECODE_PORT:-8200}" ;;
    *) fail "Usage: bash serve-shared.sh a|b" ;;
esac

if [ "$SINGLE_MACHINE" = "1" ]; then
    MODEL="${SMALL_MODEL:-Qwen/Qwen3-1.7B}"
    MEM_FRACTION="${SPLIT_MEM_FRACTION:-0.40}"
else
    MODEL="${MODEL:-Qwen/Qwen3-8B}"
    MEM_FRACTION="${MEM_FRACTION:-0.85}"
fi

[ -n "$KV_SHARED_PATH" ] || fail "KV_SHARED_PATH is not set. Both instances need the SAME directory."
mkdir -p "$KV_SHARED_PATH"
[ -w "$KV_SHARED_PATH" ] || fail "KV_SHARED_PATH (${KV_SHARED_PATH}) is not writable by this account."

# kv_both, not kv_producer or kv_consumer: this instance both writes blocks to the store
# and reads blocks from it, which is what sharing means as opposed to handing off.
KV_CONFIG="{\"kv_connector\":\"ExampleConnector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{\"shared_storage_path\":\"${KV_SHARED_PATH}\"}}"

cat <<INFO
==> vLLM instance ${INSTANCE}, sharing a key-value store
    model             ${MODEL}
    served as         ${SERVED_NAME}
    listening on      http://${SERVE_HOST}:${PORT}
    shared store      ${KV_SHARED_PATH}
    kv role           kv_both
    max model length  ${CTX}
    memory fraction   ${MEM_FRACTION}

    Both instances must load the SAME model at the SAME context length. Two instances
    that disagree about either produce blocks the other cannot use, and the symptom is
    not an error: it is that nothing is ever reused and the store fills up.

    The measurement is measure-reuse.py with BOTH endpoints given, so that turns
    alternate. A fast second turn on the other endpoint is the store doing its job.

    The store is a directory of key-value blocks: it is the users' prompts in the
    model's own representation. Put it where you would put the transcripts, and delete
    it in cleanup.

INFO

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