@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
sh 140 lines 3.6 KB
Raw
1 #!/bin/bash
2
3 echo "Warning: LMCache disaggregated prefill support for vLLM v1 is experimental and subject to change."
4
5
6 PIDS=()
7
8 # Switch to the directory of the current script
9 cd "$(dirname "${BASH_SOURCE[0]}")"
10
11 check_hf_token() {
12 if [ -z "$HF_TOKEN" ]; then
13 echo "HF_TOKEN is not set. Please set it to your Hugging Face token."
14 exit 1
15 fi
16 if [[ "$HF_TOKEN" != hf_* ]]; then
17 echo "HF_TOKEN is not a valid Hugging Face token. Please set it to your Hugging Face token."
18 exit 1
19 fi
20 echo "HF_TOKEN is set and valid."
21 }
22
23 check_num_gpus() {
24 # can you check if the number of GPUs are >=2 via nvidia-smi/rocm-smi?
25 if ! which rocm-smi > /dev/null 2>&1; then
26 num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
27 else
28 num_gpus=$(rocm-smi --showid | grep -c Instinct)
29 fi
30
31 if [ "$num_gpus" -lt 2 ]; then
32 echo "You need at least 2 GPUs to run disaggregated prefill."
33 exit 1
34 else
35 echo "Found $num_gpus GPUs."
36 fi
37 }
38
39 ensure_python_library_installed() {
40 echo "Checking if $1 is installed..."
41 if ! python3 -c "import $1" > /dev/null 2>&1; then
42 if [ "$1" == "nixl" ]; then
43 echo "$1 is not installed. Please refer to https://github.com/ai-dynamo/nixl for installation."
44 else
45 echo "$1 is not installed. Please install it via pip install $1."
46 fi
47 exit 1
48 else
49 echo "$1 is installed."
50 fi
51 }
52
53 cleanup() {
54 echo "Stopping everything…"
55 trap - INT TERM # prevent re-entrancy
56 kill -- -$$ # negative PID == “this whole process-group”
57 wait # reap children so we don't leave zombies
58 exit 0
59 }
60
61 wait_for_server() {
62 local port=$1
63 local timeout_seconds=1200
64 local start_time=$(date +%s)
65
66 echo "Waiting for server on port $port..."
67
68 while true; do
69 if curl -s "localhost:${port}/v1/completions" > /dev/null; then
70 return 0
71 fi
72
73 local now=$(date +%s)
74 if (( now - start_time >= timeout_seconds )); then
75 echo "Timeout waiting for server"
76 return 1
77 fi
78
79 sleep 1
80 done
81 }
82
83
84 main() {
85 check_hf_token
86 check_num_gpus
87 ensure_python_library_installed lmcache
88 ensure_python_library_installed nixl
89 ensure_python_library_installed pandas
90 ensure_python_library_installed datasets
91 ensure_python_library_installed vllm
92
93 trap cleanup INT
94 trap cleanup USR1
95 trap cleanup TERM
96
97 echo "Launching prefiller, decoder and proxy..."
98 echo "Please check prefiller.log, decoder.log and proxy.log for logs."
99
100 bash disagg_vllm_launcher.sh prefiller \
101 > >(tee prefiller.log) 2>&1 &
102 prefiller_pid=$!
103 PIDS+=("$prefiller_pid")
104
105 bash disagg_vllm_launcher.sh decoder \
106 > >(tee decoder.log) 2>&1 &
107 decoder_pid=$!
108 PIDS+=("$decoder_pid")
109
110 python3 disagg_proxy_server.py \
111 --host localhost \
112 --port 9000 \
113 --prefiller-host localhost \
114 --prefiller-port 8100 \
115 --decoder-host localhost \
116 --decoder-port 8200 \
117 > >(tee proxy.log) 2>&1 &
118 proxy_pid=$!
119 PIDS+=("$proxy_pid")
120
121 wait_for_server 8100
122 wait_for_server 8200
123 wait_for_server 9000
124
125 echo "All servers are up. Starting benchmark..."
126
127 # begin benchmark
128 cd ../../../../benchmarks/
129 vllm bench serve --port 9000 --seed "$(date +%s)" \
130 --model meta-llama/Llama-3.1-8B-Instruct \
131 --dataset-name random --random-input-len 7500 --random-output-len 200 \
132 --num-prompts 200 --burstiness 100 --request-rate 3.6 | tee benchmark.log
133
134 echo "Benchmarking done. Cleaning up..."
135
136 cleanup
137
138 }
139
140 main