139 lines
3.9 KiB
Bash
Executable File
139 lines
3.9 KiB
Bash
Executable File
#!/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"
|