@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 316b1bf706
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 23:45 UTC

[Tests] Add tests for headless internal DP LB (#21450)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Jul 23, 2025 at 15:49 UTC 316b1bf706f874d68d72c1bc1ba2b4e1f3b491bd
4 files changed +768 -120
.buildkite/test-pipeline.yaml
+2
@@ -165,6 +165,7 @@ steps:
165 - tests/examples/offline_inference/data_parallel.py
166 - tests/v1/test_async_llm_dp.py
167 - tests/v1/test_external_lb_dp.py
168 + - tests/v1/test_internal_lb_dp.py
169 - tests/v1/engine/test_engine_core_client.py
170 commands:
171 # test with tp=2 and external_dp=2
@@ -176,6 +177,7 @@ steps:
177 - python3 ../examples/offline_inference/data_parallel.py --enforce-eager
178 - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py
179 - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_external_lb_dp.py
180 + - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/test_internal_lb_dp.py
181 - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
182 - pytest -v -s distributed/test_utils.py
183 - pytest -v -s compile/test_basic_correctness.py
tests/v1/entrypoints/openai/test_multi_api_servers.py
+3 -120
@@ -2,136 +2,19 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import asyncio
4 import os
5 -import re
5
6 import openai # use the official client for correctness check
7 import pytest
8 import pytest_asyncio
10 -import requests
9
10 from tests.utils import RemoteOpenAIServer
11 +from tests.v1.test_utils import check_request_balancing
12
13 MODEL_NAME = "ibm-research/PowerMoE-3b"
14
15 DP_SIZE = os.getenv("DP_SIZE", "1")
16
17
19 -def get_prometheus_metrics(
20 - server: RemoteOpenAIServer) -> dict[str, dict[str, float]]:
21 - """Fetch and parse Prometheus metrics from the /metrics endpoint.
22 -
23 - Returns:
24 - Dict mapping metric names to their values grouped by labels.
25 - For example: {"vllm:request_success": {
26 - "engine=0": 5.0, "engine=1": 3.0}
27 - }
28 - """
29 - try:
30 - response = requests.get(server.url_for("metrics"), timeout=10)
31 - response.raise_for_status()
32 -
33 - metrics: dict[str, dict[str, float]] = {}
34 -
35 - # Regex patterns for Prometheus metrics
36 - metric_with_labels = re.compile(
37 - r'^([a-zA-Z_:][a-zA-Z0-9_:]*)\{([^}]*)\}\s+([\d\.\-\+e]+)$')
38 - metric_simple = re.compile(
39 - r'^([a-zA-Z_:][a-zA-Z0-9_:]*)\s+([\d\.\-\+e]+)$')
40 -
41 - for line in response.text.split('\n'):
42 - line = line.strip()
43 - # Skip comments and empty lines
44 - if not line or line.startswith('#'):
45 - continue
46 -
47 - # Try to match metric with labels first
48 - match = metric_with_labels.match(line)
49 - if match:
50 - metric_name, labels_part, value_str = match.groups()
51 - try:
52 - value = float(value_str)
53 - if metric_name not in metrics:
54 - metrics[metric_name] = {}
55 - metrics[metric_name][f'{{{labels_part}}}'] = value
56 - except ValueError:
57 - continue
58 - else:
59 - # Try simple metric without labels
60 - match = metric_simple.match(line)
61 - if match:
62 - metric_name, value_str = match.groups()
63 - try:
64 - value = float(value_str)
65 - if metric_name not in metrics:
66 - metrics[metric_name] = {}
67 - metrics[metric_name][''] = value
68 - except ValueError:
69 - continue
70 -
71 - return metrics
72 - except Exception as e:
73 - pytest.fail(f"Failed to fetch Prometheus metrics: {e}")
74 - return {}
75 -
76 -
77 -def get_engine_request_counts(
78 - metrics: dict[str, dict[str, float]]) -> dict[str, float]:
79 - """Extract request counts per engine from Prometheus metrics.
80 -
81 - Returns:
82 - Dict mapping engine indices to request counts.
83 - For example: {"0": 15.0, "1": 12.0}
84 - """
85 - engine_counts = {}
86 -
87 - # Look for request success metrics with engine labels
88 - success_metrics = metrics.get("vllm:request_success_total", {})
89 - engine_pattern = re.compile(r'engine="([^"]*)"')
90 -
91 - for labels, count in success_metrics.items():
92 - # Extract engine ID from labels using regex
93 - match = engine_pattern.search(labels)
94 - if match:
95 - engine_id = match.group(1)
96 - if engine_id not in engine_counts:
97 - engine_counts[engine_id] = 0.0
98 - engine_counts[engine_id] += count
99 -
100 - return engine_counts
101 -
102 -
103 -def check_request_balancing(server: RemoteOpenAIServer):
104 - """Check request balancing via Prometheus metrics if DP_SIZE > 1.
105 -
106 - Args:
107 - server: The RemoteOpenAIServer instance
108 - """
109 - dp_size = int(DP_SIZE)
110 - if dp_size <= 1:
111 - return
112 -
113 - # Get metrics after all requests are completed
114 - metrics = get_prometheus_metrics(server)
115 - engine_counts = get_engine_request_counts(metrics)
116 -
117 - # Check that multiple engines received requests
118 - engines_with_requests = [
119 - engine for engine, count in engine_counts.items() if count > 0
120 - ]
121 - assert len(engines_with_requests) == dp_size, (
122 - f"Expected requests to be distributed across multiple engines,"
123 - f" but only engine(s) {engines_with_requests} received "
124 - f"requests. Engine counts: {engine_counts}")
125 -
126 - # Verify that the load is reasonably balanced
127 - # (no engine should handle all requests)
128 - total_requests = sum(engine_counts.values())
129 -
130 - for count in engine_counts.values():
131 - assert count > total_requests // (dp_size + 1), (
132 - f"requests are imbalanced: {engine_counts}")
133 -
134 -
18 @pytest.fixture(scope="module")
19 def default_server_args():
20 return [
@@ -217,7 +100,7 @@ async def test_single_completion(client: openai.AsyncOpenAI,
100 assert all(completion is not None for completion in results)
101
102 # Check request balancing via Prometheus metrics if DP_SIZE > 1
220 - check_request_balancing(server)
103 + check_request_balancing(server, int(DP_SIZE))
104
105
106 @pytest.mark.asyncio
@@ -295,4 +178,4 @@ async def test_completion_streaming(client: openai.AsyncOpenAI,
178 assert all(results), "Not all streaming requests completed successfully."
179
180 # Check request balancing via Prometheus metrics if DP_SIZE > 1
298 - check_request_balancing(server)
181 + check_request_balancing(server, int(DP_SIZE))
tests/v1/test_internal_lb_dp.py new
+639
@@ -0,0 +1,639 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import asyncio
4 +import os
5 +import threading
6 +import time
7 +
8 +import openai # use the official client for correctness check
9 +import pytest
10 +import pytest_asyncio
11 +
12 +from tests.utils import RemoteOpenAIServer
13 +from tests.v1.test_utils import check_request_balancing
14 +from vllm.platforms import Platform
15 +
16 +MODEL_NAME = "ibm-research/PowerMoE-3b"
17 +
18 +# Number of data parallel ranks for multi-node internal LB testing
19 +DP_SIZE = int(os.getenv("DP_SIZE", "2"))
20 +# Default tensor parallel size to use
21 +TP_SIZE = int(os.getenv("TP_SIZE", "1"))
22 +
23 +# Number of nodes to simulate
24 +NUM_NODES = 2
25 +
26 +
27 +class MultinodeInternalLBServerManager:
28 + """Manages multi-node data parallel vLLM server instances for internal
29 + load balancer testing using --headless mode."""
30 +
31 + def __init__(self,
32 + model_name: str,
33 + dp_size: int,
34 + api_server_count: int,
35 + base_server_args: list,
36 + dp_per_node: int = 1,
37 + tp_size: int = TP_SIZE):
38 + self.model_name = model_name
39 + self.dp_size = dp_size
40 + self.dp_per_node = dp_per_node
41 + self.tp_size = tp_size
42 + self.api_server_count = api_server_count
43 + self.base_server_args = base_server_args
44 + self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = []
45 + self.server_threads: list[threading.Thread] = []
46 +
47 + def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
48 + """Start all server instances for multi-node internal LB mode."""
49 + for rank in range(0, self.dp_size, self.dp_per_node):
50 + # Create server args for this specific rank
51 + server_args = self.base_server_args.copy()
52 +
53 + if rank == 0:
54 + # Head node - runs API server and first DP rank
55 + server_args.extend([
56 + "--data-parallel-size",
57 + str(self.dp_size),
58 + "--data-parallel-size-local",
59 + str(self.dp_per_node),
60 + "--tensor-parallel-size",
61 + str(self.tp_size),
62 + "--port",
63 + "8000", # Single endpoint for all requests
64 + "--api-server-count",
65 + str(self.api_server_count),
66 + "--data-parallel-address",
67 + "127.0.0.1",
68 + "--data-parallel-rpc-port",
69 + "13345",
70 + ])
71 + else:
72 + # Secondary nodes - run in headless mode
73 + server_args.extend([
74 + "--headless",
75 + "--data-parallel-size",
76 + str(self.dp_size),
77 + "--data-parallel-size-local",
78 + str(self.dp_per_node),
79 + "--data-parallel-start-rank",
80 + str(rank),
81 + "--tensor-parallel-size",
82 + str(self.tp_size),
83 + "--data-parallel-address",
84 + "127.0.0.1",
85 + "--data-parallel-rpc-port",
86 + "13345",
87 + ])
88 +
89 + # Use a thread to start each server to allow parallel initialization
90 + def start_server(r: int, sargs: list[str]):
91 + gpus_per_node = self.tp_size * self.dp_per_node
92 + try:
93 + # Start the server
94 + server = RemoteOpenAIServer(
95 + self.model_name,
96 + sargs,
97 + auto_port=False,
98 + env_dict={
99 + "CUDA_VISIBLE_DEVICES":
100 + ",".join(
101 + str(Platform.device_id_to_physical_device_id(
102 + i)) for i in range(r, r + gpus_per_node))
103 + })
104 + server.__enter__()
105 + if r == 0:
106 + print(
107 + f"Head node (rank {r}) started successfully with "
108 + f"{self.api_server_count} API servers")
109 + else:
110 + print(f"Headless node (rank {r}) started successfully")
111 + self.servers.append((server, sargs))
112 + except Exception as e:
113 + print(f"Failed to start server rank {r}: {e}")
114 + raise
115 +
116 + thread = threading.Thread(target=start_server,
117 + args=(rank, server_args))
118 + thread.start()
119 +
120 + self.server_threads.append(thread)
121 +
122 + # Wait for all servers to start
123 + for thread in self.server_threads:
124 + thread.join()
125 +
126 + # Give servers additional time to fully initialize and coordinate
127 + time.sleep(3)
128 +
129 + if len(self.servers) != self.dp_size // self.dp_per_node:
130 + raise Exception("Servers failed to start")
131 +
132 + return self.servers
133 +
134 + def __exit__(self, exc_type, exc_val, exc_tb):
135 + """Stop all server instances."""
136 + while self.servers:
137 + try:
138 + self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
139 + except Exception as e:
140 + print(f"Error stopping server: {e}")
141 +
142 +
143 +class APIOnlyServerManager:
144 + """Manages API-only server (Node 0) and headless engines server (Node 1)
145 + for testing separated API server and engine configuration."""
146 +
147 + def __init__(self,
148 + model_name: str,
149 + dp_size: int,
150 + api_server_count: int,
151 + base_server_args: list,
152 + tp_size: int = TP_SIZE):
153 + self.model_name = model_name
154 + self.dp_size = dp_size
155 + self.tp_size = tp_size
156 + self.api_server_count = api_server_count
157 + self.base_server_args = base_server_args
158 + self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = []
159 + self.server_threads: list[threading.Thread] = []
160 +
161 + def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
162 + """Start API-only server and headless engines server."""
163 +
164 + # Start API-only server (Node 0) - no engines, only API server
165 + api_server_args = self.base_server_args.copy()
166 + api_server_args.extend([
167 + "--data-parallel-size",
168 + str(self.dp_size),
169 + "--data-parallel-size-local",
170 + "0", # No engines on this node
171 + "--tensor-parallel-size",
172 + str(self.tp_size),
173 + "--port",
174 + "8000",
175 + "--api-server-count",
176 + str(self.api_server_count),
177 + "--data-parallel-address",
178 + "127.0.0.1",
179 + "--data-parallel-rpc-port",
180 + "13345",
181 + ])
182 +
183 + # Start headless engines server (Node 1) - all engines, no API server
184 + engines_server_args = self.base_server_args.copy()
185 + engines_server_args.extend([
186 + "--headless",
187 + "--data-parallel-size",
188 + str(self.dp_size),
189 + "--data-parallel-size-local",
190 + str(self.dp_size), # All engines on this node
191 + "--tensor-parallel-size",
192 + str(self.tp_size),
193 + "--data-parallel-address",
194 + "127.0.0.1",
195 + "--data-parallel-rpc-port",
196 + "13345",
197 + ])
198 +
199 + # Use threads to start both servers in parallel
200 + def start_api_server():
201 + try:
202 + server = RemoteOpenAIServer(
203 + self.model_name,
204 + api_server_args,
205 + auto_port=False,
206 + env_dict={}) # No GPUs needed for API-only server
207 + server.__enter__()
208 + print(f"API-only server started successfully with "
209 + f"{self.api_server_count} API servers")
210 + self.servers.append((server, api_server_args))
211 + except Exception as e:
212 + print(f"Failed to start API-only server: {e}")
213 + raise
214 +
215 + def start_engines_server():
216 + try:
217 + server = RemoteOpenAIServer(
218 + self.model_name,
219 + engines_server_args,
220 + auto_port=False,
221 + env_dict={
222 + "CUDA_VISIBLE_DEVICES":
223 + ",".join(
224 + str(Platform.device_id_to_physical_device_id(i))
225 + for i in range(self.dp_size * self.tp_size))
226 + })
227 + server.__enter__()
228 + print(f"Headless engines server started successfully with "
229 + f"{self.dp_size} engines")
230 + self.servers.append((server, engines_server_args))
231 + except Exception as e:
232 + print(f"Failed to start headless engines server: {e}")
233 + raise
234 +
235 + # Start API server first
236 + api_thread = threading.Thread(target=start_api_server)
237 + api_thread.start()
238 + self.server_threads.append(api_thread)
239 +
240 + # Start engines server second
241 + engines_thread = threading.Thread(target=start_engines_server)
242 + engines_thread.start()
243 + self.server_threads.append(engines_thread)
244 +
245 + # Wait for both servers to start
246 + for thread in self.server_threads:
247 + thread.join()
248 +
249 + # Give servers additional time to fully initialize and coordinate
250 + time.sleep(3)
251 +
252 + if len(self.servers) != 2:
253 + raise Exception("Both servers failed to start")
254 +
255 + return self.servers
256 +
257 + def __exit__(self, exc_type, exc_val, exc_tb):
258 + """Stop both server instances."""
259 + while self.servers:
260 + try:
261 + self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
262 + except Exception as e:
263 + print(f"Error stopping server: {e}")
264 +
265 +
266 +@pytest.fixture(scope="module")
267 +def default_server_args():
268 + return [
269 + # use half precision for speed and memory savings in CI environment
270 + "--dtype",
271 + "bfloat16",
272 + "--max-model-len",
273 + "2048",
274 + "--max-num-seqs",
275 + "128",
276 + "--enforce-eager",
277 + ]
278 +
279 +
280 +@pytest.fixture(scope="module", params=[1, 4])
281 +def servers(request, default_server_args):
282 + api_server_count = request.param
283 + with MultinodeInternalLBServerManager(MODEL_NAME, DP_SIZE,
284 + api_server_count,
285 + default_server_args,
286 + DP_SIZE // NUM_NODES,
287 + TP_SIZE) as server_list:
288 + yield server_list
289 +
290 +
291 +@pytest.fixture(scope="module", params=[1, 4])
292 +def api_only_servers(request, default_server_args):
293 + """Fixture for API-only server + headless engines configuration."""
294 + api_server_count = request.param
295 + with APIOnlyServerManager(MODEL_NAME, DP_SIZE, api_server_count,
296 + default_server_args, TP_SIZE) as server_list:
297 + yield server_list
298 +
299 +
300 +@pytest_asyncio.fixture
301 +async def client(servers: list[tuple[RemoteOpenAIServer, list[str]]]):
302 + # For internal LB, we only connect to the head node (rank 0)
303 + # which provides the single API endpoint
304 + head_server = servers[0][0]
305 + async with head_server.get_async_client() as client:
306 + yield client
307 +
308 +
309 +@pytest_asyncio.fixture
310 +async def api_only_client(api_only_servers: list[tuple[RemoteOpenAIServer,
311 + list[str]]]):
312 + """Client fixture for API-only server configuration."""
313 + # Connect to the API-only server (first server in the list)
314 + api_server = api_only_servers[0][0]
315 + async with api_server.get_async_client() as client:
316 + yield client
317 +
318 +
319 +@pytest.mark.asyncio
320 +@pytest.mark.parametrize(
321 + "model_name",
322 + [MODEL_NAME],
323 +)
324 +async def test_multinode_dp_completion(client: openai.AsyncOpenAI,
325 + servers: list[tuple[RemoteOpenAIServer,
326 + list[str]]],
327 + model_name: str) -> None:
328 +
329 + async def make_request():
330 + completion = await client.completions.create(
331 + model=model_name,
332 + prompt="Hello, my name is",
333 + max_tokens=10,
334 + temperature=1.0)
335 +
336 + assert completion.id is not None
337 + assert completion.choices is not None and len(completion.choices) == 1
338 +
339 + choice = completion.choices[0]
340 + # The exact number of tokens can vary slightly with temperature=1.0,
341 + # so we check for a reasonable minimum length.
342 + assert len(choice.text) >= 1
343 + # Finish reason might not always be 'length' if the model finishes early
344 + # or due to other reasons, especially with high temperature.
345 + # So, we'll accept 'length' or 'stop'.
346 + assert choice.finish_reason in ("length", "stop")
347 +
348 + # Token counts can also vary, so we check they are positive.
349 + assert completion.usage.completion_tokens > 0
350 + assert completion.usage.prompt_tokens > 0
351 + assert completion.usage.total_tokens > 0
352 + return completion
353 +
354 + # Test single request
355 + result = await make_request()
356 + assert result is not None
357 + print(
358 + "Multi-node internal LB handled single completion request successfully"
359 + )
360 +
361 + await asyncio.sleep(0.5)
362 +
363 + # Send multiple requests - internal LB should distribute across DP ranks
364 + num_requests = 50
365 + all_tasks = [make_request() for _ in range(num_requests)]
366 +
367 + results = await asyncio.gather(*all_tasks)
368 + assert len(results) == num_requests
369 + assert all(completion is not None for completion in results)
370 +
371 + await asyncio.sleep(0.5)
372 +
373 + # Second burst of requests
374 + all_tasks = [make_request() for _ in range(num_requests)]
375 +
376 + results = await asyncio.gather(*all_tasks)
377 + assert len(results) == num_requests
378 + assert all(completion is not None for completion in results)
379 +
380 + _, server_args = servers[0]
381 + api_server_count = (
382 + server_args.count('--api-server-count')
383 + and server_args[server_args.index('--api-server-count') + 1] or 1)
384 + print(f"Successfully completed multi-node internal LB test with "
385 + f"{len(servers)} DP ranks (API server count: {api_server_count})")
386 +
387 + # Check request balancing via Prometheus metrics
388 + head_server = servers[0][0]
389 + check_request_balancing(head_server, DP_SIZE)
390 +
391 +
392 +@pytest.mark.asyncio
393 +@pytest.mark.parametrize(
394 + "model_name",
395 + [MODEL_NAME],
396 +)
397 +async def test_multinode_dp_completion_streaming(client: openai.AsyncOpenAI,
398 + servers: list[
399 + tuple[RemoteOpenAIServer,
400 + list[str]]],
401 + model_name: str) -> None:
402 + prompt = "What is an LLM?"
403 +
404 + async def make_streaming_request():
405 + # Perform a non-streaming request to get the expected full output
406 + single_completion = await client.completions.create(
407 + model=model_name,
408 + prompt=prompt,
409 + max_tokens=5,
410 + temperature=0.0,
411 + )
412 + single_output = single_completion.choices[0].text
413 +
414 + # Perform the streaming request
415 + stream = await client.completions.create(model=model_name,
416 + prompt=prompt,
417 + max_tokens=5,
418 + temperature=0.0,
419 + stream=True)
420 + chunks: list[str] = []
421 + finish_reason_count = 0
422 + last_chunk = None
423 + async for chunk in stream:
424 + chunks.append(chunk.choices[0].text)
425 + if chunk.choices[0].finish_reason is not None:
426 + finish_reason_count += 1
427 + last_chunk = chunk # Keep track of the last chunk
428 +
429 + # finish reason should only return in the last block for OpenAI API
430 + assert finish_reason_count == 1, (
431 + "Finish reason should appear exactly once.")
432 + assert last_chunk is not None, (
433 + "Stream should have yielded at least one chunk.")
434 + assert last_chunk.choices[
435 + 0].finish_reason == "length", "Finish reason should be 'length'."
436 + # Check that the combined text matches the non-streamed version.
437 + assert "".join(
438 + chunks
439 + ) == single_output, "Streamed output should match non-streamed output."
440 + return True # Indicate success for this request
441 +
442 + # Test single streaming request
443 + result = await make_streaming_request()
444 + assert result is not None
445 + print(
446 + "Multi-node internal LB handled single streaming request successfully")
447 +
448 + await asyncio.sleep(0.5)
449 +
450 + # Send multiple streaming requests - internal LB should distribute across
451 + # DP ranks
452 + num_requests = 50
453 + all_tasks = [make_streaming_request() for _ in range(num_requests)]
454 +
455 + results = await asyncio.gather(*all_tasks)
456 + assert len(results) == num_requests
457 + assert all(results), "Not all streaming requests completed successfully."
458 +
459 + await asyncio.sleep(0.5)
460 +
461 + # Second burst of streaming requests
462 + all_tasks = [make_streaming_request() for _ in range(num_requests)]
463 +
464 + results = await asyncio.gather(*all_tasks)
465 + assert len(results) == num_requests
466 + assert all(results), "Not all streaming requests completed successfully."
467 +
468 + _, server_args = servers[0]
469 + api_server_count = (
470 + server_args.count('--api-server-count')
471 + and server_args[server_args.index('--api-server-count') + 1] or 1)
472 + print(f"Successfully completed multi-node internal LB streaming test with "
473 + f"{len(servers)} DP ranks (API server count: {api_server_count})")
474 +
475 + # Check request balancing via Prometheus metrics
476 + head_server = servers[0][0]
477 + check_request_balancing(head_server, DP_SIZE)
478 +
479 +
480 +@pytest.mark.asyncio
481 +@pytest.mark.parametrize(
482 + "model_name",
483 + [MODEL_NAME],
484 +)
485 +async def test_api_only_multinode_dp_completion(
486 + api_only_client: openai.AsyncOpenAI,
487 + api_only_servers: list[tuple[RemoteOpenAIServer,
488 + list[str]]], model_name: str) -> None:
489 + """Test API-only server with all engines on separate headless server."""
490 +
491 + async def make_request():
492 + completion = await api_only_client.completions.create(
493 + model=model_name,
494 + prompt="Hello, my name is",
495 + max_tokens=10,
496 + temperature=1.0)
497 +
498 + assert completion.id is not None
499 + assert completion.choices is not None and len(completion.choices) == 1
500 +
501 + choice = completion.choices[0]
502 + # The exact number of tokens can vary slightly with temperature=1.0,
503 + # so we check for a reasonable minimum length.
504 + assert len(choice.text) >= 1
505 + # Finish reason might not always be 'length' if the model finishes
506 + # early or due to other reasons, especially with high temperature.
507 + # So, we'll accept 'length' or 'stop'.
508 + assert choice.finish_reason in ("length", "stop")
509 +
510 + # Token counts can also vary, so we check they are positive.
511 + assert completion.usage.completion_tokens > 0
512 + assert completion.usage.prompt_tokens > 0
513 + assert completion.usage.total_tokens > 0
514 + return completion
515 +
516 + # Test single request
517 + result = await make_request()
518 + assert result is not None
519 + print("API-only server handled single completion request successfully")
520 +
521 + await asyncio.sleep(0.5)
522 +
523 + # Send multiple requests - should be distributed across engines on
524 + # headless server
525 + num_requests = 50
526 + all_tasks = [make_request() for _ in range(num_requests)]
527 +
528 + results = await asyncio.gather(*all_tasks)
529 + assert len(results) == num_requests
530 + assert all(completion is not None for completion in results)
531 +
532 + await asyncio.sleep(0.5)
533 +
534 + # Second burst of requests
535 + all_tasks = [make_request() for _ in range(num_requests)]
536 +
537 + results = await asyncio.gather(*all_tasks)
538 + assert len(results) == num_requests
539 + assert all(completion is not None for completion in results)
540 +
541 + _, api_server_args = api_only_servers[0]
542 + api_server_count = (
543 + api_server_args.count('--api-server-count')
544 + and api_server_args[api_server_args.index('--api-server-count') + 1]
545 + or 1)
546 + print(f"Successfully completed API-only multi-node test with {DP_SIZE} "
547 + f"engines on headless server (API server count: {api_server_count})")
548 +
549 + # Check request balancing via Prometheus metrics
550 + api_server = api_only_servers[0][0]
551 + check_request_balancing(api_server, DP_SIZE)
552 +
553 +
554 +@pytest.mark.asyncio
555 +@pytest.mark.parametrize(
556 + "model_name",
557 + [MODEL_NAME],
558 +)
559 +async def test_api_only_multinode_dp_completion_streaming(
560 + api_only_client: openai.AsyncOpenAI,
561 + api_only_servers: list[tuple[RemoteOpenAIServer,
562 + list[str]]], model_name: str) -> None:
563 + """Test API-only server streaming with all engines on separate
564 + headless server."""
565 + prompt = "What is an LLM?"
566 +
567 + async def make_streaming_request():
568 + # Perform a non-streaming request to get the expected full output
569 + single_completion = await api_only_client.completions.create(
570 + model=model_name,
571 + prompt=prompt,
572 + max_tokens=5,
573 + temperature=0.0,
574 + )
575 + single_output = single_completion.choices[0].text
576 +
577 + # Perform the streaming request
578 + stream = await api_only_client.completions.create(model=model_name,
579 + prompt=prompt,
580 + max_tokens=5,
581 + temperature=0.0,
582 + stream=True)
583 + chunks: list[str] = []
584 + finish_reason_count = 0
585 + last_chunk = None
586 + async for chunk in stream:
587 + chunks.append(chunk.choices[0].text)
588 + if chunk.choices[0].finish_reason is not None:
589 + finish_reason_count += 1
590 + last_chunk = chunk # Keep track of the last chunk
591 +
592 + # finish reason should only return in the last block for OpenAI API
593 + assert finish_reason_count == 1, (
594 + "Finish reason should appear exactly once.")
595 + assert last_chunk is not None, (
596 + "Stream should have yielded at least one chunk.")
597 + assert last_chunk.choices[
598 + 0].finish_reason == "length", "Finish reason should be 'length'."
599 + # Check that the combined text matches the non-streamed version.
600 + assert "".join(
601 + chunks
602 + ) == single_output, "Streamed output should match non-streamed output."
603 + return True # Indicate success for this request
604 +
605 + # Test single streaming request
606 + result = await make_streaming_request()
607 + assert result is not None
608 + print("API-only server handled single streaming request successfully")
609 +
610 + await asyncio.sleep(0.5)
611 +
612 + # Send multiple streaming requests - should be distributed across engines
613 + num_requests = 50
614 + all_tasks = [make_streaming_request() for _ in range(num_requests)]
615 +
616 + results = await asyncio.gather(*all_tasks)
617 + assert len(results) == num_requests
618 + assert all(results), "Not all streaming requests completed successfully."
619 +
620 + await asyncio.sleep(0.5)
621 +
622 + # Second burst of streaming requests
623 + all_tasks = [make_streaming_request() for _ in range(num_requests)]
624 +
625 + results = await asyncio.gather(*all_tasks)
626 + assert len(results) == num_requests
627 + assert all(results), "Not all streaming requests completed successfully."
628 +
629 + _, api_server_args = api_only_servers[0]
630 + api_server_count = (
631 + api_server_args.count('--api-server-count')
632 + and api_server_args[api_server_args.index('--api-server-count') + 1]
633 + or 1)
634 + print(f"Successfully completed API-only streaming test with {DP_SIZE} "
635 + f"engines on headless server (API server count: {api_server_count})")
636 +
637 + # Check request balancing via Prometheus metrics
638 + api_server = api_only_servers[0][0]
639 + check_request_balancing(api_server, DP_SIZE)
tests/v1/test_utils.py
+124
@@ -1,8 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import re
5 +
6 +import pytest
7 +import requests
8 import torch
9
10 +from tests.utils import RemoteOpenAIServer
11 from vllm.v1.worker.utils import bind_kv_cache
12
13
@@ -61,3 +66,122 @@ def test_bind_kv_cache_non_attention():
66
67 assert runner_kv_caches[0] is kv_cache['model.layers.20.attn']
68 assert runner_kv_caches[1] is kv_cache['model.layers.28.attn']
69 +
70 +
71 +# Prometheus metrics utilities for testing
72 +
73 +
74 +def get_prometheus_metrics(
75 + server: RemoteOpenAIServer) -> dict[str, dict[str, float]]:
76 + """Fetch and parse Prometheus metrics from the /metrics endpoint.
77 +
78 + Returns:
79 + Dict mapping metric names to their values grouped by labels.
80 + For example: {"vllm:request_success": {
81 + "engine=0": 5.0, "engine=1": 3.0}
82 + }
83 + """
84 + try:
85 + response = requests.get(server.url_for("metrics"), timeout=10)
86 + response.raise_for_status()
87 +
88 + metrics: dict[str, dict[str, float]] = {}
89 +
90 + # Regex patterns for Prometheus metrics
91 + metric_with_labels = re.compile(
92 + r'^([a-zA-Z_:][a-zA-Z0-9_:]*)\{([^}]*)\}\s+([\d\.\-\+e]+)$')
93 + metric_simple = re.compile(
94 + r'^([a-zA-Z_:][a-zA-Z0-9_:]*)\s+([\d\.\-\+e]+)$')
95 +
96 + for line in response.text.split('\n'):
97 + line = line.strip()
98 + # Skip comments and empty lines
99 + if not line or line.startswith('#'):
100 + continue
101 +
102 + # Try to match metric with labels first
103 + match = metric_with_labels.match(line)
104 + if match:
105 + metric_name, labels_part, value_str = match.groups()
106 + try:
107 + value = float(value_str)
108 + if metric_name not in metrics:
109 + metrics[metric_name] = {}
110 + metrics[metric_name][f'{{{labels_part}}}'] = value
111 + except ValueError:
112 + continue
113 + else:
114 + # Try simple metric without labels
115 + match = metric_simple.match(line)
116 + if match:
117 + metric_name, value_str = match.groups()
118 + try:
119 + value = float(value_str)
120 + if metric_name not in metrics:
121 + metrics[metric_name] = {}
122 + metrics[metric_name][''] = value
123 + except ValueError:
124 + continue
125 +
126 + return metrics
127 + except Exception as e:
128 + pytest.fail(f"Failed to fetch Prometheus metrics: {e}")
129 + return {}
130 +
131 +
132 +def get_engine_request_counts(
133 + metrics: dict[str, dict[str, float]]) -> dict[str, float]:
134 + """Extract request counts per engine from Prometheus metrics.
135 +
136 + Returns:
137 + Dict mapping engine indices to request counts.
138 + For example: {"0": 15.0, "1": 12.0}
139 + """
140 + engine_counts = {}
141 +
142 + # Look for request success metrics with engine labels
143 + success_metrics = metrics.get("vllm:request_success_total", {})
144 + engine_pattern = re.compile(r'engine="([^"]*)"')
145 +
146 + for labels, count in success_metrics.items():
147 + # Extract engine ID from labels using regex
148 + match = engine_pattern.search(labels)
149 + if match:
150 + engine_id = match.group(1)
151 + if engine_id not in engine_counts:
152 + engine_counts[engine_id] = 0.0
153 + engine_counts[engine_id] += count
154 +
155 + return engine_counts
156 +
157 +
158 +def check_request_balancing(server: RemoteOpenAIServer, dp_size: int):
159 + """Check request balancing via Prometheus metrics if dp_size > 1.
160 +
161 + Args:
162 + server: The RemoteOpenAIServer instance
163 + dp_size: Number of data parallel ranks
164 + """
165 + if dp_size <= 1:
166 + return
167 +
168 + # Get metrics after all requests are completed
169 + metrics = get_prometheus_metrics(server)
170 + engine_counts = get_engine_request_counts(metrics)
171 +
172 + # Check that multiple engines received requests
173 + engines_with_requests = [
174 + engine for engine, count in engine_counts.items() if count > 0
175 + ]
176 + assert len(engines_with_requests) == dp_size, (
177 + f"Expected requests to be distributed across multiple engines,"
178 + f" but only engine(s) {engines_with_requests} received "
179 + f"requests. Engine counts: {engine_counts}")
180 +
181 + # Verify that the load is reasonably balanced
182 + # (no engine should handle all requests)
183 + total_requests = sum(engine_counts.values())
184 +
185 + for count in engine_counts.values():
186 + assert count > total_requests // (dp_size + 1), (
187 + f"requests are imbalanced: {engine_counts}")