#!/bin/bash
# PD-Disaggregated serving via Mooncake (RDMA + DRAM KV pool).
#
# Architecture:
#   Client → Proxy (port 8000)
#     → Prefill (port 8010, TP=4, GPUs 0-3, bootstrap 8998)
#       [prefill + store KV to DRAM pool via RDMA]
#     → Decode  (port 8020, TP=4, GPUs 4-7)
#       [pull KV from DRAM pool via RDMA + decode]
#
# Usage: bash scripts/launch_pd_mooncake.sh

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
VENV="$PROJECT_DIR/.venv/bin"
VLLM="$VENV/vllm"

MODEL_PATH="${MODEL_PATH:-$HOME/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}"
PROXY_PORT=8000
PREFILL_PORT=8010
DECODE_PORT=8020
BOOTSTRAP_PORT=8998

PROXY_SCRIPT="$PROJECT_DIR/third_party/vllm/examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py"

trap 'echo "Cleaning up..."; kill $(jobs -p) 2>/dev/null; wait 2>/dev/null' EXIT INT TERM

echo "=== PD-Disaggregated vLLM 0.18.1 (Mooncake/RDMA) ==="
echo "  Model:     $MODEL_PATH"
echo "  Prefill:   GPUs 0-3 (TP=4), port $PREFILL_PORT, bootstrap $BOOTSTRAP_PORT"
echo "  Decode:    GPUs 4-7 (TP=4), port $DECODE_PORT"
echo "  Proxy:     port $PROXY_PORT"
echo ""

# Step 1: Start prefill instance (KV producer)
echo "[1/3] Starting prefill instance..."
VLLM_MOONCAKE_BOOTSTRAP_PORT=$BOOTSTRAP_PORT \
CUDA_VISIBLE_DEVICES=0,1,2,3 \
$VLLM serve "$MODEL_PATH" \
    --host 0.0.0.0 \
    --port $PREFILL_PORT \
    --tensor-parallel-size 4 \
    --trust-remote-code \
    --enable-prefix-caching \
    --dtype auto \
    --gpu-memory-utilization 0.9 \
    --kv-transfer-config \
    '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer"}' &
PREFILL_PID=$!
echo "  Prefill PID=$PREFILL_PID"

# Step 2: Start decode instance (KV consumer)
echo "[2/3] Starting decode instance..."
CUDA_VISIBLE_DEVICES=4,5,6,7 \
$VLLM serve "$MODEL_PATH" \
    --host 0.0.0.0 \
    --port $DECODE_PORT \
    --tensor-parallel-size 4 \
    --trust-remote-code \
    --enable-prefix-caching \
    --dtype auto \
    --gpu-memory-utilization 0.8 \
    --kv-transfer-config \
    '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer"}' &
DECODE_PID=$!
echo "  Decode PID=$DECODE_PID"

# Wait for both instances
echo ""
echo "Waiting for instances..."
timeout 1200 bash -c "until curl -s localhost:$PREFILL_PORT/v1/models > /dev/null 2>&1; do sleep 5; done"
echo "  Prefill ready!"
timeout 1200 bash -c "until curl -s localhost:$DECODE_PORT/v1/models > /dev/null 2>&1; do sleep 5; done"
echo "  Decode ready!"

# Step 3: Start proxy (after instances are ready)
echo "[3/3] Starting proxy..."
$VENV/python "$PROXY_SCRIPT" \
    --prefill "http://127.0.0.1:$PREFILL_PORT" "$BOOTSTRAP_PORT" \
    --decode "http://127.0.0.1:$DECODE_PORT" \
    --host 0.0.0.0 \
    --port $PROXY_PORT &
PROXY_PID=$!
echo "  Proxy PID=$PROXY_PID"

sleep 5
echo ""
echo "=== All ready ==="
echo "  Send requests to: http://localhost:$PROXY_PORT"
echo ""

wait