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 14, 05:30 UTC
[BugFix] Threadsafe close async zmq sockets (#22877)
Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Nick Hill committed
Aug 14, 2025 at 03:44 UTC
eb08487b18f4be3cc8dbe6776c2d69d223b3737c
2 files changed
+77
-26
vllm/utils/__init__.py
+22
-2
@@ -709,8 +709,28 @@ class AsyncMicrobatchTokenizer:
709
710
711
def cancel_task_threadsafe(task: Task):
712
- if task and not task.done() and not (loop := task.get_loop()).is_closed():
713
- loop.call_soon_threadsafe(task.cancel)
712
+ if task and not task.done():
713
+ run_in_loop(task.get_loop(), task.cancel)
714
+
715
+
716
+def close_sockets(sockets: Sequence[Union[zmq.Socket, zmq.asyncio.Socket]]):
717
+ for sock in sockets:
718
+ if sock is not None:
719
+ sock.close(linger=0)
720
+
721
+
722
+def run_in_loop(loop: AbstractEventLoop, function: Callable, *args):
723
+ if in_loop(loop):
724
+ function(*args)
725
+ elif not loop.is_closed():
726
+ loop.call_soon_threadsafe(function, *args)
727
+
728
+
729
+def in_loop(event_loop: AbstractEventLoop) -> bool:
730
+ try:
731
+ return asyncio.get_running_loop() == event_loop
732
+ except RuntimeError:
733
+ return False
734
735
736
def make_async(
vllm/v1/engine/core_client.py
+55
-24
@@ -23,8 +23,8 @@ from vllm.config import VllmConfig
23
from vllm.logger import init_logger
24
from vllm.lora.request import LoRARequest
25
from vllm.tasks import SupportedTask
26
-from vllm.utils import (cancel_task_threadsafe, get_open_port,
27
- get_open_zmq_inproc_path, make_zmq_socket)
26
+from vllm.utils import (close_sockets, get_open_port, get_open_zmq_inproc_path,
27
+ in_loop, make_zmq_socket)
28
from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
29
EngineCoreRequestType,
30
ReconfigureDistributedRequest, ReconfigureRankType,
@@ -317,7 +317,7 @@ class BackgroundResources:
317
"""Used as a finalizer for clean shutdown, avoiding
318
circular reference back to the client object."""
319
320
- ctx: Union[zmq.Context]
320
+ ctx: zmq.Context
321
# If CoreEngineProcManager, it manages local engines;
322
# if CoreEngineActorManager, it manages all engines.
323
engine_manager: Optional[Union[CoreEngineProcManager,
@@ -326,6 +326,8 @@ class BackgroundResources:
326
output_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None
327
input_socket: Optional[Union[zmq.Socket, zmq.asyncio.Socket]] = None
328
first_req_send_socket: Optional[zmq.asyncio.Socket] = None
329
+ first_req_rcv_socket: Optional[zmq.asyncio.Socket] = None
330
+ stats_update_socket: Optional[zmq.asyncio.Socket] = None
331
output_queue_task: Optional[asyncio.Task] = None
332
stats_update_task: Optional[asyncio.Task] = None
333
shutdown_path: Optional[str] = None
@@ -343,23 +345,47 @@ class BackgroundResources:
345
if self.coordinator is not None:
346
self.coordinator.close()
347
346
- cancel_task_threadsafe(self.output_queue_task)
347
- cancel_task_threadsafe(self.stats_update_task)
348
+ if isinstance(self.output_socket, zmq.asyncio.Socket):
349
+ # Async case.
350
+ loop = self.output_socket._get_loop()
351
+ asyncio.get_running_loop()
352
+ sockets = (self.output_socket, self.input_socket,
353
+ self.first_req_send_socket, self.first_req_rcv_socket,
354
+ self.stats_update_socket)
355
+
356
+ tasks = (self.output_queue_task, self.stats_update_task)
357
+
358
+ def close_sockets_and_tasks():
359
+ close_sockets(sockets)
360
+ for task in tasks:
361
+ if task is not None and not task.done():
362
+ task.cancel()
363
+
364
+ if in_loop(loop):
365
+ close_sockets_and_tasks()
366
+ elif not loop.is_closed():
367
+ loop.call_soon_threadsafe(close_sockets_and_tasks)
368
+ else:
369
+ # Loop has been closed, try to clean up directly.
370
+ del tasks
371
+ del close_sockets_and_tasks
372
+ close_sockets(sockets)
373
+ del self.output_queue_task
374
+ del self.stats_update_task
375
+ else:
376
+ # Sync case.
377
349
- # ZMQ context termination can hang if the sockets
350
- # aren't explicitly closed first.
351
- for socket in (self.output_socket, self.input_socket,
352
- self.first_req_send_socket):
353
- if socket is not None:
354
- socket.close(linger=0)
378
+ # ZMQ context termination can hang if the sockets
379
+ # aren't explicitly closed first.
380
+ close_sockets((self.output_socket, self.input_socket))
381
356
- if self.shutdown_path is not None:
357
- # We must ensure that the sync output socket is
358
- # closed cleanly in its own thread.
359
- with self.ctx.socket(zmq.PAIR) as shutdown_sender:
360
- shutdown_sender.connect(self.shutdown_path)
361
- # Send shutdown signal.
362
- shutdown_sender.send(b'')
382
+ if self.shutdown_path is not None:
383
+ # We must ensure that the sync output socket is
384
+ # closed cleanly in its own thread.
385
+ with self.ctx.socket(zmq.PAIR) as shutdown_sender:
386
+ shutdown_sender.connect(self.shutdown_path)
387
+ # Send shutdown signal.
388
+ shutdown_sender.send(b'')
389
390
def validate_alive(self, frames: Sequence[zmq.Frame]):
391
if len(frames) == 1 and (frames[0].buffer
@@ -969,14 +995,19 @@ class DPAsyncMPClient(AsyncMPClient):
995
self.engine_ranks_managed[-1] + 1)
996
997
async def run_engine_stats_update_task():
972
- with make_zmq_socket(self.ctx, self.stats_update_address,
973
- zmq.XSUB) as socket, make_zmq_socket(
974
- self.ctx,
975
- self.first_req_sock_addr,
976
- zmq.PAIR,
977
- bind=False) as first_req_rcv_socket:
998
+ with (make_zmq_socket(self.ctx,
999
+ self.stats_update_address,
1000
+ zmq.XSUB,
1001
+ linger=0) as socket,
1002
+ make_zmq_socket(self.ctx,
1003
+ self.first_req_sock_addr,
1004
+ zmq.PAIR,
1005
+ bind=False,
1006
+ linger=0) as first_req_rcv_socket):
1007
assert isinstance(socket, zmq.asyncio.Socket)
1008
assert isinstance(first_req_rcv_socket, zmq.asyncio.Socket)
1009
+ self.resources.stats_update_socket = socket
1010
+ self.resources.first_req_rcv_socket = first_req_rcv_socket
1011
# Send subscription message.
1012
await socket.send(b'\x01')
1013