#!/bin/bash
set -e

# sglang PD 分离单节点启动脚本
# 104 上前 4 张卡做 prefill，后 4 张卡做 decode
# 使用 NIXL/UCX 作为传输后端（Mooncake 无法安装）

MODEL_PATH="/data/models/DeepSeek-V4-Flash"
BOOTSTRAP_PORT=30010
PREFILL_PORT=30001
DECODE_PORT=30002
LB_PORT=30000
HOST="0.0.0.0"

VENV="/data/user1/yy/envs/sglang"
PYTHON="$VENV/bin/python"
LOG_DIR="/data/user1/yy/pd_logs"
mkdir -p "$LOG_DIR"

# NIXL/UCX 配置：单节点内优先用 CUDA IPC / shared memory，禁用 IB
export SGLANG_DISAGGREGATION_NIXL_BACKEND="UCX"
export UCX_TLS="self,sm,cuda_copy,cuda_ipc,tcp"
export UCX_NET_DEVICES=""

echo "=== Starting sglang 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 "Load balancer port: $LB_PORT"
echo "Logs: $LOG_DIR"
echo ""

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

# Common args for both prefill and decode
COMMON_ARGS=(
    --trust-remote-code
    --model-path "$MODEL_PATH"
    --tp 4
    --moe-runner-backend marlin
    --disable-cuda-graph
    --disaggregation-bootstrap-port "$BOOTSTRAP_PORT"
    --disaggregation-transfer-backend nixl
    --host "$HOST"
)

# Start prefill server (GPUs 0-3)
echo "[1/3] Starting prefill server on GPUs 0-3..."
CUDA_VISIBLE_DEVICES=0,1,2,3 nohup "$PYTHON" -m sglang.launch_server \
    "${COMMON_ARGS[@]}" \
    --disaggregation-mode prefill \
    --port "$PREFILL_PORT" \
    > "$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 -50 "$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=0,1,2,3 nohup "$PYTHON" -m sglang.launch_server \
    "${COMMON_ARGS[@]}" \
    --disaggregation-mode decode \
    --base-gpu-id 4 \
    --port "$DECODE_PORT" \
    > "$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 -50 "$LOG_DIR/decode.log"
        exit 1
    fi
    echo "Waiting for decode server... ($i/120)"
    sleep 5
done

# Start PD load balancer
echo "[3/3] Starting PD load balancer on port $LB_PORT..."
nohup "$PYTHON" -m sglang_router.launch_router \
    --pd-disaggregation \
    --mini-lb \
    --prefill "http://127.0.0.1:$PREFILL_PORT" \
    --decode "http://127.0.0.1:$DECODE_PORT" \
    --host "$HOST" \
    --port "$LB_PORT" \
    > "$LOG_DIR/lb.log" 2>&1 &
LB_PID=$!
echo "LB PID: $LB_PID"

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

echo ""
echo "=== PD disaggregation is ready ==="
echo "Load balancer URL: http://127.0.0.1:$LB_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:$LB_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 $LB_PID"