@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 27066d1b2b
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 7, 04:00 UTC

[Frontend][Core] Add shutdown timeout - allowing in-flight requests to finish (#34730)

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

Mark McLoughlin committed Mar 6, 2026 at 06:04 UTC 27066d1b2bd0dea89d617afa24da611d9a32e36a
15 files changed +762 -90
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
@@ -608,6 +608,8 @@ class EngineArgs:
608 kv_offloading_backend: KVOffloadingBackend = CacheConfig.kv_offloading_backend
609 tokens_only: bool = False
610
611 + shutdown_timeout: int = 0
612 +
613 weight_transfer_config: WeightTransferConfig | None = get_field(
614 VllmConfig,
615 "weight_transfer_config",
@@ -1311,6 +1313,14 @@ class EngineArgs:
1313 default=False,
1314 action=argparse.BooleanOptionalAction,
1315 )
1316 +
1317 + parser.add_argument(
1318 + "--shutdown-timeout",
1319 + type=int,
1320 + default=0,
1321 + help="Shutdown timeout in seconds. 0 = abort, >0 = wait.",
1322 + )
1323 +
1324 return parser
1325
1326 @classmethod
@@ -1920,6 +1930,7 @@ class EngineArgs:
1930 optimization_level=self.optimization_level,
1931 performance_mode=self.performance_mode,
1932 weight_transfer_config=self.weight_transfer_config,
1933 + shutdown_timeout=self.shutdown_timeout,
1934 )
1935
1936 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
@@ -211,8 +212,12 @@ def run_headless(args: argparse.Namespace):
212 try:
213 engine_manager.join_first()
214 finally:
215 + timeout = None
216 + if shutdown_requested:
217 + timeout = vllm_config.shutdown_timeout
218 + logger.info("Waiting up to %d seconds for processes to exit", timeout)
219 + engine_manager.shutdown(timeout=timeout)
220 logger.info("Shutting down.")
215 - engine_manager.close()
221
222
223 def run_multi_api_server(args: argparse.Namespace):
@@ -229,6 +234,19 @@ def run_multi_api_server(args: argparse.Namespace):
234 if num_api_servers > 1:
235 setup_multiprocess_prometheus()
236
237 + shutdown_requested = False
238 +
239 + # Catch SIGTERM and SIGINT to allow graceful shutdown.
240 + def signal_handler(signum, frame):
241 + nonlocal shutdown_requested
242 + logger.debug("Received %d signal.", signum)
243 + if not shutdown_requested:
244 + shutdown_requested = True
245 + raise SystemExit
246 +
247 + signal.signal(signal.SIGTERM, signal_handler)
248 + signal.signal(signal.SIGINT, signal_handler)
249 +
250 listen_address, sock = setup_server(args)
251
252 engine_args = vllm.AsyncEngineArgs.from_cli_args(args)
@@ -290,11 +308,29 @@ def run_multi_api_server(args: argparse.Namespace):
308 api_server_manager = APIServerProcessManager(**api_server_manager_kwargs)
309
310 # Wait for API servers
293 - wait_for_completion_or_failure(
294 - api_server_manager=api_server_manager,
295 - engine_manager=local_engine_manager,
296 - coordinator=coordinator,
297 - )
311 + try:
312 + wait_for_completion_or_failure(
313 + api_server_manager=api_server_manager,
314 + engine_manager=local_engine_manager,
315 + coordinator=coordinator,
316 + )
317 + finally:
318 + timeout = shutdown_by = None
319 + if shutdown_requested:
320 + timeout = vllm_config.shutdown_timeout
321 + shutdown_by = time.monotonic() + timeout
322 + logger.info("Waiting up to %d seconds for processes to exit", timeout)
323 +
324 + def to_timeout(deadline: float | None) -> float | None:
325 + return (
326 + deadline if deadline is None else max(deadline - time.monotonic(), 0.0)
327 + )
328 +
329 + api_server_manager.shutdown(timeout=timeout)
330 + if local_engine_manager:
331 + local_engine_manager.shutdown(timeout=to_timeout(shutdown_by))
332 + if coordinator:
333 + coordinator.shutdown(timeout=to_timeout(shutdown_by))
334
335
336 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
@@ -238,6 +238,8 @@ class EngineCoreRequestType(enum.Enum):
238 UTILITY = b"\x03"
239 # Sentinel used within EngineCoreProc.
240 EXECUTOR_FAILED = b"\x04"
241 + # Sentinel to wake up input_queue.get() during shutdown.
242 + WAKEUP = b"\x05"
243
244
245 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
@@ -10,6 +10,7 @@ from collections import defaultdict, deque
10 from collections.abc import Callable, Generator
11 from concurrent.futures import Future
12 from contextlib import ExitStack, contextmanager
13 +from enum import IntEnum
14 from functools import partial
15 from inspect import isclass, signature
16 from logging import DEBUG
@@ -62,6 +63,7 @@ from vllm.v1.engine import (
63 from vllm.v1.engine.utils import (
64 EngineHandshakeMetadata,
65 EngineZmqAddresses,
66 + SignalCallback,
67 get_device_indices,
68 )
69 from vllm.v1.executor import Executor
@@ -776,6 +778,12 @@ class EngineCore:
778 raise NotImplementedError
779
780
781 +class EngineShutdownState(IntEnum):
782 + RUNNING = 0
783 + REQUESTED = 1
784 + SHUTTING_DOWN = 2
785 +
786 +
787 class EngineCoreProc(EngineCore):
788 """ZMQ-wrapper for running EngineCore in background process."""
789
@@ -803,6 +811,7 @@ class EngineCoreProc(EngineCore):
811 self.engine_index = engine_index
812 identity = self.engine_index.to_bytes(length=2, byteorder="little")
813 self.engines_running = False
814 + self.shutdown_state = EngineShutdownState.RUNNING
815
816 with self._perform_handshakes(
817 handshake_address,
@@ -1053,25 +1062,11 @@ class EngineCoreProc(EngineCore):
1062 def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs):
1063 """Launch EngineCore busy loop in background process."""
1064
1056 - # Signal handler used for graceful termination.
1057 - # SystemExit exception is only raised once to allow this and worker
1058 - # processes to terminate without error
1059 - shutdown_requested = False
1060 -
1065 # Ensure we can serialize transformer config after spawning
1066 maybe_register_config_serialize_by_value()
1067
1064 - def signal_handler(signum, frame):
1065 - nonlocal shutdown_requested
1066 - if not shutdown_requested:
1067 - shutdown_requested = True
1068 - raise SystemExit()
1069 -
1070 - # Either SIGTERM or SIGINT will terminate the engine_core
1071 - signal.signal(signal.SIGTERM, signal_handler)
1072 - signal.signal(signal.SIGINT, signal_handler)
1073 -
1068 engine_core: EngineCoreProc | None = None
1069 + signal_callback: SignalCallback | None = None
1070 try:
1071 vllm_config: VllmConfig = kwargs["vllm_config"]
1072 parallel_config: ParallelConfig = vllm_config.parallel_config
@@ -1119,6 +1114,22 @@ class EngineCoreProc(EngineCore):
1114 engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs)
1115
1116 assert engine_core is not None
1117 +
1118 + def wakeup_engine():
1119 + # Wakes up idle engine via input_queue when shutdown is requested
1120 + # Not safe in a signal handler - we may interrupt the main thread
1121 + # while it is holding the non-reentrant input_queue.mutex
1122 + engine_core.input_queue.put_nowait((EngineCoreRequestType.WAKEUP, None))
1123 +
1124 + signal_callback = SignalCallback(wakeup_engine)
1125 +
1126 + def signal_handler(signum, frame):
1127 + engine_core.shutdown_state = EngineShutdownState.REQUESTED
1128 + signal_callback.trigger()
1129 +
1130 + signal.signal(signal.SIGTERM, signal_handler)
1131 + signal.signal(signal.SIGINT, signal_handler)
1132 +
1133 engine_core.run_busy_loop()
1134
1135 except SystemExit:
@@ -1132,6 +1143,10 @@ class EngineCoreProc(EngineCore):
1143 engine_core._send_engine_dead()
1144 raise e
1145 finally:
1146 + signal.signal(signal.SIGTERM, signal.SIG_DFL)
1147 + signal.signal(signal.SIGINT, signal.SIG_DFL)
1148 + if signal_callback is not None:
1149 + signal_callback.stop()
1150 if engine_core is not None:
1151 engine_core.shutdown()
1152
@@ -1146,21 +1161,25 @@ class EngineCoreProc(EngineCore):
1161 or bool(self.batch_queue)
1162 )
1163
1164 + def is_running(self) -> bool:
1165 + """Returns true if shutdown has not been requested."""
1166 + return self.shutdown_state == EngineShutdownState.RUNNING
1167 +
1168 def run_busy_loop(self):
1169 """Core busy loop of the EngineCore."""
1151 -
1152 - # Loop until process is sent a SIGINT or SIGTERM
1153 - while True:
1170 + while self._handle_shutdown():
1171 # 1) Poll the input queue until there is work to do.
1172 self._process_input_queue()
1173 # 2) Step the engine core and return the outputs.
1174 self._process_engine_step()
1175
1176 + raise SystemExit
1177 +
1178 def _process_input_queue(self):
1179 """Exits when an engine step needs to be performed."""
1180
1181 waited = False
1163 - while not self.has_work():
1182 + while not self.has_work() and self.is_running():
1183 # Notify callbacks waiting for engine to become idle.
1184 self._notify_idle_state_callbacks()
1185 if self.input_queue.empty():
@@ -1212,18 +1231,60 @@ class EngineCoreProc(EngineCore):
1231 callback = self._idle_state_callbacks.pop()
1232 callback(self)
1233
1234 + def _handle_shutdown(self) -> bool:
1235 + # Check if shutdown was requested and handle it
1236 + if self.shutdown_state == EngineShutdownState.RUNNING:
1237 + return True
1238 +
1239 + if self.shutdown_state == EngineShutdownState.REQUESTED:
1240 + shutdown_timeout = self.vllm_config.shutdown_timeout
1241 +
1242 + logger.info("Shutdown initiated (timeout=%d)", shutdown_timeout)
1243 +
1244 + if shutdown_timeout == 0:
1245 + num_requests = self.scheduler.get_num_unfinished_requests()
1246 + if num_requests > 0:
1247 + logger.info("Aborting %d requests", num_requests)
1248 + aborted_reqs = self.scheduler.finish_requests(
1249 + None, RequestStatus.FINISHED_ABORTED
1250 + )
1251 + self._send_abort_outputs(aborted_reqs)
1252 + else:
1253 + num_requests = self.scheduler.get_num_unfinished_requests()
1254 + if num_requests > 0:
1255 + logger.info(
1256 + "Draining %d in-flight requests (timeout=%ds)",
1257 + num_requests,
1258 + shutdown_timeout,
1259 + )
1260 +
1261 + self.shutdown_state = EngineShutdownState.SHUTTING_DOWN
1262 +
1263 + # Exit when no work remaining
1264 + if not self.has_work():
1265 + logger.info("Shutdown complete")
1266 + return False
1267 +
1268 + return True
1269 +
1270 def _handle_client_request(
1271 self, request_type: EngineCoreRequestType, request: Any
1272 ) -> None:
1273 """Dispatch request from client."""
1274
1220 - if request_type == EngineCoreRequestType.ADD:
1275 + if request_type == EngineCoreRequestType.WAKEUP:
1276 + return
1277 + elif request_type == EngineCoreRequestType.ADD:
1278 req, request_wave = request
1279 + if self._reject_add_in_shutdown(req):
1280 + return
1281 self.add_request(req, request_wave)
1282 elif request_type == EngineCoreRequestType.ABORT:
1283 self.abort_requests(request)
1284 elif request_type == EngineCoreRequestType.UTILITY:
1285 client_idx, call_id, method_name, args = request
1286 + if self._reject_utility_in_shutdown(client_idx, call_id, method_name):
1287 + return
1288 output = UtilityOutput(call_id)
1289 # Lazily look-up utility method so that failure will be handled/returned.
1290 get_result = lambda: (method := getattr(self, method_name)) and method(
@@ -1240,6 +1301,27 @@ class EngineCoreProc(EngineCore):
1301 "Unrecognized input request type encountered: %s", request_type
1302 )
1303
1304 + def _reject_add_in_shutdown(self, request: Request) -> bool:
1305 + if self.shutdown_state == EngineShutdownState.RUNNING:
1306 + return False
1307 +
1308 + logger.info("Rejecting request %s (server shutting down)", request.request_id)
1309 + self._send_abort_outputs_to_client([request.request_id], request.client_index)
1310 + return True
1311 +
1312 + def _reject_utility_in_shutdown(
1313 + self, client_idx: int, call_id: int, method_name: str
1314 + ) -> bool:
1315 + if self.shutdown_state == EngineShutdownState.RUNNING:
1316 + return False
1317 +
1318 + logger.warning("Rejecting utility call %s (server shutting down)", method_name)
1319 + output = UtilityOutput(call_id, failure_message="Server shutting down")
1320 + self.output_queue.put_nowait(
1321 + (client_idx, EngineCoreOutputs(utility_output=output))
1322 + )
1323 + return True
1324 +
1325 @staticmethod
1326 def _invoke_utility_method(
1327 name: str, get_result: Callable, output: UtilityOutput, enqueue_output: Callable
@@ -1453,22 +1535,7 @@ class EngineCoreProc(EngineCore):
1535 logger.exception(
1536 "Unexpected error pre-processing request %s", request.request_id
1537 )
1456 - self.output_queue.put_nowait(
1457 - (
1458 - request.client_index,
1459 - EngineCoreOutputs(
1460 - engine_index=self.engine_index,
1461 - finished_requests={request.request_id},
1462 - outputs=[
1463 - EngineCoreOutput(
1464 - request_id=request.request_id,
1465 - new_token_ids=[],
1466 - finish_reason=FinishReason.ERROR,
1467 - )
1468 - ],
1469 - ),
1470 - )
1471 - )
1538 + self._send_error_outputs_to_client([request.request_id], request.client_index)
1539
1540 def pause_scheduler(
1541 self, mode: PauseMode = "abort", clear_cache: bool = True
@@ -1511,6 +1578,26 @@ class EngineCoreProc(EngineCore):
1578 self._idle_state_callbacks.append(partial(engine_idle_callback, future=future))
1579 return future
1580
1581 + def _send_finish_outputs_to_client(
1582 + self, req_ids: list[str], client_index: int, finish_reason: FinishReason
1583 + ) -> None:
1584 + outputs = [
1585 + EngineCoreOutput(req_id, [], finish_reason=finish_reason)
1586 + for req_id in req_ids
1587 + ]
1588 + eco = EngineCoreOutputs(finished_requests=req_ids, outputs=outputs)
1589 + self.output_queue.put_nowait((client_index, eco))
1590 +
1591 + def _send_abort_outputs_to_client(
1592 + self, req_ids: list[str], client_index: int
1593 + ) -> None:
1594 + self._send_finish_outputs_to_client(req_ids, client_index, FinishReason.ABORT)
1595 +
1596 + def _send_error_outputs_to_client(
1597 + self, req_ids: list[str], client_index: int
1598 + ) -> None:
1599 + self._send_finish_outputs_to_client(req_ids, client_index, FinishReason.ERROR)
1600 +
1601 def _send_abort_outputs(self, aborted_reqs: list[tuple[str, int]]) -> None:
1602 # TODO(nick) this will be moved inside the scheduler
1603 if aborted_reqs:
@@ -1519,12 +1606,7 @@ class EngineCoreProc(EngineCore):
1606 for req_id, client_index in aborted_reqs:
1607 by_client[client_index].add(req_id)
1608 for client_index, req_ids in by_client.items():
1522 - outputs = [
1523 - EngineCoreOutput(req_id, [], finish_reason=FinishReason.ABORT)
1524 - for req_id in req_ids
1525 - ]
1526 - eco = EngineCoreOutputs(finished_requests=req_ids, outputs=outputs)
1527 - self.output_queue.put_nowait((client_index, eco))
1609 + self._send_abort_outputs_to_client(list(req_ids), client_index)
1610
1611
1612 class DPEngineCoreProc(EngineCoreProc):
@@ -1642,7 +1724,7 @@ class DPEngineCoreProc(EngineCoreProc):
1724 """Core busy loop of the EngineCore for data parallel case."""
1725
1726 # Loop until process is sent a SIGINT or SIGTERM
1645 - while True:
1727 + while self._handle_shutdown():
1728 # 1) Poll the input queue until there is work to do.
1729 self._process_input_queue()
1730
@@ -1690,6 +1772,8 @@ class DPEngineCoreProc(EngineCoreProc):
1772 self.current_wave += 1
1773 self.step_counter = 0
1774
1775 + raise SystemExit
1776 +
1777 def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:
1778 # Optimization - only perform finish-sync all-reduce every 32 steps.
1779 self.step_counter += 1
vllm/v1/engine/core_client.py
+10 -7
@@ -127,7 +127,7 @@ class EngineCoreClient(ABC):
127 return AsyncMPClient(*client_args)
128
129 @abstractmethod
130 - def shutdown(self): ...
130 + def shutdown(self, timeout: float | None = None) -> None: ...
131
132 def get_output(self) -> EngineCoreOutputs:
133 raise NotImplementedError
@@ -297,7 +297,7 @@ class InprocClient(EngineCoreClient):
297 if len(request_ids) > 0:
298 self.engine_core.abort_requests(request_ids)
299
300 - def shutdown(self) -> None:
300 + def shutdown(self, timeout: float | None = None) -> None:
301 self.engine_core.shutdown()
302
303 def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
@@ -389,9 +389,9 @@ class BackgroundResources:
389
390 self.engine_dead = True
391 if self.engine_manager is not None:
392 - self.engine_manager.close()
392 + self.engine_manager.shutdown()
393 if self.coordinator is not None:
394 - self.coordinator.close()
394 + self.coordinator.shutdown()
395
396 if isinstance(self.output_socket, zmq.asyncio.Socket):
397 # Async case.
@@ -636,9 +636,12 @@ class MPClient(EngineCoreClient):
636 if not success:
637 self._finalizer()
638
639 - def shutdown(self):
640 - # Terminate background resources.
641 - self._finalizer()
639 + def shutdown(self, timeout: float | None = None) -> None:
640 + """Shutdown engine manager under timeout and clean up resources."""
641 + self._finalizer.detach()
642 + if self.resources.engine_manager is not None:
643 + self.resources.engine_manager.shutdown(timeout=timeout)
644 + self.resources()
645
646 def _format_exception(self, e: Exception) -> Exception:
647 """If errored, use EngineDeadError so root cause is clear."""
vllm/v1/engine/launch.py
+3
@@ -119,6 +119,9 @@ class LaunchEngineClient(EngineClient):
119 async def is_paused(self) -> bool:
120 return False
121
122 + def shutdown(self, timeout: float | None = None) -> None:
123 + pass
124 +
125 async def encode(
126 self,
127 prompt: PromptType | ProcessorInputs,
vllm/v1/engine/utils.py
+34 -5
@@ -3,6 +3,7 @@
3
4 import contextlib
5 import os
6 +import threading
7 import weakref
8 from collections.abc import Callable, Iterator
9 from dataclasses import dataclass
@@ -151,11 +152,12 @@ class CoreEngineProcManager:
152 finally:
153 # Kill other procs if not all are running.
154 if self.finished_procs():
154 - self.close()
155 + self.shutdown()
156
156 - def close(self):
157 - """Shutdown all procs."""
158 - self._finalizer()
157 + def shutdown(self, timeout: float | None = None) -> None:
158 + """Shutdown engine core processes with configurable timeout."""
159 + if self._finalizer.detach() is not None:
160 + shutdown(self.processes, timeout=timeout)
161
162 def join_first(self):
163 """Wait for any process to exit."""
@@ -173,6 +175,33 @@ class CoreEngineProcManager:
175 }
176
177
178 +class SignalCallback:
179 + """Safely trigger a callback from signal handler context via a dedicated thread."""
180 +
181 + def __init__(self, callback: Callable[[], None]):
182 + self._callback = callback
183 + self._event = threading.Event()
184 + self._stopped = False
185 + self._thread = threading.Thread(
186 + target=self._run,
187 + daemon=True,
188 + name="signal-callback",
189 + )
190 + self._thread.start()
191 +
192 + def _run(self):
193 + self._event.wait()
194 + if not self._stopped:
195 + self._callback()
196 +
197 + def trigger(self):
198 + self._event.set()
199 +
200 + def stop(self):
201 + self._stopped = True
202 + self._event.set()
203 +
204 +
205 @contextlib.contextmanager
206 def set_device_control_env_var(
207 vllm_config: VllmConfig, local_dp_rank: int
@@ -768,7 +797,7 @@ class CoreEngineActorManager:
797 def get_run_refs(self):
798 return self.run_refs
799
771 - def close(self):
800 + def shutdown(self, timeout: float | None = None) -> None:
801 import ray
802
803 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: