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, 21:00 UTC
[V1][BugFix] Fix remaining sync engine client shutdown errors/hangs (#13869)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Mar 4, 2025 at 07:06 UTC
5db6b2c9614d70246d3386b6fb0cc655c72aeb85
3 files changed
+68
-40
tests/v1/engine/test_llm_engine.py
-2
@@ -15,8 +15,6 @@ DTYPE = "half"
15
def _vllm_model(apc: bool, vllm_runner, monkeypatch):
16
"""Set up VllmRunner instance."""
17
monkeypatch.setenv("VLLM_USE_V1", "1")
18
- # TODO(nick): Single-proc to work around a ZMQ shutdown hang for now.
19
- monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
18
return vllm_runner(
19
MODEL,
20
dtype=DTYPE,
vllm/utils.py
+12
-10
@@ -500,6 +500,10 @@ def get_open_zmq_ipc_path() -> str:
500
return f"ipc://{base_rpc_path}/{uuid4()}"
501
502
503
+def get_open_zmq_inproc_path() -> str:
504
+ return f"inproc://{uuid4()}"
505
+
506
+
507
def get_open_port() -> int:
508
"""
509
Get an open port for the vLLM process to listen on.
@@ -2108,12 +2112,12 @@ def get_exception_traceback():
2112
def make_zmq_socket(
2113
ctx: Union[zmq.asyncio.Context, zmq.Context], # type: ignore[name-defined]
2114
path: str,
2111
- type: Any,
2115
+ socket_type: Any,
2116
) -> Union[zmq.Socket, zmq.asyncio.Socket]: # type: ignore[name-defined]
2117
"""Make a ZMQ socket with the proper bind/connect semantics."""
2118
2119
mem = psutil.virtual_memory()
2116
- socket = ctx.socket(type)
2120
+ socket = ctx.socket(socket_type)
2121
2122
# Calculate buffer size based on system memory
2123
total_mem = mem.total / 1024**3
@@ -2127,29 +2131,27 @@ def make_zmq_socket(
2131
else:
2132
buf_size = -1 # Use system default buffer size
2133
2130
- if type == zmq.constants.PULL:
2134
+ if socket_type == zmq.constants.PULL:
2135
socket.setsockopt(zmq.constants.RCVHWM, 0)
2136
socket.setsockopt(zmq.constants.RCVBUF, buf_size)
2137
socket.connect(path)
2134
- elif type == zmq.constants.PUSH:
2138
+ elif socket_type == zmq.constants.PUSH:
2139
socket.setsockopt(zmq.constants.SNDHWM, 0)
2140
socket.setsockopt(zmq.constants.SNDBUF, buf_size)
2141
socket.bind(path)
2142
else:
2139
- raise ValueError(f"Unknown Socket Type: {type}")
2143
+ raise ValueError(f"Unknown Socket Type: {socket_type}")
2144
2145
return socket
2146
2147
2148
@contextlib.contextmanager
2145
-def zmq_socket_ctx(
2146
- path: str,
2147
- type: Any) -> Iterator[zmq.Socket]: # type: ignore[name-defined]
2149
+def zmq_socket_ctx(path: str, socket_type: Any) -> Iterator[zmq.Socket]:
2150
"""Context manager for a ZMQ socket"""
2151
2150
- ctx = zmq.Context(io_threads=2) # type: ignore[attr-defined]
2152
+ ctx = zmq.Context() # type: ignore[attr-defined]
2153
try:
2152
- yield make_zmq_socket(ctx, path, type)
2154
+ yield make_zmq_socket(ctx, path, socket_type)
2155
2156
except KeyboardInterrupt:
2157
logger.debug("Got Keyboard Interrupt.")
vllm/v1/engine/core_client.py
+56
-28
@@ -18,8 +18,8 @@ import zmq.asyncio
18
from vllm.config import VllmConfig
19
from vllm.logger import init_logger
20
from vllm.lora.request import LoRARequest
21
-from vllm.utils import (get_open_zmq_ipc_path, kill_process_tree,
22
- make_zmq_socket)
21
+from vllm.utils import (get_open_zmq_inproc_path, get_open_zmq_ipc_path,
22
+ kill_process_tree, make_zmq_socket)
23
from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
24
EngineCoreRequestType, UtilityOutput)
25
from vllm.v1.engine.core import EngineCore, EngineCoreProc
@@ -202,10 +202,11 @@ class BackgroundResources:
202
"""Used as a finalizer for clean shutdown, avoiding
203
circular reference back to the client object."""
204
205
- ctx: Union[zmq.Context, zmq.asyncio.Context] = None
205
+ ctx: Union[zmq.Context] = None
206
output_socket: Union[zmq.Socket, zmq.asyncio.Socket] = None
207
input_socket: Union[zmq.Socket, zmq.asyncio.Socket] = None
208
proc_handle: Optional[BackgroundProcHandle] = None
209
+ shutdown_path: Optional[str] = None
210
211
def __call__(self):
212
"""Clean up background resources."""
@@ -218,8 +219,13 @@ class BackgroundResources:
219
self.output_socket.close(linger=0)
220
if self.input_socket is not None:
221
self.input_socket.close(linger=0)
221
- if self.ctx is not None:
222
- self.ctx.destroy(linger=0)
222
+ if self.shutdown_path is not None:
223
+ # We must ensure that the sync output socket is
224
+ # closed cleanly in its own thread.
225
+ with self.ctx.socket(zmq.PAIR) as shutdown_sender:
226
+ shutdown_sender.connect(self.shutdown_path)
227
+ # Send shutdown signal.
228
+ shutdown_sender.send(b'')
229
230
231
class MPClient(EngineCoreClient):
@@ -261,28 +267,23 @@ class MPClient(EngineCoreClient):
267
self.decoder = MsgpackDecoder(EngineCoreOutputs)
268
269
# ZMQ setup.
264
- self.ctx = (
265
- zmq.asyncio.Context() # type: ignore[attr-defined]
266
- if asyncio_mode else zmq.Context()) # type: ignore[attr-defined]
270
+ sync_ctx = zmq.Context()
271
+ self.ctx = zmq.asyncio.Context(sync_ctx) if asyncio_mode else sync_ctx
272
273
# This will ensure resources created so far are closed
274
# when the client is garbage collected, even if an
275
# exception is raised mid-construction.
271
- resources = BackgroundResources(ctx=self.ctx)
272
- self._finalizer = weakref.finalize(self, resources)
276
+ self.resources = BackgroundResources(ctx=sync_ctx)
277
+ self._finalizer = weakref.finalize(self, self.resources)
278
274
- # Paths and sockets for IPC.
275
- output_path = get_open_zmq_ipc_path()
279
+ # Paths for IPC.
280
+ self.output_path = get_open_zmq_ipc_path()
281
input_path = get_open_zmq_ipc_path()
277
- resources.output_socket = make_zmq_socket(self.ctx, output_path,
278
- zmq.constants.PULL)
279
- resources.input_socket = make_zmq_socket(self.ctx, input_path,
280
- zmq.constants.PUSH)
282
283
# Start EngineCore in background process.
283
- resources.proc_handle = BackgroundProcHandle(
284
+ self.resources.proc_handle = BackgroundProcHandle(
285
input_path=input_path,
285
- output_path=output_path,
286
+ output_path=self.output_path,
287
process_name="EngineCore",
288
target_fn=EngineCoreProc.run_engine_core,
289
process_kwargs={
@@ -291,8 +292,10 @@ class MPClient(EngineCoreClient):
292
"log_stats": log_stats,
293
})
294
294
- self.output_socket = resources.output_socket
295
- self.input_socket = resources.input_socket
295
+ # Create input socket.
296
+ self.resources.input_socket = make_zmq_socket(self.ctx, input_path,
297
+ zmq.constants.PUSH)
298
+ self.input_socket = self.resources.input_socket
299
self.utility_results: dict[int, AnyFuture] = {}
300
301
def shutdown(self):
@@ -325,27 +328,48 @@ class SyncMPClient(MPClient):
328
329
# Ensure that the outputs socket processing thread does not have
330
# a ref to the client which prevents gc.
328
- output_socket = self.output_socket
331
+ ctx = self.ctx
332
+ output_path = self.output_path
333
decoder = self.decoder
334
utility_results = self.utility_results
335
outputs_queue = self.outputs_queue
336
337
+ shutdown_path = get_open_zmq_inproc_path()
338
+ self.resources.shutdown_path = shutdown_path
339
+
340
def process_outputs_socket():
341
+ shutdown_socket = ctx.socket(zmq.PAIR)
342
+ shutdown_socket.bind(shutdown_path)
343
+ out_socket = make_zmq_socket(ctx, output_path, zmq.constants.PULL)
344
try:
345
+ poller = zmq.Poller()
346
+ poller.register(shutdown_socket)
347
+ poller.register(out_socket)
348
while True:
336
- (frame, ) = output_socket.recv_multipart(copy=False)
349
+ socks = poller.poll()
350
+ if not socks:
351
+ continue
352
+ if len(socks) == 2 or socks[0][0] == shutdown_socket:
353
+ # shutdown signal, exit thread.
354
+ break
355
+
356
+ (frame, ) = out_socket.recv_multipart(copy=False)
357
outputs = decoder.decode(frame.buffer)
358
if outputs.utility_output:
359
_process_utility_output(outputs.utility_output,
360
utility_results)
361
else:
362
outputs_queue.put_nowait(outputs)
343
- except zmq.error.ContextTerminated:
344
- # Expected when the class is GC'd / during process termination.
345
- pass
363
+ finally:
364
+ # Close sockets.
365
+ shutdown_socket.close(linger=0)
366
+ out_socket.close(linger=0)
367
368
# Process outputs from engine in separate thread.
348
- Thread(target=process_outputs_socket, daemon=True).start()
369
+ self.output_queue_thread = Thread(target=process_outputs_socket,
370
+ name="EngineCoreOutputQueueThread",
371
+ daemon=True)
372
+ self.output_queue_thread.start()
373
374
def get_output(self) -> EngineCoreOutputs:
375
return self.outputs_queue.get()
@@ -424,10 +448,13 @@ class AsyncMPClient(MPClient):
448
# Perform IO in separate task to parallelize as much as possible.
449
# Avoid task having direct reference back to the client.
450
self.outputs_queue = asyncio.Queue()
427
- output_socket = self.output_socket
451
decoder = self.decoder
452
utility_results = self.utility_results
453
outputs_queue = self.outputs_queue
454
+ output_path = self.output_path
455
+ output_socket = make_zmq_socket(self.ctx, output_path,
456
+ zmq.constants.PULL)
457
+ self.resources.output_socket = output_socket
458
459
async def process_outputs_socket():
460
while True:
@@ -439,7 +466,8 @@ class AsyncMPClient(MPClient):
466
else:
467
outputs_queue.put_nowait(outputs)
468
442
- self.queue_task = asyncio.create_task(process_outputs_socket())
469
+ self.queue_task = asyncio.create_task(process_outputs_socket(),
470
+ name="EngineCoreOutputQueueTask")
471
472
async def get_output_async(self) -> EngineCoreOutputs:
473
if self.outputs_queue is None: