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 6, 14:30 UTC
[CI] Stabilize multinode DP internal LB completion tests (#36356)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 16, 2026 at 17:40 UTC
4f9b14c21cd4eb4b56c972b3280be41d341056d1
1 file changed
+89
-94
tests/v1/distributed/test_internal_lb_dp.py
+89
-94
@@ -12,7 +12,7 @@ import pytest
12
import pytest_asyncio
13
import requests
14
15
-from tests.utils import RemoteOpenAIServer
15
+from tests.utils import ROCM_ENV_OVERRIDES, RemoteOpenAIServer
16
from tests.v1.utils import check_request_balancing
17
from vllm.platforms import current_platform
18
@@ -27,6 +27,84 @@ TP_SIZE = int(os.getenv("TP_SIZE", "1"))
27
NUM_NODES = 2
28
29
30
+async def _make_completion_request(
31
+ client: openai.AsyncOpenAI,
32
+ model_name: str,
33
+) -> openai.types.Completion:
34
+ """Make a single completion request and validate the response.
35
+
36
+ Uses temperature=1.0 to ensure diverse outputs across concurrent
37
+ requests for realistic load balancer testing.
38
+ """
39
+ completion = await client.completions.create(
40
+ model=model_name,
41
+ prompt="Hello, my name is",
42
+ max_tokens=5,
43
+ temperature=1.0,
44
+ )
45
+
46
+ assert completion.id is not None, (
47
+ f"Expected non-None completion id. usage={completion.usage!r}"
48
+ )
49
+ assert completion.choices is not None and len(completion.choices) == 1, (
50
+ f"Expected 1 choice, got "
51
+ f"{len(completion.choices) if completion.choices else 'None'}"
52
+ )
53
+
54
+ choice = completion.choices[0]
55
+ # With temperature=1.0, the model may emit a stop token immediately,
56
+ # producing empty text with finish_reason='stop'. This is valid
57
+ # model behavior - the test's purpose is load balancing, not output
58
+ # quality.
59
+ assert choice.finish_reason in ("length", "stop"), (
60
+ f"Expected finish_reason 'length' or 'stop', "
61
+ f"got {choice.finish_reason!r}. text={choice.text!r}"
62
+ )
63
+ if choice.finish_reason == "length":
64
+ assert len(choice.text) >= 1, (
65
+ f"Expected non-empty text with finish_reason='length', got {choice.text!r}"
66
+ )
67
+
68
+ assert completion.usage.prompt_tokens > 0, (
69
+ f"Expected positive prompt_tokens, got {completion.usage.prompt_tokens}"
70
+ )
71
+ assert completion.usage.total_tokens > 0, (
72
+ f"Expected positive total_tokens, got {completion.usage.total_tokens}"
73
+ )
74
+ return completion
75
+
76
+
77
+async def _run_request_bursts(
78
+ client: openai.AsyncOpenAI,
79
+ model_name: str,
80
+ num_requests: int = 200,
81
+ num_bursts: int = 2,
82
+):
83
+ """Send multiple bursts of completion requests and validate all succeed."""
84
+ for burst in range(num_bursts):
85
+ all_tasks = []
86
+ for _ in range(num_requests):
87
+ all_tasks.append(
88
+ asyncio.create_task(_make_completion_request(client, model_name))
89
+ )
90
+ await asyncio.sleep(0.01)
91
+
92
+ results = await asyncio.gather(*all_tasks, return_exceptions=True)
93
+ assert len(results) == num_requests, (
94
+ f"Burst {burst}: expected {num_requests} results, got {len(results)}"
95
+ )
96
+
97
+ for result in results:
98
+ if isinstance(result, BaseException):
99
+ raise result
100
+
101
+ assert all(completion is not None for completion in results), (
102
+ f"Burst {burst}: some completions were None"
103
+ )
104
+
105
+ await asyncio.sleep(0.5)
106
+
107
+
108
class MultinodeInternalLBServerManager:
109
"""Manages multi-node data parallel vLLM server instances for internal
110
load balancer testing using --headless mode."""
@@ -108,6 +186,7 @@ class MultinodeInternalLBServerManager:
186
auto_port=False,
187
env_dict={
188
"VLLM_SERVER_DEV_MODE": "1",
189
+ **ROCM_ENV_OVERRIDES,
190
current_platform.device_control_env_var: ",".join(
191
str(current_platform.device_id_to_physical_device_id(i))
192
for i in range(r, r + gpus_per_node)
@@ -229,6 +308,7 @@ class APIOnlyServerManager:
308
auto_port=False,
309
env_dict={
310
"VLLM_SERVER_DEV_MODE": "1",
311
+ **ROCM_ENV_OVERRIDES,
312
# No GPUs needed for API-only server
313
},
314
)
@@ -249,10 +329,11 @@ class APIOnlyServerManager:
329
engines_server_args,
330
auto_port=False,
331
env_dict={
332
+ **ROCM_ENV_OVERRIDES,
333
current_platform.device_control_env_var: ",".join(
334
str(current_platform.device_id_to_physical_device_id(i))
335
for i in range(self.dp_size * self.tp_size)
255
- )
336
+ ),
337
},
338
)
339
server.__enter__()
@@ -395,58 +476,15 @@ async def test_multinode_dp_completion(
476
servers: list[tuple[RemoteOpenAIServer, list[str]]],
477
model_name: str,
478
) -> None:
398
- async def make_request():
399
- completion = await client.completions.create(
400
- model=model_name, prompt="Hello, my name is", max_tokens=5, temperature=1.0
401
- )
402
-
403
- assert completion.id is not None
404
- assert completion.choices is not None and len(completion.choices) == 1
405
-
406
- choice = completion.choices[0]
407
- # The exact number of tokens can vary slightly with temperature=1.0,
408
- # so we check for a reasonable minimum length.
409
- assert len(choice.text) >= 1
410
- # Finish reason might not always be 'length' if the model finishes early
411
- # or due to other reasons, especially with high temperature.
412
- # So, we'll accept 'length' or 'stop'.
413
- assert choice.finish_reason in ("length", "stop")
414
-
415
- # Token counts can also vary, so we check they are positive.
416
- assert completion.usage.completion_tokens > 0
417
- assert completion.usage.prompt_tokens > 0
418
- assert completion.usage.total_tokens > 0
419
- return completion
420
-
479
# Test single request
422
- result = await make_request()
480
+ result = await _make_completion_request(client, model_name)
481
assert result is not None
482
print("Multi-node internal LB handled single completion request successfully")
483
484
await asyncio.sleep(0.5)
485
428
- # Send multiple requests - internal LB should distribute across DP ranks
429
- num_requests = 200
430
- all_tasks = []
431
- for _ in range(num_requests):
432
- all_tasks.append(asyncio.create_task(make_request()))
433
- await asyncio.sleep(0.01)
434
-
435
- results = await asyncio.gather(*all_tasks)
436
- assert len(results) == num_requests
437
- assert all(completion is not None for completion in results)
438
-
439
- await asyncio.sleep(0.5)
440
-
441
- # Second burst of requests
442
- all_tasks = []
443
- for _ in range(num_requests):
444
- all_tasks.append(asyncio.create_task(make_request()))
445
- await asyncio.sleep(0.01)
446
-
447
- results = await asyncio.gather(*all_tasks)
448
- assert len(results) == num_requests
449
- assert all(completion is not None for completion in results)
486
+ # Send multiple bursts - internal LB should distribute across DP ranks
487
+ await _run_request_bursts(client, model_name)
488
489
_, server_args = servers[0]
490
api_server_count = (
@@ -570,59 +608,16 @@ async def test_api_only_multinode_dp_completion(
608
) -> None:
609
"""Test API-only server with all engines on separate headless server."""
610
573
- async def make_request():
574
- completion = await api_only_client.completions.create(
575
- model=model_name, prompt="Hello, my name is", max_tokens=5, temperature=1.0
576
- )
577
-
578
- assert completion.id is not None
579
- assert completion.choices is not None and len(completion.choices) == 1
580
-
581
- choice = completion.choices[0]
582
- # The exact number of tokens can vary slightly with temperature=1.0,
583
- # so we check for a reasonable minimum length.
584
- assert len(choice.text) >= 1
585
- # Finish reason might not always be 'length' if the model finishes
586
- # early or due to other reasons, especially with high temperature.
587
- # So, we'll accept 'length' or 'stop'.
588
- assert choice.finish_reason in ("length", "stop")
589
-
590
- # Token counts can also vary, so we check they are positive.
591
- assert completion.usage.completion_tokens > 0
592
- assert completion.usage.prompt_tokens > 0
593
- assert completion.usage.total_tokens > 0
594
- return completion
595
-
611
# Test single request
597
- result = await make_request()
612
+ result = await _make_completion_request(api_only_client, model_name)
613
assert result is not None
614
print("API-only server handled single completion request successfully")
615
616
await asyncio.sleep(0.5)
617
603
- # Send multiple requests - should be distributed across engines on
618
+ # Send multiple bursts - should be distributed across engines on
619
# headless server
605
- num_requests = 200
606
- all_tasks = []
607
- for _ in range(num_requests):
608
- all_tasks.append(asyncio.create_task(make_request()))
609
- await asyncio.sleep(0.01)
610
-
611
- results = await asyncio.gather(*all_tasks)
612
- assert len(results) == num_requests
613
- assert all(completion is not None for completion in results)
614
-
615
- await asyncio.sleep(0.5)
616
-
617
- # Second burst of requests
618
- all_tasks = []
619
- for _ in range(num_requests):
620
- all_tasks.append(asyncio.create_task(make_request()))
621
- await asyncio.sleep(0.01)
622
-
623
- results = await asyncio.gather(*all_tasks)
624
- assert len(results) == num_requests
625
- assert all(completion is not None for completion in results)
620
+ await _run_request_bursts(api_only_client, model_name)
621
622
api_server, api_server_args = api_only_servers[0]
623
api_server_count = (