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, 08:45 UTC
[CI] Clean up remote servers on pytest parent exit (#41570)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
May 3, 2026 at 22:11 UTC
67058ca326ac566db1aac6fc20ff0c43510422f8
1 file changed
+80
-3
tests/utils.py
+80
-3
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import asyncio
5
+import atexit
6
import contextlib
7
import copy
8
import functools
@@ -134,6 +135,11 @@ class RemoteVLLMServer:
135
"""
136
137
DUMMY_API_KEY = "token-abc123" # vLLM's OpenAI server does not need API key
138
+ _active_servers: set["RemoteVLLMServer"] = set()
139
+ _active_servers_lock = threading.RLock()
140
+ _cleanup_hooks_registered = False
141
+ _signal_hooks_registered = False
142
+ _previous_signal_handlers: dict[int, Any] = {}
143
proc: subprocess.Popen
144
145
def _create_cli_subcommand(self):
@@ -209,6 +215,7 @@ class RemoteVLLMServer:
215
)
216
217
self._pre_download_model(model, args)
218
+ self._shutdown_complete = False
219
220
# Record GPU memory before server start so we know what
221
# "released" looks like.
@@ -221,6 +228,7 @@ class RemoteVLLMServer:
228
)
229
230
self._start_server(model, vllm_serve_args, env_dict)
231
+ self._register_active_server()
232
max_wait_seconds = max_wait_seconds or 480
233
try:
234
self._wait_for_server(url=self.url_for("health"), timeout=max_wait_seconds)
@@ -246,8 +254,70 @@ class RemoteVLLMServer:
254
(when the server fails to start). Must be safe to call even if
255
the process is already dead.
256
"""
249
- self._terminate_process_tree()
250
- self._wait_for_gpu_memory_release()
257
+ if self._shutdown_complete:
258
+ return
259
+
260
+ self._shutdown_complete = True
261
+ try:
262
+ self._terminate_process_tree()
263
+ self._wait_for_gpu_memory_release()
264
+ finally:
265
+ self._unregister_active_server()
266
+
267
+ @classmethod
268
+ def _ensure_cleanup_hooks_registered(cls) -> None:
269
+ """Register process-exit cleanup for detached server subprocesses."""
270
+ root_cls = RemoteVLLMServer
271
+ with root_cls._active_servers_lock:
272
+ if not root_cls._cleanup_hooks_registered:
273
+ atexit.register(root_cls._shutdown_active_servers)
274
+ root_cls._cleanup_hooks_registered = True
275
+
276
+ if (
277
+ threading.current_thread() is threading.main_thread()
278
+ and not root_cls._signal_hooks_registered
279
+ ):
280
+ for signum in (signal.SIGTERM, signal.SIGINT):
281
+ root_cls._previous_signal_handlers[signum] = signal.getsignal(
282
+ signum
283
+ )
284
+ signal.signal(signum, root_cls._handle_parent_signal)
285
+ root_cls._signal_hooks_registered = True
286
+
287
+ def _register_active_server(self) -> None:
288
+ """Track this server so parent-process exits still clean it up."""
289
+ RemoteVLLMServer._ensure_cleanup_hooks_registered()
290
+ with RemoteVLLMServer._active_servers_lock:
291
+ RemoteVLLMServer._active_servers.add(self)
292
+
293
+ def _unregister_active_server(self) -> None:
294
+ with RemoteVLLMServer._active_servers_lock:
295
+ RemoteVLLMServer._active_servers.discard(self)
296
+
297
+ @classmethod
298
+ def _shutdown_active_servers(cls) -> None:
299
+ """Best-effort shutdown for all live RemoteVLLMServer instances."""
300
+ with cls._active_servers_lock:
301
+ servers = list(cls._active_servers)
302
+
303
+ for server in servers:
304
+ with contextlib.suppress(Exception):
305
+ server._shutdown()
306
+
307
+ @classmethod
308
+ def _handle_parent_signal(cls, signum, frame) -> None:
309
+ """Clean up detached servers before letting the signal terminate pytest."""
310
+ cls._shutdown_active_servers()
311
+
312
+ previous_handler = cls._previous_signal_handlers.get(signum, signal.SIG_DFL)
313
+ if callable(previous_handler):
314
+ previous_handler(signum, frame)
315
+ elif previous_handler == signal.SIG_IGN:
316
+ return
317
+ elif signum == signal.SIGINT:
318
+ raise KeyboardInterrupt
319
+ else:
320
+ raise SystemExit(128 + signum)
321
322
def _terminate_process_tree(self) -> None:
323
"""Kill the server process tree without waiting for GPU memory release.
@@ -315,6 +385,9 @@ class RemoteVLLMServer:
385
if not servers:
386
return
387
388
+ for server in servers:
389
+ server._shutdown_complete = True
390
+
391
threads = [
392
threading.Thread(
393
target=s._terminate_process_tree,
@@ -339,7 +412,11 @@ class RemoteVLLMServer:
412
else s._pre_server_gpu_memory
413
),
414
)
342
- earliest._wait_for_gpu_memory_release()
415
+ try:
416
+ earliest._wait_for_gpu_memory_release()
417
+ finally:
418
+ for server in servers:
419
+ server._unregister_active_server()
420
421
def _kill_process_group_survivors(
422
self, pgid: int | None, timeout: float = 15.0