@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2b05b8ce69
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 15, 14:45 UTC

[V1][Frontend] Improve Shutdown And Logs (#11737)

Signed-off-by: rshaw@neuralmagic.com <rshaw@neuralmagic.com> Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com> Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: rshaw@neuralmagic.com <rshaw@neuralmagic.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com> Co-authored-by: Russell Bryant <rbryant@redhat.com> Co-authored-by: Andrew Feldman <afeldman@neuralmagic.com> Co-authored-by: afeldman-nm <156691304+afeldman-nm@users.noreply.github.com> Co-authored-by: Nick Hill <nhill@redhat.com>

Robert Shaw committed Apr 16, 2025 at 22:48 UTC 2b05b8ce69fff90b7f1e3285dcd05ced2b79fd97
16 files changed +1047 -363
.buildkite/test-pipeline.yaml
+1
@@ -552,6 +552,7 @@ steps:
552 # - pytest -v -s spec_decode/e2e/test_integration_dist_tp2.py
553 - VLLM_USE_V1=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s test_sharded_state_loader.py
554 - VLLM_USE_V1=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s kv_transfer/test_disagg.py
555 + - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
556
557 - label: Plugin Tests (2 GPUs) # 40min
558 working_dir: "/vllm-workspace/tests"
tests/v1/shutdown/test_delete.py new
+97
@@ -0,0 +1,97 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Test that we handle a startup Error and shutdown."""
3 +
4 +import pytest
5 +
6 +from tests.utils import wait_for_gpu_memory_to_clear
7 +from tests.v1.shutdown.utils import (SHUTDOWN_TEST_THRESHOLD_BYTES,
8 + SHUTDOWN_TEST_TIMEOUT_SEC)
9 +from vllm import LLM, SamplingParams
10 +from vllm.engine.arg_utils import AsyncEngineArgs
11 +from vllm.sampling_params import RequestOutputKind
12 +from vllm.utils import cuda_device_count_stateless
13 +from vllm.v1.engine.async_llm import AsyncLLM
14 +
15 +MODELS = ["meta-llama/Llama-3.2-1B"]
16 +
17 +
18 +@pytest.mark.asyncio
19 +@pytest.mark.timeout(SHUTDOWN_TEST_TIMEOUT_SEC)
20 +@pytest.mark.parametrize("model", MODELS)
21 +@pytest.mark.parametrize("tensor_parallel_size", [2, 1])
22 +@pytest.mark.parametrize("send_one_request", [False, True])
23 +async def test_async_llm_delete(model: str, tensor_parallel_size: int,
24 + send_one_request: bool) -> None:
25 + """Test that AsyncLLM frees GPU memory upon deletion.
26 + AsyncLLM always uses an MP client.
27 +
28 + Args:
29 + model: model under test
30 + tensor_parallel_size: degree of tensor parallelism
31 + send_one_request: send one request to engine before deleting
32 + """
33 + if cuda_device_count_stateless() < tensor_parallel_size:
34 + pytest.skip(reason="Not enough CUDA devices")
35 +
36 + engine_args = AsyncEngineArgs(model=model,
37 + enforce_eager=True,
38 + tensor_parallel_size=tensor_parallel_size)
39 +
40 + # Instantiate AsyncLLM; make request to complete any deferred
41 + # initialization; then delete instance
42 + async_llm = AsyncLLM.from_engine_args(engine_args)
43 + if send_one_request:
44 + async for _ in async_llm.generate(
45 + "Hello my name is",
46 + request_id="abc",
47 + sampling_params=SamplingParams(
48 + max_tokens=1, output_kind=RequestOutputKind.DELTA)):
49 + pass
50 + del async_llm
51 +
52 + # Confirm all the processes are cleaned up.
53 + wait_for_gpu_memory_to_clear(
54 + devices=list(range(tensor_parallel_size)),
55 + threshold_bytes=SHUTDOWN_TEST_THRESHOLD_BYTES,
56 + )
57 +
58 +
59 +@pytest.mark.timeout(SHUTDOWN_TEST_TIMEOUT_SEC)
60 +@pytest.mark.parametrize("model", MODELS)
61 +@pytest.mark.parametrize("tensor_parallel_size", [2, 1])
62 +@pytest.mark.parametrize("enable_multiprocessing", [True])
63 +@pytest.mark.parametrize("send_one_request", [False, True])
64 +def test_llm_delete(monkeypatch, model: str, tensor_parallel_size: int,
65 + enable_multiprocessing: bool,
66 + send_one_request: bool) -> None:
67 + """Test that LLM frees GPU memory upon deletion.
68 + TODO(andy) - LLM without multiprocessing.
69 +
70 + Args:
71 + model: model under test
72 + tensor_parallel_size: degree of tensor parallelism
73 + enable_multiprocessing: enable workers in separate process(es)
74 + send_one_request: send one request to engine before deleting
75 + """
76 + if cuda_device_count_stateless() < tensor_parallel_size:
77 + pytest.skip(reason="Not enough CUDA devices")
78 +
79 + with monkeypatch.context() as m:
80 + MP_VALUE = "1" if enable_multiprocessing else "0"
81 + m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", MP_VALUE)
82 +
83 + # Instantiate LLM; make request to complete any deferred
84 + # initialization; then delete instance
85 + llm = LLM(model=model,
86 + enforce_eager=True,
87 + tensor_parallel_size=tensor_parallel_size)
88 + if send_one_request:
89 + llm.generate("Hello my name is",
90 + sampling_params=SamplingParams(max_tokens=1))
91 + del llm
92 +
93 + # Confirm all the processes are cleaned up.
94 + wait_for_gpu_memory_to_clear(
95 + devices=list(range(tensor_parallel_size)),
96 + threshold_bytes=SHUTDOWN_TEST_THRESHOLD_BYTES,
97 + )
tests/v1/shutdown/test_forward_error.py new
+129
@@ -0,0 +1,129 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Test that we handle an Error in model forward and shutdown."""
3 +
4 +import asyncio
5 +
6 +import pytest
7 +
8 +from tests.utils import wait_for_gpu_memory_to_clear
9 +from tests.v1.shutdown.utils import (SHUTDOWN_TEST_THRESHOLD_BYTES,
10 + SHUTDOWN_TEST_TIMEOUT_SEC)
11 +from vllm import LLM, AsyncEngineArgs, SamplingParams
12 +from vllm.distributed import get_tensor_model_parallel_rank
13 +from vllm.model_executor.models.llama import LlamaForCausalLM
14 +from vllm.utils import cuda_device_count_stateless
15 +from vllm.v1.engine.async_llm import AsyncLLM
16 +from vllm.v1.engine.exceptions import EngineDeadError
17 +
18 +MODELS = ["meta-llama/Llama-3.2-1B"]
19 +
20 +
21 +def evil_forward(self, *args, **kwargs):
22 + """Evil forward method that raise an exception after 10 calls."""
23 + NUMBER_OF_GOOD_PASSES = 10
24 +
25 + if not hasattr(self, "num_calls"):
26 + self.num_calls = 0
27 +
28 + if (self.num_calls == NUMBER_OF_GOOD_PASSES
29 + and get_tensor_model_parallel_rank() == 0):
30 + raise Exception("Simulated illegal memory access on Rank 0!")
31 + self.num_calls += 1
32 +
33 + return self.model(*args, **kwargs)
34 +
35 +
36 +@pytest.mark.asyncio
37 +@pytest.mark.parametrize("tensor_parallel_size", [2, 1])
38 +@pytest.mark.parametrize("model", MODELS)
39 +async def test_async_llm_model_error(monkeypatch, tensor_parallel_size: int,
40 + model: str) -> None:
41 + """Test that AsyncLLM propagates a forward pass error and frees memory.
42 +
43 + AsyncLLM always uses an MP client.
44 + """
45 + if cuda_device_count_stateless() < tensor_parallel_size:
46 + pytest.skip(reason="Not enough CUDA devices")
47 +
48 + # Monkeypatch an error in the model.
49 + monkeypatch.setattr(LlamaForCausalLM, "forward", evil_forward)
50 +
51 + engine_args = AsyncEngineArgs(model=model,
52 + enforce_eager=True,
53 + tensor_parallel_size=tensor_parallel_size)
54 + async_llm = AsyncLLM.from_engine_args(engine_args)
55 +
56 + async def generate(request_id: str):
57 + generator = async_llm.generate("Hello my name is",
58 + request_id=request_id,
59 + sampling_params=SamplingParams())
60 + try:
61 + async for _ in generator:
62 + pass
63 + except Exception as e:
64 + return e
65 +
66 + NUM_REQS = 3
67 + tasks = [generate(f"request-{idx}") for idx in range(NUM_REQS)]
68 + outputs = await asyncio.gather(*tasks)
69 +
70 + # Every request should get an EngineDeadError.
71 + for output in outputs:
72 + assert isinstance(output, EngineDeadError)
73 +
74 + # AsyncLLM should be errored.
75 + assert async_llm.errored
76 +
77 + # We should not be able to make another request.
78 + with pytest.raises(EngineDeadError):
79 + async for _ in async_llm.generate("Hello my name is",
80 + request_id="abc",
81 + sampling_params=SamplingParams()):
82 + raise Exception("We should not get here.")
83 +
84 + # Confirm all the processes are cleaned up.
85 + wait_for_gpu_memory_to_clear(
86 + devices=list(range(tensor_parallel_size)),
87 + threshold_bytes=2 * 2**30,
88 + timeout_s=60,
89 + )
90 +
91 + # NOTE: shutdown is handled by the API Server if an exception
92 + # occurs, so it is expected that we would need to call this.
93 + async_llm.shutdown()
94 +
95 +
96 +@pytest.mark.timeout(SHUTDOWN_TEST_TIMEOUT_SEC)
97 +@pytest.mark.parametrize("enable_multiprocessing", [True])
98 +@pytest.mark.parametrize("tensor_parallel_size", [2, 1])
99 +@pytest.mark.parametrize("model", MODELS)
100 +def test_llm_model_error(monkeypatch, tensor_parallel_size: int,
101 + enable_multiprocessing: bool, model: str) -> None:
102 + """Test that LLM propagates a forward pass error and frees memory.
103 + TODO(andy) - LLM without multiprocessing; LLM with multiprocessing
104 + and >1 rank
105 + """
106 + if cuda_device_count_stateless() < tensor_parallel_size:
107 + pytest.skip(reason="Not enough CUDA devices")
108 +
109 + with monkeypatch.context() as m:
110 +
111 + MP_VALUE = "1" if enable_multiprocessing else "0"
112 + m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", MP_VALUE)
113 +
114 + # Monkeypatch an error in the model.
115 + m.setattr(LlamaForCausalLM, "forward", evil_forward)
116 +
117 + llm = LLM(model=model,
118 + enforce_eager=True,
119 + tensor_parallel_size=tensor_parallel_size)
120 +
121 + with pytest.raises(
122 + EngineDeadError if enable_multiprocessing else Exception):
123 + llm.generate("Hello my name is Robert and I")
124 +
125 + # Confirm all the processes are cleaned up.
126 + wait_for_gpu_memory_to_clear(
127 + devices=list(range(tensor_parallel_size)),
128 + threshold_bytes=SHUTDOWN_TEST_THRESHOLD_BYTES,
129 + )
tests/v1/shutdown/test_processor_error.py new
+69
@@ -0,0 +1,69 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Test error handling in Processor. Should not impact other reqs."""
3 +
4 +import asyncio
5 +
6 +import pytest
7 +
8 +from tests.v1.shutdown.utils import SHUTDOWN_TEST_TIMEOUT_SEC
9 +from vllm import SamplingParams
10 +from vllm.engine.arg_utils import AsyncEngineArgs
11 +from vllm.inputs.data import TokensPrompt
12 +from vllm.sampling_params import RequestOutputKind
13 +from vllm.v1.engine.async_llm import AsyncLLM
14 +from vllm.v1.engine.exceptions import EngineGenerateError
15 +
16 +MODELS = ["meta-llama/Llama-3.2-1B"]
17 +
18 +
19 +@pytest.mark.asyncio
20 +@pytest.mark.timeout(SHUTDOWN_TEST_TIMEOUT_SEC)
21 +@pytest.mark.parametrize("model", MODELS)
22 +async def test_async_llm_processor_error(model: str) -> None:
23 + """Test that AsyncLLM propagates a processor error.
24 + Test empty tokens prompt (failure) and non-empty prompt (no failure.)
25 + AsyncLLM always uses an MP client.
26 + """
27 + engine_args = AsyncEngineArgs(model=model, enforce_eager=True)
28 + async_llm = AsyncLLM.from_engine_args(engine_args)
29 +
30 + async def generate(request_id: str):
31 + # [] is not allowed and will raise a ValueError in Processor.
32 + generator = async_llm.generate(TokensPrompt([]),
33 + request_id=request_id,
34 + sampling_params=SamplingParams())
35 + try:
36 + async for _ in generator:
37 + pass
38 + except Exception as e:
39 + return e
40 +
41 + NUM_REQS = 3
42 + tasks = [generate(f"request-{idx}") for idx in range(NUM_REQS)]
43 + outputs = await asyncio.gather(*tasks)
44 +
45 + # Every request should have get an EngineGenerateError.
46 + for output in outputs:
47 + with pytest.raises(EngineGenerateError):
48 + raise output
49 +
50 + # AsyncLLM should be errored.
51 + assert not async_llm.errored
52 +
53 + # This should be no problem.
54 + EXPECTED_TOKENS = 5
55 + outputs = []
56 + async for out in async_llm.generate(
57 + "Hello my name is",
58 + request_id="abc",
59 + sampling_params=SamplingParams(
60 + max_tokens=EXPECTED_TOKENS,
61 + output_kind=RequestOutputKind.DELTA)):
62 + outputs.append(out)
63 +
64 + generated_tokens = []
65 + for out in outputs:
66 + generated_tokens.extend(out.outputs[0].token_ids)
67 + assert len(generated_tokens) == EXPECTED_TOKENS
68 +
69 + async_llm.shutdown()
tests/v1/shutdown/test_startup_error.py new
+97
@@ -0,0 +1,97 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Test that we handle a startup Error and shutdown."""
3 +
4 +import pytest
5 +
6 +from tests.utils import wait_for_gpu_memory_to_clear
7 +from tests.v1.shutdown.utils import (SHUTDOWN_TEST_THRESHOLD_BYTES,
8 + SHUTDOWN_TEST_TIMEOUT_SEC)
9 +from vllm import LLM
10 +from vllm.distributed import get_tensor_model_parallel_rank
11 +from vllm.engine.arg_utils import AsyncEngineArgs
12 +from vllm.model_executor.models.llama import LlamaForCausalLM
13 +from vllm.utils import cuda_device_count_stateless
14 +from vllm.v1.engine.async_llm import AsyncLLM
15 +
16 +MODELS = ["meta-llama/Llama-3.2-1B"]
17 +
18 +
19 +def evil_method(self, *args, **kwargs):
20 + """Evil method that raises an exception."""
21 +
22 + if get_tensor_model_parallel_rank() == 0:
23 + raise Exception("Simulated Error in startup!")
24 +
25 + return self.model(*args, **kwargs, intermediate_tensors=None)
26 +
27 +
28 +@pytest.mark.timeout(SHUTDOWN_TEST_TIMEOUT_SEC)
29 +@pytest.mark.parametrize("model", MODELS)
30 +@pytest.mark.parametrize("tensor_parallel_size", [2, 1])
31 +@pytest.mark.parametrize("failing_method", ["forward", "load_weights"])
32 +def test_async_llm_startup_error(monkeypatch, model: str,
33 + tensor_parallel_size: int,
34 + failing_method: str) -> None:
35 + """Test that AsyncLLM propagates an __init__ error & frees memory.
36 + Test profiling (forward()) and load weights failures.
37 + AsyncLLM always uses an MP client.
38 + """
39 + if cuda_device_count_stateless() < tensor_parallel_size:
40 + pytest.skip(reason="Not enough CUDA devices")
41 +
42 + # Monkeypatch an error in the model.
43 + monkeypatch.setattr(LlamaForCausalLM, failing_method, evil_method)
44 +
45 + engine_args = AsyncEngineArgs(model=model,
46 + enforce_eager=True,
47 + tensor_parallel_size=tensor_parallel_size)
48 +
49 + # Confirm we get an exception.
50 + with pytest.raises(Exception, match="initialization failed"):
51 + _ = AsyncLLM.from_engine_args(engine_args)
52 +
53 + # Confirm all the processes are cleaned up.
54 + wait_for_gpu_memory_to_clear(
55 + devices=list(range(tensor_parallel_size)),
56 + threshold_bytes=SHUTDOWN_TEST_THRESHOLD_BYTES,
57 + )
58 +
59 +
60 +@pytest.mark.timeout(SHUTDOWN_TEST_TIMEOUT_SEC)
61 +@pytest.mark.parametrize("model", MODELS)
62 +@pytest.mark.parametrize("tensor_parallel_size", [2, 1])
63 +@pytest.mark.parametrize("enable_multiprocessing", [True])
64 +@pytest.mark.parametrize("failing_method", ["forward", "load_weights"])
65 +def test_llm_startup_error(monkeypatch, model: str, tensor_parallel_size: int,
66 + enable_multiprocessing: bool,
67 + failing_method: str) -> None:
68 + """Test that LLM propagates an __init__ error and frees memory.
69 + Test profiling (forward()) and load weights failures.
70 + TODO(andy) - LLM without multiprocessing.
71 + """
72 + if model != "meta-llama/Llama-3.2-1B":
73 + pytest.skip(reason="Only test meta-llama/Llama-3.2-1B")
74 + if cuda_device_count_stateless() < tensor_parallel_size:
75 + pytest.skip(reason="Not enough CUDA devices")
76 +
77 + with monkeypatch.context() as m:
78 +
79 + MP_VALUE = "1" if enable_multiprocessing else "0"
80 + m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", MP_VALUE)
81 +
82 + # Monkeypatch an error in the model.
83 + monkeypatch.setattr(LlamaForCausalLM, failing_method, evil_method)
84 +
85 + with pytest.raises(
86 + Exception,
87 + match="initialization failed"
88 + if enable_multiprocessing else "Simulated Error in startup!"):
89 + _ = LLM(model=model,
90 + enforce_eager=True,
91 + tensor_parallel_size=tensor_parallel_size)
92 +
93 + # Confirm all the processes are cleaned up.
94 + wait_for_gpu_memory_to_clear(
95 + devices=list(range(tensor_parallel_size)),
96 + threshold_bytes=SHUTDOWN_TEST_THRESHOLD_BYTES,
97 + )
tests/v1/shutdown/utils.py new
+5
@@ -0,0 +1,5 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Shutdown test utils"""
3 +
4 +SHUTDOWN_TEST_TIMEOUT_SEC = 120
5 +SHUTDOWN_TEST_THRESHOLD_BYTES = 2 * 2**30
vllm/distributed/device_communicators/shm_broadcast.py
+23 -8
@@ -7,11 +7,13 @@ import time
7 from contextlib import contextmanager
8 from dataclasses import dataclass, field
9 from multiprocessing import shared_memory
10 -from typing import List, Optional, Tuple, Union
10 +from threading import Event
11 +from typing import Any, List, Optional, Tuple, Union
12 from unittest.mock import patch
13
14 import torch
15 import torch.distributed as dist
16 +import zmq
17 from torch.distributed import ProcessGroup
18 from zmq import IPV6 # type: ignore
19 from zmq import SUB, SUBSCRIBE, XPUB, XPUB_VERBOSE, Context # type: ignore
@@ -400,7 +402,9 @@ class MessageQueue:
402 break
403
404 @contextmanager
403 - def acquire_read(self, timeout: Optional[float] = None):
405 + def acquire_read(self,
406 + timeout: Optional[float] = None,
407 + cancel: Optional[Event] = None):
408 assert self._is_local_reader, "Only readers can acquire read"
409 start_time = time.monotonic()
410 n_warning = 1
@@ -430,6 +434,9 @@ class MessageQueue:
434 )
435 n_warning += 1
436
437 + if cancel is not None and cancel.is_set():
438 + raise RuntimeError("cancelled")
439 +
440 # if we time out, raise an exception
441 if (timeout is not None
442 and time.monotonic() - start_time > timeout):
@@ -464,10 +471,12 @@ class MessageQueue:
471 if self.n_remote_reader > 0:
472 self.remote_socket.send(serialized_obj)
473
467 - def dequeue(self, timeout: Optional[float] = None):
474 + def dequeue(self,
475 + timeout: Optional[float] = None,
476 + cancel: Optional[Event] = None):
477 """ Read from message queue with optional timeout (in seconds) """
478 if self._is_local_reader:
470 - with self.acquire_read(timeout) as buf:
479 + with self.acquire_read(timeout, cancel) as buf:
480 overflow = buf[0] == 1
481 if not overflow:
482 # no need to know the size of serialized object
@@ -475,15 +484,21 @@ class MessageQueue:
484 # see https://docs.python.org/3/library/pickle.html
485 obj = pickle.loads(buf[1:])
486 if overflow:
478 - recv = self.local_socket.recv()
479 - obj = pickle.loads(recv)
487 + obj = MessageQueue.recv(self.local_socket, timeout)
488 elif self._is_remote_reader:
481 - recv = self.remote_socket.recv()
482 - obj = pickle.loads(recv)
489 + obj = MessageQueue.recv(self.remote_socket, timeout)
490 else:
491 raise RuntimeError("Only readers can dequeue")
492 return obj
493
494 + @staticmethod
495 + def recv(socket: zmq.Socket, timeout: Optional[float]) -> Any:
496 + timeout_ms = None if timeout is None else int(timeout * 1000)
497 + if not socket.poll(timeout=timeout_ms):
498 + raise TimeoutError
499 + recv = socket.recv(copy=False)
500 + return pickle.loads(recv.buffer)
501 +
502 def broadcast_object(self, obj=None):
503 if self._is_writer:
504 self.enqueue(obj)
vllm/entrypoints/launcher.py
+62 -36
@@ -12,9 +12,11 @@ from fastapi import FastAPI, Request, Response
12 from vllm import envs
13 from vllm.engine.async_llm_engine import AsyncEngineDeadError
14 from vllm.engine.multiprocessing import MQEngineDeadError
15 +from vllm.engine.protocol import EngineClient
16 from vllm.entrypoints.ssl import SSLCertRefresher
17 from vllm.logger import init_logger
18 from vllm.utils import find_process_using_port
19 +from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
20
21 logger = init_logger(__name__)
22
@@ -40,6 +42,8 @@ async def serve_http(app: FastAPI,
42
43 loop = asyncio.get_running_loop()
44
45 + watchdog_task = loop.create_task(
46 + watchdog_loop(server, app.state.engine_client))
47 server_task = loop.create_task(
48 server.serve(sockets=[sock] if sock else None))
49
@@ -52,6 +56,7 @@ async def serve_http(app: FastAPI,
56 def signal_handler() -> None:
57 # prevents the uvicorn signal handler to exit early
58 server_task.cancel()
59 + watchdog_task.cancel()
60 if ssl_cert_refresher:
61 ssl_cert_refresher.stop()
62
@@ -73,48 +78,69 @@ async def serve_http(app: FastAPI,
78 port, process, " ".join(process.cmdline()))
79 logger.info("Shutting down FastAPI HTTP server.")
80 return server.shutdown()
81 + finally:
82 + watchdog_task.cancel()
83 +
84 +
85 +async def watchdog_loop(server: uvicorn.Server, engine: EngineClient):
86 + """
87 + # Watchdog task that runs in the background, checking
88 + # for error state in the engine. Needed to trigger shutdown
89 + # if an exception arises is StreamingResponse() generator.
90 + """
91 + VLLM_WATCHDOG_TIME_S = 5.0
92 + while True:
93 + await asyncio.sleep(VLLM_WATCHDOG_TIME_S)
94 + terminate_if_errored(server, engine)
95 +
96 +
97 +def terminate_if_errored(server: uvicorn.Server, engine: EngineClient):
98 + """
99 + See discussions here on shutting down a uvicorn server
100 + https://github.com/encode/uvicorn/discussions/1103
101 + In this case we cannot await the server shutdown here
102 + because handler must first return to close the connection
103 + for this request.
104 + """
105 + engine_errored = engine.errored and not engine.is_running
106 + if not envs.VLLM_KEEP_ALIVE_ON_ENGINE_DEATH and engine_errored:
107 + server.should_exit = True
108
109
110 def _add_shutdown_handlers(app: FastAPI, server: uvicorn.Server) -> None:
79 - """Adds handlers for fatal errors that should crash the server"""
111 + """
112 + VLLM V1 AsyncLLM catches exceptions and returns
113 + only two types: EngineGenerateError and EngineDeadError.
114 +
115 + EngineGenerateError is raised by the per request generate()
116 + method. This error could be request specific (and therefore
117 + recoverable - e.g. if there is an error in input processing).
118 +
119 + EngineDeadError is raised by the background output_handler
120 + method. This error is global and therefore not recoverable.
121 +
122 + We register these @app.exception_handlers to return nice
123 + responses to the end user if they occur and shut down if needed.
124 + See https://fastapi.tiangolo.com/tutorial/handling-errors/
125 + for more details on how exception handlers work.
126 +
127 + If an exception is encountered in a StreamingResponse
128 + generator, the exception is not raised, since we already sent
129 + a 200 status. Rather, we send an error message as the next chunk.
130 + Since the exception is not raised, this means that the server
131 + will not automatically shut down. Instead, we use the watchdog
132 + background task for check for errored state.
133 + """
134
135 @app.exception_handler(RuntimeError)
82 - async def runtime_error_handler(request: Request, __):
83 - """On generic runtime error, check to see if the engine has died.
84 - It probably has, in which case the server will no longer be able to
85 - handle requests. Trigger a graceful shutdown with a SIGTERM."""
86 - engine = request.app.state.engine_client
87 - if (not envs.VLLM_KEEP_ALIVE_ON_ENGINE_DEATH and engine.errored
88 - and not engine.is_running):
89 - logger.fatal("AsyncLLMEngine has failed, terminating server "
90 - "process")
91 - # See discussions here on shutting down a uvicorn server
92 - # https://github.com/encode/uvicorn/discussions/1103
93 - # In this case we cannot await the server shutdown here because
94 - # this handler must first return to close the connection for
95 - # this request.
96 - server.should_exit = True
97 -
98 - return Response(status_code=HTTPStatus.INTERNAL_SERVER_ERROR)
99 -
136 @app.exception_handler(AsyncEngineDeadError)
101 - async def async_engine_dead_handler(_, __):
102 - """Kill the server if the async engine is already dead. It will
103 - not handle any further requests."""
104 - if not envs.VLLM_KEEP_ALIVE_ON_ENGINE_DEATH:
105 - logger.fatal("AsyncLLMEngine is already dead, terminating server "
106 - "process")
107 - server.should_exit = True
108 -
109 - return Response(status_code=HTTPStatus.INTERNAL_SERVER_ERROR)
110 -
137 @app.exception_handler(MQEngineDeadError)
112 - async def mq_engine_dead_handler(_, __):
113 - """Kill the server if the mq engine is already dead. It will
114 - not handle any further requests."""
115 - if not envs.VLLM_KEEP_ALIVE_ON_ENGINE_DEATH:
116 - logger.fatal("MQLLMEngine is already dead, terminating server "
117 - "process")
118 - server.should_exit = True
138 + @app.exception_handler(EngineDeadError)
139 + @app.exception_handler(EngineGenerateError)
140 + async def runtime_exception_handler(request: Request, __):
141 + terminate_if_errored(
142 + server=server,
143 + engine=request.app.state.engine_client,
144 + )
145
146 return Response(status_code=HTTPStatus.INTERNAL_SERVER_ERROR)
vllm/v1/engine/__init__.py
+2
@@ -156,3 +156,5 @@ class EngineCoreRequestType(enum.Enum):
156 ABORT = b'\x01'
157 START_DP = b'\x02'
158 UTILITY = b'\x03'
159 + # Sentinel used within EngineCoreProc.
160 + EXECUTOR_FAILED = b'\x04'
vllm/v1/engine/async_llm.py
+115 -73
@@ -1,8 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -
2 import asyncio
3 import logging
5 -import os
4 from collections.abc import AsyncGenerator, Mapping
5 from copy import copy
6 from typing import Optional, Union
@@ -26,9 +24,10 @@ from vllm.sampling_params import SamplingParams
24 from vllm.transformers_utils.tokenizer import AnyTokenizer
25 from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
26 from vllm.usage.usage_lib import UsageContext
29 -from vllm.utils import Device, cdiv, kill_process_tree
27 +from vllm.utils import Device, cdiv
28 from vllm.v1.engine import EngineCoreRequest
31 -from vllm.v1.engine.core_client import EngineCoreClient
29 +from vllm.v1.engine.core_client import AsyncMPClient, DPAsyncMPClient
30 +from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
31 from vllm.v1.engine.output_processor import (OutputProcessor,
32 RequestOutputCollector)
33 from vllm.v1.engine.parallel_sampling import ParentRequest
@@ -61,8 +60,6 @@ class AsyncLLM(EngineClient):
60 "AsyncLLMEngine.from_vllm_config(...) or explicitly set "
61 "VLLM_USE_V1=0 or 1 and report this issue on Github.")
62
64 - assert start_engine_loop
65 -
63 self.model_config = vllm_config.model_config
64 self.vllm_config = vllm_config
65 self.log_requests = log_requests
@@ -99,15 +96,23 @@ class AsyncLLM(EngineClient):
96 log_stats=self.log_stats)
97
98 # EngineCore (starts the engine in background process).
102 - self.engine_core = EngineCoreClient.make_client(
103 - multiprocess_mode=True,
104 - asyncio_mode=True,
99 + core_client_class = AsyncMPClient if (
100 + vllm_config.parallel_config.data_parallel_size
101 + == 1) else DPAsyncMPClient
102 +
103 + self.engine_core = core_client_class(
104 vllm_config=vllm_config,
105 executor_class=executor_class,
106 log_stats=self.log_stats,
107 )
108
109 self.output_handler: Optional[asyncio.Task] = None
110 + try:
111 + # Start output handler eagerly if we are in the asyncio eventloop.
112 + asyncio.get_running_loop()
113 + self._run_output_handler()
114 + except RuntimeError:
115 + pass
116
117 @classmethod
118 def from_vllm_config(
@@ -165,6 +170,9 @@ class AsyncLLM(EngineClient):
170 usage_context=usage_context,
171 )
172
173 + def __del__(self):
174 + self.shutdown()
175 +
176 def shutdown(self):
177 """Shutdown, cleaning up the background proc and IPC."""
178
@@ -187,6 +195,9 @@ class AsyncLLM(EngineClient):
195 ) -> RequestOutputCollector:
196 """Add new request to the AsyncLLM."""
197
198 + if self.errored:
199 + raise EngineDeadError()
200 +
201 assert isinstance(params, SamplingParams), \
202 "Pooling is not supported in V1"
203
@@ -261,9 +272,7 @@ class AsyncLLM(EngineClient):
272 # We start the output_handler on the first call to generate() so
273 # we can call __init__ before the event loop, which enables us
274 # to handle startup failure gracefully in the OpenAI server.
264 - if self.output_handler is None:
265 - self.output_handler = asyncio.create_task(
266 - self._run_output_handler())
275 + self._run_output_handler()
276
277 q = await self.add_request(
278 request_id,
@@ -288,62 +297,96 @@ class AsyncLLM(EngineClient):
297 finished = out.finished
298 yield out
299
291 - # If the request is disconnected by the client, the
292 - # generate() task will be canceled. So, we abort the
293 - # request if we end up here.
300 + # If the request is disconnected by the client, generate()
301 + # is cancelled. So, we abort the request if we end up here.
302 except asyncio.CancelledError:
303 await self.abort(request_id)
304 + if self.log_requests:
305 + logger.info("Request %s aborted.", request_id)
306 raise
307
298 - async def _run_output_handler(self):
299 - """Background loop: pulls from EngineCore and pushes to AsyncStreams."""
308 + # Engine is dead. Do not abort since we shut down.
309 + except EngineDeadError:
310 + if self.log_requests:
311 + logger.info("Request %s failed (engine dead).", request_id)
312 + raise
313
301 - try:
302 - while True:
303 - # 1) Pull EngineCoreOutputs from the EngineCore.
304 - outputs = await self.engine_core.get_output_async()
305 - num_outputs = len(outputs.outputs)
306 -
307 - iteration_stats = IterationStats() if (
308 - self.log_stats and num_outputs) else None
309 -
310 - # Split outputs into chunks of at most
311 - # VLLM_V1_OUTPUT_PROC_CHUNK_SIZE, so that we don't block the
312 - # event loop for too long.
313 - if num_outputs <= VLLM_V1_OUTPUT_PROC_CHUNK_SIZE:
314 - slices = (outputs.outputs, )
315 - else:
316 - slices = np.array_split(
317 - outputs.outputs,
318 - cdiv(num_outputs, VLLM_V1_OUTPUT_PROC_CHUNK_SIZE))
319 -
320 - for i, outputs_slice in enumerate(slices):
321 - # 2) Process EngineCoreOutputs.
322 - processed_outputs = self.output_processor.process_outputs(
323 - outputs_slice, outputs.timestamp, iteration_stats)
324 - # NOTE: RequestOutputs are pushed to their queues.
325 - assert not processed_outputs.request_outputs
326 -
327 - # Allow other asyncio tasks to run between chunks
328 - if i + 1 < len(slices):
329 - await asyncio.sleep(0)
330 -
331 - # 3) Abort any reqs that finished due to stop strings.
332 - await self.engine_core.abort_requests_async(
333 - processed_outputs.reqs_to_abort)
334 -
335 - # 4) Logging.
336 - # TODO(rob): make into a coroutine and launch it in
337 - # background thread once Prometheus overhead is non-trivial.
338 - self._record_stats(
339 - engine_index=outputs.engine_index,
340 - scheduler_stats=outputs.scheduler_stats,
341 - iteration_stats=iteration_stats,
342 - )
314 + # Request validation error.
315 + except ValueError:
316 + if self.log_requests:
317 + logger.info("Request %s failed (bad request).", request_id)
318 + raise
319
320 + # Unexpected error in the generate() task (possibly recoverable).
321 except Exception as e:
345 - logger.exception("EngineCore output handler hit an error: %s", e)
346 - kill_process_tree(os.getpid())
322 + await self.abort(request_id)
323 + if self.log_requests:
324 + logger.info("Request %s failed.", request_id)
325 + raise EngineGenerateError() from e
326 +
327 + def _run_output_handler(self):
328 + """Background loop: pulls from EngineCore and pushes to AsyncStreams."""
329 +
330 + if self.output_handler is not None:
331 + return
332 +
333 + # Ensure that the task doesn't have a circular ref back to the AsyncLLM
334 + # object, or else it won't be garbage collected and cleaned up properly.
335 + engine_core = self.engine_core
336 + output_processor = self.output_processor
337 + log_stats = self.log_stats
338 + stat_loggers = self.stat_loggers if log_stats else None
339 +
340 + async def output_handler():
341 + try:
342 + while True:
343 + # 1) Pull EngineCoreOutputs from the EngineCore.
344 + outputs = await engine_core.get_output_async()
345 + num_outputs = len(outputs.outputs)
346 +
347 + iteration_stats = IterationStats() if (
348 + log_stats and num_outputs) else None
349 +
350 + # Split outputs into chunks of at most
351 + # VLLM_V1_OUTPUT_PROC_CHUNK_SIZE, so that we don't block the
352 + # event loop for too long.
353 + if num_outputs <= VLLM_V1_OUTPUT_PROC_CHUNK_SIZE:
354 + slices = (outputs.outputs, )
355 + else:
356 + slices = np.array_split(
357 + outputs.outputs,
358 + cdiv(num_outputs, VLLM_V1_OUTPUT_PROC_CHUNK_SIZE))
359 +
360 + for i, outputs_slice in enumerate(slices):
361 + # 2) Process EngineCoreOutputs.
362 + processed_outputs = output_processor.process_outputs(
363 + outputs_slice, outputs.timestamp, iteration_stats)
364 + # NOTE: RequestOutputs are pushed to their queues.
365 + assert not processed_outputs.request_outputs
366 +
367 + # Allow other asyncio tasks to run between chunks
368 + if i + 1 < len(slices):
369 + await asyncio.sleep(0)
370 +
371 + # 3) Abort any reqs that finished due to stop strings.
372 + await engine_core.abort_requests_async(
373 + processed_outputs.reqs_to_abort)
374 +
375 + # 4) Logging.
376 + # TODO(rob): make into a coroutine and launch it in
377 + # background thread once Prometheus overhead is non-trivial.
378 + if stat_loggers:
379 + assert outputs.scheduler_stats is not None
380 + AsyncLLM._record_stats(
381 + stat_loggers[outputs.engine_index],
382 + scheduler_stats=outputs.scheduler_stats,
383 + iteration_stats=iteration_stats,
384 + )
385 + except Exception as e:
386 + logger.exception("AsyncLLM output_handler failed.")
387 + output_processor.propagate_error(e)
388 +
389 + self.output_handler = asyncio.create_task(output_handler())
390
391 async def abort(self, request_id: str) -> None:
392 """Abort RequestId in OutputProcessor and EngineCore."""
@@ -354,17 +397,15 @@ class AsyncLLM(EngineClient):
397 if self.log_requests:
398 logger.info("Aborted request %s.", request_id)
399
400 + @staticmethod
401 def _record_stats(
358 - self,
359 - scheduler_stats: Optional[SchedulerStats],
402 + stat_loggers: list[StatLoggerBase],
403 + scheduler_stats: SchedulerStats,
404 iteration_stats: Optional[IterationStats],
361 - engine_index: int = 0,
405 ):
363 - if not self.log_stats:
364 - return
365 -
366 - assert scheduler_stats is not None
367 - for stat_logger in self.stat_loggers[engine_index]:
406 + """static so that it can be used from the output_handler task
407 + without a circular ref to AsyncLLM."""
408 + for stat_logger in stat_loggers:
409 stat_logger.record(scheduler_stats=scheduler_stats,
410 iteration_stats=iteration_stats)
411
@@ -451,16 +492,17 @@ class AsyncLLM(EngineClient):
492
493 @property
494 def is_running(self) -> bool:
454 - return True
495 + # Is None before the loop is started.
496 + return self.output_handler is None or not self.output_handler.done()
497
498 @property
499 def is_stopped(self) -> bool:
458 - return False
500 + return self.errored
501
502 @property
503 def errored(self) -> bool:
462 - return False
504 + return self.engine_core.resources.engine_dead or not self.is_running
505
506 @property
507 def dead_error(self) -> BaseException:
466 - return Exception() # TODO: implement
508 + return EngineDeadError()
vllm/v1/engine/core.py
+60 -27
@@ -11,9 +11,7 @@ from logging import DEBUG
11 from typing import Any, Callable, Optional, TypeVar, Union
12
13 import msgspec
14 -import psutil
14 import zmq
16 -import zmq.asyncio
15
16 from vllm.config import ParallelConfig, VllmConfig
17 from vllm.distributed import stateless_destroy_torch_distributed_process_group
@@ -22,8 +20,7 @@ from vllm.logger import init_logger
20 from vllm.lora.request import LoRARequest
21 from vllm.transformers_utils.config import (
22 maybe_register_config_serialize_by_value)
25 -from vllm.utils import (get_exception_traceback, resolve_obj_by_qualname,
26 - zmq_socket_ctx)
23 +from vllm.utils import resolve_obj_by_qualname, zmq_socket_ctx
24 from vllm.v1.core.kv_cache_utils import (get_kv_cache_config,
25 unify_kv_cache_configs)
26 from vllm.v1.core.sched.interface import SchedulerInterface
@@ -50,12 +47,11 @@ _R = TypeVar('_R') # Return type for collective_rpc
47 class EngineCore:
48 """Inner loop of vLLM's Engine."""
49
53 - def __init__(
54 - self,
55 - vllm_config: VllmConfig,
56 - executor_class: type[Executor],
57 - log_stats: bool,
58 - ):
50 + def __init__(self,
51 + vllm_config: VllmConfig,
52 + executor_class: type[Executor],
53 + log_stats: bool,
54 + executor_fail_callback: Optional[Callable] = None):
55 assert vllm_config.model_config.runner_type != "pooling"
56
57 logger.info("Initializing a V1 LLM engine (v%s) with config: %s",
@@ -65,6 +61,9 @@ class EngineCore:
61
62 # Setup Model.
63 self.model_executor = executor_class(vllm_config)
64 + if executor_fail_callback is not None:
65 + self.model_executor.register_failure_callback(
66 + executor_fail_callback)
67
68 # Setup KV Caches and update CacheConfig after profiling.
69 num_gpu_blocks, num_cpu_blocks, kv_cache_config = \
@@ -254,7 +253,8 @@ class EngineCore:
253 return engine_core_outputs
254
255 def shutdown(self):
257 - self.model_executor.shutdown()
256 + if self.model_executor:
257 + self.model_executor.shutdown()
258
259 def profile(self, is_start: bool = True):
260 self.model_executor.profile(is_start)
@@ -308,6 +308,8 @@ class EngineCore:
308 class EngineCoreProc(EngineCore):
309 """ZMQ-wrapper for running EngineCore in background process."""
310
311 + ENGINE_CORE_DEAD = b'ENGINE_CORE_DEAD'
312 +
313 def __init__(
314 self,
315 input_path: str,
@@ -317,11 +319,16 @@ class EngineCoreProc(EngineCore):
319 log_stats: bool,
320 engine_index: int = 0,
321 ):
320 - super().__init__(vllm_config, executor_class, log_stats)
322 + input_queue = queue.Queue[tuple[EngineCoreRequestType, Any]]()
323 +
324 + executor_fail_callback = lambda: input_queue.put_nowait(
325 + (EngineCoreRequestType.EXECUTOR_FAILED, b''))
326 +
327 + super().__init__(vllm_config, executor_class, log_stats,
328 + executor_fail_callback)
329
330 self.step_fn = (self.step if self.batch_queue is None else
331 self.step_with_batch_queue)
324 -
332 self.global_unfinished_reqs = False
333
334 # Background Threads and Queues for IO. These enable us to
@@ -329,15 +336,16 @@ class EngineCoreProc(EngineCore):
336 # and to overlap some serialization/deserialization with the
337 # model forward pass.
338 # Threads handle Socket <-> Queues and core_busy_loop uses Queue.
332 - self.input_queue: queue.Queue[tuple[EngineCoreRequestType,
333 - Any]] = queue.Queue()
334 - self.output_queue: queue.Queue[EngineCoreOutputs] = queue.Queue()
339 + self.input_queue = input_queue
340 + self.output_queue = queue.Queue[Union[EngineCoreOutputs, bytes]]()
341 threading.Thread(target=self.process_input_socket,
342 args=(input_path, engine_index),
343 daemon=True).start()
338 - threading.Thread(target=self.process_output_socket,
339 - args=(output_path, engine_index),
340 - daemon=True).start()
344 + self.output_thread = threading.Thread(
345 + target=self.process_output_socket,
346 + args=(output_path, engine_index),
347 + daemon=True)
348 + self.output_thread.start()
349
350 @staticmethod
351 def run_engine_core(*args,
@@ -364,7 +372,6 @@ class EngineCoreProc(EngineCore):
372 signal.signal(signal.SIGTERM, signal_handler)
373 signal.signal(signal.SIGINT, signal_handler)
374
367 - parent_process = psutil.Process().parent()
375 engine_core: Optional[EngineCoreProc] = None
376 try:
377 parallel_config: ParallelConfig = kwargs[
@@ -380,13 +387,15 @@ class EngineCoreProc(EngineCore):
387 engine_core.run_busy_loop()
388
389 except SystemExit:
383 - logger.debug("EngineCore interrupted.")
384 -
385 - except Exception:
386 - traceback = get_exception_traceback()
387 - logger.error("EngineCore hit an exception: %s", traceback)
388 - parent_process.send_signal(signal.SIGUSR1)
390 + logger.debug("EngineCore exiting.")
391
392 + except Exception as e:
393 + if engine_core is None:
394 + logger.exception("EngineCore failed to start.")
395 + else:
396 + logger.exception("EngineCore encountered a fatal error.")
397 + engine_core._send_engine_dead()
398 + raise e
399 finally:
400 if engine_core is not None:
401 engine_core.shutdown()
@@ -458,6 +467,11 @@ class EngineCoreProc(EngineCore):
467 f" failed: {str(e)}")
468 self.output_queue.put_nowait(
469 EngineCoreOutputs(utility_output=output))
470 + elif request_type == EngineCoreRequestType.EXECUTOR_FAILED:
471 + raise RuntimeError("Executor failed.")
472 + else:
473 + logger.error("Unrecognized input request type encountered: %s",
474 + request_type)
475
476 @staticmethod
477 def _convert_msgspec_args(method, args):
@@ -473,6 +487,18 @@ class EngineCoreProc(EngineCore):
487 and not isinstance(v, p.annotation) else v
488 for v, p in zip(args, arg_types))
489
490 + def _send_engine_dead(self):
491 + """Send EngineDead status to the EngineCoreClient."""
492 +
493 + # Put ENGINE_CORE_DEAD in the queue.
494 + self.output_queue.put_nowait(EngineCoreProc.ENGINE_CORE_DEAD)
495 +
496 + # Wait until msg sent by the daemon before shutdown.
497 + self.output_thread.join(timeout=5.0)
498 + if self.output_thread.is_alive():
499 + logger.fatal("vLLM shutdown signal from EngineCore failed "
500 + "to send. Please report this issue.")
501 +
502 def process_input_socket(self, input_path: str, engine_index: int):
503 """Input socket IO thread."""
504
@@ -511,9 +537,16 @@ class EngineCoreProc(EngineCore):
537 # Reuse send buffer.
538 buffer = bytearray()
539
514 - with zmq_socket_ctx(output_path, zmq.constants.PUSH) as socket:
540 + # We must set linger to ensure the ENGINE_CORE_DEAD
541 + # message is sent prior to closing the socket.
542 + with zmq_socket_ctx(output_path, zmq.constants.PUSH,
543 + linger=4000) as socket:
544 while True:
545 outputs = self.output_queue.get()
546 + if outputs == EngineCoreProc.ENGINE_CORE_DEAD:
547 + socket.send(outputs, copy=False)
548 + break
549 + assert not isinstance(outputs, bytes)
550 outputs.engine_index = engine_index
551 buffers = encoder.encode_into(outputs, buffer)
552 socket.send_multipart(buffers, copy=False)
vllm/v1/engine/core_client.py
+126 -90
@@ -1,14 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -
2 import asyncio
4 -import os
3 import queue
6 -import signal
7 -import threading
4 import uuid
5 import weakref
6 from abc import ABC, abstractmethod
11 -from collections.abc import Awaitable
7 +from collections.abc import Awaitable, Sequence
8 from concurrent.futures import Future
9 from dataclasses import dataclass, field
10 from threading import Thread
@@ -21,10 +17,11 @@ from vllm.config import VllmConfig
17 from vllm.logger import init_logger
18 from vllm.lora.request import LoRARequest
19 from vllm.utils import (get_open_zmq_inproc_path, get_open_zmq_ipc_path,
24 - kill_process_tree, make_zmq_socket)
20 + make_zmq_socket)
21 from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
22 EngineCoreRequestType, UtilityOutput)
23 from vllm.v1.engine.core import EngineCore, EngineCoreProc
24 +from vllm.v1.engine.exceptions import EngineDeadError
25 from vllm.v1.executor.abstract import Executor
26 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr
27 from vllm.v1.utils import BackgroundProcHandle
@@ -305,14 +302,22 @@ class BackgroundResources:
302 core_engines: list[CoreEngine] = field(default_factory=list)
303 output_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None
304 input_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None
305 + output_queue_task: Optional[asyncio.Task] = None
306 shutdown_path: Optional[str] = None
307
308 + # Set if any of the engines are dead. Here so that the output
309 + # processing threads can access it without holding a ref to the client.
310 + engine_dead: bool = False
311 +
312 def __call__(self):
313 """Clean up background resources."""
314
315 for core_engine in self.core_engines:
316 core_engine.close()
317
318 + if self.output_queue_task is not None:
319 + self.output_queue_task.cancel()
320 +
321 # ZMQ context termination can hang if the sockets
322 # aren't explicitly closed first.
323 if self.output_socket is not None:
@@ -327,6 +332,12 @@ class BackgroundResources:
332 # Send shutdown signal.
333 shutdown_sender.send(b'')
334
335 + def validate_alive(self, frames: Sequence[zmq.Frame]):
336 + if len(frames) == 1 and (frames[0].buffer
337 + == EngineCoreProc.ENGINE_CORE_DEAD):
338 + self.engine_dead = True
339 + raise EngineDeadError()
340 +
341
342 class MPClient(EngineCoreClient):
343 """
@@ -348,27 +359,6 @@ class MPClient(EngineCoreClient):
359 executor_class: type[Executor],
360 log_stats: bool,
361 ):
351 - # The child processes will send SIGUSR1 when unrecoverable
352 - # errors happen. We kill the process tree here so that the
353 - # stack trace is very evident.
354 - # TODO(rob): rather than killing the main process, we should
355 - # figure out how to raise an AsyncEngineDeadError and
356 - # handle at the API server level so we can return a better
357 - # error code to the clients calling vLLM.
358 - def sigusr1_handler(signum, frame):
359 - logger.fatal("Got fatal signal from worker processes, shutting "
360 - "down. See stack trace above for root cause issue.")
361 - kill_process_tree(os.getpid())
362 -
363 - if threading.current_thread() == threading.main_thread():
364 - signal.signal(signal.SIGUSR1, sigusr1_handler)
365 - else:
366 - logger.warning("SIGUSR1 handler not installed because we are not "
367 - "running in the main thread. In this case the "
368 - "forked engine process may not be killed when "
369 - "an exception is raised, and you need to handle "
370 - "the engine process shutdown manually.")
371 -
362 # Serialization setup.
363 self.encoder = MsgpackEncoder()
364 self.decoder = MsgpackDecoder(EngineCoreOutputs)
@@ -378,32 +368,37 @@ class MPClient(EngineCoreClient):
368 self.ctx = zmq.asyncio.Context(sync_ctx) if asyncio_mode else sync_ctx
369
370 # This will ensure resources created so far are closed
381 - # when the client is garbage collected, even if an
371 + # when the client is garbage collected, even if an
372 # exception is raised mid-construction.
373 self.resources = BackgroundResources(ctx=sync_ctx)
374 self._finalizer = weakref.finalize(self, self.resources)
385 -
386 - # Paths and sockets for IPC.
387 - self.output_path = get_open_zmq_ipc_path()
388 - input_path = get_open_zmq_ipc_path()
389 - self.input_socket = make_zmq_socket(self.ctx,
390 - input_path,
391 - zmq.ROUTER,
392 - bind=True)
393 - self.resources.input_socket = self.input_socket
394 -
395 - new_core_engine = lambda index, local_dp_rank=None: CoreEngine(
396 - vllm_config, executor_class, log_stats, input_path, self.
397 - output_path, index, local_dp_rank)
398 -
399 - # Start engine core process(es).
400 - self._init_core_engines(vllm_config, new_core_engine,
401 - self.resources.core_engines)
402 -
403 - # Wait for engine core process(es) to start.
404 - self._wait_for_engine_startup()
405 -
406 - self.utility_results: dict[int, AnyFuture] = {}
375 + success = False
376 + try:
377 + # Paths and sockets for IPC.
378 + self.output_path = get_open_zmq_ipc_path()
379 + input_path = get_open_zmq_ipc_path()
380 + self.input_socket = make_zmq_socket(self.ctx,
381 + input_path,
382 + zmq.ROUTER,
383 + bind=True)
384 + self.resources.input_socket = self.input_socket
385 +
386 + new_core_engine = lambda index, local_dp_rank=None: CoreEngine(
387 + vllm_config, executor_class, log_stats, input_path, self.
388 + output_path, index, local_dp_rank)
389 +
390 + # Start engine core process(es).
391 + self._init_core_engines(vllm_config, new_core_engine,
392 + self.resources.core_engines)
393 +
394 + # Wait for engine core process(es) to start.
395 + self._wait_for_engine_startup()
396 +
397 + self.utility_results: dict[int, AnyFuture] = {}
398 + success = True
399 + finally:
400 + if not success:
401 + self._finalizer()
402
403 def _wait_for_engine_startup(self):
404 # Get a sync handle to the socket which can be sync or async.
@@ -451,8 +446,18 @@ class MPClient(EngineCoreClient):
446 self.core_engine = core_engine
447
448 def shutdown(self):
449 + # Terminate background resources.
450 self._finalizer()
451
452 + def _format_exception(self, e: Exception) -> Exception:
453 + """If errored, use EngineDeadError so root cause is clear."""
454 + return EngineDeadError(
455 + suppress_context=True) if self.resources.engine_dead else e
456 +
457 + def ensure_alive(self):
458 + if self.resources.engine_dead:
459 + raise EngineDeadError()
460 +
461
462 def _process_utility_output(output: UtilityOutput,
463 utility_results: dict[int, AnyFuture]):
@@ -476,7 +481,7 @@ class SyncMPClient(MPClient):
481 log_stats=log_stats,
482 )
483
479 - self.outputs_queue: queue.Queue[EngineCoreOutputs] = queue.Queue()
484 + self.outputs_queue = queue.Queue[Union[EngineCoreOutputs, Exception]]()
485
486 # Ensure that the outputs socket processing thread does not have
487 # a ref to the client which prevents gc.
@@ -487,7 +492,8 @@ class SyncMPClient(MPClient):
492 outputs_queue = self.outputs_queue
493
494 shutdown_path = get_open_zmq_inproc_path()
490 - self.resources.shutdown_path = shutdown_path
495 + resources = self.resources
496 + resources.shutdown_path = shutdown_path
497
498 def process_outputs_socket():
499 shutdown_socket = ctx.socket(zmq.PAIR)
@@ -506,12 +512,15 @@ class SyncMPClient(MPClient):
512 break
513
514 frames = out_socket.recv_multipart(copy=False)
515 + resources.validate_alive(frames)
516 outputs = decoder.decode(frames)
517 if outputs.utility_output:
518 _process_utility_output(outputs.utility_output,
519 utility_results)
520 else:
521 outputs_queue.put_nowait(outputs)
522 + except Exception as e:
523 + outputs_queue.put_nowait(e)
524 finally:
525 # Close sockets.
526 shutdown_socket.close(linger=0)
@@ -524,9 +533,16 @@ class SyncMPClient(MPClient):
533 self.output_queue_thread.start()
534
535 def get_output(self) -> EngineCoreOutputs:
527 - return self.outputs_queue.get()
536 + # If an exception arises in process_outputs_socket task,
537 + # it is forwarded to the outputs_queue so we can raise it
538 + # from this (run_output_handler) task to shut down the server.
539 + outputs = self.outputs_queue.get()
540 + if isinstance(outputs, Exception):
541 + raise self._format_exception(outputs) from None
542 + return outputs
543
544 def _send_input(self, request_type: EngineCoreRequestType, request: Any):
545 + self.ensure_alive()
546 # (Identity, RequestType, SerializedRequest)
547 msg = (self.core_engine.identity, request_type.value,
548 *self.encoder.encode(request))
@@ -608,61 +624,81 @@ class AsyncMPClient(MPClient):
624 log_stats=log_stats,
625 )
626
611 - self.outputs_queue: Optional[asyncio.Queue[EngineCoreOutputs]] = None
612 - self.queue_task: Optional[asyncio.Task] = None
613 -
614 - self.outputs_handler: Optional[Callable[
615 - [AsyncMPClient, EngineCoreOutputs], Awaitable[None]]] = None
627 + self.outputs_queue = asyncio.Queue[Union[EngineCoreOutputs,
628 + Exception]]()
629 + try:
630 + # If we are running in an asyncio event loop, start the queue task.
631 + # Otherwise, it will be started lazily. If it is not started here,
632 + # we could miss EXECUTOR_FAILED messages from engine core if they
633 + # occur prior to any requests being sent.
634 + asyncio.get_running_loop()
635 + self._ensure_output_queue_task()
636 + except RuntimeError:
637 + pass
638
639 def _ensure_output_queue_task(self):
618 - if self.outputs_queue is not None:
640 + resources = self.resources
641 + if resources.output_queue_task is not None:
642 return
643
644 # Perform IO in separate task to parallelize as much as possible.
645 # Avoid task having direct reference back to the client.
623 - self.outputs_queue = asyncio.Queue()
646 decoder = self.decoder
647 utility_results = self.utility_results
648 outputs_queue = self.outputs_queue
627 - output_handler = self.outputs_handler
649 + output_handler: Optional[Callable[[AsyncMPClient, EngineCoreOutputs],
650 + Awaitable[None]]] = getattr(
651 + self.__class__,
652 + "process_engine_outputs", None)
653 _self_ref = weakref.ref(self) if output_handler else None
654 output_path = self.output_path
655 output_socket = make_zmq_socket(self.ctx, output_path,
656 zmq.constants.PULL)
632 - self.resources.output_socket = output_socket
657 + resources.output_socket = output_socket
658
659 async def process_outputs_socket():
635 - while True:
636 - frames = await output_socket.recv_multipart(copy=False)
637 - outputs: EngineCoreOutputs = decoder.decode(frames)
638 - if outputs.utility_output:
639 - _process_utility_output(outputs.utility_output,
640 - utility_results)
641 - continue
642 -
643 - if output_handler is not None:
644 - assert _self_ref is not None
645 - _self = _self_ref()
646 - if not _self:
647 - # Client has been garbage collected, abort.
648 - return
649 - await output_handler(_self, outputs)
650 -
651 - if outputs.outputs or outputs.scheduler_stats:
652 - outputs_queue.put_nowait(outputs)
653 -
654 - self.queue_task = asyncio.create_task(process_outputs_socket(),
655 - name="EngineCoreOutputQueueTask")
660 + try:
661 + while True:
662 + frames = await output_socket.recv_multipart(copy=False)
663 + resources.validate_alive(frames)
664 + outputs: EngineCoreOutputs = decoder.decode(frames)
665 + if outputs.utility_output:
666 + _process_utility_output(outputs.utility_output,
667 + utility_results)
668 + continue
669 +
670 + if output_handler is not None:
671 + assert _self_ref is not None
672 + _self = _self_ref()
673 + if not _self:
674 + # Client has been garbage collected, abort.
675 + return
676 + await output_handler(_self, outputs)
677 +
678 + if outputs.outputs or outputs.scheduler_stats:
679 + outputs_queue.put_nowait(outputs)
680 + except Exception as e:
681 + outputs_queue.put_nowait(e)
682 +
683 + resources.output_queue_task = asyncio.create_task(
684 + process_outputs_socket(), name="EngineCoreOutputQueueTask")
685
686 async def get_output_async(self) -> EngineCoreOutputs:
687 self._ensure_output_queue_task()
688 + # If an exception arises in process_outputs_socket task,
689 + # it is forwarded to the outputs_queue so we can raise it
690 + # from this (run_output_handler) task to shut down the server.
691 assert self.outputs_queue is not None
660 - return await self.outputs_queue.get()
692 + outputs = await self.outputs_queue.get()
693 + if isinstance(outputs, Exception):
694 + raise self._format_exception(outputs) from None
695 + return outputs
696
697 def _send_input(self,
698 request_type: EngineCoreRequestType,
699 request: Any,
700 engine: Optional[CoreEngine] = None) -> Awaitable[None]:
701 + self.ensure_alive()
702 if engine is None:
703 engine = self.core_engine
704
@@ -671,6 +707,7 @@ class AsyncMPClient(MPClient):
707
708 def _send_input_message(self, message: tuple[bytestr, ...],
709 engine: CoreEngine) -> Awaitable[None]:
710 + self.ensure_alive()
711 message = (engine.identity, ) + message
712 return self.input_socket.send_multipart(message, copy=False)
713
@@ -754,18 +791,17 @@ class DPAsyncMPClient(AsyncMPClient):
791
792 def __init__(self, vllm_config: VllmConfig, executor_class: type[Executor],
793 log_stats: bool):
757 - super().__init__(vllm_config, executor_class, log_stats)
794
759 - assert len(self.core_engines) > 1
795 + self.num_engines_running = 0
796 + self.reqs_in_flight: dict[str, CoreEngine] = {}
797 +
798 + super().__init__(vllm_config, executor_class, log_stats)
799
800 # Control message used for triggering dp idle mode loop.
801 self.start_dp_msg = (EngineCoreRequestType.START_DP.value,
802 *self.encoder.encode(None))
803
765 - self.num_engines_running = 0
766 - self.reqs_in_flight: dict[str, CoreEngine] = {}
767 -
768 - self.outputs_handler = DPAsyncMPClient.process_engine_outputs # type: ignore[assignment]
804 + assert len(self.core_engines) > 1
805
806 def _init_core_engines(
807 self,
vllm/v1/engine/exceptions.py new
+16
@@ -0,0 +1,16 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +class EngineGenerateError(Exception):
3 + """Raised when a AsyncLLM.generate() fails. Recoverable."""
4 + pass
5 +
6 +
7 +class EngineDeadError(Exception):
8 + """Raised when the EngineCore dies. Unrecoverable."""
9 +
10 + def __init__(self, *args, suppress_context: bool = False, **kwargs):
11 + ENGINE_DEAD_MESSAGE = "EngineCore encountered an issue. See stack trace (above) for the root cause." # noqa: E501
12 +
13 + super().__init__(ENGINE_DEAD_MESSAGE, *args, **kwargs)
14 + # Make stack trace clearer when using with LLMEngine by
15 + # silencing irrelevant ZMQError.
16 + self.__suppress_context__ = suppress_context
vllm/v1/engine/output_processor.py
+24 -9
@@ -28,32 +28,40 @@ class RequestOutputCollector:
28
29 def __init__(self, output_kind: RequestOutputKind):
30 self.aggregate = output_kind == RequestOutputKind.DELTA
31 - self.output: Optional[RequestOutput] = None
31 + self.output: Optional[Union[RequestOutput, Exception]] = None
32 self.ready = asyncio.Event()
33
34 - def put(self, output: RequestOutput) -> None:
35 - if self.output is None:
34 + def put(self, output: Union[RequestOutput, Exception]) -> None:
35 + """Non-blocking put operation."""
36 + if self.output is None or isinstance(output, Exception):
37 self.output = output
38 self.ready.set()
38 - elif self.aggregate:
39 - # Coalesce the outputs in delta case.
40 - self.output.add(output)
41 - else:
42 - # Just replace latest in non-delta case.
43 - self.output = output
39 + elif isinstance(self.output, RequestOutput):
40 + if self.aggregate:
41 + # Coalesce the outputs in delta case.
42 + self.output.add(output)
43 + else:
44 + # Just replace latest in non-delta case.
45 + self.output = output
46
47 async def get(self) -> RequestOutput:
48 + """Get operation blocks on put event."""
49 while (output := self.output) is None:
50 await self.ready.wait()
51 self.output = None
52 self.ready.clear()
53 + if isinstance(output, Exception):
54 + raise output
55 return output
56
57 def get_nowait(self) -> Optional[RequestOutput]:
58 + """Non-blocking get operation."""
59 output = self.output
60 if output is not None:
61 self.output = None
62 self.ready.clear()
63 + if isinstance(output, Exception):
64 + raise output
65 return output
66
67
@@ -235,6 +243,13 @@ class OutputProcessor:
243 def has_unfinished_requests(self) -> bool:
244 return len(self.request_states) > 0
245
246 + def propagate_error(self, e: Exception):
247 + """Propagate error to all generate() tasks."""
248 +
249 + for _, state in self.request_states.items():
250 + assert state.queue is not None
251 + state.queue.put(e)
252 +
253 def abort_requests(
254 self,
255 request_ids: Iterable[str],
vllm/v1/executor/abstract.py
+10 -1
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 from concurrent.futures import Future
4 -from typing import Union
4 +from typing import Callable, Union
5
6 import torch
7 import torch.distributed as dist
@@ -15,6 +15,8 @@ from vllm.executor.uniproc_executor import ( # noqa
15 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
16 from vllm.v1.outputs import ModelRunnerOutput
17
18 +FailureCallback = Callable[[], None]
19 +
20
21 class Executor(ExecutorBase):
22 """
@@ -62,6 +64,13 @@ class Executor(ExecutorBase):
64 args=(kv_cache_configs, ))
65 self.collective_rpc("compile_or_warm_up_model")
66
67 + def register_failure_callback(self, callback: FailureCallback):
68 + """
69 + Register a function to be called if the executor enters a permanent
70 + failed state.
71 + """
72 + pass
73 +
74 def determine_available_memory(self) -> list[int]: # in bytes
75 output = self.collective_rpc("determine_available_memory")
76 return output
vllm/v1/executor/multiproc_executor.py
+211 -119
@@ -1,21 +1,23 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -
2 +import multiprocessing
3 import os
4 import pickle
5 import signal
6 import sys
7 +import threading
8 import time
9 import traceback
10 import weakref
11 +from concurrent.futures import Future
12 from dataclasses import dataclass
13 from enum import Enum, auto
14 from functools import partial
15 +from multiprocessing.connection import Connection
16 from multiprocessing.process import BaseProcess
14 -from typing import Any, Callable, Optional, Union
17 +from threading import Thread
18 +from typing import Any, Callable, Optional, Union, cast
19
20 import cloudpickle
17 -import psutil
18 -import zmq
21
22 from vllm.config import VllmConfig
23 from vllm.distributed import (destroy_distributed_environment,
@@ -26,8 +28,9 @@ from vllm.executor.multiproc_worker_utils import (
28 _add_prefix, set_multiprocessing_worker_envs)
29 from vllm.logger import init_logger
30 from vllm.utils import (get_distributed_init_method, get_mp_context,
29 - get_open_port, get_open_zmq_ipc_path, zmq_socket_ctx)
30 -from vllm.v1.executor.abstract import Executor
31 + get_open_port)
32 +from vllm.v1.executor.abstract import Executor, FailureCallback
33 +from vllm.v1.outputs import ModelRunnerOutput
34 from vllm.worker.worker_base import WorkerWrapperBase
35
36 logger = init_logger(__name__)
@@ -35,6 +38,8 @@ logger = init_logger(__name__)
38 POLLING_TIMEOUT_MS = 5000
39 POLLING_TIMEOUT_S = POLLING_TIMEOUT_MS // 1000
40
41 +EXECUTE_MODEL_TIMEOUT_S = 30
42 +
43
44 class MultiprocExecutor(Executor):
45
@@ -42,19 +47,9 @@ class MultiprocExecutor(Executor):
47 # Call self.shutdown at exit to clean up
48 # and ensure workers will be terminated.
49 self._finalizer = weakref.finalize(self, self.shutdown)
45 -
46 - # The child processes will send SIGUSR1 when unrecoverable
47 - # errors happen.
48 - def sigusr1_handler(signum, frame):
49 - logger.fatal(
50 - "MulitprocExecutor got fatal signal from worker processes, "
51 - "shutting down. See stack trace above for root cause issue.")
52 - # Propagate error up to parent process.
53 - parent_process = psutil.Process().parent()
54 - parent_process.send_signal(signal.SIGUSR1)
55 - self.shutdown()
56 -
57 - signal.signal(signal.SIGUSR1, sigusr1_handler)
50 + self.is_failed = False
51 + self.shutdown_event = threading.Event()
52 + self.failure_callback: Optional[FailureCallback] = None
53
54 self.world_size = self.parallel_config.world_size
55 tensor_parallel_size = self.parallel_config.tensor_parallel_size
@@ -78,28 +73,94 @@ class MultiprocExecutor(Executor):
73 scheduler_output_handle = self.rpc_broadcast_mq.export_handle()
74
75 # Create workers
81 - self.workers: list[WorkerProcHandle] = []
82 - for rank in range(self.world_size):
83 - worker = WorkerProc.make_worker_process(self.vllm_config, rank,
84 - rank,
85 - distributed_init_method,
86 - scheduler_output_handle)
87 - self.workers.append(worker)
88 -
89 - # Ensure message queues are ready. Will deadlock if re-ordered
90 - # Must be kept consistent with the WorkerProc
91 - self.rpc_broadcast_mq.wait_until_ready()
92 - for w in self.workers:
93 - w.worker_response_mq.wait_until_ready()
76 + unready_workers: list[UnreadyWorkerProcHandle] = []
77 + success = False
78 + try:
79 + for rank in range(self.world_size):
80 + unready_workers.append(
81 + WorkerProc.make_worker_process(
82 + vllm_config=self.vllm_config,
83 + local_rank=rank,
84 + rank=rank,
85 + distributed_init_method=distributed_init_method,
86 + input_shm_handle=scheduler_output_handle,
87 + ))
88 +
89 + # Workers must be created before wait_for_ready to avoid
90 + # deadlock, since worker.init_device() does a device sync.
91 + self.workers = WorkerProc.wait_for_ready(unready_workers)
92 +
93 + # Ensure message queues are ready. Will deadlock if re-ordered
94 + # Must be kept consistent with the WorkerProc.
95 + self.rpc_broadcast_mq.wait_until_ready()
96 + for w in self.workers:
97 + w.worker_response_mq.wait_until_ready()
98 +
99 + self.start_worker_monitor()
100 + success = True
101 + finally:
102 + if not success:
103 + # Clean up the worker procs if there was a failure.
104 + self._ensure_worker_termination(
105 + [w.proc for w in unready_workers])
106 +
107 + def start_worker_monitor(self):
108 + workers = self.workers
109 + self_ref = weakref.ref(self)
110 +
111 + # Monitors worker process liveness. If any die unexpectedly,
112 + # logs an error, shuts down the executor and invokes the failure
113 + # callback to inform the engine.
114 + def monitor_workers():
115 + sentinels = [h.proc.sentinel for h in workers]
116 + died = multiprocessing.connection.wait(sentinels)
117 + _self = self_ref()
118 + if not _self or getattr(_self, 'shutting_down', False):
119 + return
120 + _self.is_failed = True
121 + proc_name = next(h.proc.name for h in workers
122 + if h.proc.sentinel == died[0])
123 + logger.error(
124 + "Worker proc %s died unexpectedly, "
125 + "shutting down executor.", proc_name)
126 + _self.shutdown()
127 + callback = _self.failure_callback
128 + if callback is not None:
129 + _self.failure_callback = None
130 + callback()
131 +
132 + Thread(target=monitor_workers,
133 + daemon=True,
134 + name="MultiprocWorkerMonitor").start()
135 +
136 + def register_failure_callback(self, callback: FailureCallback):
137 + if self.is_failed:
138 + callback()
139 + else:
140 + self.failure_callback = callback
141 +
142 + def execute_model(
143 + self,
144 + scheduler_output,
145 + ) -> Union[ModelRunnerOutput, Future[ModelRunnerOutput]]:
146 + (output, ) = self.collective_rpc("execute_model",
147 + args=(scheduler_output, ),
148 + rank0_reply_only=True,
149 + timeout=EXECUTE_MODEL_TIMEOUT_S)
150 + return output
151
152 def collective_rpc(self,
153 method: Union[str, Callable],
97 - timeout: Optional[float] = None,
154 + timeout: Optional[float] = 180.0,
155 args: tuple = (),
99 - kwargs: Optional[dict] = None) -> list[Any]:
156 + kwargs: Optional[dict] = None,
157 + rank0_reply_only: bool = False) -> list[Any]:
158 start_time = time.monotonic()
159 kwargs = kwargs or {}
160
161 + if self.is_failed:
162 + raise RuntimeError("Executor failed.")
163 +
164 # NOTE: If the args are heterogeneous, then we pack them into a list,
165 # and unpack them in the method of every worker, because every worker
166 # knows their own rank.
@@ -109,30 +170,30 @@ class MultiprocExecutor(Executor):
170 else:
171 send_method = cloudpickle.dumps(
172 method, protocol=pickle.HIGHEST_PROTOCOL)
112 - self.rpc_broadcast_mq.enqueue((send_method, args, kwargs))
173 + self.rpc_broadcast_mq.enqueue(
174 + (send_method, args, kwargs, rank0_reply_only))
175
114 - responses = [None] * self.world_size
115 - for w in self.workers:
176 + workers = (self.workers[0], ) if rank0_reply_only else self.workers
177 + responses = [None] * len(workers)
178 + for w in workers:
179 dequeue_timeout = timeout - (time.monotonic() - start_time
180 ) if timeout is not None else None
181 status, result = w.worker_response_mq.dequeue(
119 - timeout=dequeue_timeout)
182 + timeout=dequeue_timeout, cancel=self.shutdown_event)
183
184 if status != WorkerProc.ResponseStatus.SUCCESS:
185 raise RuntimeError(
123 - "Worker failed with error %s, please check the"
124 - " stack trace above for the root cause", result)
186 + f"Worker failed with error '{result}', please check the"
187 + " stack trace above for the root cause")
188
189 responses[w.rank] = result
190
191 return responses
192 except TimeoutError as e:
193 raise TimeoutError(f"RPC call to {method} timed out.") from e
131 - except Exception as e:
132 - # Re-raise any other exceptions
133 - raise e
194
135 - def _ensure_worker_termination(self):
195 + @staticmethod
196 + def _ensure_worker_termination(worker_procs: list[BaseProcess]):
197 """Ensure that all worker processes are terminated. Assumes workers have
198 received termination requests. Waits for processing, then sends
199 termination and kill signals if needed."""
@@ -150,7 +211,7 @@ class MultiprocExecutor(Executor):
211 return False
212
213 # Send SIGTERM if still running
153 - active_procs = [w.proc for w in self.workers if w.proc.is_alive()]
214 + active_procs = [proc for proc in worker_procs if proc.is_alive()]
215 for p in active_procs:
216 p.terminate()
217 if not wait_for_termination(active_procs, 4):
@@ -159,22 +220,14 @@ class MultiprocExecutor(Executor):
220 for p in active_procs:
221 p.kill()
222
162 - self._cleanup_sockets()
163 -
164 - def _cleanup_sockets(self):
165 - for w in self.workers:
166 - # Remove the zmq ipc socket file
167 - socket_path = w.ready_path.replace("ipc://", "")
168 - if os and os.path.exists(socket_path):
169 - os.remove(socket_path)
170 -
223 def shutdown(self):
224 """Properly shut down the executor and its workers"""
225 if not getattr(self, 'shutting_down', False):
226 self.shutting_down = True
227 + self.shutdown_event.set()
228 for w in self.workers:
229 w.worker_response_mq = None
177 - self._ensure_worker_termination()
230 + self._ensure_worker_termination([w.proc for w in self.workers])
231
232 self.rpc_broadcast_mq = None
233
@@ -183,13 +236,30 @@ class MultiprocExecutor(Executor):
236 return
237
238
239 +@dataclass
240 +class UnreadyWorkerProcHandle:
241 + """WorkerProcess handle before READY."""
242 + proc: BaseProcess
243 + rank: int
244 + ready_pipe: Connection
245 +
246 +
247 @dataclass
248 class WorkerProcHandle:
249 proc: BaseProcess
250 rank: int
190 - ready_path: str
251 worker_response_mq: MessageQueue # The worker process writes to this MQ
252
253 + @classmethod
254 + def from_unready_handle(
255 + cls, unready_handle: UnreadyWorkerProcHandle,
256 + worker_response_mq: MessageQueue) -> "WorkerProcHandle":
257 + return cls(
258 + proc=unready_handle.proc,
259 + rank=unready_handle.rank,
260 + worker_response_mq=worker_response_mq,
261 + )
262 +
263
264 class WorkerProc:
265 """Wrapper that runs one Worker in a separate process."""
@@ -203,7 +273,6 @@ class WorkerProc:
273 rank: int,
274 distributed_init_method: str,
275 input_shm_handle: Handle,
206 - ready_path: str,
276 ):
277 self.rank = rank
278 wrapper = WorkerWrapperBase(vllm_config=vllm_config, rpc_rank=rank)
@@ -231,18 +300,8 @@ class WorkerProc:
300
301 # Initializes a message queue for sending the model output
302 self.worker_response_mq = MessageQueue(1, 1)
234 - worker_response_mq_handle = self.worker_response_mq.export_handle()
235 -
236 - # Send Readiness signal to EngineCore process.
237 - # Set linger here because we want to ensure the message has
238 - # been sent before the context is closed.
239 - with zmq_socket_ctx(ready_path, zmq.constants.PUSH,
240 - linger=10000) as ready_socket:
241 - payload = pickle.dumps(worker_response_mq_handle,
242 - protocol=pickle.HIGHEST_PROTOCOL)
243 - ready_socket.send_string(WorkerProc.READY_STR)
244 - ready_socket.send(payload)
303
304 + # Initialize device and loads weights
305 self.worker.init_device()
306 self.worker.load_model()
307
@@ -253,12 +312,10 @@ class WorkerProc:
312 rank: int,
313 distributed_init_method: str,
314 input_shm_handle, # Receive SchedulerOutput
256 - ) -> WorkerProcHandle:
315 + ) -> UnreadyWorkerProcHandle:
316 context = get_mp_context()
258 -
259 - # ZMQ path for worker to send ready message and shm_broadcast handle
260 - # back to core process.
261 - ready_path = get_open_zmq_ipc_path()
317 + # (reader, writer)
318 + reader, writer = context.Pipe(duplex=False)
319
320 process_kwargs = {
321 "vllm_config": vllm_config,
@@ -266,24 +323,57 @@ class WorkerProc:
323 "rank": rank,
324 "distributed_init_method": distributed_init_method,
325 "input_shm_handle": input_shm_handle,
269 - "ready_path": ready_path,
326 + "ready_pipe": (reader, writer),
327 }
328 # Run EngineCore busy loop in background process.
329 proc = context.Process(target=WorkerProc.worker_main,
330 kwargs=process_kwargs,
331 + name=f"VllmWorker-{rank}",
332 daemon=True)
333
276 - with zmq_socket_ctx(ready_path, zmq.constants.PULL) as ready_socket:
277 - proc.start()
278 -
279 - # Wait for startup
280 - worker_response_mq_handle = WorkerProc.wait_for_startup(
281 - proc, ready_socket)
282 -
283 - worker_response_mq = MessageQueue.create_from_handle(
284 - worker_response_mq_handle, 0)
334 + proc.start()
335 + writer.close()
336 + return UnreadyWorkerProcHandle(proc, rank, reader)
337
286 - return WorkerProcHandle(proc, rank, ready_path, worker_response_mq)
338 + @staticmethod
339 + def wait_for_ready(
340 + unready_proc_handles: list[UnreadyWorkerProcHandle]
341 + ) -> list[WorkerProcHandle]:
342 +
343 + e = Exception("WorkerProc initialization failed due to "
344 + "an exception in a background process. "
345 + "See stack trace for root cause.")
346 +
347 + pipes = {handle.ready_pipe: handle for handle in unready_proc_handles}
348 + ready_proc_handles: list[Optional[WorkerProcHandle]] = (
349 + [None] * len(unready_proc_handles))
350 + while pipes:
351 + ready = multiprocessing.connection.wait(pipes.keys())
352 + for pipe in ready:
353 + assert isinstance(pipe, Connection)
354 + try:
355 + # Wait until the WorkerProc is ready.
356 + unready_proc_handle = pipes.pop(pipe)
357 + response: dict[str, Any] = pipe.recv()
358 + if response["status"] != "READY":
359 + raise e
360 +
361 + # Extract the message queue handle.
362 + worker_response_mq = MessageQueue.create_from_handle(
363 + response["handle"], 0)
364 + ready_proc_handles[unready_proc_handle.rank] = (
365 + WorkerProcHandle.from_unready_handle(
366 + unready_proc_handle, worker_response_mq))
367 +
368 + except EOFError:
369 + e.__suppress_context__ = True
370 + raise e from None
371 +
372 + finally:
373 + # Close connection.
374 + pipe.close()
375 +
376 + return cast(list[WorkerProcHandle], ready_proc_handles)
377
378 def shutdown(self):
379 self.rpc_broadcast_mq = None
@@ -312,51 +402,51 @@ class WorkerProc:
402 signal.signal(signal.SIGINT, signal_handler)
403
404 worker = None
405 + # tuple[Connection, Connection]
406 + reader, ready_writer = kwargs.pop("ready_pipe")
407 try:
408 + reader.close()
409 worker = WorkerProc(*args, **kwargs)
410
411 + # Send READY once we know everything is loaded
412 + ready_writer.send({
413 + "status":
414 + WorkerProc.READY_STR,
415 + "handle":
416 + worker.worker_response_mq.export_handle(),
417 + })
418 +
419 # Ensure message queues are ready. Will deadlock if re-ordered.
420 # Must be kept consistent with the Executor
421 worker.rpc_broadcast_mq.wait_until_ready()
422 worker.worker_response_mq.wait_until_ready()
423 + ready_writer.close()
424 + ready_writer = None
425
426 worker.worker_busy_loop()
427
325 - except SystemExit:
326 - logger.debug("Worker interrupted.")
327 -
428 except Exception:
329 - # worker_busy_loop sends exceptions to Executor
330 - # for shutdown, but if there is an error in startup or an
331 - # error with IPC itself, we need to alert the parent.
332 - psutil.Process().parent().send_signal(signal.SIGUSR1)
333 - raise
429 + # NOTE: if an Exception arises in busy_loop, we send
430 + # a FAILURE message over the MQ RPC to notify the Executor,
431 + # which triggers system shutdown.
432 + # TODO(rob): handle case where the MQ itself breaks.
433 +
434 + if ready_writer is not None:
435 + logger.exception("WorkerProc failed to start.")
436 + else:
437 + logger.exception("WorkerProc failed.")
438 +
439 + # The parent sends a SIGTERM to all worker processes if
440 + # any worker dies. Set this value so we don't re-throw
441 + # SystemExit() to avoid zmq exceptions in __del__.
442 + shutdown_requested = True
443
444 finally:
445 + if ready_writer is not None:
446 + ready_writer.close()
447 # Clean up once worker exits busy loop
448 if worker is not None:
449 worker.shutdown()
339 - worker = None
340 -
341 - @staticmethod
342 - def wait_for_startup(
343 - proc: BaseProcess,
344 - ready_socket: zmq.Socket,
345 - ) -> Optional[Handle]:
346 - """Wait until the Worker is ready."""
347 -
348 - # Wait for Worker to send READY.
349 - while ready_socket.poll(timeout=POLLING_TIMEOUT_MS) == 0:
350 - logger.debug("Waiting for WorkerProc to startup.")
351 -
352 - if not proc.is_alive():
353 - raise RuntimeError("WorkerProc failed to start.")
354 -
355 - message = ready_socket.recv_string()
356 - assert message == WorkerProc.READY_STR
357 - handle_frame = ready_socket.recv(copy=False)
358 - handle = pickle.loads(handle_frame.buffer)
359 - return handle
450
451 class ResponseStatus(Enum):
452 SUCCESS = auto()
@@ -365,7 +455,7 @@ class WorkerProc:
455 def worker_busy_loop(self):
456 """Main busy loop for Multiprocessing Workers"""
457 while True:
368 - method, args, kwargs = self.rpc_broadcast_mq.dequeue()
458 + method, args, kwargs, rank0_only = self.rpc_broadcast_mq.dequeue()
459
460 try:
461 if isinstance(method, str):
@@ -377,12 +467,14 @@ class WorkerProc:
467 # Notes have been introduced in python 3.11
468 if hasattr(e, "add_note"):
469 e.add_note(traceback.format_exc())
380 - logger.exception("WorkerProc hit an exception: %s", exc_info=e)
470 + logger.exception("WorkerProc hit an exception.")
471 # exception might not be serializable, so we convert it to
472 # string, only for logging purpose.
383 - self.worker_response_mq.enqueue(
384 - (WorkerProc.ResponseStatus.FAILURE, str(e)))
473 + if not rank0_only or self.rank == 0:
474 + self.worker_response_mq.enqueue(
475 + (WorkerProc.ResponseStatus.FAILURE, str(e)))
476 continue
477
387 - self.worker_response_mq.enqueue(
388 - (WorkerProc.ResponseStatus.SUCCESS, output))
478 + if not rank0_only or self.rank == 0:
479 + self.worker_response_mq.enqueue(
480 + (WorkerProc.ResponseStatus.SUCCESS, output))