@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9ffb8cea96

[CI] Add e2e test for scale-out EC connector flow (#54973)

Signed-off-by: zhouyou9505 <zhouyou9505@gmail.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: zhouyou9505 <zhouyou9505@gmail.com> Co-authored-by: Cursor Grok 4.6 <noreply@cursor.com> Co-authored-by: OpenAI Codex <noreply@openai.com>

Nicolò Lucchesi committed Sep 9, 2026 at 16:38 UTC 9ffb8cea96369e2164af184a10faaa618eaf6b92
5 files changed +617 -2
.buildkite/test_areas/disaggregated_ec.yaml new
+16
@@ -0,0 +1,16 @@
1 +group: Disaggregated EC
2 +depends_on:
3 + - image-build
4 +steps:
5 +- label: ":nvidia: (L4) Scale-out EC Connector E2E"
6 + key: scale-out-ec-e2e-2-gpus
7 + timeout_in_minutes: 40
8 + working_dir: "/vllm-workspace/tests"
9 + device: l4
10 + num_devices: 2
11 + source_file_dependencies:
12 + - vllm/distributed/ec_transfer/
13 + - vllm/entrypoints/scale_out/
14 + - tests/entrypoints/scale_out/ec_integration/
15 + commands:
16 + - bash entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh
tests/entrypoints/scale_out/ec_integration/README.md new
+43
@@ -0,0 +1,43 @@
1 +# Scale-out EC connector e2e tests
2 +
3 +End-to-end tests for the disaggregated multimodal (encoder-cache) flow over
4 +the scale-out endpoints, in the bash-script style of
5 +`tests/v1/kv_connector/nixl_integration`.
6 +
7 +## Topology
8 +
9 +```text
10 +render (GPU-less) encode (EC producer)
11 + /v1/chat/completions/render /inference/v1/generate (kwargs_data)
12 + /v1/chat/completions/derender |
13 + prefill (EC consumer)
14 + /inference/v1/generate (mm_metadata + ec_transfer_params)
15 +```
16 +
17 +- The **render** server (`vllm launch render`) preprocesses the chat request
18 + into token ids and multimodal features (`kwargs_data` + `mm_metadata`).
19 +- The **encode** instance (`ec_role: ec_producer`, encode-only) runs the
20 + vision encoder on `kwargs_data` and publishes embeddings through the
21 + `ECExampleConnector` shared storage. Its response carries
22 + `ec_transfer_params`.
23 +- The **prefill** instance (`ec_role: ec_consumer`) receives metadata-only
24 + features (`mm_metadata`, no `kwargs_data`) plus `ec_transfer_params`, loads
25 + the embeddings through the EC connector, and generates output token ids.
26 +- The **derender** endpoint turns the output token ids back into a chat
27 + completion response.
28 +
29 +## Running
30 +
31 +```bash
32 +bash tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh
33 +```
34 +
35 +The script first collects baseline outputs from a single `vllm serve`
36 +instance, then runs the scale-out topology and compares the derendered
37 +outputs for exact equality. It also checks that metadata-only features are
38 +rejected when `ec_transfer_params` is missing.
39 +
40 +Defaults: `Qwen/Qwen3-VL-2B-Instruct`, encode on GPU 0, prefill on GPU 1.
41 +Override via environment variables (`MODEL`, `GPU_E`, `GPU_PD`,
42 +`EC_SHARED_STORAGE_PATH`, `PYTHON=.venv/bin/python`, ...); see the header of
43 +the script for the full list.
tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh new
+275
@@ -0,0 +1,275 @@
1 +#!/bin/bash
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +#
5 +# Scale-out EC connector E2E test
6 +#
7 +# Tests the disaggregated multimodal flow over the scale-out endpoints,
8 +# following the bash-script style of tests/v1/kv_connector/nixl_integration:
9 +#
10 +# 1. Baseline: a single vLLM instance serving /v1/chat/completions
11 +# 2. Scale-out:
12 +# render (GPU-less, `vllm launch render`)
13 +# encode (EC producer, encode-only)
14 +# prefill (EC consumer)
15 +#
16 +# The Python client renders each multimodal request once, sends the full
17 +# kwargs_data to the encode instance, then sends metadata-only features
18 +# (mm_metadata, no kwargs_data) plus ec_transfer_params to the prefill
19 +# instance, and derenders the output tokens. Outputs are compared against
20 +# the baseline for exact equality.
21 +#
22 +# Usage:
23 +# bash tests/entrypoints/scale_out/ec_integration/run_scale_out_ec_e2e_test.sh
24 +#
25 +# Requires 2 GPUs by default (encode on GPU 0, prefill on GPU 1); the
26 +# baseline reuses GPU 0 after the encode instance is stopped.
27 +
28 +set -euo pipefail
29 +
30 +# Resolve the repository root from the script location instead of `.git`.
31 +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
32 +GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
33 +
34 +# Model and engine configuration
35 +MODEL="${MODEL:-Qwen/Qwen3-VL-2B-Instruct}"
36 +MAX_MODEL_LEN="${MAX_MODEL_LEN:-8192}"
37 +GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.7}"
38 +MAX_NUM_SEQS="${MAX_NUM_SEQS:-128}"
39 +
40 +# GPU configuration. encode + prefill run concurrently on separate GPUs;
41 +# the baseline reuses GPU_E after the scale-out instances are stopped.
42 +GPU_E="${GPU_E:-0}"
43 +GPU_PD="${GPU_PD:-1}"
44 +GPU_SINGLE="${GPU_SINGLE:-$GPU_E}"
45 +
46 +# Ports
47 +RENDER_PORT="${RENDER_PORT:-19600}"
48 +ENCODE_PORT="${ENCODE_PORT:-19601}"
49 +PREFILL_PORT="${PREFILL_PORT:-19602}"
50 +BASELINE_PORT="${BASELINE_PORT:-19603}"
51 +
52 +# Storage path for the ECExampleConnector encoder cache
53 +EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/tmp/ec_scale_out_test}"
54 +
55 +# Python interpreter (override to .venv/bin/python for local runs)
56 +PYTHON="${PYTHON:-python3}"
57 +
58 +# Output files and logs
59 +LOG_PATH="${LOG_PATH:-/tmp}"
60 +BASELINE_FILE="${BASELINE_FILE:-/tmp/vllm_scale_out_ec_baseline.txt}"
61 +TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-600}"
62 +BASELINE_LOG="${LOG_PATH}/scale_out_ec_baseline.log"
63 +RENDER_LOG="${LOG_PATH}/scale_out_ec_render.log"
64 +ENCODE_LOG="${LOG_PATH}/scale_out_ec_encode.log"
65 +PREFILL_LOG="${LOG_PATH}/scale_out_ec_prefill.log"
66 +
67 +mkdir -p "$LOG_PATH"
68 +
69 +# Trap the SIGINT signal (triggered by Ctrl+C)
70 +trap 'kill $(jobs -pr) 2>/dev/null || true' SIGINT SIGTERM EXIT
71 +
72 +wait_for_server() {
73 + local port=$1
74 + local pid=$2
75 + local name=$3
76 + local deadline=$((SECONDS + TIMEOUT_SECONDS))
77 +
78 + while ((SECONDS < deadline)); do
79 + if ! kill -0 "$pid" 2>/dev/null; then
80 + echo "$name instance exited before becoming healthy"
81 + return 1
82 + fi
83 + if curl --max-time 2 -fsS "http://localhost:${port}/health" \
84 + > /dev/null 2>&1; then
85 + return 0
86 + fi
87 + sleep 1
88 + done
89 +
90 + echo "Timed out waiting for $name instance on port $port"
91 + return 1
92 +}
93 +
94 +dump_log() {
95 + local name=$1
96 + local log=$2
97 +
98 + echo "--- $name log (last 200 lines)"
99 + if [[ -f "$log" ]]; then
100 + tail -n 200 "$log" || true
101 + else
102 + echo "Log file not found: $log"
103 + fi
104 +}
105 +
106 +dump_scale_out_logs() {
107 + dump_log "render" "$RENDER_LOG"
108 + dump_log "encode" "$ENCODE_LOG"
109 + dump_log "prefill" "$PREFILL_LOG"
110 +}
111 +
112 +cleanup_instances() {
113 + echo "Cleaning up any running vLLM instances..."
114 + pkill -TERM -f "vllm serve" || true
115 + pkill -TERM -f "vllm launch render" || true
116 + sleep 3
117 + pkill -9 -f "vllm serve" || true
118 + pkill -9 -f "vllm launch render" || true
119 + sleep 2
120 +}
121 +
122 +# Step 1: baseline single instance
123 +run_baseline() {
124 + echo "================================"
125 + echo "Running BASELINE (single instance)"
126 + echo "================================"
127 +
128 + cleanup_instances
129 + rm -rf "$EC_SHARED_STORAGE_PATH"
130 +
131 + echo "Starting baseline instance on GPU $GPU_SINGLE, port $BASELINE_PORT"
132 + env CUDA_VISIBLE_DEVICES="$GPU_SINGLE" vllm serve "$MODEL" \
133 + --port "$BASELINE_PORT" \
134 + --max-model-len "$MAX_MODEL_LEN" \
135 + --enforce-eager \
136 + --gpu-memory-utilization 0.9 \
137 + --max-num-seqs "$MAX_NUM_SEQS" \
138 + > "$BASELINE_LOG" 2>&1 &
139 + local BASELINE_PID=$!
140 +
141 + echo "Waiting for baseline instance to start..."
142 + if ! wait_for_server "$BASELINE_PORT" "$BASELINE_PID" "baseline"; then
143 + dump_log "baseline" "$BASELINE_LOG"
144 + return 1
145 + fi
146 +
147 + if ! "$PYTHON" "${SCRIPT_DIR}/test_scale_out_ec_e2e.py" \
148 + --mode baseline \
149 + --service_url "http://localhost:$BASELINE_PORT" \
150 + --model_name "$MODEL" \
151 + --baseline_file "$BASELINE_FILE"; then
152 + dump_log "baseline" "$BASELINE_LOG"
153 + return 1
154 + fi
155 +
156 + echo "Stopping baseline instance..."
157 + kill "$BASELINE_PID" 2>/dev/null || true
158 + sleep 2
159 + cleanup_instances
160 +}
161 +
162 +# Step 2: render + encode (EC producer) + prefill (EC consumer)
163 +run_scale_out_ec() {
164 + echo "================================"
165 + echo "Running SCALE-OUT EC (render + encode + prefill)"
166 + echo "================================"
167 +
168 + cleanup_instances
169 + rm -rf "$EC_SHARED_STORAGE_PATH"
170 + mkdir -p "$EC_SHARED_STORAGE_PATH"
171 +
172 + declare -a PIDS=()
173 +
174 + # GPU-less render/derender server
175 + echo "Starting render server on port $RENDER_PORT"
176 + vllm launch render "$MODEL" \
177 + --port "$RENDER_PORT" \
178 + > "$RENDER_LOG" 2>&1 &
179 + PIDS+=("$!")
180 +
181 + # Encode-only EC producer instance. It runs the vision encoder and
182 + # publishes embeddings through the ECExampleConnector shared storage.
183 + echo "Starting encode instance on GPU $GPU_E, port $ENCODE_PORT"
184 + env CUDA_VISIBLE_DEVICES="$GPU_E" VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \
185 + vllm serve "$MODEL" \
186 + --port "$ENCODE_PORT" \
187 + --max-model-len "$MAX_MODEL_LEN" \
188 + --enforce-eager \
189 + --gpu-memory-utilization 0.01 \
190 + --no-enable-prefix-caching \
191 + --max-num-batched-tokens 114688 \
192 + --max-num-seqs "$MAX_NUM_SEQS" \
193 + --ec-transfer-config '{
194 + "ec_connector": "ECExampleConnector",
195 + "ec_role": "ec_producer",
196 + "ec_connector_extra_config": {
197 + "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
198 + }
199 + }' \
200 + > "$ENCODE_LOG" 2>&1 &
201 + PIDS+=("$!")
202 +
203 + # Prefill/decode EC consumer instance. It receives metadata-only
204 + # features plus ec_transfer_params and loads the embeddings that the
205 + # encode instance published.
206 + echo "Starting prefill instance on GPU $GPU_PD, port $PREFILL_PORT"
207 + env CUDA_VISIBLE_DEVICES="$GPU_PD" VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 \
208 + vllm serve "$MODEL" \
209 + --port "$PREFILL_PORT" \
210 + --max-model-len "$MAX_MODEL_LEN" \
211 + --enforce-eager \
212 + --enable-mm-embeds \
213 + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
214 + --max-num-seqs "$MAX_NUM_SEQS" \
215 + --ec-transfer-config '{
216 + "ec_connector": "ECExampleConnector",
217 + "ec_role": "ec_consumer",
218 + "ec_connector_extra_config": {
219 + "shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
220 + }
221 + }' \
222 + > "$PREFILL_LOG" 2>&1 &
223 + PIDS+=("$!")
224 +
225 + echo "Waiting for render instance..."
226 + if ! wait_for_server "$RENDER_PORT" "${PIDS[0]}" "render"; then
227 + dump_scale_out_logs
228 + return 1
229 + fi
230 + echo "Waiting for encode instance..."
231 + if ! wait_for_server "$ENCODE_PORT" "${PIDS[1]}" "encode"; then
232 + dump_scale_out_logs
233 + return 1
234 + fi
235 + echo "Waiting for prefill instance..."
236 + if ! wait_for_server "$PREFILL_PORT" "${PIDS[2]}" "prefill"; then
237 + dump_scale_out_logs
238 + return 1
239 + fi
240 +
241 + echo "All scale-out EC services are up!"
242 +
243 + if ! "$PYTHON" "${SCRIPT_DIR}/test_scale_out_ec_e2e.py" \
244 + --mode disagg \
245 + --render_url "http://localhost:$RENDER_PORT" \
246 + --encode_url "http://localhost:$ENCODE_PORT" \
247 + --prefill_url "http://localhost:$PREFILL_PORT" \
248 + --model_name "$MODEL" \
249 + --baseline_file "$BASELINE_FILE"; then
250 + dump_scale_out_logs
251 + return 1
252 + fi
253 +
254 + echo "Stopping scale-out EC instances..."
255 + for pid in "${PIDS[@]}"; do
256 + kill "$pid" 2>/dev/null || true
257 + done
258 + sleep 2
259 + cleanup_instances
260 +}
261 +
262 +# Main execution
263 +echo "================================"
264 +echo "Scale-out EC E2E Test Suite"
265 +echo "Model: $MODEL"
266 +echo "================================"
267 +
268 +run_baseline
269 +run_scale_out_ec
270 +
271 +rm -f "$BASELINE_FILE"
272 +
273 +echo "================================"
274 +echo "All scale-out EC e2e tests finished!"
275 +echo "================================"
tests/entrypoints/scale_out/ec_integration/test_scale_out_ec_e2e.py new
+276
@@ -0,0 +1,276 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""E2E test for the scale-out EC connector flow (render -> encode -> prefill).
4 +
5 +Mimics the bash-script style of ``tests/v1/kv_connector/nixl_integration``:
6 +the instances are started by ``run_scale_out_ec_e2e_test.sh`` and this file is
7 +the HTTP client that drives the disaggregated request path and compares its
8 +output against a single-instance baseline.
9 +
10 +Topology under test (started by the bash script):
11 +
12 + render (GPU-less) -- /v1/chat/completions/render, /derender
13 + encode (EC producer) -- /inference/v1/generate
14 + prefill (EC consumer) -- /inference/v1/generate
15 +
16 +The client renders each multimodal chat request once, sends the full
17 +``kwargs_data`` to the encode instance, then sends metadata-only features
18 +(``mm_metadata`` without ``kwargs_data``) plus the encode response's
19 +``ec_transfer_params`` to the prefill instance, and finally derenders the
20 +output token ids. The derendered text must match the baseline exactly.
21 +
22 +Usage (from the bash script, or manually):
23 + # Baseline mode (saves outputs):
24 + python test_scale_out_ec_e2e.py \
25 + --mode baseline \
26 + --service_url http://localhost:19603 \
27 + --model_name Qwen/Qwen3-VL-2B-Instruct \
28 + --baseline_file /tmp/vllm_scale_out_ec_baseline.txt
29 +
30 + # Disagg mode (compares outputs):
31 + python test_scale_out_ec_e2e.py \
32 + --mode disagg \
33 + --render_url http://localhost:19600 \
34 + --encode_url http://localhost:19601 \
35 + --prefill_url http://localhost:19602 \
36 + --model_name Qwen/Qwen3-VL-2B-Instruct \
37 + --baseline_file /tmp/vllm_scale_out_ec_baseline.txt
38 +"""
39 +
40 +import argparse
41 +import json
42 +import sys
43 +import time
44 +
45 +import requests
46 +
47 +from vllm.assets.image import ImageAsset
48 +from vllm.multimodal.utils import encode_image_url
49 +
50 +MAX_OUTPUT_LEN = 128
51 +REQUEST_TIMEOUT = 300
52 +
53 +image_1 = ImageAsset("stop_sign").pil_image.resize((1280, 720))
54 +image_2 = ImageAsset("cherry_blossom").pil_image.resize((1280, 720))
55 +image_1_url = encode_image_url(image_1)
56 +image_2_url = encode_image_url(image_2)
57 +
58 +# Each prompt is a chat completion request body; the same body is sent to the
59 +# baseline instance, to /render, and to /derender (as chat_request).
60 +SAMPLE_PROMPTS: list[dict] = [
61 + {
62 + "messages": [
63 + {
64 + "role": "user",
65 + "content": [
66 + {"type": "image_url", "image_url": {"url": image_1_url}},
67 + {"type": "text", "text": "What's in this image?"},
68 + ],
69 + }
70 + ],
71 + "description": "single image",
72 + },
73 + {
74 + "messages": [
75 + {
76 + "role": "user",
77 + "content": [
78 + {"type": "image_url", "image_url": {"url": image_2_url}},
79 + {"type": "image_url", "image_url": {"url": image_1_url}},
80 + {"type": "text", "text": "Describe these 2 images in detail."},
81 + ],
82 + }
83 + ],
84 + "description": "two images",
85 + },
86 + {
87 + "messages": [{"role": "user", "content": "What is the capital of France?"}],
88 + "description": "text only",
89 + },
90 +]
91 +
92 +
93 +def chat_request(prompt: dict) -> dict:
94 + return {
95 + "model": None, # filled by the caller
96 + "messages": prompt["messages"],
97 + "max_tokens": MAX_OUTPUT_LEN,
98 + "temperature": 0.0,
99 + "seed": 42,
100 + }
101 +
102 +
103 +def post_json(url: str, payload: dict) -> requests.Response:
104 + response = requests.post(url, json=payload, timeout=REQUEST_TIMEOUT)
105 + if response.status_code != 200:
106 + raise RuntimeError(
107 + f"POST {url} failed with {response.status_code}: {response.text}"
108 + )
109 + return response
110 +
111 +
112 +def run_baseline(args) -> None:
113 + outputs = []
114 + for prompt in SAMPLE_PROMPTS:
115 + request = chat_request(prompt) | {"model": args.model_name}
116 + response = post_json(f"{args.service_url}/v1/chat/completions", request).json()
117 + outputs.append(response["choices"][0]["message"]["content"])
118 + print(f"[baseline] {prompt['description']}: {outputs[-1][:80]!r}")
119 +
120 + with open(args.baseline_file, "w") as f:
121 + json.dump(outputs, f)
122 + print(f"Baseline saved to {args.baseline_file}")
123 +
124 +
125 +def _assert_mm_features(features: dict) -> None:
126 + kwargs_data = features.get("kwargs_data") or {}
127 + mm_metadata = features.get("mm_metadata")
128 + assert any(item is not None for items in kwargs_data.values() for item in items), (
129 + f"render response has no kwargs_data items: {features}"
130 + )
131 + assert mm_metadata, "render response is missing mm_metadata"
132 + for modality, hashes in features["mm_hashes"].items():
133 + assert modality in mm_metadata, f"mm_metadata missing modality {modality}"
134 + assert len(mm_metadata[modality]) == len(hashes)
135 +
136 +
137 +def run_disagg(args) -> None:
138 + with open(args.baseline_file) as f:
139 + baseline = json.load(f)
140 +
141 + failures = []
142 + for prompt, expected in zip(SAMPLE_PROMPTS, baseline, strict=True):
143 + request = chat_request(prompt) | {"model": args.model_name}
144 + got = run_disagg_one(args, request)
145 + if got != expected:
146 + failures.append(
147 + f"{prompt['description']}:\n"
148 + f" baseline: {expected!r}\n disagg: {got!r}"
149 + )
150 + else:
151 + print(f"[disagg] {prompt['description']}: matched baseline")
152 +
153 + if failures:
154 + print("Output mismatch:")
155 + for failure in failures:
156 + print(failure)
157 + sys.exit(1)
158 + print("All scale-out EC outputs matched the baseline")
159 +
160 +
161 +def run_disagg_one(args, request: dict) -> str:
162 + # 1. Render: chat request -> token ids + multimodal features (GPU-less).
163 + rendered = post_json(
164 + f"{args.render_url}/v1/chat/completions/render", request
165 + ).json()
166 + token_ids = rendered["token_ids"]
167 + features = rendered.get("features")
168 + request_id = rendered["request_id"]
169 +
170 + # 2. Encode: full kwargs_data so the vision encoder can run. The
171 + # response carries ec_transfer_params for the prefill instance.
172 + ec_transfer_params = None
173 + if features is not None:
174 + _assert_mm_features(features)
175 + encode_payload = {
176 + "request_id": request_id,
177 + "token_ids": token_ids,
178 + "features": {
179 + "mm_hashes": features["mm_hashes"],
180 + "mm_placeholders": features["mm_placeholders"],
181 + "kwargs_data": features["kwargs_data"],
182 + },
183 + "sampling_params": {"max_tokens": 1, "temperature": 0.0},
184 + }
185 + encode_response = post_json(
186 + f"{args.encode_url}/inference/v1/generate", encode_payload
187 + ).json()
188 + ec_transfer_params = encode_response.get("ec_transfer_params")
189 + assert ec_transfer_params, (
190 + f"encode response has no ec_transfer_params: {encode_response}"
191 + )
192 +
193 + # 3. Prefill: metadata-only features. kwargs_data is dropped; the
194 + # embeddings arrive through the EC connector.
195 + if features is not None:
196 + prefill_features = {
197 + "mm_hashes": features["mm_hashes"],
198 + "mm_placeholders": features["mm_placeholders"],
199 + "mm_metadata": features["mm_metadata"],
200 + }
201 + _check_metadata_only_rejected_without_ec(args, token_ids, prefill_features)
202 + else:
203 + prefill_features = None
204 +
205 + prefill_payload = {
206 + "request_id": request_id,
207 + "token_ids": token_ids,
208 + "sampling_params": {"max_tokens": MAX_OUTPUT_LEN, "temperature": 0.0},
209 + }
210 + if prefill_features is not None:
211 + prefill_payload["features"] = prefill_features
212 + prefill_payload["ec_transfer_params"] = ec_transfer_params
213 + prefill_response = post_json(
214 + f"{args.prefill_url}/inference/v1/generate", prefill_payload
215 + ).json()
216 + output_token_ids = prefill_response["choices"][0]["token_ids"]
217 + assert output_token_ids, f"prefill returned no token ids: {prefill_response}"
218 +
219 + # 4. Derender: token ids -> chat completion response (GPU-less).
220 + derendered = post_json(
221 + f"{args.render_url}/v1/chat/completions/derender",
222 + {
223 + "model": args.model_name,
224 + "generate_response": prefill_response,
225 + "prompt_tokens": len(token_ids),
226 + "chat_request": request,
227 + },
228 + ).json()
229 + return derendered["choices"][0]["message"]["content"]
230 +
231 +
232 +def _check_metadata_only_rejected_without_ec(
233 + args, token_ids: list[int], prefill_features: dict
234 +) -> None:
235 + """mm_metadata without kwargs_data must be rejected without EC params."""
236 + payload = {
237 + "token_ids": token_ids,
238 + "features": prefill_features,
239 + "sampling_params": {"max_tokens": 1, "temperature": 0.0},
240 + }
241 + response = requests.post(
242 + f"{args.prefill_url}/inference/v1/generate",
243 + json=payload,
244 + timeout=REQUEST_TIMEOUT,
245 + )
246 + assert response.status_code in (400, 422), (
247 + "metadata-only features without ec_transfer_params should be rejected, "
248 + f"got {response.status_code}: {response.text}"
249 + )
250 +
251 +
252 +def main() -> None:
253 + parser = argparse.ArgumentParser(description=__doc__)
254 + parser.add_argument("--mode", choices=["baseline", "disagg"], required=True)
255 + parser.add_argument("--service_url", help="baseline single-instance URL")
256 + parser.add_argument("--render_url", help="GPU-less render server URL")
257 + parser.add_argument("--encode_url", help="EC producer (encode) URL")
258 + parser.add_argument("--prefill_url", help="EC consumer (prefill) URL")
259 + parser.add_argument("--model_name", required=True)
260 + parser.add_argument("--baseline_file", required=True)
261 + args = parser.parse_args()
262 +
263 + start = time.time()
264 + if args.mode == "baseline":
265 + assert args.service_url, "--service_url is required in baseline mode"
266 + run_baseline(args)
267 + else:
268 + assert args.render_url, "--render_url is required in disagg mode"
269 + assert args.encode_url, "--encode_url is required in disagg mode"
270 + assert args.prefill_url, "--prefill_url is required in disagg mode"
271 + run_disagg(args)
272 + print(f"Mode {args.mode} finished in {time.time() - start:.1f}s")
273 +
274 +
275 +if __name__ == "__main__":
276 + main()
vllm/distributed/ec_transfer/ec_connector/example_connector.py
+7 -2
@@ -17,6 +17,7 @@ from vllm.distributed.ec_transfer.ec_connector.utils import (
17 collect_ec_item_metadata,
18 )
19 from vllm.logger import init_logger
20 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed
21 from vllm.v1.core.sched.output import SchedulerOutput
22
23 if TYPE_CHECKING:
@@ -99,7 +100,10 @@ class ECExampleConnector(ECConnectorBase):
100 if mm_data.mm_hash in encoder_cache:
101 continue
102 filename = self._generate_filename_debug(mm_data.mm_hash)
102 - ec_cache = safetensors.torch.load_file(filename, device=device)["ec_cache"]
103 + with gpu_sync_allowed():
104 + ec_cache = safetensors.torch.load_file(filename, device=device)[
105 + "ec_cache"
106 + ]
107 encoder_cache[mm_data.mm_hash] = ec_cache
108 logger.debug("Success load encoder cache for hash %s", mm_data.mm_hash)
109
@@ -121,7 +125,8 @@ class ECExampleConnector(ECConnectorBase):
125 return
126 filename = self._generate_filename_debug(mm_hash)
127 ec_cache = encoder_cache[mm_hash]
124 - tensors = {"ec_cache": ec_cache.detach().cpu()}
128 + with gpu_sync_allowed():
129 + tensors = {"ec_cache": ec_cache.detach().cpu()}
130 safetensors.torch.save_file(tensors, filename)
131 logger.debug("Save cache successful for mm_hash %s", mm_hash)
132