Skip to content

Commit 95f5574

Browse files
MengqingCaoMengqingCao
committed
[CI/UT][PD Disaggreate] Initialize PD Disaggreate UT
Signed-off-by: MengqingCao <cmq0113@163.com>
1 parent a0c3e9b commit 95f5574

File tree

9 files changed

+341
-16
lines changed

9 files changed

+341
-16
lines changed

‎.github/workflows/vllm_ascend_test_pd.yaml‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -50,12 +50,6 @@ jobs:
5050
options: >-
5151
--device /dev/davinci0
5252
--device /dev/davinci1
53-
--device /dev/davinci2
54-
--device /dev/davinci3
55-
--device /dev/davinci4
56-
--device /dev/davinci5
57-
--device /dev/davinci6
58-
--device /dev/davinci7
5953
--device /dev/davinci_manager
6054
--device /dev/devmm_svm
6155
--device /dev/hisi_hdc
@@ -100,3 +94,8 @@ jobs:
10094
run: |
10195
pip install -r requirements-dev.txt
10296
pip install -v -e .
97+
98+
- name: Run vllm-project/vllm-ascend PD Disaggregation test
99+
if: github.event_name == 'schedule' || github.event.label.name == 'module:pd'
100+
run: |
101+
pytest -sv tests/e2e/pd_disaggreate/test_pd_e2e.py

‎examples/disaggregated_prefill/disaggregated_prefill_offline.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from multiprocessing import Event, Process
1414

1515
kv_connector_extra_config = {
16-
"prompt_device_ips": ["1.2.3.1", "1.2.3.2"],
16+
"prefill_device_ips": ["1.2.3.1", "1.2.3.2"],
1717
"decode_device_ips": ["1.2.3.9", "1.2.3.10"],
1818
"llmdatadist_comm_port": 26000,
1919
}

‎examples/disaggregated_prefill/p2p_disaggrefated_prefill_proxy.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,13 @@ async def handle_request():
181181

182182

183183
if __name__ == "__main__":
184-
t = start_service_discovery("0.0.0.0", 30001)
185-
app.run(host="0.0.0.0", port=10001)
184+
import argparse
185+
parser = argparse.ArgumentParser(
186+
description="args of disaggregated-prefill proxy")
187+
parser.add_argument("--http-port", type=int, default=10001)
188+
parser.add_argument("--register-port", type=int, default=10002)
189+
args = parser.parse_args()
190+
191+
t = start_service_discovery("0.0.0.0", args.register_port)
192+
app.run(host="0.0.0.0", port=args.http_port)
186193
t.join()

‎packages.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
11
git
22
vim
33
wget
4+
jq
5+
curl

‎requirements-dev.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,3 +10,5 @@ types-jsonschema
1010
xgrammar
1111
zmq
1212
numba
13+
quart
14+
types-psutil

‎tests/e2e/pd_disaggreate/setup_pd.sh‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
#!/bin/bash
2+
3+
#
4+
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
5+
#
6+
# Licensed under the Apache License, Version 2.0 (the "License");
7+
# you may not use this file except in compliance with the License.
8+
# You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing, software
13+
# distributed under the License is distributed on an "AS IS" BASIS,
14+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
# See the License for the specific language governing permissions and
16+
# limitations under the License.
17+
# This file is a part of the vllm-ascend project.
18+
#
19+
20+
21+
function cleanup_instances() {
22+
VLLM_PID=$(pgrep -f "vllm serve")
23+
_info "===> Try kill -2 ${VLLM_PID} to exit."
24+
kill -2 "$VLLM_PID"
25+
wait_for_exit "$VLLM_PID"
26+
}
27+
28+
29+
function run_prefill_instance() {
30+
local model_name=$1
31+
local tp_size=$2
32+
local prefill_port=$3
33+
local register_port=$4
34+
local prefill_device_ips=$5
35+
local decode_device_ips=$6
36+
37+
echo "================================"
38+
echo "Testing model: $model_name"
39+
echo "================================"
40+
# Start prefill instance
41+
42+
KV_CONFIG=$(jq -n \
43+
--arg kv_connector "AscendSimpleConnector" \
44+
--arg kv_buffer_device "npu" \
45+
--arg kv_role "kv_producer" \
46+
--argjson kv_parallel_size 8 \
47+
--arg kv_port "11001" \
48+
--argjson prefill_device_ips "$prefill_device_ips" \
49+
--argjson decode_device_ips "$decode_device_ips" \
50+
--argjson llmdatadist_comm_port "26000" \
51+
--arg proxy_ip "0.0.0.0" \
52+
--argjson proxy_port "$register_port" \
53+
--argjson http_port "$prefill_port" \
54+
'{
55+
"kv_connector": $kv_connector,
56+
"kv_buffer_device": $kv_buffer_device,
57+
"kv_role": $kv_role,
58+
"kv_parallel_size": $kv_parallel_size,
59+
"kv_port": $kv_port,
60+
"kv_connector_extra_config": {
61+
"prefill_device_ips": $prefill_device_ips,
62+
"decode_device_ips": $decode_device_ips,
63+
"llmdatadist_comm_port": $llmdatadist_comm_port,
64+
"proxy_ip": $proxy_ip,
65+
"proxy_port": $proxy_port,
66+
"http_port": $http_port
67+
}
68+
}')
69+
70+
# start prefill instance
71+
ASCEND_RT_VISIBLE_DEVICES=0 vllm serve $model_name \
72+
--host 0.0.0.0 \
73+
--port $prefill_port \
74+
--tensor-parallel-size $tp_size \
75+
--served-model-name Deepseek \
76+
--max-model-len 2000 \
77+
--trust-remote-code \
78+
--kv-transfer-config "$KV_CONFIG" &
79+
}
80+
81+
82+
83+
function run_decode_instance() {
84+
# Start decode instance
85+
local model_name=$1
86+
local tp_size=$2
87+
local decode_port=$3
88+
local register_port=$4
89+
local prefill_device_ips=$5
90+
local decode_device_ips=$6
91+
92+
KV_CONFIG=$(jq -n \
93+
--arg kv_connector "AscendSimpleConnector" \
94+
--arg kv_buffer_device "npu" \
95+
--arg kv_role "kv_consumer" \
96+
--argjson kv_parallel_size 8 \
97+
--arg kv_port "21001" \
98+
--argjson prefill_device_ips "$prefill_device_ips" \
99+
--argjson decode_device_ips "$decode_device_ips" \
100+
--argjson llmdatadist_comm_port "26000" \
101+
--arg proxy_ip "0.0.0.0" \
102+
--argjson proxy_port "$register_port" \
103+
--argjson http_port "$decode_port" \
104+
'{
105+
"kv_connector": $kv_connector,
106+
"kv_buffer_device": $kv_buffer_device,
107+
"kv_role": $kv_role,
108+
"kv_parallel_size": $kv_parallel_size,
109+
"kv_port": $kv_port,
110+
"kv_connector_extra_config": {
111+
"prefill_device_ips": $prefill_device_ips,
112+
"decode_device_ips": $decode_device_ips,
113+
"llmdatadist_comm_port": $llmdatadist_comm_port,
114+
"proxy_ip": $proxy_ip,
115+
"proxy_port": $proxy_port,
116+
"http_port": $http_port
117+
}
118+
}')
119+
120+
# start decode instance
121+
ASCEND_RT_VISIBLE_DEVICES=1 vllm serve $model_name \
122+
--host 0.0.0.0 \
123+
--port $decode_port \
124+
--tensor-parallel-size $tp_size \
125+
--seed 1024 \
126+
--served-model-name Deepseek \
127+
--max-model-len 2000 \
128+
--max-num-batched-tokens 2000 \
129+
--trust-remote-code \
130+
--gpu-memory-utilization 0.9 \
131+
--kv-transfer-config "$KV_CONFIG" &
132+
}
133+
134+
function run_proxy_server() {
135+
# Build the command for the proxy server with all the hosts and ports
136+
register_port=$1
137+
proxy_port=$2
138+
PROXY_CMD="python examples/disaggregated_prefill/p2p_disaggrefated_prefill_proxy.py --http-port $proxy_port --register-port $register_port"
139+
140+
# Start the proxy server
141+
echo "Starting proxy server with command: $PROXY_CMD"
142+
$PROXY_CMD &
143+
144+
# Wait for the proxy to start
145+
sleep 3
146+
}

‎tests/e2e/pd_disaggreate/test_pd_e2e.py‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
#!/bin/bash
2+
3+
#
4+
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
5+
#
6+
# Licensed under the Apache License, Version 2.0 (the "License");
7+
# you may not use this file except in compliance with the License.
8+
# You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing, software
13+
# distributed under the License is distributed on an "AS IS" BASIS,
14+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
# See the License for the specific language governing permissions and
16+
# limitations under the License.
17+
# This file is a part of the vllm-ascend project.
18+
#
19+
20+
import os
21+
import signal
22+
import subprocess
23+
import time
24+
25+
import psutil
26+
import requests
27+
28+
29+
def kill_process_and_children(pid):
30+
try:
31+
parent = psutil.Process(pid)
32+
children = parent.children(recursive=True)
33+
for child in children:
34+
print(f"Killing child process {child.pid}")
35+
child.kill()
36+
print(f"Killing parent process {pid}")
37+
parent.kill()
38+
except psutil.NoSuchProcess:
39+
pass
40+
41+
42+
def kill_all_vllm_related():
43+
current_pid = os.getpid()
44+
45+
for proc in psutil.process_iter(['pid', 'cmdline']):
46+
try:
47+
if proc.pid == current_pid:
48+
continue
49+
cmd = ' '.join(proc.info['cmdline'])
50+
if "vllm" in cmd or "proxy" in cmd or "engine_worker" in cmd:
51+
kill_process_and_children(proc.pid)
52+
except Exception:
53+
continue
54+
55+
56+
PROXY_PORT = 10102
57+
DECODE_PORT = 8002
58+
59+
SCRIPT_PATH = os.path.abspath("./tests/e2e/run_disagg_pd.sh")
60+
61+
62+
def wait_for_port(port, timeout=30):
63+
import socket
64+
start = time.time()
65+
while time.time() - start < timeout:
66+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
67+
if sock.connect_ex(("127.0.0.1", port)) == 0:
68+
return True
69+
time.sleep(1)
70+
raise TimeoutError(f"Port {port} not ready after {timeout}s")
71+
72+
73+
def start_and_test_pipeline():
74+
print("Launching bash script to run vLLM PD setup...")
75+
proc = subprocess.Popen(["bash", SCRIPT_PATH])
76+
try:
77+
print("Waiting for proxy port to be available...")
78+
wait_for_port(PROXY_PORT, 1200)
79+
wait_for_port(DECODE_PORT, 1200)
80+
81+
# request
82+
payload = {
83+
"model": "Deepseek",
84+
"prompt": "The future of AI is",
85+
"max_tokens": 64,
86+
"temperature": 0,
87+
}
88+
response = requests.post(
89+
f"http://localhost:{PROXY_PORT}/v1/completions",
90+
headers={"Content-Type": "application/json"},
91+
json=payload,
92+
timeout=10)
93+
assert response.status_code == 200, f"HTTP failed: {response.status_code}"
94+
result = response.json()
95+
print("Response:", result)
96+
assert "text" in result["choices"][0]
97+
assert len(result["choices"][0]["text"].strip()) > 0
98+
99+
finally:
100+
# clean up subprocesses
101+
print("Cleaning up subprocess...")
102+
proc.send_signal(signal.SIGINT)
103+
try:
104+
proc.wait(timeout=10)
105+
except subprocess.TimeoutExpired:
106+
proc.kill()
107+
kill_all_vllm_related()
108+
109+
110+
def test_disaggregated_pd_pipeline():
111+
start_and_test_pipeline()

‎tests/e2e/run_disagg_pd.sh‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
#!/bin/bash
2+
3+
#
4+
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
5+
#
6+
# Licensed under the Apache License, Version 2.0 (the "License");
7+
# you may not use this file except in compliance with the License.
8+
# You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing, software
13+
# distributed under the License is distributed on an "AS IS" BASIS,
14+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
# See the License for the specific language governing permissions and
16+
# limitations under the License.
17+
# This file is a part of the vllm-ascend project.
18+
#
19+
20+
set -eo errexit
21+
22+
. $(dirname "$0")/common.sh
23+
. $(dirname "$0")/pd_disaggreate/setup_pd.sh
24+
25+
export VLLM_USE_MODELSCOPE="True"
26+
27+
MODEL_NAME="deepseek-ai/DeepSeek-V2-Lite"
28+
# TODO: add tp case
29+
TP_SIZE=1
30+
31+
# TODO: support multi-card
32+
prefill_ip=$(/usr/local/Ascend/driver/tools/hccn_tool -i 0 -ip -g | grep "ipaddr" | awk -F: '{print $2}' | xargs)
33+
PREFILL_DEVICE_IPS="[\"$prefill_ip\"]"
34+
35+
decode_ip=$(/usr/local/Ascend/driver/tools/hccn_tool -i 1 -ip -g | grep "ipaddr" | awk -F: '{print $2}' | xargs)
36+
DECODE_DEVICE_IPS="[\"$decode_ip\"]"
37+
38+
_info "====> Start pd disaggregated test"
39+
REGISTER_PORT=10101
40+
PREOXY_PORT=10102
41+
run_proxy_server $REGISTER_PORT $PREOXY_PORT
42+
_info "Started pd disaggregated proxy server"
43+
44+
PREFILL_PROC_NAME="Prefill-instance"
45+
PREFILL_PORT=8001
46+
run_prefill_instance $MODEL_NAME $TP_SIZE $PREFILL_PORT $REGISTER_PORT $PREFILL_DEVICE_IPS $DECODE_DEVICE_IPS
47+
_info "Starting prefill instance"
48+
49+
wait_url_ready $PREFILL_PROC_NAME "http://localhost:${PREFILL_PORT}/v1/completions"
50+
51+
DECODE_PROC_NAME="Decode-instance"
52+
DECODE_PORT=8002
53+
run_decode_instance $MODEL_NAME $TP_SIZE $DECODE_PORT $REGISTER_PORT $PREFILL_DEVICE_IPS $DECODE_DEVICE_IPS
54+
_info "Starting decode instance"
55+
56+
wait_url_ready $DECODE_PROC_NAME "http://localhost:${DECODE_PORT}/v1/completions"
57+
58+
_info "pd disaggregated system is ready for handling request"

0 commit comments

Comments
 (0)