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, 05:30 UTC
[BugFix] Fix port lookup in internal DP LB tests (#22252)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Aug 14, 2025 at 20:17 UTC
ae05a6d83d8c7aab9bdfb98f92f915a3a1f852bd
1 file changed
+32
-22
tests/v1/test_internal_lb_dp.py
+32
-22
@@ -4,6 +4,8 @@ import asyncio
4
import os
5
import threading
6
import time
7
+import traceback
8
+from typing import Optional, cast
9
10
import openai # use the official client for correctness check
11
import pytest
@@ -41,12 +43,15 @@ class MultinodeInternalLBServerManager:
43
self.tp_size = tp_size
44
self.api_server_count = api_server_count
45
self.base_server_args = base_server_args
44
- self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = []
46
+ self.servers: list[Optional[tuple[RemoteOpenAIServer,
47
+ list[str]]]] = [None] * (dp_size //
48
+ dp_per_node)
49
self.server_threads: list[threading.Thread] = []
50
51
def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
52
"""Start all server instances for multi-node internal LB mode."""
49
- for rank in range(0, self.dp_size, self.dp_per_node):
53
+ for server_idx, rank in enumerate(
54
+ range(0, self.dp_size, self.dp_per_node)):
55
# Create server args for this specific rank
56
server_args = self.base_server_args.copy()
57
@@ -87,7 +92,7 @@ class MultinodeInternalLBServerManager:
92
])
93
94
# Use a thread to start each server to allow parallel initialization
90
- def start_server(r: int, sargs: list[str]):
95
+ def start_server(sidx: int, r: int, sargs: list[str]):
96
gpus_per_node = self.tp_size * self.dp_per_node
97
try:
98
# Start the server
@@ -110,13 +115,14 @@ class MultinodeInternalLBServerManager:
115
f"{self.api_server_count} API servers")
116
else:
117
print(f"Headless node (rank {r}) started successfully")
113
- self.servers.append((server, sargs))
118
+ self.servers[sidx] = (server, sargs)
119
except Exception as e:
120
print(f"Failed to start server rank {r}: {e}")
121
+ traceback.print_exc()
122
raise
123
124
thread = threading.Thread(target=start_server,
119
- args=(rank, server_args))
125
+ args=(server_idx, rank, server_args))
126
thread.start()
127
128
self.server_threads.append(thread)
@@ -128,18 +134,20 @@ class MultinodeInternalLBServerManager:
134
# Give servers additional time to fully initialize and coordinate
135
time.sleep(3)
136
131
- if len(self.servers) != self.dp_size // self.dp_per_node:
137
+ if not all(self.servers):
138
raise Exception("Servers failed to start")
139
134
- return self.servers
140
+ return cast(list[tuple[RemoteOpenAIServer, list[str]]], self.servers)
141
142
def __exit__(self, exc_type, exc_val, exc_tb):
143
"""Stop all server instances."""
144
while self.servers:
139
- try:
140
- self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
141
- except Exception as e:
142
- print(f"Error stopping server: {e}")
145
+ if server := self.servers.pop():
146
+ try:
147
+ server[0].__exit__(exc_type, exc_val, exc_tb)
148
+ except Exception as e:
149
+ print(f"Error stopping server: {e}")
150
+ traceback.print_exc()
151
152
153
class APIOnlyServerManager:
@@ -157,7 +165,8 @@ class APIOnlyServerManager:
165
self.tp_size = tp_size
166
self.api_server_count = api_server_count
167
self.base_server_args = base_server_args
160
- self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = []
168
+ self.servers: list[Optional[tuple[RemoteOpenAIServer,
169
+ list[str]]]] = [None] * 2
170
self.server_threads: list[threading.Thread] = []
171
172
def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
@@ -209,7 +218,7 @@ class APIOnlyServerManager:
218
server.__enter__()
219
print(f"API-only server started successfully with "
220
f"{self.api_server_count} API servers")
212
- self.servers.append((server, api_server_args))
221
+ self.servers[0] = (server, api_server_args)
222
except Exception as e:
223
print(f"Failed to start API-only server: {e}")
224
raise
@@ -231,7 +240,7 @@ class APIOnlyServerManager:
240
server.__enter__()
241
print(f"Headless engines server started successfully with "
242
f"{self.dp_size} engines")
234
- self.servers.append((server, engines_server_args))
243
+ self.servers[1] = (server, engines_server_args)
244
except Exception as e:
245
print(f"Failed to start headless engines server: {e}")
246
raise
@@ -253,18 +262,20 @@ class APIOnlyServerManager:
262
# Give servers additional time to fully initialize and coordinate
263
time.sleep(3)
264
256
- if len(self.servers) != 2:
265
+ if not all(self.servers):
266
raise Exception("Both servers failed to start")
267
259
- return self.servers
268
+ return cast(list[tuple[RemoteOpenAIServer, list[str]]], self.servers)
269
270
def __exit__(self, exc_type, exc_val, exc_tb):
271
"""Stop both server instances."""
272
while self.servers:
264
- try:
265
- self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
266
- except Exception as e:
267
- print(f"Error stopping server: {e}")
273
+ if server := self.servers.pop():
274
+ try:
275
+ server[0].__exit__(exc_type, exc_val, exc_tb)
276
+ except Exception as e:
277
+ print(f"Error stopping server: {e}")
278
+ traceback.print_exc()
279
280
281
@pytest.fixture(scope="module")
@@ -560,7 +571,7 @@ async def test_api_only_multinode_dp_completion(
571
assert len(results) == num_requests
572
assert all(completion is not None for completion in results)
573
563
- _, api_server_args = api_only_servers[0]
574
+ api_server, api_server_args = api_only_servers[0]
575
api_server_count = (
576
api_server_args.count('--api-server-count')
577
and api_server_args[api_server_args.index('--api-server-count') + 1]
@@ -569,7 +580,6 @@ async def test_api_only_multinode_dp_completion(
580
f"engines on headless server (API server count: {api_server_count})")
581
582
# Check request balancing via Prometheus metrics
572
- api_server = api_only_servers[0][0]
583
check_request_balancing(api_server, DP_SIZE)
584
585