[CI/UT][PD Disaggreate] Initialize PD Disaggreate UT

requirements-dev.txt

@@ -9,3 +9,4 @@ ray
 types-jsonschema
 xgrammar
 zmq
+quart

tests/e2e/pd_disaggreate/setup_pd.sh

@@ -0,0 +1,138 @@
+#!/bin/bash
+
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+
+function cleanup_instances() {
+  VLLM_PID=$(pgrep -f "vllm serve")
+  _info "===> Try kill -2 ${VLLM_PID} to exit."
+  kill -2 "$VLLM_PID"
+  wait_for_exit "$VLLM_PID"
+}
+
+
+function run_prefill_instance() {
+  local model_name=$1
+  local tp_size=$2
+  local prefill_port=$3
+  local register_port=$4
+
+  echo "================================"
+  echo "Testing model: $model_name"
+  echo "================================"
+  # Start prefill instance
+
+  KV_CONFIG=$(jq -n \
+    --arg kv_connector "AscendSimpleConnector" \
+    --arg kv_buffer_device "npu" \
+    --arg kv_role "kv_producer" \
+    --argjson kv_parallel_size 8 \
+    --arg kv_port "11001" \
+    --argjson llmdatadist_comm_port "26000" \
+    --arg proxy_ip "0.0.0.0" \
+    --argjson proxy_port "$register_port" \
+    --argjson http_port "$prefill_port" \
+    '{
+      "kv_connector": $kv_connector,
+      "kv_buffer_device": $kv_buffer_device,
+      "kv_role": $kv_role,
+      "kv_parallel_size": $kv_parallel_size,
+      "kv_port": $kv_port,
+      "kv_connector_extra_config": {
+        "prompt_device_ips": ["29.7.130.29"],
+        "decode_device_ips": ["29.7.186.66"],
+        "llmdatadist_comm_port": $llmdatadist_comm_port,
+        "proxy_ip": $proxy_ip,
+        "proxy_port": $proxy_port,
+        "http_port": $http_port
+      }
+    }')
+
+  # start prefill instance
+  ASCEND_RT_VISIBLE_DEVICES=0 vllm serve $model_name \
+  --host 0.0.0.0 \
+  --port $prefill_port \
+  --tensor-parallel-size $tp_size \
+  --served-model-name Deepseek \
+  --max-model-len 2000 \
+  --trust-remote-code \
+  --kv-transfer-config "$KV_CONFIG" &
+}
+
+
+
+function run_decode_instance() {
+  # Start decode instance
+  local model_name=$1
+  local tp_size=$2
+  local decode_port=$3
+  local register_port=$4
+
+  KV_CONFIG=$(jq -n \
+    --arg kv_connector "AscendSimpleConnector" \
+    --arg kv_buffer_device "npu" \
+    --arg kv_role "kv_consumer" \
+    --argjson kv_parallel_size 8 \
+    --arg kv_port "21001" \
+    --argjson llmdatadist_comm_port "26000" \
+    --arg proxy_ip "0.0.0.0" \
+    --argjson proxy_port "$register_port" \
+    --argjson http_port "$decode_port" \
+    '{
+      "kv_connector": $kv_connector,
+      "kv_buffer_device": $kv_buffer_device,
+      "kv_role": $kv_role,
+      "kv_parallel_size": $kv_parallel_size,
+      "kv_port": $kv_port,
+      "kv_connector_extra_config": {
+        "prompt_device_ips": ["29.7.130.29"],
+        "decode_device_ips": ["29.7.186.66"],
+        "llmdatadist_comm_port": $llmdatadist_comm_port,
+        "proxy_ip": $proxy_ip,
+        "proxy_port": $proxy_port,
+        "http_port": $http_port
+      }
+    }')
+
+  # start decode instance
+  ASCEND_RT_VISIBLE_DEVICES=1 vllm serve $model_name \
+    --host 0.0.0.0 \
+    --port $decode_port \
+    --tensor-parallel-size $tp_size \
+    --seed 1024 \
+    --served-model-name Deepseek \
+    --max-model-len 2000 \
+    --max-num-batched-tokens 2000 \
+    --trust-remote-code \
+    --gpu-memory-utilization 0.9 \
+    --kv-transfer-config "$KV_CONFIG" &
+}
+
+function run_proxy_server() {
+  # Build the command for the proxy server with all the hosts and ports
+  register_port=$1
+  proxy_port=$2
+  PROXY_CMD="python examples/disaggregated_prefill/p2p_disaggrefated_prefill_proxy.py --http-port $proxy_port --register-port $register_port"
+
+  # Start the proxy server
+  echo "Starting proxy server with command: $PROXY_CMD"
+  $PROXY_CMD &
+
+  # Wait for the proxy to start
+  sleep 3
+}

tests/e2e/pd_disaggreate/test_pd_e2e.py

@@ -0,0 +1,60 @@
+import os
+import signal
+import subprocess
+import time
+
+import requests
+
+PROXY_PORT = 8192
+REGISTER_PORT = 8193
+
+SCRIPT_PATH = os.path.abspath("./tests/e2e/run_disagg_pd.sh")
+
+
+def wait_for_port(port, timeout=30):
+    import socket
+    start = time.time()
+    while time.time() - start < timeout:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            if sock.connect_ex(("127.0.0.1", port)) == 0:
+                return True
+        time.sleep(1)
+    raise TimeoutError(f"Port {port} not ready after {timeout}s")
+
+
+def start_and_test_pipeline():
+    print("Launching bash script to run vLLM PD setup...")
+    proc = subprocess.Popen(["bash", SCRIPT_PATH])
+    try:
+        print("Waiting for proxy port to be available...")
+        wait_for_port(PROXY_PORT, 1200)
+
+        # request
+        prompt = "The future of AI is"
+        payload = {
+            "model": "Deepseek/DeepSeek-V2-Lite-Chat",
+            "prompt": prompt,
+            "max_tokens": 64,
+            "temperature": 0,
+        }
+        response = requests.post(f"http://localhost:{PROXY_PORT}/generate",
+                                 json=payload,
+                                 timeout=10)
+        assert response.status_code == 200, f"HTTP failed: {response.status_code}"
+        result = response.json()
+        print("Response:", result)
+        assert "text" in result["choices"][0]
+        assert len(result["choices"][0]["text"].strip()) > 0
+
+    finally:
+        # clean up subprocesses
+        print("Cleaning up subprocess...")
+        proc.send_signal(signal.SIGINT)
+        try:
+            proc.wait(timeout=10)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+
+
+def test_disaggregated_pd_pipeline():
+    start_and_test_pipeline()

tests/e2e/run_disagg_pd.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+set -eo errexit
+
+. $(dirname "$0")/common.sh
+. $(dirname "$0")/pd_disaggreate/setup_pd.sh
+
+MODEL_NAME="/home/cmq/vllm-workspace/DeepSeek-V2-Lite"
+TP_SIZE=1
+
+_info "====> Start pd disaggregated test"
+REGISTER_PORT=10101
+PREOXY_PORT=10102
+run_proxy_server $REGISTER_PORT $PREOXY_PORT
+_info "Started pd disaggregated proxy server"
+
+PREFILL_PROC_NAME="Prefill-instance"
+PREFILL_PORT=8001
+run_prefill_instance $MODEL_NAME $TP_SIZE $PREFILL_PORT $REGISTER_PORT
+_info "Startting prefill instance"
+
+wait_url_ready $PREFILL_PROC_NAME "http://localhost:${PREFILL_PORT}/v1/completions"
+
+DECODE_PROC_NAME="Decode-instance"
+DECODE_PORT=8002
+run_decode_instance  $MODEL_NAME $TP_SIZE $DECODE_PORT $REGISTER_PORT
+_info "Startting decode instance"
+
+wait_url_ready $DECODE_PROC_NAME "http://localhost:${DECODE_PORT}/v1/completions"
+
+_info "pd disaggregated system is ready for handling request"

vllm_ascend/distributed/kv_transfer/simple_pipe.py

@@ -98,7 +98,7 @@ def __init__(
         if proxy_ip == "" or proxy_port == "":
             self.proxy_address = ""
         else:
-            self.proxy_address = proxy_ip + ":" + proxy_port
+            self.proxy_address = proxy_ip + ":" + str(proxy_port)
 
         self._register_thread = None
         if port_offset == 0 and self.proxy_address != "":