@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b07bf83c7d

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, 13:30 UTC

[BugFix] Avoid race conditions in zero-copy tensor transmission (#17203)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Apr 25, 2025 at 23:00 UTC b07bf83c7d610ee88176b70eb58f66627b4576a2

3 files changed +77 -12

tests/v1/test_serial_utils.py

@@ -32,6 +32,7 @@ class MyType:

32 large_f_contig_tensor: torch.Tensor

33 small_non_contig_tensor: torch.Tensor

34 large_non_contig_tensor: torch.Tensor

35 + empty_tensor: torch.Tensor

38 def test_encode_decode():

@@ -58,6 +59,7 @@ def test_encode_decode():

59 large_f_contig_tensor=torch.rand(1024, 4).t(),

60 small_non_contig_tensor=torch.rand(2, 4)[:, 1:3],

61 large_non_contig_tensor=torch.rand(1024, 512)[:, 10:20],

62 + empty_tensor=torch.empty(0),

63 )

65 encoder = MsgpackEncoder(size_threshold=256)

@@ -193,3 +195,4 @@ def assert_equal(obj1: MyType, obj2: MyType):

195 obj2.small_non_contig_tensor)

196 assert torch.equal(obj1.large_non_contig_tensor,

197 obj2.large_non_contig_tensor)

198 + assert torch.equal(obj1.empty_tensor, obj2.empty_tensor)

vllm/v1/engine/core.py

+22 -3

@@ -5,6 +5,7 @@ import signal

5 import sys

6 import threading

7 import time

8 +from collections import deque

9 from concurrent.futures import Future

10 from inspect import isclass, signature

11 from logging import DEBUG

@@ -527,8 +528,12 @@ class EngineCoreProc(EngineCore):

528

529 # Msgpack serialization encoding.

530 encoder = MsgpackEncoder()

530 - # Reuse send buffer.

531 - buffer = bytearray()

531 + # Send buffers to reuse.

532 + reuse_buffers: list[bytearray] = []

533 + # Keep references to outputs and buffers until zmq is finished

534 + # with them (outputs may contain tensors/np arrays whose

535 + # backing buffers were extracted for zero-copy send).

536 + pending = deque[tuple[zmq.MessageTracker, Any, bytearray]]()

537

538 # We must set linger to ensure the ENGINE_CORE_DEAD

539 # message is sent prior to closing the socket.

@@ -541,8 +546,22 @@ class EngineCoreProc(EngineCore):

546 break

547 assert not isinstance(outputs, bytes)

548 outputs.engine_index = engine_index

549 +

550 + # Reclaim buffers that zmq is finished with.

551 + while pending and pending[-1][0].done:

552 + reuse_buffers.append(pending.pop()[2])

553 +

554 + buffer = reuse_buffers.pop() if reuse_buffers else bytearray()

555 buffers = encoder.encode_into(outputs, buffer)

545 - socket.send_multipart(buffers, copy=False)

556 + tracker = socket.send_multipart(buffers,

557 + copy=False,

558 + track=True)

559 + if not tracker.done:

560 + ref = outputs if len(buffers) > 1 else None

561 + pending.appendleft((tracker, ref, buffer))

562 + elif len(reuse_buffers) < 2:

563 + # Keep at most 2 buffers to reuse.

564 + reuse_buffers.append(buffer)

565

566

567 class DPEngineCoreProc(EngineCoreProc):

vllm/v1/engine/core_client.py

+52 -9

@@ -1,9 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

2 import asyncio

3 +import contextlib

4 import queue

5 import uuid

6 import weakref

7 from abc import ABC, abstractmethod

8 +from collections import deque

9 from collections.abc import Awaitable, Sequence

10 from concurrent.futures import Future

11 from dataclasses import dataclass, field

@@ -396,6 +398,12 @@ class MPClient(EngineCoreClient):

398 self._wait_for_engine_startup()

399

400 self.utility_results: dict[int, AnyFuture] = {}

401 +

402 + # Request objects which may contain pytorch-allocated tensors

403 + # that we need to keep references to until zmq is done with the

404 + # underlying data.

405 + self.pending_messages = deque[tuple[zmq.MessageTracker, Any]]()

406 +

407 success = True

408 finally:

409 if not success:

@@ -459,6 +467,14 @@ class MPClient(EngineCoreClient):

467 if self.resources.engine_dead:

468 raise EngineDeadError()

469

470 + def add_pending_message(self, tracker: zmq.MessageTracker, msg: Any):

471 + if not tracker.done:

472 + self.pending_messages.appendleft((tracker, msg))

473 +

474 + def free_pending_messages(self):

475 + while self.pending_messages and self.pending_messages[-1][0].done:

476 + self.pending_messages.pop()

477 +

478

479 def _process_utility_output(output: UtilityOutput,

480 utility_results: dict[int, AnyFuture]):

@@ -544,10 +560,18 @@ class SyncMPClient(MPClient):

560

561 def _send_input(self, request_type: EngineCoreRequestType, request: Any):

562 self.ensure_alive()

563 + self.free_pending_messages()

564 # (Identity, RequestType, SerializedRequest)

565 msg = (self.core_engine.identity, request_type.value,

566 *self.encoder.encode(request))

550 - self.input_socket.send_multipart(msg, copy=False)

567 +

568 + if len(msg) <= 3:

569 + # No auxiliary buffers => no tensor backing buffers in request.

570 + self.input_socket.send_multipart(msg, copy=False)

571 + return

572 +

573 + tracker = self.input_socket.send_multipart(msg, copy=False, track=True)

574 + self.add_pending_message(tracker, request)

575

576 def call_utility(self, method: str, *args) -> Any:

577 call_id = uuid.uuid1().int >> 64

@@ -698,19 +722,38 @@ class AsyncMPClient(MPClient):

722 def _send_input(self,

723 request_type: EngineCoreRequestType,

724 request: Any,

701 - engine: Optional[CoreEngine] = None) -> Awaitable[None]:

725 + engine: Optional[CoreEngine] = None) -> Awaitable[Any]:

726 self.ensure_alive()

727 if engine is None:

728 engine = self.core_engine

729

730 message = (request_type.value, *self.encoder.encode(request))

707 - return self._send_input_message(message, engine)

708 -

709 - def _send_input_message(self, message: tuple[bytestr, ...],

710 - engine: CoreEngine) -> Awaitable[None]:

731 + return self._send_input_message(message, engine, request)

732 +

733 + def _send_input_message(self, message: tuple[bytestr,

734 + ...], engine: CoreEngine,

735 + objects: Any) -> Awaitable[Any]:

736 + """

737 + objects is a reference to retain until zmq is finished with the

738 + buffers, in case they were extracted from tensors in the request.

739 + """

740 self.ensure_alive()

712 - message = (engine.identity, ) + message

713 - return self.input_socket.send_multipart(message, copy=False)

741 + self.free_pending_messages()

742 +

743 + msg = (engine.identity, ) + message

744 + if not objects or len(msg) <= 3:

745 + # No auxiliary buffers => no tensor backing buffers in request.

746 + return self.input_socket.send_multipart(msg, copy=False)

747 +

748 + future: asyncio.Future[zmq.MessageTracker]

749 + future = self.input_socket.send_multipart(msg, copy=False, track=True)

750 +

751 + def add_pending(f: asyncio.Future[zmq.MessageTracker]):

752 + with contextlib.suppress(BaseException):

753 + self.add_pending_message(f.result(), objects)

754 +

755 + future.add_done_callback(add_pending)

756 + return future

757

758 async def call_utility_async(self, method: str, *args) -> Any:

759 return await self._call_utility_async(method,

@@ -724,7 +767,7 @@ class AsyncMPClient(MPClient):

767 self.utility_results[call_id] = future

768 message = (EngineCoreRequestType.UTILITY.value, *self.encoder.encode(

769 (call_id, method, args)))

727 - await self._send_input_message(message, engine)

770 + await self._send_input_message(message, engine, args)

771 self._ensure_output_queue_task()

772 return await future

773