| 1 | #!/bin/bash |
| 2 | |
| 3 | # ============================================================================= |
| 4 | # vLLM Disaggregated Serving Script for Mooncake Connector |
| 5 | # ============================================================================= |
| 6 | # This script demonstrates disaggregated prefill and decode serving using |
| 7 | # Mooncake Connector. |
| 8 | # |
| 9 | # Configuration can be customized via environment variables: |
| 10 | # MODEL: Model to serve |
| 11 | # PREFILL_GPUS: Comma-separated GPU IDs for prefill servers |
| 12 | # DECODE_GPUS: Comma-separated GPU IDs for decode servers |
| 13 | # PREFILL_PORTS: Comma-separated ports for prefill servers |
| 14 | # BOOTSTRAP_PORTS: Bootstrap server port launched by prefill servers |
| 15 | # DECODE_PORTS: Comma-separated ports for decode servers |
| 16 | # PROXY_PORT: Proxy server port used to setup P/D disaggregated connection. |
| 17 | # TIMEOUT_SECONDS: Server startup timeout |
| 18 | # ============================================================================= |
| 19 | |
| 20 | # Configuration - can be overridden via environment variables |
| 21 | MODEL=${MODEL:-Qwen/Qwen2.5-7B-Instruct} |
| 22 | TIMEOUT_SECONDS=${TIMEOUT_SECONDS:-1200} |
| 23 | PROXY_PORT=${PROXY_PORT:-8000} |
| 24 | |
| 25 | PREFILL_GPUS=${PREFILL_GPUS:-0} |
| 26 | DECODE_GPUS=${DECODE_GPUS:-1} |
| 27 | PREFILL_PORTS=${PREFILL_PORTS:-8010} |
| 28 | BOOTSTRAP_PORTS=${BOOTSTRAP_PORTS:-8998} |
| 29 | DECODE_PORTS=${DECODE_PORTS:-8020} |
| 30 | |
| 31 | echo "Warning: Mooncake Connector support for vLLM v1 is experimental and subject to change." |
| 32 | echo "" |
| 33 | echo "Architecture Configuration:" |
| 34 | echo " Model: $MODEL" |
| 35 | echo " Prefill GPUs: $PREFILL_GPUS, Ports: $PREFILL_PORTS, Bootstrap Port:$BOOTSTRAP_PORTS" |
| 36 | echo " Decode GPUs: $DECODE_GPUS, Ports: $DECODE_PORTS" |
| 37 | echo " Proxy Port: $PROXY_PORT" |
| 38 | echo " Timeout: ${TIMEOUT_SECONDS}s" |
| 39 | echo "" |
| 40 | |
| 41 | PIDS=() |
| 42 | |
| 43 | # Switch to the directory of the current script |
| 44 | cd "$(dirname "${BASH_SOURCE[0]}")" |
| 45 | |
| 46 | check_required_files() { |
| 47 | local files=("mooncake_connector_proxy.py") |
| 48 | for file in "${files[@]}"; do |
| 49 | if [[ ! -f "$file" ]]; then |
| 50 | echo "Required file $file not found in $(pwd)" |
| 51 | exit 1 |
| 52 | fi |
| 53 | done |
| 54 | } |
| 55 | |
| 56 | check_hf_token() { |
| 57 | if [ -z "$HF_TOKEN" ]; then |
| 58 | echo "HF_TOKEN is not set. Please set it to your Hugging Face token." |
| 59 | echo "Example: export HF_TOKEN=your_token_here" |
| 60 | exit 1 |
| 61 | fi |
| 62 | if [[ "$HF_TOKEN" != hf_* ]]; then |
| 63 | echo "HF_TOKEN is not a valid Hugging Face token. Please set it to your Hugging Face token." |
| 64 | exit 1 |
| 65 | fi |
| 66 | echo "HF_TOKEN is set and valid." |
| 67 | } |
| 68 | |
| 69 | check_num_gpus() { |
| 70 | # Check if the number of GPUs are >=2 via nvidia-smi |
| 71 | num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l) |
| 72 | if [ "$num_gpus" -lt 2 ]; then |
| 73 | echo "You need at least 2 GPUs to run disaggregated prefill." |
| 74 | exit 1 |
| 75 | else |
| 76 | echo "Found $num_gpus GPUs." |
| 77 | fi |
| 78 | } |
| 79 | |
| 80 | ensure_python_library_installed() { |
| 81 | echo "Checking if $1 is installed..." |
| 82 | if ! python3 -c "import $1" > /dev/null 2>&1; then |
| 83 | echo "$1 is not installed. Please install it via pip install $1." |
| 84 | exit 1 |
| 85 | else |
| 86 | echo "$1 is installed." |
| 87 | fi |
| 88 | } |
| 89 | |
| 90 | cleanup() { |
| 91 | echo "Stopping everything…" |
| 92 | trap - INT TERM # prevent re-entrancy |
| 93 | pkill -9 -f "mooncake_connector_proxy.py" |
| 94 | kill -- -$$ # negative PID == "this whole process-group" |
| 95 | wait # reap children so we don't leave zombies |
| 96 | exit 0 |
| 97 | } |
| 98 | |
| 99 | wait_for_server() { |
| 100 | local port=$1 |
| 101 | local timeout_seconds=$TIMEOUT_SECONDS |
| 102 | local start_time=$(date +%s) |
| 103 | |
| 104 | echo "Waiting for server on port $port..." |
| 105 | |
| 106 | while true; do |
| 107 | if curl -s "localhost:${port}/v1/completions" > /dev/null; then |
| 108 | echo "Server on port $port is ready." |
| 109 | return 0 |
| 110 | fi |
| 111 | |
| 112 | local now=$(date +%s) |
| 113 | if (( now - start_time >= timeout_seconds )); then |
| 114 | echo "Timeout waiting for server on port $port" |
| 115 | return 1 |
| 116 | fi |
| 117 | |
| 118 | sleep 1 |
| 119 | done |
| 120 | } |
| 121 | |
| 122 | main() { |
| 123 | check_required_files |
| 124 | check_hf_token |
| 125 | check_num_gpus |
| 126 | ensure_python_library_installed vllm |
| 127 | ensure_python_library_installed mooncake.engine |
| 128 | |
| 129 | trap cleanup INT |
| 130 | trap cleanup USR1 |
| 131 | trap cleanup TERM |
| 132 | |
| 133 | echo "Launching disaggregated serving components..." |
| 134 | echo "Please check the log files for detailed output:" |
| 135 | echo " - prefill*.log: Prefill server logs" |
| 136 | echo " - decode*.log: Decode server logs" |
| 137 | echo " - proxy.log: Proxy server log" |
| 138 | |
| 139 | # Parse GPU and port arrays |
| 140 | IFS=',' read -ra PREFILL_GPU_ARRAY <<< "$PREFILL_GPUS" |
| 141 | IFS=',' read -ra DECODE_GPU_ARRAY <<< "$DECODE_GPUS" |
| 142 | IFS=',' read -ra PREFILL_PORT_ARRAY <<< "$PREFILL_PORTS" |
| 143 | IFS=',' read -ra BOOTSTRAP_PORT_ARRAY <<< "$BOOTSTRAP_PORTS" |
| 144 | IFS=',' read -ra DECODE_PORT_ARRAY <<< "$DECODE_PORTS" |
| 145 | |
| 146 | proxy_args=() |
| 147 | |
| 148 | # ============================================================================= |
| 149 | # Launch Prefill Servers (X Producers) |
| 150 | # ============================================================================= |
| 151 | echo "" |
| 152 | echo "Starting ${#PREFILL_GPU_ARRAY[@]} prefill server(s)..." |
| 153 | for i in "${!PREFILL_GPU_ARRAY[@]}"; do |
| 154 | local gpu_id=${PREFILL_GPU_ARRAY[$i]} |
| 155 | local port=${PREFILL_PORT_ARRAY[$i]} |
| 156 | local bootstrap_port=${BOOTSTRAP_PORT_ARRAY[$i]} |
| 157 | |
| 158 | echo " Prefill server $((i+1)): GPU $gpu_id, Port $port, Bootstrap Port $bootstrap_port" |
| 159 | VLLM_MOONCAKE_BOOTSTRAP_PORT=$bootstrap_port CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \ |
| 160 | --port "$port" \ |
| 161 | --kv-transfer-config \ |
| 162 | "{\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_producer\"}" > prefill$((i+1)).log 2>&1 & |
| 163 | PIDS+=($!) |
| 164 | proxy_args+=(--prefill "http://0.0.0.0:${port}" "$bootstrap_port") |
| 165 | done |
| 166 | |
| 167 | # ============================================================================= |
| 168 | # Launch Decode Servers (Y Decoders) |
| 169 | # ============================================================================= |
| 170 | echo "" |
| 171 | echo "Starting ${#DECODE_GPU_ARRAY[@]} decode server(s)..." |
| 172 | for i in "${!DECODE_GPU_ARRAY[@]}"; do |
| 173 | local gpu_id=${DECODE_GPU_ARRAY[$i]} |
| 174 | local port=${DECODE_PORT_ARRAY[$i]} |
| 175 | |
| 176 | echo " Decode server $((i+1)): GPU $gpu_id, Port $port" |
| 177 | CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \ |
| 178 | --port "$port" \ |
| 179 | --kv-transfer-config \ |
| 180 | "{\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_consumer\"}" > decode$((i+1)).log 2>&1 & |
| 181 | PIDS+=($!) |
| 182 | proxy_args+=(--decode "http://0.0.0.0:${port}") |
| 183 | done |
| 184 | |
| 185 | # ============================================================================= |
| 186 | # Launch Proxy Server |
| 187 | # ============================================================================= |
| 188 | echo "" |
| 189 | echo "Starting proxy server on port $PROXY_PORT..." |
| 190 | python3 mooncake_connector_proxy.py "${proxy_args[@]}" --port "$PROXY_PORT" > proxy.log 2>&1 & |
| 191 | PIDS+=($!) |
| 192 | |
| 193 | # ============================================================================= |
| 194 | # Wait for All Servers to Start |
| 195 | # ============================================================================= |
| 196 | echo "" |
| 197 | echo "Waiting for all servers to start..." |
| 198 | for port in "${PREFILL_PORT_ARRAY[@]}" "${DECODE_PORT_ARRAY[@]}"; do |
| 199 | if ! wait_for_server "$port"; then |
| 200 | echo "Failed to start server on port $port" |
| 201 | cleanup |
| 202 | # shellcheck disable=SC2317 |
| 203 | exit 1 |
| 204 | fi |
| 205 | done |
| 206 | |
| 207 | echo "" |
| 208 | echo "All servers are up. Starting benchmark..." |
| 209 | |
| 210 | # ============================================================================= |
| 211 | # Run Benchmark |
| 212 | # ============================================================================= |
| 213 | vllm bench serve --port "$PROXY_PORT" --seed "$(date +%s)" \ |
| 214 | --backend vllm --model "$MODEL" \ |
| 215 | --dataset-name random --random-input-len 7500 --random-output-len 200 \ |
| 216 | --num-prompts 200 --burstiness 100 --request-rate 2 | tee benchmark.log |
| 217 | |
| 218 | echo "Benchmarking done. Cleaning up..." |
| 219 | |
| 220 | cleanup |
| 221 | } |
| 222 | |
| 223 | main |