@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7afe0faab1
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 6, 19:15 UTC

[Frontend][Core] Re-add shutdown timeout - allowing in-flight requests to finish (#36666)

Signed-off-by: Mark McLoughlin <markmc@redhat.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Mark McLoughlin committed Mar 13, 2026 at 19:10 UTC 7afe0faab1eb2ab84cda5cab29b24046e516f7b8
14 files changed +762 -96
tests/entrypoints/openai/test_shutdown.py
+459
@@ -1,14 +1,20 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Integration tests for shutdown behavior, timeout, and signal handling."""
4
5 +import asyncio
6 import signal
7 import subprocess
8 import sys
9 import time
10 +from dataclasses import dataclass, field
11
12 +import httpx
13 import openai
14 +import psutil
15 import pytest
16
17 +from tests.utils import RemoteOpenAIServer
18 from vllm.platforms import current_platform
19 from vllm.utils.network_utils import get_open_port
20
@@ -18,6 +24,101 @@ MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
24 _IS_ROCM = current_platform.is_rocm()
25 _SERVER_STARTUP_TIMEOUT = 120
26 _PROCESS_EXIT_TIMEOUT = 15
27 +_SHUTDOWN_DETECTION_TIMEOUT = 10
28 +_CHILD_CLEANUP_TIMEOUT = 10
29 +
30 +
31 +def _get_child_pids(parent_pid: int) -> list[int]:
32 + try:
33 + parent = psutil.Process(parent_pid)
34 + return [c.pid for c in parent.children(recursive=True)]
35 + except psutil.NoSuchProcess:
36 + return []
37 +
38 +
39 +async def _assert_children_cleaned_up(
40 + child_pids: list[int],
41 + timeout: float = _CHILD_CLEANUP_TIMEOUT,
42 +):
43 + """Wait for child processes to exit and fail if any remain."""
44 + if not child_pids:
45 + return
46 +
47 + deadline = time.time() + timeout
48 + while time.time() < deadline:
49 + still_alive = []
50 + for pid in child_pids:
51 + try:
52 + p = psutil.Process(pid)
53 + if p.is_running() and p.status() != psutil.STATUS_ZOMBIE:
54 + still_alive.append(pid)
55 + except psutil.NoSuchProcess:
56 + pass
57 + if not still_alive:
58 + return
59 + await asyncio.sleep(0.5)
60 +
61 + pytest.fail(
62 + f"Child processes {still_alive} still alive after {timeout}s. "
63 + f"Process cleanup may not be working correctly."
64 + )
65 +
66 +
67 +@dataclass
68 +class ShutdownState:
69 + got_503: bool = False
70 + got_500: bool = False
71 + requests_after_sigterm: int = 0
72 + aborted_requests: int = 0
73 + connection_errors: int = 0
74 + stop_requesting: bool = False
75 + errors: list[str] = field(default_factory=list)
76 +
77 +
78 +async def _concurrent_request_loop(
79 + client: openai.AsyncOpenAI,
80 + state: ShutdownState,
81 + sigterm_sent: asyncio.Event | None = None,
82 + concurrency: int = 10,
83 +):
84 + """Run multiple concurrent requests to keep the server busy."""
85 +
86 + async def single_request():
87 + while not state.stop_requesting:
88 + try:
89 + response = await client.completions.create(
90 + model=MODEL_NAME,
91 + prompt="Write a story: ",
92 + max_tokens=200,
93 + )
94 + if sigterm_sent is not None and sigterm_sent.is_set():
95 + state.requests_after_sigterm += 1
96 + # Check if any choice has finish_reason='abort'
97 + if any(choice.finish_reason == "abort" for choice in response.choices):
98 + state.aborted_requests += 1
99 + except openai.APIStatusError as e:
100 + if e.status_code == 503:
101 + state.got_503 = True
102 + elif e.status_code == 500:
103 + state.got_500 = True
104 + else:
105 + state.errors.append(f"API error: {e}")
106 + except (openai.APIConnectionError, httpx.RemoteProtocolError):
107 + state.connection_errors += 1
108 + if sigterm_sent is not None and sigterm_sent.is_set():
109 + break
110 + except Exception as e:
111 + state.errors.append(f"Unexpected error: {e}")
112 + break
113 + await asyncio.sleep(0.01)
114 +
115 + tasks = [asyncio.create_task(single_request()) for _ in range(concurrency)]
116 + try:
117 + await asyncio.gather(*tasks, return_exceptions=True)
118 + finally:
119 + for t in tasks:
120 + if not t.done():
121 + t.cancel()
122
123
124 @pytest.mark.asyncio
@@ -103,3 +204,361 @@ async def test_shutdown_on_engine_failure():
204
205 return_code = proc.wait(timeout=_PROCESS_EXIT_TIMEOUT)
206 assert return_code is not None
207 +
208 +
209 +@pytest.mark.asyncio
210 +async def test_wait_timeout_completes_requests():
211 + """Verify wait timeout: new requests rejected, in-flight requests complete."""
212 + server_args = [
213 + "--dtype",
214 + "bfloat16",
215 + "--max-model-len",
216 + "256",
217 + "--enforce-eager",
218 + "--gpu-memory-utilization",
219 + "0.05",
220 + "--max-num-seqs",
221 + "4",
222 + "--shutdown-timeout",
223 + "30",
224 + ]
225 +
226 + with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
227 + client = remote_server.get_async_client()
228 + proc = remote_server.proc
229 + child_pids = _get_child_pids(proc.pid)
230 +
231 + state = ShutdownState()
232 + sigterm_sent = asyncio.Event()
233 +
234 + request_task = asyncio.create_task(
235 + _concurrent_request_loop(client, state, sigterm_sent, concurrency=10)
236 + )
237 +
238 + await asyncio.sleep(0.5)
239 + proc.send_signal(signal.SIGTERM)
240 + sigterm_sent.set()
241 +
242 + try:
243 + await asyncio.wait_for(request_task, timeout=_SHUTDOWN_DETECTION_TIMEOUT)
244 + except asyncio.TimeoutError:
245 + pass
246 + finally:
247 + state.stop_requesting = True
248 + if not request_task.done():
249 + request_task.cancel()
250 + await asyncio.gather(request_task, return_exceptions=True)
251 +
252 + # wait timeout should complete in-flight requests
253 + assert state.requests_after_sigterm > 0, (
254 + f"Wait timeout should complete in-flight requests. "
255 + f"503: {state.got_503}, 500: {state.got_500}, "
256 + f"conn_errors: {state.connection_errors}, errors: {state.errors}"
257 + )
258 + # server must stop accepting new requests (503, 500, or connection close)
259 + assert state.got_503 or state.got_500 or state.connection_errors > 0, (
260 + f"Server should stop accepting requests. "
261 + f"completed: {state.requests_after_sigterm}, errors: {state.errors}"
262 + )
263 +
264 + await _assert_children_cleaned_up(child_pids)
265 +
266 +
267 +@pytest.mark.asyncio
268 +@pytest.mark.parametrize("wait_for_engine_idle", [0.0, 2.0])
269 +async def test_abort_timeout_exits_quickly(wait_for_engine_idle: float):
270 + server_args = [
271 + "--dtype",
272 + "bfloat16",
273 + "--max-model-len",
274 + "256",
275 + "--enforce-eager",
276 + "--gpu-memory-utilization",
277 + "0.05",
278 + "--max-num-seqs",
279 + "4",
280 + "--shutdown-timeout",
281 + "0",
282 + ]
283 +
284 + with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
285 + proc = remote_server.proc
286 + child_pids = _get_child_pids(proc.pid)
287 +
288 + if wait_for_engine_idle > 0:
289 + client = remote_server.get_async_client()
290 + # Send requests to ensure engine is fully initialized
291 + for _ in range(2):
292 + await client.completions.create(
293 + model=MODEL_NAME,
294 + prompt="Test request: ",
295 + max_tokens=10,
296 + )
297 + # Wait for engine to become idle
298 + await asyncio.sleep(wait_for_engine_idle)
299 +
300 + start_time = time.time()
301 + proc.send_signal(signal.SIGTERM)
302 +
303 + # abort timeout (0) should exit promptly
304 + for _ in range(20):
305 + if proc.poll() is not None:
306 + break
307 + time.sleep(0.1)
308 +
309 + if proc.poll() is None:
310 + proc.kill()
311 + proc.wait(timeout=5)
312 + pytest.fail("Process did not exit after SIGTERM with abort timeout")
313 +
314 + exit_time = time.time() - start_time
315 + assert exit_time < 2, f"Default shutdown took too long: {exit_time:.1f}s"
316 + assert proc.returncode in (0, -15, None), f"Unexpected: {proc.returncode}"
317 +
318 + await _assert_children_cleaned_up(child_pids)
319 +
320 +
321 +@pytest.mark.asyncio
322 +async def test_wait_timeout_with_short_duration():
323 + """Verify server exits cleanly with a short wait timeout."""
324 + wait_timeout = 3
325 + server_args = [
326 + "--dtype",
327 + "bfloat16",
328 + "--max-model-len",
329 + "256",
330 + "--enforce-eager",
331 + "--gpu-memory-utilization",
332 + "0.05",
333 + "--max-num-seqs",
334 + "4",
335 + "--shutdown-timeout",
336 + str(wait_timeout),
337 + ]
338 +
339 + with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
340 + client = remote_server.get_async_client()
341 + proc = remote_server.proc
342 + child_pids = _get_child_pids(proc.pid)
343 +
344 + state = ShutdownState()
345 + request_task = asyncio.create_task(
346 + _concurrent_request_loop(client, state, concurrency=3)
347 + )
348 +
349 + await asyncio.sleep(0.5)
350 +
351 + start_time = time.time()
352 + proc.send_signal(signal.SIGTERM)
353 +
354 + # server should exit within wait_timeout + buffer
355 + max_wait = wait_timeout + 15
356 + for _ in range(int(max_wait * 10)):
357 + if proc.poll() is not None:
358 + break
359 + time.sleep(0.1)
360 +
361 + exit_time = time.time() - start_time
362 +
363 + state.stop_requesting = True
364 + if not request_task.done():
365 + request_task.cancel()
366 + await asyncio.gather(request_task, return_exceptions=True)
367 +
368 + if proc.poll() is None:
369 + proc.kill()
370 + proc.wait(timeout=5)
371 + pytest.fail(f"Process did not exit within {max_wait}s after SIGTERM")
372 +
373 + assert exit_time < wait_timeout + 10, (
374 + f"Took too long to exit ({exit_time:.1f}s), expected <{wait_timeout + 10}s"
375 + )
376 + assert proc.returncode in (0, -15, None), f"Unexpected: {proc.returncode}"
377 +
378 + await _assert_children_cleaned_up(child_pids)
379 +
380 +
381 +@pytest.mark.asyncio
382 +async def test_abort_timeout_fails_inflight_requests():
383 + """Verify abort timeout (0) immediately aborts in-flight requests."""
384 + server_args = [
385 + "--dtype",
386 + "bfloat16",
387 + "--max-model-len",
388 + "256",
389 + "--enforce-eager",
390 + "--gpu-memory-utilization",
391 + "0.05",
392 + "--max-num-seqs",
393 + "4",
394 + "--shutdown-timeout",
395 + "0",
396 + ]
397 +
398 + with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
399 + client = remote_server.get_async_client()
400 + proc = remote_server.proc
401 + child_pids = _get_child_pids(proc.pid)
402 +
403 + state = ShutdownState()
404 + sigterm_sent = asyncio.Event()
405 +
406 + request_task = asyncio.create_task(
407 + _concurrent_request_loop(client, state, sigterm_sent, concurrency=10)
408 + )
409 +
410 + await asyncio.sleep(0.5)
411 +
412 + proc.send_signal(signal.SIGTERM)
413 + sigterm_sent.set()
414 +
415 + try:
416 + await asyncio.wait_for(request_task, timeout=5)
417 + except asyncio.TimeoutError:
418 + pass
419 + finally:
420 + state.stop_requesting = True
421 + if not request_task.done():
422 + request_task.cancel()
423 + await asyncio.gather(request_task, return_exceptions=True)
424 +
425 + # With abort timeout (0), requests should be aborted (finish_reason='abort')
426 + # or rejected (connection errors or API errors)
427 + assert (
428 + state.aborted_requests > 0
429 + or state.connection_errors > 0
430 + or state.got_500
431 + or state.got_503
432 + ), (
433 + f"Abort timeout should cause request aborts or failures. "
434 + f"aborted: {state.aborted_requests}, "
435 + f"503: {state.got_503}, 500: {state.got_500}, "
436 + f"conn_errors: {state.connection_errors}, "
437 + f"completed: {state.requests_after_sigterm}"
438 + )
439 +
440 + # Verify fast shutdown
441 + start_time = time.time()
442 + for _ in range(100):
443 + if proc.poll() is not None:
444 + break
445 + time.sleep(0.1)
446 +
447 + exit_time = time.time() - start_time
448 + assert exit_time < 10, f"Abort timeout shutdown took too long: {exit_time:.1f}s"
449 +
450 + await _assert_children_cleaned_up(child_pids)
451 +
452 +
453 +@pytest.mark.asyncio
454 +async def test_request_rejection_during_shutdown():
455 + """Verify new requests are rejected with error during shutdown."""
456 + server_args = [
457 + "--dtype",
458 + "bfloat16",
459 + "--max-model-len",
460 + "256",
461 + "--enforce-eager",
462 + "--gpu-memory-utilization",
463 + "0.05",
464 + "--max-num-seqs",
465 + "4",
466 + "--shutdown-timeout",
467 + "30",
468 + ]
469 +
470 + with RemoteOpenAIServer(MODEL_NAME, server_args) as remote_server:
471 + client = remote_server.get_async_client()
472 + proc = remote_server.proc
473 + child_pids = _get_child_pids(proc.pid)
474 +
475 + proc.send_signal(signal.SIGTERM)
476 +
477 + await asyncio.sleep(1.0)
478 +
479 + # Try to send new requests - they should be rejected
480 + rejected_count = 0
481 + for _ in range(10):
482 + try:
483 + await client.completions.create(
484 + model=MODEL_NAME, prompt="Hello", max_tokens=10
485 + )
486 + except (
487 + openai.APIStatusError,
488 + openai.APIConnectionError,
489 + httpx.RemoteProtocolError,
490 + ):
491 + rejected_count += 1
492 + await asyncio.sleep(0.1)
493 +
494 + assert rejected_count > 0, (
495 + f"Expected requests to be rejected during shutdown, "
496 + f"but {rejected_count} were rejected out of 10"
497 + )
498 +
499 + await _assert_children_cleaned_up(child_pids)
500 +
501 +
502 +@pytest.mark.asyncio
503 +async def test_multi_api_server_shutdown():
504 + """Verify shutdown works with multiple API servers."""
505 + server_args = [
506 + "--dtype",
507 + "bfloat16",
508 + "--max-model-len",
509 + "256",
510 + "--enforce-eager",
511 + "--gpu-memory-utilization",
512 + "0.05",
513 + "--max-num-seqs",
514 + "4",
515 + "--shutdown-timeout",
516 + "30",
517 + "--api-server-count",
518 + "2",
519 + ]
520 +
521 + with RemoteOpenAIServer(MODEL_NAME, server_args, auto_port=True) as remote_server:
522 + client = remote_server.get_async_client()
523 + proc = remote_server.proc
524 + child_pids = _get_child_pids(proc.pid)
525 +
526 + assert len(child_pids) >= 2, (
527 + f"Expected at least 2 child processes, got {len(child_pids)}"
528 + )
529 +
530 + state = ShutdownState()
531 + sigterm_sent = asyncio.Event()
532 +
533 + # Start concurrent requests across both API servers
534 + request_task = asyncio.create_task(
535 + _concurrent_request_loop(client, state, sigterm_sent, concurrency=8)
536 + )
537 +
538 + await asyncio.sleep(0.5)
539 +
540 + # Send SIGTERM to parent - should propagate to all children
541 + proc.send_signal(signal.SIGTERM)
542 + sigterm_sent.set()
543 +
544 + try:
545 + await asyncio.wait_for(request_task, timeout=_SHUTDOWN_DETECTION_TIMEOUT)
546 + except asyncio.TimeoutError:
547 + pass
548 + finally:
549 + state.stop_requesting = True
550 + if not request_task.done():
551 + request_task.cancel()
552 + await asyncio.gather(request_task, return_exceptions=True)
553 +
554 + for _ in range(300): # up to 30 seconds
555 + if proc.poll() is not None:
556 + break
557 + time.sleep(0.1)
558 +
559 + if proc.poll() is None:
560 + proc.kill()
561 + proc.wait(timeout=5)
562 + pytest.fail("Process did not exit after SIGTERM")
563 +
564 + await _assert_children_cleaned_up(child_pids)
tests/entrypoints/test_api_server_process_manager.py
+15 -7
@@ -79,7 +79,7 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update):
79 finally:
80 # Always clean up the processes
81 print("Cleaning up processes...")
82 - manager.close()
82 + manager.shutdown()
83
84 # Give processes time to terminate
85 time.sleep(0.2)
@@ -111,6 +111,8 @@ def test_wait_for_completion_or_failure(api_server_args):
111 wait_for_completion_or_failure(api_server_manager=manager)
112 except Exception as e:
113 result["exception"] = e
114 + finally:
115 + manager.shutdown()
116
117 # Start a thread to run wait_for_completion_or_failure
118 wait_thread = threading.Thread(target=run_with_exception_capture, daemon=True)
@@ -143,7 +145,7 @@ def test_wait_for_completion_or_failure(api_server_args):
145 assert not proc.is_alive(), f"Process {i} should not be alive"
146
147 finally:
146 - manager.close()
148 + manager.shutdown()
149 time.sleep(0.2)
150
151
@@ -174,11 +176,14 @@ def test_normal_completion(api_server_args):
176 # since all processes have already
177 # terminated, it should return immediately
178 # with no error
177 - wait_for_completion_or_failure(api_server_manager=manager)
179 + try:
180 + wait_for_completion_or_failure(api_server_manager=manager)
181 + finally:
182 + manager.shutdown()
183
184 finally:
185 # Clean up just in case
181 - manager.close()
186 + manager.shutdown()
187 time.sleep(0.2)
188
189
@@ -201,7 +206,7 @@ def test_external_process_monitoring(api_server_args):
206 def __init__(self, proc):
207 self.proc = proc
208
204 - def close(self):
209 + def shutdown(self):
210 if self.proc.is_alive():
211 self.proc.terminate()
212 self.proc.join(timeout=0.5)
@@ -226,6 +231,9 @@ def test_external_process_monitoring(api_server_args):
231 )
232 except Exception as e:
233 result["exception"] = e
234 + finally:
235 + manager.shutdown()
236 + mock_coordinator.shutdown()
237
238 # Start a thread to run wait_for_completion_or_failure
239 wait_thread = threading.Thread(target=run_with_exception_capture, daemon=True)
@@ -259,6 +267,6 @@ def test_external_process_monitoring(api_server_args):
267
268 finally:
269 # Clean up
262 - manager.close()
263 - mock_coordinator.close()
270 + manager.shutdown()
271 + mock_coordinator.shutdown()
272 time.sleep(0.2)
vllm/config/vllm.py
+6
@@ -327,6 +327,12 @@ class VllmConfig:
327 weight_transfer_config: WeightTransferConfig | None = None
328 """The configurations for weight transfer during RL training."""
329
330 + shutdown_timeout: int = Field(default=0, ge=0)
331 + """Shutdown grace period for in-flight requests. Shutdown will be delayed for
332 + up to this amount of time to allow already-running requests to complete. Any
333 + remaining requests are aborted once the timeout is reached.
334 + """
335 +
336 def compute_hash(self) -> str:
337 """
338 WARNING: Whenever a new field is added to this config,
vllm/engine/arg_utils.py
+11
@@ -606,6 +606,8 @@ class EngineArgs:
606 kv_offloading_backend: KVOffloadingBackend = CacheConfig.kv_offloading_backend
607 tokens_only: bool = False
608
609 + shutdown_timeout: int = 0
610 +
611 weight_transfer_config: WeightTransferConfig | None = get_field(
612 VllmConfig,
613 "weight_transfer_config",
@@ -1308,6 +1310,14 @@ class EngineArgs:
1310 default=False,
1311 action=argparse.BooleanOptionalAction,
1312 )
1313 +
1314 + parser.add_argument(
1315 + "--shutdown-timeout",
1316 + type=int,
1317 + default=0,
1318 + help="Shutdown timeout in seconds. 0 = abort, >0 = wait.",
1319 + )
1320 +
1321 return parser
1322
1323 @classmethod
@@ -1916,6 +1926,7 @@ class EngineArgs:
1926 optimization_level=self.optimization_level,
1927 performance_mode=self.performance_mode,
1928 weight_transfer_config=self.weight_transfer_config,
1929 + shutdown_timeout=self.shutdown_timeout,
1930 )
1931
1932 return config
vllm/engine/protocol.py
+5
@@ -200,6 +200,11 @@ class EngineClient(ABC):
200 """Return whether the engine is currently paused."""
201 ...
202
203 + @abstractmethod
204 + def shutdown(self, timeout: float | None = None) -> None:
205 + """Shutdown the engine with optional timeout."""
206 + ...
207 +
208 async def scale_elastic_ep(
209 self, new_data_parallel_size: int, drain_timeout: int = 300
210 ) -> None:
vllm/entrypoints/cli/serve.py
+42 -6
@@ -3,6 +3,7 @@
3
4 import argparse
5 import signal
6 +import time
7
8 import uvloop
9
@@ -222,8 +223,12 @@ def run_headless(args: argparse.Namespace):
223 try:
224 engine_manager.join_first()
225 finally:
226 + timeout = None
227 + if shutdown_requested:
228 + timeout = vllm_config.shutdown_timeout
229 + logger.info("Waiting up to %d seconds for processes to exit", timeout)
230 + engine_manager.shutdown(timeout=timeout)
231 logger.info("Shutting down.")
226 - engine_manager.close()
232
233
234 def run_multi_api_server(args: argparse.Namespace):
@@ -234,6 +239,19 @@ def run_multi_api_server(args: argparse.Namespace):
239 if num_api_servers > 1:
240 setup_multiprocess_prometheus()
241
242 + shutdown_requested = False
243 +
244 + # Catch SIGTERM and SIGINT to allow graceful shutdown.
245 + def signal_handler(signum, frame):
246 + nonlocal shutdown_requested
247 + logger.debug("Received %d signal.", signum)
248 + if not shutdown_requested:
249 + shutdown_requested = True
250 + raise SystemExit
251 +
252 + signal.signal(signal.SIGTERM, signal_handler)
253 + signal.signal(signal.SIGINT, signal_handler)
254 +
255 listen_address, sock = setup_server(args)
256
257 engine_args = vllm.AsyncEngineArgs.from_cli_args(args)
@@ -295,11 +313,29 @@ def run_multi_api_server(args: argparse.Namespace):
313 api_server_manager = APIServerProcessManager(**api_server_manager_kwargs)
314
315 # Wait for API servers
298 - wait_for_completion_or_failure(
299 - api_server_manager=api_server_manager,
300 - engine_manager=local_engine_manager,
301 - coordinator=coordinator,
302 - )
316 + try:
317 + wait_for_completion_or_failure(
318 + api_server_manager=api_server_manager,
319 + engine_manager=local_engine_manager,
320 + coordinator=coordinator,
321 + )
322 + finally:
323 + timeout = shutdown_by = None
324 + if shutdown_requested:
325 + timeout = vllm_config.shutdown_timeout
326 + shutdown_by = time.monotonic() + timeout
327 + logger.info("Waiting up to %d seconds for processes to exit", timeout)
328 +
329 + def to_timeout(deadline: float | None) -> float | None:
330 + return (
331 + deadline if deadline is None else max(deadline - time.monotonic(), 0.0)
332 + )
333 +
334 + api_server_manager.shutdown(timeout=timeout)
335 + if local_engine_manager:
336 + local_engine_manager.shutdown(timeout=to_timeout(shutdown_by))
337 + if coordinator:
338 + coordinator.shutdown(timeout=to_timeout(shutdown_by))
339
340
341 def run_api_server_worker_proc(
vllm/entrypoints/launcher.py
+23 -5
@@ -4,6 +4,7 @@
4 import asyncio
5 import signal
6 import socket
7 +from functools import partial
8 from typing import Any
9
10 import uvicorn
@@ -91,12 +92,10 @@ async def serve_http(
92 )
93 )
94
95 + shutdown_event = asyncio.Event()
96 +
97 def signal_handler() -> None:
95 - # prevents the uvicorn signal handler to exit early
96 - server_task.cancel()
97 - watchdog_task.cancel()
98 - if ssl_cert_refresher:
99 - ssl_cert_refresher.stop()
98 + shutdown_event.set()
99
100 async def dummy_shutdown() -> None:
101 pass
@@ -104,6 +103,24 @@ async def serve_http(
103 loop.add_signal_handler(signal.SIGINT, signal_handler)
104 loop.add_signal_handler(signal.SIGTERM, signal_handler)
105
106 + async def handle_shutdown() -> None:
107 + await shutdown_event.wait()
108 +
109 + engine_client = app.state.engine_client
110 + timeout = engine_client.vllm_config.shutdown_timeout
111 +
112 + await loop.run_in_executor(
113 + None, partial(engine_client.shutdown, timeout=timeout)
114 + )
115 +
116 + server.should_exit = True
117 + server_task.cancel()
118 + watchdog_task.cancel()
119 + if ssl_cert_refresher:
120 + ssl_cert_refresher.stop()
121 +
122 + shutdown_task = loop.create_task(handle_shutdown())
123 +
124 try:
125 await server_task
126 return dummy_shutdown()
@@ -120,6 +137,7 @@ async def serve_http(
137 logger.info("Shutting down FastAPI HTTP server.")
138 return server.shutdown()
139 finally:
140 + shutdown_task.cancel()
141 watchdog_task.cancel()
142
143
vllm/v1/engine/__init__.py
+2
@@ -226,6 +226,8 @@ class EngineCoreRequestType(enum.Enum):
226 UTILITY = b"\x03"
227 # Sentinel used within EngineCoreProc.
228 EXECUTOR_FAILED = b"\x04"
229 + # Sentinel to wake up input_queue.get() during shutdown.
230 + WAKEUP = b"\x05"
231
232
233 class ReconfigureDistributedRequest(msgspec.Struct):
vllm/v1/engine/async_llm.py
+2 -3
@@ -264,16 +264,15 @@ class AsyncLLM(EngineClient):
264 def __del__(self):
265 self.shutdown()
266
267 - def shutdown(self):
267 + def shutdown(self, timeout: float | None = None) -> None:
268 """Shutdown, cleaning up the background proc and IPC."""
269 -
269 shutdown_prometheus()
270
271 if renderer := getattr(self, "renderer", None):
272 renderer.shutdown()
273
274 if engine_core := getattr(self, "engine_core", None):
276 - engine_core.shutdown()
275 + engine_core.shutdown(timeout=timeout)
276
277 handler = getattr(self, "output_handler", None)
278 if handler is not None:
vllm/v1/engine/coordinator.py
+4 -2
@@ -104,8 +104,10 @@ class DPCoordinator:
104 """Returns tuple of ZMQ input address, output address."""
105 return self.coord_in_address, self.coord_out_address
106
107 - def close(self):
108 - self._finalizer()
107 + def shutdown(self, timeout: float | None = None) -> None:
108 + """Shutdown coordinator process with configurable timeout."""
109 + if self._finalizer.detach() is not None:
110 + shutdown([self.proc], timeout=timeout)
111
112
113 class EngineState:
vllm/v1/engine/core.py
+127 -43
@@ -9,6 +9,7 @@ from collections import defaultdict, deque
9 from collections.abc import Callable, Generator
10 from concurrent.futures import Future
11 from contextlib import ExitStack, contextmanager
12 +from enum import IntEnum
13 from functools import partial
14 from inspect import isclass, signature
15 from logging import DEBUG
@@ -61,6 +62,7 @@ from vllm.v1.engine import (
62 from vllm.v1.engine.utils import (
63 EngineHandshakeMetadata,
64 EngineZmqAddresses,
65 + SignalCallback,
66 get_device_indices,
67 )
68 from vllm.v1.executor import Executor
@@ -765,6 +767,12 @@ class EngineCore:
767 raise NotImplementedError
768
769
770 +class EngineShutdownState(IntEnum):
771 + RUNNING = 0
772 + REQUESTED = 1
773 + SHUTTING_DOWN = 2
774 +
775 +
776 class EngineCoreProc(EngineCore):
777 """ZMQ-wrapper for running EngineCore in background process."""
778
@@ -792,6 +800,7 @@ class EngineCoreProc(EngineCore):
800 self.engine_index = engine_index
801 identity = self.engine_index.to_bytes(length=2, byteorder="little")
802 self.engines_running = False
803 + self.shutdown_state = EngineShutdownState.RUNNING
804
805 with self._perform_handshakes(
806 handshake_address,
@@ -1020,25 +1029,11 @@ class EngineCoreProc(EngineCore):
1029 def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs):
1030 """Launch EngineCore busy loop in background process."""
1031
1023 - # Signal handler used for graceful termination.
1024 - # SystemExit exception is only raised once to allow this and worker
1025 - # processes to terminate without error
1026 - shutdown_requested = False
1027 -
1032 # Ensure we can serialize transformer config after spawning
1033 maybe_register_config_serialize_by_value()
1034
1031 - def signal_handler(signum, frame):
1032 - nonlocal shutdown_requested
1033 - if not shutdown_requested:
1034 - shutdown_requested = True
1035 - raise SystemExit()
1036 -
1037 - # Either SIGTERM or SIGINT will terminate the engine_core
1038 - signal.signal(signal.SIGTERM, signal_handler)
1039 - signal.signal(signal.SIGINT, signal_handler)
1040 -
1035 engine_core: EngineCoreProc | None = None
1036 + signal_callback: SignalCallback | None = None
1037 try:
1038 vllm_config: VllmConfig = kwargs["vllm_config"]
1039 parallel_config: ParallelConfig = vllm_config.parallel_config
@@ -1078,6 +1073,22 @@ class EngineCoreProc(EngineCore):
1073 engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs)
1074
1075 assert engine_core is not None
1076 +
1077 + def wakeup_engine():
1078 + # Wakes up idle engine via input_queue when shutdown is requested
1079 + # Not safe in a signal handler - we may interrupt the main thread
1080 + # while it is holding the non-reentrant input_queue.mutex
1081 + engine_core.input_queue.put_nowait((EngineCoreRequestType.WAKEUP, None))
1082 +
1083 + signal_callback = SignalCallback(wakeup_engine)
1084 +
1085 + def signal_handler(signum, frame):
1086 + engine_core.shutdown_state = EngineShutdownState.REQUESTED
1087 + signal_callback.trigger()
1088 +
1089 + signal.signal(signal.SIGTERM, signal_handler)
1090 + signal.signal(signal.SIGINT, signal_handler)
1091 +
1092 engine_core.run_busy_loop()
1093
1094 except SystemExit:
@@ -1091,6 +1102,10 @@ class EngineCoreProc(EngineCore):
1102 engine_core._send_engine_dead()
1103 raise e
1104 finally:
1105 + signal.signal(signal.SIGTERM, signal.SIG_DFL)
1106 + signal.signal(signal.SIGINT, signal.SIG_DFL)
1107 + if signal_callback is not None:
1108 + signal_callback.stop()
1109 if engine_core is not None:
1110 engine_core.shutdown()
1111
@@ -1105,21 +1120,25 @@ class EngineCoreProc(EngineCore):
1120 or bool(self.batch_queue)
1121 )
1122
1123 + def is_running(self) -> bool:
1124 + """Returns true if shutdown has not been requested."""
1125 + return self.shutdown_state == EngineShutdownState.RUNNING
1126 +
1127 def run_busy_loop(self):
1128 """Core busy loop of the EngineCore."""
1110 -
1111 - # Loop until process is sent a SIGINT or SIGTERM
1112 - while True:
1129 + while self._handle_shutdown():
1130 # 1) Poll the input queue until there is work to do.
1131 self._process_input_queue()
1132 # 2) Step the engine core and return the outputs.
1133 self._process_engine_step()
1134
1135 + raise SystemExit
1136 +
1137 def _process_input_queue(self):
1138 """Exits when an engine step needs to be performed."""
1139
1140 waited = False
1122 - while not self.has_work():
1141 + while not self.has_work() and self.is_running():
1142 # Notify callbacks waiting for engine to become idle.
1143 self._notify_idle_state_callbacks()
1144 if self.input_queue.empty():
@@ -1171,18 +1190,60 @@ class EngineCoreProc(EngineCore):
1190 callback = self._idle_state_callbacks.pop()
1191 callback(self)
1192
1193 + def _handle_shutdown(self) -> bool:
1194 + # Check if shutdown was requested and handle it
1195 + if self.shutdown_state == EngineShutdownState.RUNNING:
1196 + return True
1197 +
1198 + if self.shutdown_state == EngineShutdownState.REQUESTED:
1199 + shutdown_timeout = self.vllm_config.shutdown_timeout
1200 +
1201 + logger.info("Shutdown initiated (timeout=%d)", shutdown_timeout)
1202 +
1203 + if shutdown_timeout == 0:
1204 + num_requests = self.scheduler.get_num_unfinished_requests()
1205 + if num_requests > 0:
1206 + logger.info("Aborting %d requests", num_requests)
1207 + aborted_reqs = self.scheduler.finish_requests(
1208 + None, RequestStatus.FINISHED_ABORTED
1209 + )
1210 + self._send_abort_outputs(aborted_reqs)
1211 + else:
1212 + num_requests = self.scheduler.get_num_unfinished_requests()
1213 + if num_requests > 0:
1214 + logger.info(
1215 + "Draining %d in-flight requests (timeout=%ds)",
1216 + num_requests,
1217 + shutdown_timeout,
1218 + )
1219 +
1220 + self.shutdown_state = EngineShutdownState.SHUTTING_DOWN
1221 +
1222 + # Exit when no work remaining
1223 + if not self.has_work():
1224 + logger.info("Shutdown complete")
1225 + return False
1226 +
1227 + return True
1228 +
1229 def _handle_client_request(
1230 self, request_type: EngineCoreRequestType, request: Any
1231 ) -> None:
1232 """Dispatch request from client."""
1233
1179 - if request_type == EngineCoreRequestType.ADD:
1234 + if request_type == EngineCoreRequestType.WAKEUP:
1235 + return
1236 + elif request_type == EngineCoreRequestType.ADD:
1237 req, request_wave = request
1238 + if self._reject_add_in_shutdown(req):
1239 + return
1240 self.add_request(req, request_wave)
1241 elif request_type == EngineCoreRequestType.ABORT:
1242 self.abort_requests(request)
1243 elif request_type == EngineCoreRequestType.UTILITY:
1244 client_idx, call_id, method_name, args = request
1245 + if self._reject_utility_in_shutdown(client_idx, call_id, method_name):
1246 + return
1247 output = UtilityOutput(call_id)
1248 # Lazily look-up utility method so that failure will be handled/returned.
1249 get_result = lambda: (method := getattr(self, method_name)) and method(
@@ -1199,6 +1260,27 @@ class EngineCoreProc(EngineCore):
1260 "Unrecognized input request type encountered: %s", request_type
1261 )
1262
1263 + def _reject_add_in_shutdown(self, request: Request) -> bool:
1264 + if self.shutdown_state == EngineShutdownState.RUNNING:
1265 + return False
1266 +
1267 + logger.info("Rejecting request %s (server shutting down)", request.request_id)
1268 + self._send_abort_outputs_to_client([request.request_id], request.client_index)
1269 + return True
1270 +
1271 + def _reject_utility_in_shutdown(
1272 + self, client_idx: int, call_id: int, method_name: str
1273 + ) -> bool:
1274 + if self.shutdown_state == EngineShutdownState.RUNNING:
1275 + return False
1276 +
1277 + logger.warning("Rejecting utility call %s (server shutting down)", method_name)
1278 + output = UtilityOutput(call_id, failure_message="Server shutting down")
1279 + self.output_queue.put_nowait(
1280 + (client_idx, EngineCoreOutputs(utility_output=output))
1281 + )
1282 + return True
1283 +
1284 @staticmethod
1285 def _invoke_utility_method(
1286 name: str, get_result: Callable, output: UtilityOutput, enqueue_output: Callable
@@ -1412,22 +1494,7 @@ class EngineCoreProc(EngineCore):
1494 logger.exception(
1495 "Unexpected error pre-processing request %s", request.request_id
1496 )
1415 - self.output_queue.put_nowait(
1416 - (
1417 - request.client_index,
1418 - EngineCoreOutputs(
1419 - engine_index=self.engine_index,
1420 - finished_requests={request.request_id},
1421 - outputs=[
1422 - EngineCoreOutput(
1423 - request_id=request.request_id,
1424 - new_token_ids=[],
1425 - finish_reason=FinishReason.ERROR,
1426 - )
1427 - ],
1428 - ),
1429 - )
1430 - )
1497 + self._send_error_outputs_to_client([request.request_id], request.client_index)
1498
1499 def pause_scheduler(
1500 self, mode: PauseMode = "abort", clear_cache: bool = True
@@ -1470,6 +1537,26 @@ class EngineCoreProc(EngineCore):
1537 self._idle_state_callbacks.append(partial(engine_idle_callback, future=future))
1538 return future
1539
1540 + def _send_finish_outputs_to_client(
1541 + self, req_ids: list[str], client_index: int, finish_reason: FinishReason
1542 + ) -> None:
1543 + outputs = [
1544 + EngineCoreOutput(req_id, [], finish_reason=finish_reason)
1545 + for req_id in req_ids
1546 + ]
1547 + eco = EngineCoreOutputs(finished_requests=req_ids, outputs=outputs)
1548 + self.output_queue.put_nowait((client_index, eco))
1549 +
1550 + def _send_abort_outputs_to_client(
1551 + self, req_ids: list[str], client_index: int
1552 + ) -> None:
1553 + self._send_finish_outputs_to_client(req_ids, client_index, FinishReason.ABORT)
1554 +
1555 + def _send_error_outputs_to_client(
1556 + self, req_ids: list[str], client_index: int
1557 + ) -> None:
1558 + self._send_finish_outputs_to_client(req_ids, client_index, FinishReason.ERROR)
1559 +
1560 def _send_abort_outputs(self, aborted_reqs: list[tuple[str, int]]) -> None:
1561 # TODO(nick) this will be moved inside the scheduler
1562 if aborted_reqs:
@@ -1478,12 +1565,7 @@ class EngineCoreProc(EngineCore):
1565 for req_id, client_index in aborted_reqs:
1566 by_client[client_index].add(req_id)
1567 for client_index, req_ids in by_client.items():
1481 - outputs = [
1482 - EngineCoreOutput(req_id, [], finish_reason=FinishReason.ABORT)
1483 - for req_id in req_ids
1484 - ]
1485 - eco = EngineCoreOutputs(finished_requests=req_ids, outputs=outputs)
1486 - self.output_queue.put_nowait((client_index, eco))
1568 + self._send_abort_outputs_to_client(list(req_ids), client_index)
1569
1570
1571 class DPEngineCoreProc(EngineCoreProc):
@@ -1601,7 +1683,7 @@ class DPEngineCoreProc(EngineCoreProc):
1683 """Core busy loop of the EngineCore for data parallel case."""
1684
1685 # Loop until process is sent a SIGINT or SIGTERM
1604 - while True:
1686 + while self._handle_shutdown():
1687 # 1) Poll the input queue until there is work to do.
1688 self._process_input_queue()
1689
@@ -1649,6 +1731,8 @@ class DPEngineCoreProc(EngineCoreProc):
1731 self.current_wave += 1
1732 self.step_counter = 0
1733
1734 + raise SystemExit
1735 +
1736 def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:
1737 # Optimization - only perform finish-sync all-reduce every 32 steps.
1738 self.step_counter += 1
vllm/v1/engine/core_client.py
+12 -12
@@ -128,7 +128,7 @@ class EngineCoreClient(ABC):
128 return AsyncMPClient(*client_args)
129
130 @abstractmethod
131 - def shutdown(self): ...
131 + def shutdown(self, timeout: float | None = None) -> None: ...
132
133 def get_output(self) -> EngineCoreOutputs:
134 raise NotImplementedError
@@ -298,7 +298,7 @@ class InprocClient(EngineCoreClient):
298 if len(request_ids) > 0:
299 self.engine_core.abort_requests(request_ids)
300
301 - def shutdown(self) -> None:
301 + def shutdown(self, timeout: float | None = None) -> None:
302 self.engine_core.shutdown()
303
304 def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
@@ -390,9 +390,9 @@ class BackgroundResources:
390
391 self.engine_dead = True
392 if self.engine_manager is not None:
393 - self.engine_manager.close()
393 + self.engine_manager.shutdown()
394 if self.coordinator is not None:
395 - self.coordinator.close()
395 + self.coordinator.shutdown()
396
397 if isinstance(self.output_socket, zmq.asyncio.Socket):
398 # Async case.
@@ -581,10 +581,7 @@ class MPClient(EngineCoreClient):
581 )
582
583 with launch_core_engines(
584 - vllm_config,
585 - executor_class,
586 - log_stats,
587 - addresses,
584 + vllm_config, executor_class, log_stats, addresses
585 ) as (engine_manager, coordinator, addresses):
586 self.resources.coordinator = coordinator
587 self.resources.engine_manager = engine_manager
@@ -649,9 +646,12 @@ class MPClient(EngineCoreClient):
646 if not success:
647 self._finalizer()
648
652 - def shutdown(self):
653 - # Terminate background resources.
654 - self._finalizer()
649 + def shutdown(self, timeout: float | None = None) -> None:
650 + """Shutdown engine manager under timeout and clean up resources."""
651 + if self._finalizer.detach() is not None:
652 + if self.resources.engine_manager is not None:
653 + self.resources.engine_manager.shutdown(timeout=timeout)
654 + self.resources()
655
656 def _format_exception(self, e: Exception) -> Exception:
657 """If errored, use EngineDeadError so root cause is clear."""
@@ -695,7 +695,7 @@ class MPClient(EngineCoreClient):
695 sentinels = [proc.sentinel for proc in engine_processes]
696 died = multiprocessing.connection.wait(sentinels)
697 _self = self_ref()
698 - if not _self or _self.resources.engine_dead:
698 + if not _self or not _self._finalizer.alive or _self.resources.engine_dead:
699 return
700 _self.resources.engine_dead = True
701 proc_name = next(
vllm/v1/engine/utils.py
+35 -6
@@ -3,8 +3,9 @@
3
4 import contextlib
5 import os
6 +import threading
7 import weakref
7 -from collections.abc import Iterator
8 +from collections.abc import Callable, Iterator
9 from dataclasses import dataclass
10 from enum import Enum, auto
11 from multiprocessing import Process, connection
@@ -146,11 +147,12 @@ class CoreEngineProcManager:
147 finally:
148 # Kill other procs if not all are running.
149 if self.finished_procs():
149 - self.close()
150 + self.shutdown()
151
151 - def close(self):
152 - """Shutdown all procs."""
153 - self._finalizer()
152 + def shutdown(self, timeout: float | None = None) -> None:
153 + """Shutdown engine core processes with configurable timeout."""
154 + if self._finalizer.detach() is not None:
155 + shutdown(self.processes, timeout=timeout)
156
157 def join_first(self):
158 """Wait for any process to exit."""
@@ -168,6 +170,33 @@ class CoreEngineProcManager:
170 }
171
172
173 +class SignalCallback:
174 + """Safely trigger a callback from signal handler context via a dedicated thread."""
175 +
176 + def __init__(self, callback: Callable[[], None]):
177 + self._callback = callback
178 + self._event = threading.Event()
179 + self._stopped = False
180 + self._thread = threading.Thread(
181 + target=self._run,
182 + daemon=True,
183 + name="signal-callback",
184 + )
185 + self._thread.start()
186 +
187 + def _run(self):
188 + self._event.wait()
189 + if not self._stopped:
190 + self._callback()
191 +
192 + def trigger(self):
193 + self._event.set()
194 +
195 + def stop(self):
196 + self._stopped = True
197 + self._event.set()
198 +
199 +
200 @contextlib.contextmanager
201 def set_device_control_env_var(
202 vllm_config: VllmConfig, local_dp_rank: int
@@ -763,7 +792,7 @@ class CoreEngineActorManager:
792 def get_run_refs(self):
793 return self.run_refs
794
766 - def close(self):
795 + def shutdown(self, timeout: float | None = None) -> None:
796 import ray
797
798 for actor in self.local_engine_actors + self.remote_engine_actors:
vllm/v1/utils.py
+19 -12
@@ -220,8 +220,10 @@ class APIServerProcessManager:
220 # The extra processes are managed by their owners
221 self._finalizer = weakref.finalize(self, shutdown, self.processes)
222
223 - def close(self) -> None:
224 - self._finalizer()
223 + def shutdown(self, timeout: float | None = None) -> None:
224 + """Shutdown API server processes with configurable timeout"""
225 + if self._finalizer.detach() is not None:
226 + shutdown(self.processes, timeout=timeout)
227
228
229 def wait_for_completion_or_failure(
@@ -288,25 +290,30 @@ def wait_for_completion_or_failure(
290 except Exception as e:
291 logger.exception("Exception occurred while running API servers: %s", str(e))
292 raise
291 - finally:
292 - logger.info("Terminating remaining processes ...")
293 - api_server_manager.close()
294 - if coordinator:
295 - coordinator.close()
296 - if engine_manager:
297 - engine_manager.close()
293
294
295 # Note(rob): shutdown function cannot be a bound method,
296 # else the gc cannot collect the object.
302 -def shutdown(procs: list[BaseProcess]):
297 +def shutdown(procs: list[BaseProcess], timeout: float | None = None) -> None:
298 + """Shutdown processes with timeout.
299 +
300 + Args:
301 + procs: List of processes to shutdown
302 + timeout: Maximum time in seconds to wait for graceful shutdown
303 + """
304 + if timeout is None:
305 + timeout = 0.0
306 +
307 + # Allow at least 5 seconds for remaining procs to terminate.
308 + timeout = max(timeout, 5.0)
309 +
310 # Shutdown the process.
311 for proc in procs:
312 if proc.is_alive():
313 proc.terminate()
314
308 - # Allow 5 seconds for remaining procs to terminate.
309 - deadline = time.monotonic() + 5
315 + # Allow time for remaining procs to terminate.
316 + deadline = time.monotonic() + timeout
317 for proc in procs:
318 remaining = deadline - time.monotonic()
319 if remaining <= 0: