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 4, 11:15 UTC
[Ci][BugFix] Fix slow DP tests due to bad teardown logic (#41166)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Apr 29, 2026 at 16:31 UTC
18599bfdf2f9dc117b004491f1eba766934310fc
4 files changed
+81
-30
tests/utils.py
+55
-6
@@ -14,9 +14,10 @@ import signal
14
import subprocess
15
import sys
16
import tempfile
17
+import threading
18
import time
19
import warnings
19
-from collections.abc import Callable, Iterable
20
+from collections.abc import Callable, Iterable, Sequence
21
from contextlib import ExitStack, contextmanager, suppress
22
from multiprocessing import Process
23
from pathlib import Path
@@ -62,8 +63,6 @@ from vllm.utils.torch_utils import (
63
FP8_DTYPE = current_platform.fp8_dtype()
64
65
if current_platform.is_rocm():
65
- import threading
66
-
66
from amdsmi import (
67
amdsmi_get_gpu_vram_usage,
68
amdsmi_get_processor_handles,
@@ -247,6 +246,16 @@ class RemoteVLLMServer:
246
(when the server fails to start). Must be safe to call even if
247
the process is already dead.
248
"""
249
+ self._terminate_process_tree()
250
+ self._wait_for_gpu_memory_release()
251
+
252
+ def _terminate_process_tree(self) -> None:
253
+ """Kill the server process tree without waiting for GPU memory release.
254
+
255
+ Split out from ``_shutdown`` so that ``shutdown_many`` can run this
256
+ phase in parallel for sibling servers and then wait for GPU memory
257
+ release once at the end.
258
+ """
259
pid = self.proc.pid
260
261
# Get the process group ID. Because we used
@@ -288,9 +297,49 @@ class RemoteVLLMServer:
297
# prevent VRAM from being reclaimed by the driver.
298
self._kill_process_group_survivors(pgid)
299
291
- # Wait for GPU memory to actually be freed, not just
292
- # "stabilized at whatever level it's at".
293
- self._wait_for_gpu_memory_release()
300
+ @classmethod
301
+ def shutdown_many(cls, servers: Sequence["RemoteVLLMServer"]) -> None:
302
+ """Shut down multiple sibling servers and wait for GPU memory once.
303
+
304
+ Test fixtures that hold several ``RemoteVLLMServer`` instances at
305
+ once must NOT shut them down by calling each server's ``__exit__``
306
+ sequentially: every server measures total GPU memory across all
307
+ visible devices in ``_wait_for_gpu_memory_release``, so the first
308
+ server's wait blocks the full timeout because later sibling
309
+ servers are still holding GPU memory.
310
+
311
+ Instead, this method terminates every server's process tree in
312
+ parallel, then runs the GPU-memory-release wait once against the
313
+ earliest recorded baseline (memory before any server started).
314
+ """
315
+ if not servers:
316
+ return
317
+
318
+ threads = [
319
+ threading.Thread(
320
+ target=s._terminate_process_tree,
321
+ name=f"shutdown-{s.proc.pid}",
322
+ daemon=True,
323
+ )
324
+ for s in servers
325
+ ]
326
+ for t in threads:
327
+ t.start()
328
+ for t in threads:
329
+ t.join()
330
+
331
+ # Use the smallest pre-server baseline so the wait targets memory
332
+ # usage before *any* of these sibling servers started, not after
333
+ # earlier siblings had already allocated.
334
+ earliest = min(
335
+ servers,
336
+ key=lambda s: (
337
+ float("inf")
338
+ if s._pre_server_gpu_memory is None
339
+ else s._pre_server_gpu_memory
340
+ ),
341
+ )
342
+ earliest._wait_for_gpu_memory_release()
343
344
def _kill_process_group_survivors(
345
self, pgid: int | None, timeout: float = 15.0
tests/v1/distributed/test_external_lb_dp.py
+6
-5
@@ -111,11 +111,12 @@ class ExternalLBServerManager:
111
112
def __exit__(self, exc_type, exc_val, exc_tb):
113
"""Stop all server instances."""
114
- while self.servers:
115
- try:
116
- self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
117
- except Exception as e:
118
- print(f"Error stopping server: {e}")
114
+ servers = [s for s, _ in self.servers]
115
+ self.servers.clear()
116
+ try:
117
+ RemoteOpenAIServer.shutdown_many(servers)
118
+ except Exception as e:
119
+ print(f"Error stopping servers: {e}")
120
121
122
@pytest.fixture(scope="module")
tests/v1/distributed/test_hybrid_lb_dp.py
+6
-5
@@ -134,11 +134,12 @@ class HybridLBServerManager:
134
135
def __exit__(self, exc_type, exc_val, exc_tb):
136
"""Stop all server instances."""
137
- while self.servers:
138
- try:
139
- self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
140
- except Exception as e:
141
- print(f"Error stopping server: {e}")
137
+ servers = [s for s, _ in self.servers]
138
+ self.servers.clear()
139
+ try:
140
+ RemoteOpenAIServer.shutdown_many(servers)
141
+ except Exception as e:
142
+ print(f"Error stopping servers: {e}")
143
144
145
@pytest.fixture(scope="module")
tests/v1/distributed/test_internal_lb_dp.py
+14
-14
@@ -228,13 +228,13 @@ class MultinodeInternalLBServerManager:
228
229
def __exit__(self, exc_type, exc_val, exc_tb):
230
"""Stop all server instances."""
231
- while self.servers:
232
- if server := self.servers.pop():
233
- try:
234
- server[0].__exit__(exc_type, exc_val, exc_tb)
235
- except Exception as e:
236
- print(f"Error stopping server: {e}")
237
- traceback.print_exc()
231
+ servers = [entry[0] for entry in self.servers if entry is not None]
232
+ self.servers.clear()
233
+ try:
234
+ RemoteOpenAIServer.shutdown_many(servers)
235
+ except Exception as e:
236
+ print(f"Error stopping servers: {e}")
237
+ traceback.print_exc()
238
239
240
class APIOnlyServerManager:
@@ -370,13 +370,13 @@ class APIOnlyServerManager:
370
371
def __exit__(self, exc_type, exc_val, exc_tb):
372
"""Stop both server instances."""
373
- while self.servers:
374
- if server := self.servers.pop():
375
- try:
376
- server[0].__exit__(exc_type, exc_val, exc_tb)
377
- except Exception as e:
378
- print(f"Error stopping server: {e}")
379
- traceback.print_exc()
373
+ servers = [entry[0] for entry in self.servers if entry is not None]
374
+ self.servers.clear()
375
+ try:
376
+ RemoteOpenAIServer.shutdown_many(servers)
377
+ except Exception as e:
378
+ print(f"Error stopping servers: {e}")
379
+ traceback.print_exc()
380
381
382
@pytest.fixture(scope="module")