#!/bin/bash
set -e

# vLLM + Mooncake PD 分离单节点启动脚本
# 104 上前 4 张卡做 prefill，后 4 张卡做 decode
# 使用 TCP 传输（因为当前 IB/RDMA 未启用）

MODEL_PATH="/data/models/DeepSeek-V4-Flash"
BOOTSTRAP_PORT=8998
PREFILL_PORT=8000
DECODE_PORT=8001
ROUTER_PORT=30000
HOST="0.0.0.0"

VENV="/data/user1/yy/envs/vllm"
PYTHON="$VENV/bin/python"
VLLM="$VENV/bin/vllm"
VLLM_ROUTER="$VENV/bin/vllm-router"
LOG_DIR="/data/user1/yy/vllm_pd_logs"
mkdir -p "$LOG_DIR"

echo "=== Starting vLLM PD disaggregation (single node) ==="
echo "Model: $MODEL_PATH"
echo "Prefill GPUs: 0-3, port $PREFILL_PORT"
echo "Decode GPUs: 4-7, port $DECODE_PORT"
echo "Router port: $ROUTER_PORT"
echo "Logs: $LOG_DIR"
echo ""

# Clean up old logs
rm -f "$LOG_DIR"/*.log

# 使用 TCP 协议（IB 没起来，不能用 rdma）
export VLLM_MOONCAKE_PROTOCOL=tcp
export VLLM_MOONCAKE_BOOTSTRAP_PORT=$BOOTSTRAP_PORT

# Common args
PREFILL_KV_CONFIG='{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"mooncake_protocol":"tcp"}}'
DECODE_KV_CONFIG='{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail","kv_buffer_device":"cuda","kv_connector_extra_config":{"enforce_handshake_compat":false,"mooncake_protocol":"tcp"}}'

# Start prefill server (GPUs 0-3)
echo "[1/3] Starting prefill server on GPUs 0-3..."
CUDA_VISIBLE_DEVICES=0,1,2,3 nohup "$VLLM" serve "$MODEL_PATH" \
    --trust-remote-code \
    --kv-cache-dtype fp8 \
    --block-size 256 \
    --port "$PREFILL_PORT" \
    --data-parallel-size 4 \
    --enable-expert-parallel \
    --tokenizer-mode deepseek_v4 \
    --reasoning-parser deepseek_v4 \
    --max-model-len auto \
    --max-num-batched-tokens 16384 \
    --max-num-seqs 8 \
    --enforce-eager \
    --no-disable-hybrid-kv-cache-manager \
    --disable-uvicorn-access-log \
    --kv-transfer-config "$PREFILL_KV_CONFIG" \
    > "$LOG_DIR/prefill.log" 2>&1 &
PREFILL_PID=$!
echo "Prefill PID: $PREFILL_PID"

# Wait for prefill health
for i in {1..120}; do
    if curl -s "http://127.0.0.1:$PREFILL_PORT/health" > /dev/null 2>&1; then
        echo "Prefill server is ready"
        break
    fi
    if ! kill -0 $PREFILL_PID 2>/dev/null; then
        echo "ERROR: Prefill server exited early"
        tail -80 "$LOG_DIR/prefill.log"
        exit 1
    fi
    echo "Waiting for prefill server... ($i/120)"
    sleep 5
done

# Start decode server (GPUs 4-7)
echo "[2/3] Starting decode server on GPUs 4-7..."
CUDA_VISIBLE_DEVICES=4,5,6,7 nohup "$VLLM" serve "$MODEL_PATH" \
    --trust-remote-code \
    --kv-cache-dtype fp8 \
    --block-size 256 \
    --port "$DECODE_PORT" \
    --data-parallel-size 4 \
    --enable-expert-parallel \
    --tokenizer-mode deepseek_v4 \
    --reasoning-parser deepseek_v4 \
    --max-model-len auto \
    --max-num-seqs 512 \
    --max-num-batched-tokens 512 \
    --compilation-config '{"mode":0,"cudagraph_mode":"FULL_DECODE_ONLY","max_cudagraph_capture_size":512,"compile_ranges_endpoints":[512]}' \
    --no-disable-hybrid-kv-cache-manager \
    --disable-uvicorn-access-log \
    --kv-transfer-config "$DECODE_KV_CONFIG" \
    > "$LOG_DIR/decode.log" 2>&1 &
DECODE_PID=$!
echo "Decode PID: $DECODE_PID"

# Wait for decode health
for i in {1..120}; do
    if curl -s "http://127.0.0.1:$DECODE_PORT/health" > /dev/null 2>&1; then
        echo "Decode server is ready"
        break
    fi
    if ! kill -0 $DECODE_PID 2>/dev/null; then
        echo "ERROR: Decode server exited early"
        tail -80 "$LOG_DIR/decode.log"
        exit 1
    fi
    echo "Waiting for decode server... ($i/120)"
    sleep 5
done

# Start vllm-router
echo "[3/3] Starting vllm-router on port $ROUTER_PORT..."
nohup "$VLLM_ROUTER" \
    --policy round_robin \
    --vllm-pd-disaggregation \
    --prefill "http://127.0.0.1:$PREFILL_PORT" \
    --decode "http://127.0.0.1:$DECODE_PORT" \
    --host 127.0.0.1 \
    --port "$ROUTER_PORT" \
    --intra-node-data-parallel-size 4 \
    --kv-connector mooncake \
    > "$LOG_DIR/router.log" 2>&1 &
ROUTER_PID=$!
echo "Router PID: $ROUTER_PID"

# Wait for router health
for i in {1..60}; do
    if curl -s "http://127.0.0.1:$ROUTER_PORT/health" > /dev/null 2>&1; then
        echo "Router is ready"
        break
    fi
    if ! kill -0 $ROUTER_PID 2>/dev/null; then
        echo "ERROR: Router exited early"
        tail -50 "$LOG_DIR/router.log"
        exit 1
    fi
    echo "Waiting for router... ($i/60)"
    sleep 2
done

echo ""
echo "=== vLLM PD disaggregation is ready ==="
echo "Router URL: http://127.0.0.1:$ROUTER_PORT"
echo "Prefill URL: http://127.0.0.1:$PREFILL_PORT"
echo "Decode URL: http://127.0.0.1:$DECODE_PORT"
echo ""
echo "Test command:"
echo "  curl http://127.0.0.1:$ROUTER_PORT/v1/completions \\"
echo "    -H \"Content-Type: application/json\" \\"
echo "    -d '{\"model\":\"DeepSeek-V4-Flash\",\"prompt\":\"Hello\",\"max_tokens\":32}'"
echo ""
echo "To stop: kill $PREFILL_PID $DECODE_PID $ROUTER_PID"