#!/bin/bash
# PD-Disaggregated serving: 1 prefill (TP=4, GPUs 0-3) + 1 decode (TP=4, GPUs 4-7)
# Uses vLLM 0.18.1's P2pNcclConnector + XpYd proxy.
#
# Architecture:
#   Client → Proxy (port 10001)
#     → Prefill (port 20003, kv_port 21001) [max_tokens=1, does prefill + KV push]
#     → Decode  (port 20005, kv_port 22001) [full generation, KV pulled from prefill]
#
# Usage: bash scripts/launch_pd_separated.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=30001      # ZMQ service discovery
CLIENT_PORT=10001     # HTTP proxy for clients
PREFILL_PORT=20003
DECODE_PORT=20005
KV_PORT_P=21001
KV_PORT_D=22001

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

echo "=== PD-Disaggregated vLLM 0.18.1 ==="
echo "  Model:   $MODEL_PATH"
echo "  Prefill: GPUs 0-3 (TP=4), port $PREFILL_PORT, kv_port $KV_PORT_P"
echo "  Decode:  GPUs 4-7 (TP=4), port $DECODE_PORT, kv_port $KV_PORT_D"
echo "  Proxy:   ZMQ=$PROXY_PORT, HTTP=$CLIENT_PORT"
echo ""

# Step 1: Start proxy FIRST (P/D instances register via ZMQ)
echo "[1/3] Starting proxy..."
PROXY_SCRIPT="$PROJECT_DIR/third_party/vllm/examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_proxy_p2p_nccl_xpyd.py"
$VENV/python "$PROXY_SCRIPT" &
PROXY_PID=$!
sleep 2
echo "  Proxy PID=$PROXY_PID"

# Step 2: Start prefill instance (KV producer)
echo "[2/3] Starting prefill instance..."
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\":\"P2pNcclConnector\",\"kv_role\":\"kv_producer\",\"kv_buffer_size\":\"1e1\",\"kv_port\":\"$KV_PORT_P\",\"kv_connector_extra_config\":{\"proxy_ip\":\"127.0.0.1\",\"proxy_port\":\"$PROXY_PORT\",\"http_port\":\"$PREFILL_PORT\",\"send_type\":\"PUT_ASYNC\",\"nccl_num_channels\":\"16\"}}" &
PREFILL_PID=$!
echo "  Prefill PID=$PREFILL_PID"

# Step 3: Start decode instance (KV consumer)
echo "[3/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\":\"P2pNcclConnector\",\"kv_role\":\"kv_consumer\",\"kv_buffer_size\":\"8e9\",\"kv_port\":\"$KV_PORT_D\",\"kv_connector_extra_config\":{\"proxy_ip\":\"127.0.0.1\",\"proxy_port\":\"$PROXY_PORT\",\"http_port\":\"$DECODE_PORT\",\"send_type\":\"PUT_ASYNC\",\"nccl_num_channels\":\"16\"}}" &
DECODE_PID=$!
echo "  Decode PID=$DECODE_PID"

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

echo ""
echo "=== All ready ==="
echo "  Send requests to: http://localhost:$CLIENT_PORT"
echo ""

wait