diff --git a/tests/v1/test_serial_utils.py b/tests/v1/test_serial_utils.py index 4ed8724e60f..9f7761c22b5 100644 --- a/tests/v1/test_serial_utils.py +++ b/tests/v1/test_serial_utils.py @@ -423,3 +423,139 @@ def test_multiple_senders_single_receiver_ipc(): assert torch.allclose(decoded.prompt_embeds, original_tensor), ( f"Value mismatch for sender {sender_idx} msg {msg_idx}" ) + + +def _logprobs_outputs(num_reqs: int, num_prompt_tokens: int): + """An EngineCoreOutputs carrying prompt logprobs, as the engine core sends + it: many requests, each with per-token tensors small enough that pyzmq + copies their frames, while the accumulated payload frame is large enough + that pyzmq sends it zero-copy.""" + from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs + from vllm.v1.outputs import LogprobsTensors + + outputs = [] + for req in range(num_reqs): + num_tokens = num_prompt_tokens + req % 4 + outputs.append( + EngineCoreOutput( + request_id=f"req-{req:08d}", + new_token_ids=[req], + new_prompt_logprobs_tensors=LogprobsTensors( + logprob_token_ids=torch.arange( + num_tokens * 2, dtype=torch.int64 + ).view(num_tokens, 2), + logprobs=torch.zeros(num_tokens, 2, dtype=torch.float32), + selected_token_ranks=torch.zeros(num_tokens, dtype=torch.int32), + ), + ) + ) + return EngineCoreOutputs(outputs=outputs) + + +def test_payload_buffer_reuse_does_not_corrupt_in_flight_messages(): + """The engine core recycles the msgpack payload buffer across messages + (`MsgpackEncoder.encode_into`). It may only do so once zmq has finished + sending that buffer, otherwise a newer payload is delivered alongside the + older message's zero-copy tensor frames. + + `Socket.send_multipart(track=True)` cannot be used to detect this: it + returns a tracker for the last frame only, and pyzmq copies frames below + `zmq.COPY_THRESHOLD` and reports them as already-sent. + """ + import zmq + + from vllm.v1.engine import EngineCoreOutputs + from vllm.v1.engine.core import EngineCoreProc + + num_msgs = 100 + encoder = MsgpackEncoder() + decoder = MsgpackDecoder(EngineCoreOutputs) + # Enough requests that the payload frame is zero-copied rather than copied + # by pyzmq, which is what makes early reuse observable. + messages = [_logprobs_outputs(300, 24 + i % 8) for i in range(num_msgs)] + assert len(encoder.encode(messages[0])[0]) >= zmq.COPY_THRESHOLD + + reuse_buffers: list[bytearray] = [] + pending: list[tuple[zmq.MessageTracker, bytearray]] = [] + with zmq.Context() as ctx: + push = ctx.socket(zmq.PUSH) + push.bind("inproc://test-payload-reuse") + pull = ctx.socket(zmq.PULL) + pull.connect("inproc://test-payload-reuse") + + for outputs in messages: + while pending and pending[0][0].done: + reuse_buffers.append(pending.pop(0)[1]) + buffer = reuse_buffers.pop() if reuse_buffers else bytearray() + buffers = encoder.encode_into(outputs, buffer) + tracker = EngineCoreProc._send_msg_tracking_payload(push, buffers) + if tracker.done: + reuse_buffers.append(buffer) + else: + pending.append((tracker, buffer)) + + for i, sent in enumerate(messages): + received = decoder.decode(pull.recv_multipart(copy=False)) + assert len(received.outputs) == len(sent.outputs), f"message {i}" + for expected, actual in zip(sent.outputs, received.outputs): + sent_ids = expected.new_prompt_logprobs_tensors.logprob_token_ids + got_ids = actual.new_prompt_logprobs_tensors.logprob_token_ids + assert actual.request_id == expected.request_id, f"message {i}" + assert torch.equal(got_ids, sent_ids), ( + f"message {i} request {actual.request_id}: corrupted " + f"prompt logprobs, {got_ids.shape} vs {sent_ids.shape}" + ) + push.close(linger=0) + pull.close(linger=0) + + +def test_zero_copy_frames_survive_without_caller_side_references(): + """Callers don't need to retain the encoded object until zmq has sent it: + for a zero-copy frame, zmq holds its own reference to the backing buffer. + + The engine core clients rely on this when sending requests that carry + tensors (e.g. prompt embeds) without tracking the messages. + + What makes that safe is that `tensor_data()` hands zmq a memoryview which + transitively references the source tensor, so refcounting - not timing - + keeps the memory from being freed and reused underneath zmq. + """ + import gc + + import zmq + + from vllm.v1.utils import tensor_data + + num_elems = 100_000 # comfortably over zmq.COPY_THRESHOLD + expected = torch.arange(num_elems, dtype=torch.int64) + encoder = MsgpackEncoder() + decoder = MsgpackDecoder(RequestWithTensor) + + # The buffer handed to zmq must keep the tensor's storage alive by itself. + holder = tensor_data(expected).obj + while getattr(holder, "base", None) is not None: + holder = holder.base + assert isinstance(holder, torch.Tensor) + assert holder.data_ptr() == expected.data_ptr() + + with zmq.Context() as ctx: + push = ctx.socket(zmq.PUSH) + push.bind("inproc://test-zero-copy-lifetime") + pull = ctx.socket(zmq.PULL) + pull.connect("inproc://test-zero-copy-lifetime") + + request = RequestWithTensor(prompt_embeds=expected.clone(), data="req") + buffers = encoder.encode(request) + assert max(len(buf) for buf in buffers) >= zmq.COPY_THRESHOLD + push.send_multipart(buffers, copy=False) + + # Drop every reference the sender holds, then churn the allocator. + del request, buffers + gc.collect() + torch.arange(num_elems * 4, dtype=torch.int64) + + decoded = decoder.decode(pull.recv_multipart(copy=False)) + assert decoded.prompt_embeds is not None + assert torch.equal(decoded.prompt_embeds, expected) + push.close(linger=0) + pull.close(linger=0) diff --git a/vllm/envs.py b/vllm/envs.py index fb54619c748..f7e01c33275 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1547,7 +1547,7 @@ environment_variables: dict[str, Callable[[], Any]] = { # tensors above will instead be sent via a separate message. # While the sending side still actually copies the tensor # in all cases, on the receiving side, tensors above this - # limit will actually be zero-copy decoded. + # limit will actually be zero-copy decoded. The unit is bytes. "VLLM_MSGPACK_ZERO_COPY_THRESHOLD": lambda: int( os.getenv("VLLM_MSGPACK_ZERO_COPY_THRESHOLD", "256") ), diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index ecac92f5fe0..66135273002 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -7,7 +7,7 @@ import signal import threading import time from collections import defaultdict, deque -from collections.abc import Callable, Generator +from collections.abc import Callable, Generator, Sequence from concurrent.futures import Future from contextlib import ExitStack, contextmanager from enum import IntEnum @@ -87,7 +87,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus -from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder +from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.utils import compute_iteration_details from vllm.version import __version__ as VLLM_VERSION @@ -1748,10 +1748,11 @@ class EngineCoreProc(EngineCore): encoder = MsgpackEncoder() # Send buffers to reuse. reuse_buffers: list[bytearray] = [] - # Keep references to outputs and buffers until zmq is finished - # with them (outputs may contain tensors/np arrays whose - # backing buffers were extracted for zero-copy send). - pending = deque[tuple[zmq.MessageTracker, Any, bytearray]]() + # Payload buffers that can't be reused yet because zmq may still be + # sending them. + # Buffers of the zero-copy tensor/ndarray frames don't need tracking + # here: zmq itself holds a reference to each until it's done with it. + pending = deque[tuple[zmq.MessageTracker, bytearray]]() # We must set linger to ensure the ENGINE_CORE_DEAD # message is sent prior to closing the socket. @@ -1792,20 +1793,38 @@ class EngineCoreProc(EngineCore): # Reclaim buffers that zmq is finished with. while pending and pending[-1][0].done: - reuse_buffers.append(pending.pop()[2]) + reclaimed = pending.pop()[1] + if len(reuse_buffers) < max_reuse_bufs: + reuse_buffers.append(reclaimed) buffer = reuse_buffers.pop() if reuse_buffers else bytearray() buffers = encoder.encode_into(outputs, buffer) - tracker = sockets[client_index].send_multipart( - buffers, copy=False, track=True + tracker = self._send_msg_tracking_payload( + sockets[client_index], buffers ) if not tracker.done: - ref = outputs if len(buffers) > 1 else None - pending.appendleft((tracker, ref, buffer)) + pending.appendleft((tracker, buffer)) elif len(reuse_buffers) < max_reuse_bufs: # Limit the number of buffers to reuse. reuse_buffers.append(buffer) + @staticmethod + def _send_msg_tracking_payload( + socket: zmq.Socket, buffers: Sequence[bytestr] + ) -> zmq.MessageTracker: + """Send `buffers` as a zero-copy multipart message, returning a tracker + for the *first* frame. + + Used instead of `Socket.send_multipart()` because we reuse the buffer + passed to `MsgpackEncoder.encode_into()`: `send_multipart()` returns a + tracker for the last frame only. + """ + more_flag = zmq.SNDMORE if len(buffers) > 1 else 0 + tracker = socket.send(buffers[0], more_flag, copy=False, track=True) + if more_flag: + socket.send_multipart(buffers[1:], copy=False) + return tracker + def _handle_request_preproc_error(self, request: EngineCoreRequest) -> None: """Log and return a request-scoped error response for exceptions raised from the add request preprocessing in the input socket processing thread. diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index 9460fdf48f7..febaa10ce61 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -7,7 +7,7 @@ import sys import uuid import weakref from abc import ABC, abstractmethod -from collections import Counter, defaultdict, deque +from collections import Counter, defaultdict from collections.abc import Awaitable, Callable, Sequence from concurrent.futures import Future from dataclasses import dataclass @@ -671,11 +671,6 @@ class MPClient(EngineCoreClient): self.core_engine: EngineIdentity = self.core_engines[0] self.utility_results: dict[int, AnyFuture] = {} - # Request objects which may contain pytorch-allocated tensors - # that we need to keep references to until zmq is done with the - # underlying data. - self.pending_messages = deque[tuple[zmq.MessageTracker, Any]]() - # Start monitoring engine core processes for unexpected failures self.start_engine_core_monitor() @@ -707,14 +702,6 @@ class MPClient(EngineCoreClient): if self.resources.engine_dead: raise EngineDeadError() - def add_pending_message(self, tracker: zmq.MessageTracker, msg: Any): - if not tracker.done: - self.pending_messages.appendleft((tracker, msg)) - - def free_pending_messages(self): - while self.pending_messages and self.pending_messages[-1][0].done: - self.pending_messages.pop() - def dp_engines_running(self) -> bool: return self.engines_running @@ -896,17 +883,12 @@ class SyncMPClient(MPClient): def _send_input(self, request_type: EngineCoreRequestType, request: Any): self.ensure_alive() - self.free_pending_messages() # (Identity, RequestType, SerializedRequest) msg = (self.core_engine, request_type.value, *self.encoder.encode(request)) - - if len(msg) <= 3: - # No auxiliary buffers => no tensor backing buffers in request. - self.input_socket.send_multipart(msg, copy=False) - return - - tracker = self.input_socket.send_multipart(msg, copy=False, track=True) - self.add_pending_message(tracker, request) + # Any zero-copy tensor/ndarray frames are kept alive by zmq itself + # until it's finished sending them (there is a ref chain from the underlying + # memoryview back to the original owning tensor/ndarray). + self.input_socket.send_multipart(msg, copy=False) def call_utility(self, method: str, *args) -> Any: call_id = uuid.uuid1().int >> 64 @@ -1129,32 +1111,16 @@ class AsyncMPClient(MPClient): engine = self.core_engine message = (request_type.value, *self.encoder.encode(request)) - return self._send_input_message(message, engine, request) + return self._send_input_message(message, engine) def _send_input_message( - self, message: tuple[bytestr, ...], engine: EngineIdentity, objects: Any + self, message: tuple[bytestr, ...], engine: EngineIdentity ) -> Awaitable[Any]: - """ - objects is a reference to retain until zmq is finished with the - buffers, in case they were extracted from tensors in the request. - """ self.ensure_alive() - self.free_pending_messages() - - msg = (engine,) + message - if not objects or len(msg) <= 3: - # No auxiliary buffers => no tensor backing buffers in request. - return self.input_socket.send_multipart(msg, copy=False) - - future: asyncio.Future[zmq.MessageTracker] - future = self.input_socket.send_multipart(msg, copy=False, track=True) - - def add_pending(f: asyncio.Future[zmq.MessageTracker]): - with contextlib.suppress(BaseException): - self.add_pending_message(f.result(), objects) - - future.add_done_callback(add_pending) - return future + # Any zero-copy tensor/ndarray frames are kept alive by zmq itself + # until it's finished sending them (there is a ref chain from the underlying + # memoryview back to the original owning tensor/ndarray). + return self.input_socket.send_multipart((engine,) + message, copy=False) async def call_utility_async(self, method: str, *args) -> Any: return await self._call_utility_async(method, *args, engine=self.core_engine) @@ -1169,7 +1135,7 @@ class AsyncMPClient(MPClient): EngineCoreRequestType.UTILITY.value, *self.encoder.encode((self.client_index, call_id, method, args)), ) - await self._send_input_message(message, engine, args) + await self._send_input_message(message, engine) self._ensure_output_queue_task() return await future