#!/usr/bin/env bash
# Purpose: start the vLLM instance that reads prompts and writes their key-value blocks
#          out through the chosen connector, printing the whole configuration first so the
#          run is reproducible from the terminal log alone
# 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; for CONNECTOR=shared, KV_SHARED_PATH exists and is writable from
#          BOTH machines; nothing else listening on PREFILL_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}"
PREFILL_PORT="${PREFILL_PORT:-8100}"
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:-}"
PREFILL_SIDE_CHANNEL_PORT="${PREFILL_SIDE_CHANNEL_PORT:-5600}"
CLUSTER_IFACE="${CLUSTER_IFACE:-}"
SPLIT="${SPLIT_MEM_FRACTION:-0.40}"
SINGLE_MACHINE="${SINGLE_MACHINE:-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"
    printf '    NOTE: single-machine path. Using %s at memory fraction %s so that two\n' "$MODEL" "$MEM_FRACTION"
    printf '          engine processes fit on one device. Record this in the notebook: the\n'
    printf '          transfer crosses the loopback interface, which measures no network.\n\n'
fi

# --- build the connector configuration ------------------------------------------------
case "$CONNECTOR" in
    shared)
        [ -n "$KV_SHARED_PATH" ] || fail "CONNECTOR=shared needs KV_SHARED_PATH set to a directory both machines can write."
        mkdir -p "$KV_SHARED_PATH"
        [ -w "$KV_SHARED_PATH" ] || fail "KV_SHARED_PATH ($KV_SHARED_PATH) is not writable by this account."
        KV_CONFIG="{\"kv_connector\":\"ExampleConnector\",\"kv_role\":\"kv_producer\",\"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_producer"}'
        export VLLM_NIXL_SIDE_CHANNEL_HOST="$SIDE_CHANNEL_ADDR"
        export VLLM_NIXL_SIDE_CHANNEL_PORT="$PREFILL_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_producer"}'
        ;;
    *)
        fail "CONNECTOR must be one of: shared, nixl, mooncake. Got '${CONNECTOR}'."
        ;;
esac

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

    Two things to copy into the notebook from the startup log: the key-value cache size
    the engine settled on, and the maximum concurrency that implies at this context
    length. The decode instance must be started with the SAME model and the SAME context
    length, or its blocks will not match these and every request will silently re-prefill.

    --kv-transfer-config is documented on vLLM's disaggregated prefilling page. If your
    build rejects it, run "vllm serve --help | grep kv" and record what it does accept.

INFO

exec vllm serve "$MODEL" \
    --host "$SERVE_HOST" \
    --port "$PREFILL_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"