| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections import UserDict |
| 4 | from dataclasses import dataclass |
| 5 | |
| 6 | import msgspec |
| 7 | import numpy as np |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from vllm.multimodal.inputs import ( |
| 12 | MultiModalBatchedField, |
| 13 | MultiModalFieldElem, |
| 14 | MultiModalFlatField, |
| 15 | MultiModalKwargsItem, |
| 16 | MultiModalKwargsItems, |
| 17 | MultiModalSharedField, |
| 18 | NestedTensors, |
| 19 | ) |
| 20 | from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder |
| 21 | |
| 22 | pytestmark = pytest.mark.cpu_test |
| 23 | |
| 24 | |
| 25 | class UnrecognizedType(UserDict): |
| 26 | def __init__(self, an_int: int): |
| 27 | super().__init__() |
| 28 | self.an_int = an_int |
| 29 | |
| 30 | |
| 31 | @dataclass |
| 32 | class MyType: |
| 33 | tensor1: torch.Tensor |
| 34 | a_string: str |
| 35 | list_of_tensors: list[torch.Tensor] |
| 36 | numpy_array: np.ndarray |
| 37 | unrecognized: UnrecognizedType |
| 38 | small_f_contig_tensor: torch.Tensor |
| 39 | large_f_contig_tensor: torch.Tensor |
| 40 | small_non_contig_tensor: torch.Tensor |
| 41 | large_non_contig_tensor: torch.Tensor |
| 42 | empty_tensor: torch.Tensor |
| 43 | |
| 44 | |
| 45 | def test_encode_decode(monkeypatch: pytest.MonkeyPatch): |
| 46 | """Test encode/decode loop with zero-copy tensors.""" |
| 47 | with monkeypatch.context() as m: |
| 48 | m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 49 | |
| 50 | obj = MyType( |
| 51 | tensor1=torch.randint(low=0, high=100, size=(1024,), dtype=torch.int32), |
| 52 | a_string="hello", |
| 53 | list_of_tensors=[ |
| 54 | torch.rand((1, 10), dtype=torch.float32), |
| 55 | torch.rand((3, 5, 4000), dtype=torch.float64), |
| 56 | torch.tensor(1984), # test scalar too |
| 57 | # Make sure to test bf16 which numpy doesn't support. |
| 58 | torch.rand((3, 5, 1000), dtype=torch.bfloat16), |
| 59 | torch.tensor( |
| 60 | [float("-inf"), float("inf")] * 1024, dtype=torch.bfloat16 |
| 61 | ), |
| 62 | ], |
| 63 | numpy_array=np.arange(512), |
| 64 | unrecognized=UnrecognizedType(33), |
| 65 | small_f_contig_tensor=torch.rand(5, 4).t(), |
| 66 | large_f_contig_tensor=torch.rand(1024, 4).t(), |
| 67 | small_non_contig_tensor=torch.rand(2, 4)[:, 1:3], |
| 68 | large_non_contig_tensor=torch.rand(1024, 512)[:, 10:20], |
| 69 | empty_tensor=torch.empty(0), |
| 70 | ) |
| 71 | |
| 72 | encoder = MsgpackEncoder(size_threshold=256) |
| 73 | decoder = MsgpackDecoder(MyType) |
| 74 | |
| 75 | encoded = encoder.encode(obj) |
| 76 | |
| 77 | # There should be the main buffer + 4 large tensor buffers |
| 78 | # + 1 large numpy array. "large" is <= 512 bytes. |
| 79 | # The two small tensors are encoded inline. |
| 80 | assert len(encoded) == 8 |
| 81 | |
| 82 | decoded: MyType = decoder.decode(encoded) |
| 83 | |
| 84 | assert_equal(decoded, obj) |
| 85 | |
| 86 | # Test encode_into case |
| 87 | |
| 88 | preallocated = bytearray() |
| 89 | |
| 90 | encoded2 = encoder.encode_into(obj, preallocated) |
| 91 | |
| 92 | assert len(encoded2) == 8 |
| 93 | assert encoded2[0] is preallocated |
| 94 | |
| 95 | decoded2: MyType = decoder.decode(encoded2) |
| 96 | |
| 97 | assert_equal(decoded2, obj) |
| 98 | |
| 99 | |
| 100 | class MyRequest(msgspec.Struct): |
| 101 | mm: list[MultiModalKwargsItems] | None |
| 102 | |
| 103 | |
| 104 | def test_multimodal_kwargs(): |
| 105 | e1 = MultiModalFieldElem( |
| 106 | torch.zeros(1000, dtype=torch.bfloat16), |
| 107 | MultiModalBatchedField(), |
| 108 | ) |
| 109 | e2 = MultiModalFieldElem( |
| 110 | [torch.zeros(1000, dtype=torch.int8) for _ in range(4)], |
| 111 | MultiModalFlatField( |
| 112 | slices=[[slice(1, 2, 3), slice(4, 5, 6)], [slice(None, 2)]], |
| 113 | dim=0, |
| 114 | ), |
| 115 | ) |
| 116 | e3 = MultiModalFieldElem( |
| 117 | torch.zeros(1000, dtype=torch.int32), |
| 118 | MultiModalSharedField(batch_size=4), |
| 119 | ) |
| 120 | e4 = MultiModalFieldElem( |
| 121 | torch.zeros(1000, dtype=torch.int32), |
| 122 | MultiModalFlatField(slices=[slice(1, 2, 3), slice(4, 5, 6)], dim=2), |
| 123 | ) |
| 124 | mm = MultiModalKwargsItems( |
| 125 | { |
| 126 | "audio": [MultiModalKwargsItem({"a0": e1})], |
| 127 | "video": [MultiModalKwargsItem({"v0": e2})], |
| 128 | "image": [MultiModalKwargsItem({"i0": e3, "i1": e4})], |
| 129 | } |
| 130 | ) |
| 131 | |
| 132 | # pack mm kwargs into a mock request so that it can be decoded properly |
| 133 | req = MyRequest([mm]) |
| 134 | |
| 135 | encoder = MsgpackEncoder() |
| 136 | decoder = MsgpackDecoder(MyRequest) |
| 137 | |
| 138 | encoded = encoder.encode(req) |
| 139 | |
| 140 | assert len(encoded) == 8 |
| 141 | |
| 142 | total_len = sum(memoryview(x).cast("B").nbytes for x in encoded) |
| 143 | |
| 144 | # expected total encoding length, should be 14319, +-20 for minor changes |
| 145 | assert 14300 <= total_len <= 14340 |
| 146 | decoded = decoder.decode(encoded).mm[0] |
| 147 | assert isinstance(decoded, MultiModalKwargsItems) |
| 148 | |
| 149 | # check all modalities were recovered and do some basic sanity checks |
| 150 | assert len(decoded) == 3 |
| 151 | images = decoded["image"] |
| 152 | assert len(images) == 1 |
| 153 | assert len(images[0].items()) == 2 |
| 154 | assert list(images[0].keys()) == ["i0", "i1"] |
| 155 | |
| 156 | # check the tensor contents and layout in the main dict |
| 157 | mm_data = mm.get_data() |
| 158 | decoded_data = decoded.get_data() |
| 159 | assert all(nested_equal(mm_data[k], decoded_data[k]) for k in mm_data) |
| 160 | |
| 161 | |
| 162 | def nested_equal(a: NestedTensors, b: NestedTensors): |
| 163 | if isinstance(a, torch.Tensor): |
| 164 | return torch.equal(a, b) |
| 165 | return all(nested_equal(x, y) for x, y in zip(a, b)) |
| 166 | |
| 167 | |
| 168 | def assert_equal(obj1: MyType, obj2: MyType): |
| 169 | assert torch.equal(obj1.tensor1, obj2.tensor1) |
| 170 | assert obj1.a_string == obj2.a_string |
| 171 | assert all( |
| 172 | torch.equal(a, b) for a, b in zip(obj1.list_of_tensors, obj2.list_of_tensors) |
| 173 | ) |
| 174 | assert np.array_equal(obj1.numpy_array, obj2.numpy_array) |
| 175 | assert obj1.unrecognized.an_int == obj2.unrecognized.an_int |
| 176 | assert torch.equal(obj1.small_f_contig_tensor, obj2.small_f_contig_tensor) |
| 177 | assert torch.equal(obj1.large_f_contig_tensor, obj2.large_f_contig_tensor) |
| 178 | assert torch.equal(obj1.small_non_contig_tensor, obj2.small_non_contig_tensor) |
| 179 | assert torch.equal(obj1.large_non_contig_tensor, obj2.large_non_contig_tensor) |
| 180 | assert torch.equal(obj1.empty_tensor, obj2.empty_tensor) |
| 181 | |
| 182 | |
| 183 | def test_dict_serialization(): |
| 184 | """Test encoding and decoding of a generic Python object using pickle.""" |
| 185 | encoder = MsgpackEncoder() |
| 186 | decoder = MsgpackDecoder() |
| 187 | |
| 188 | # Create a sample Python object |
| 189 | obj = {"key": "value", "number": 42} |
| 190 | |
| 191 | # Encode the object |
| 192 | encoded = encoder.encode(obj) |
| 193 | |
| 194 | # Decode the object |
| 195 | decoded = decoder.decode(encoded) |
| 196 | |
| 197 | # Verify the decoded object matches the original |
| 198 | assert obj == decoded, "Decoded object does not match the original object." |
| 199 | |
| 200 | |
| 201 | def test_tensor_serialization(): |
| 202 | """Test encoding and decoding of a torch.Tensor.""" |
| 203 | encoder = MsgpackEncoder() |
| 204 | decoder = MsgpackDecoder(torch.Tensor) |
| 205 | |
| 206 | # Create a sample tensor |
| 207 | tensor = torch.rand(10, 10) |
| 208 | |
| 209 | # Encode the tensor |
| 210 | encoded = encoder.encode(tensor) |
| 211 | |
| 212 | # Decode the tensor |
| 213 | decoded = decoder.decode(encoded) |
| 214 | |
| 215 | # Verify the decoded tensor matches the original |
| 216 | assert torch.allclose(tensor, decoded), ( |
| 217 | "Decoded tensor does not match the original tensor." |
| 218 | ) |
| 219 | |
| 220 | |
| 221 | def test_numpy_array_serialization(): |
| 222 | """Test encoding and decoding of a numpy array.""" |
| 223 | encoder = MsgpackEncoder() |
| 224 | decoder = MsgpackDecoder(np.ndarray) |
| 225 | |
| 226 | # Create a sample numpy array |
| 227 | array = np.random.rand(10, 10) |
| 228 | |
| 229 | # Encode the numpy array |
| 230 | encoded = encoder.encode(array) |
| 231 | |
| 232 | # Decode the numpy array |
| 233 | decoded = decoder.decode(encoded) |
| 234 | |
| 235 | # Verify the decoded array matches the original |
| 236 | assert np.allclose(array, decoded), ( |
| 237 | "Decoded numpy array does not match the original array." |
| 238 | ) |
| 239 | |
| 240 | |
| 241 | class CustomClass: |
| 242 | def __init__(self, value): |
| 243 | self.value = value |
| 244 | |
| 245 | def __eq__(self, other): |
| 246 | return isinstance(other, CustomClass) and self.value == other.value |
| 247 | |
| 248 | |
| 249 | def test_custom_class_serialization_allowed_with_pickle( |
| 250 | monkeypatch: pytest.MonkeyPatch, |
| 251 | ): |
| 252 | """Test that serializing a custom class succeeds when allow_pickle=True.""" |
| 253 | with monkeypatch.context() as m: |
| 254 | m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 255 | encoder = MsgpackEncoder() |
| 256 | decoder = MsgpackDecoder(CustomClass) |
| 257 | |
| 258 | obj = CustomClass("test_value") |
| 259 | |
| 260 | # Encode the custom class |
| 261 | encoded = encoder.encode(obj) |
| 262 | |
| 263 | # Decode the custom class |
| 264 | decoded = decoder.decode(encoded) |
| 265 | |
| 266 | # Verify the decoded object matches the original |
| 267 | assert obj == decoded, "Decoded object does not match the original object." |
| 268 | |
| 269 | |
| 270 | def test_custom_class_serialization_disallowed_without_pickle(): |
| 271 | """Test that serializing a custom class fails when allow_pickle=False.""" |
| 272 | encoder = MsgpackEncoder() |
| 273 | |
| 274 | obj = CustomClass("test_value") |
| 275 | |
| 276 | with pytest.raises(TypeError): |
| 277 | # Attempt to encode the custom class |
| 278 | encoder.encode(obj) |
| 279 | |
| 280 | |
| 281 | @dataclass |
| 282 | class RequestWithTensor: |
| 283 | """Mock request with non-multimodal tensor field like EngineCoreRequest.""" |
| 284 | |
| 285 | prompt_embeds: torch.Tensor | None |
| 286 | data: str |
| 287 | |
| 288 | |
| 289 | def test_non_multimodal_tensor_with_ipc(): |
| 290 | """Test that non-multimodal tensor fields work correctly with IPC enabled. |
| 291 | |
| 292 | This reproduces the bug where fields like prompt_embeds: torch.Tensor | None |
| 293 | would fail to decode when IPC is enabled because _decode_tensor expected a |
| 294 | raw tensor tuple but received a msgpack-decoded TensorIpcHandle list. |
| 295 | """ |
| 296 | import torch.multiprocessing as torch_mp |
| 297 | |
| 298 | from vllm.v1.engine.tensor_ipc import TensorIpcReceiver, TensorIpcSender |
| 299 | |
| 300 | # Create tensor queues for IPC |
| 301 | tensor_queues = [torch_mp.Queue()] |
| 302 | |
| 303 | # Create encoder with IPC sender |
| 304 | sender = TensorIpcSender(tensor_queues[0]) |
| 305 | encoder = MsgpackEncoder(oob_tensor_consumer=sender) |
| 306 | |
| 307 | # Create decoder with IPC receiver |
| 308 | receiver = TensorIpcReceiver(tensor_queues[0]) |
| 309 | decoder = MsgpackDecoder(RequestWithTensor, oob_tensor_provider=receiver) |
| 310 | |
| 311 | # Create a request with a non-multimodal tensor |
| 312 | original_tensor = torch.randn(5, 10, dtype=torch.float32) |
| 313 | request = RequestWithTensor(prompt_embeds=original_tensor, data="test_data") |
| 314 | |
| 315 | # Encode the request - this should send the tensor via IPC |
| 316 | encoded = encoder.encode(request) |
| 317 | |
| 318 | # Verify encoding succeeded |
| 319 | assert len(encoded) > 0 |
| 320 | |
| 321 | # Decode the request - this should retrieve the tensor from IPC queue |
| 322 | # Previously this would fail because the decoder tried to unpack the |
| 323 | # handle list as raw tensor bytes metadata. |
| 324 | decoded = decoder.decode(encoded) |
| 325 | |
| 326 | # Verify the decoded request matches the original |
| 327 | assert isinstance(decoded, RequestWithTensor) |
| 328 | assert decoded.data == "test_data" |
| 329 | assert decoded.prompt_embeds is not None |
| 330 | assert torch.allclose(decoded.prompt_embeds, original_tensor), ( |
| 331 | "Decoded tensor does not match the original tensor." |
| 332 | ) |
| 333 | |
| 334 | |
| 335 | def test_non_multimodal_tensor_with_ipc_none_value(): |
| 336 | """Test that None values for tensor fields work correctly with IPC enabled.""" |
| 337 | import torch.multiprocessing as torch_mp |
| 338 | |
| 339 | from vllm.v1.engine.tensor_ipc import TensorIpcReceiver, TensorIpcSender |
| 340 | |
| 341 | # Create tensor queues for IPC |
| 342 | tensor_queues = [torch_mp.Queue()] |
| 343 | |
| 344 | # Create encoder with IPC sender |
| 345 | sender = TensorIpcSender(tensor_queues[0]) |
| 346 | encoder = MsgpackEncoder(oob_tensor_consumer=sender) |
| 347 | |
| 348 | # Create decoder with IPC receiver |
| 349 | receiver = TensorIpcReceiver(tensor_queues[0]) |
| 350 | decoder = MsgpackDecoder(RequestWithTensor, oob_tensor_provider=receiver) |
| 351 | |
| 352 | # Create a request with None for the tensor field |
| 353 | request = RequestWithTensor(prompt_embeds=None, data="test_data_with_none") |
| 354 | |
| 355 | # Encode and decode the request |
| 356 | encoded = encoder.encode(request) |
| 357 | decoded = decoder.decode(encoded) |
| 358 | |
| 359 | # Verify the decoded request matches the original |
| 360 | assert isinstance(decoded, RequestWithTensor) |
| 361 | assert decoded.data == "test_data_with_none" |
| 362 | assert decoded.prompt_embeds is None |
| 363 | |
| 364 | |
| 365 | def test_multiple_senders_single_receiver_ipc(): |
| 366 | """Test N senders sharing a queue with a single receiver via msgpack. |
| 367 | |
| 368 | Simulates the real vLLM topology where multiple API server frontends |
| 369 | each have their own MsgpackEncoder + TensorIpcSender, all putting |
| 370 | tensors onto the same torch.mp queue, and a single engine core |
| 371 | decodes them with one MsgpackDecoder + TensorIpcReceiver. |
| 372 | """ |
| 373 | import torch.multiprocessing as torch_mp |
| 374 | |
| 375 | from vllm.v1.engine.tensor_ipc import TensorIpcReceiver, TensorIpcSender |
| 376 | |
| 377 | num_senders = 3 |
| 378 | num_messages_per_sender = 2 |
| 379 | tensor_queue = torch_mp.Queue() |
| 380 | |
| 381 | # Create N independent senders (each gets its own uuid-based sender_id) |
| 382 | senders = [] |
| 383 | encoders = [] |
| 384 | for _ in range(num_senders): |
| 385 | s = TensorIpcSender(tensor_queue) |
| 386 | senders.append(s) |
| 387 | encoders.append(MsgpackEncoder(oob_tensor_consumer=s)) |
| 388 | |
| 389 | # Single receiver |
| 390 | receiver = TensorIpcReceiver(tensor_queue) |
| 391 | decoder = MsgpackDecoder(RequestWithTensor, oob_tensor_provider=receiver) |
| 392 | |
| 393 | # Encode messages from all senders, interleaving the order |
| 394 | # so that tensors from different senders land on the queue interleaved. |
| 395 | encoded_payloads: list[tuple[int, int, torch.Tensor, list]] = [] |
| 396 | for msg_idx in range(num_messages_per_sender): |
| 397 | for sender_idx in range(num_senders): |
| 398 | tensor = torch.full( |
| 399 | (sender_idx + 1, msg_idx + 2), |
| 400 | float(sender_idx * 100 + msg_idx), |
| 401 | dtype=torch.float32, |
| 402 | ) |
| 403 | req = RequestWithTensor( |
| 404 | prompt_embeds=tensor, |
| 405 | data=f"s{sender_idx}_m{msg_idx}", |
| 406 | ) |
| 407 | encoded = encoders[sender_idx].encode(req) |
| 408 | encoded_payloads.append((sender_idx, msg_idx, tensor, encoded)) |
| 409 | |
| 410 | # Decode all messages — the receiver must correctly match each |
| 411 | # tensor handle to the right TensorIpcData from the shared queue. |
| 412 | for sender_idx, msg_idx, original_tensor, encoded in encoded_payloads: |
| 413 | decoded = decoder.decode(encoded) |
| 414 | assert isinstance(decoded, RequestWithTensor) |
| 415 | assert decoded.data == f"s{sender_idx}_m{msg_idx}" |
| 416 | assert decoded.prompt_embeds is not None |
| 417 | assert decoded.prompt_embeds.shape == original_tensor.shape, ( |
| 418 | f"Shape mismatch for sender {sender_idx} msg {msg_idx}: " |
| 419 | f"{decoded.prompt_embeds.shape} != {original_tensor.shape}" |
| 420 | ) |
| 421 | assert torch.allclose(decoded.prompt_embeds, original_tensor), ( |
| 422 | f"Value mismatch for sender {sender_idx} msg {msg_idx}" |
| 423 | ) |
| 424 | |
| 425 | |
| 426 | def _logprobs_outputs(num_reqs: int, num_prompt_tokens: int): |
| 427 | """An EngineCoreOutputs carrying prompt logprobs, as the engine core sends |
| 428 | it: many requests, each with per-token tensors small enough that pyzmq |
| 429 | copies their frames, while the accumulated payload frame is large enough |
| 430 | that pyzmq sends it zero-copy.""" |
| 431 | from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs |
| 432 | from vllm.v1.outputs import LogprobsTensors |
| 433 | |
| 434 | outputs = [] |
| 435 | for req in range(num_reqs): |
| 436 | num_tokens = num_prompt_tokens + req % 4 |
| 437 | outputs.append( |
| 438 | EngineCoreOutput( |
| 439 | request_id=f"req-{req:08d}", |
| 440 | new_token_ids=[req], |
| 441 | new_prompt_logprobs_tensors=LogprobsTensors( |
| 442 | logprob_token_ids=torch.arange( |
| 443 | num_tokens * 2, dtype=torch.int64 |
| 444 | ).view(num_tokens, 2), |
| 445 | logprobs=torch.zeros(num_tokens, 2, dtype=torch.float32), |
| 446 | selected_token_ranks=torch.zeros(num_tokens, dtype=torch.int32), |
| 447 | ), |
| 448 | ) |
| 449 | ) |
| 450 | return EngineCoreOutputs(outputs=outputs) |
| 451 | |
| 452 | |
| 453 | def test_payload_buffer_reuse_does_not_corrupt_in_flight_messages(): |
| 454 | """The engine core recycles the msgpack payload buffer across messages |
| 455 | (`MsgpackEncoder.encode_into`). It may only do so once zmq has finished |
| 456 | sending that buffer, otherwise a newer payload is delivered alongside the |
| 457 | older message's zero-copy tensor frames. |
| 458 | |
| 459 | `Socket.send_multipart(track=True)` cannot be used to detect this: it |
| 460 | returns a tracker for the last frame only, and pyzmq copies frames below |
| 461 | `zmq.COPY_THRESHOLD` and reports them as already-sent. |
| 462 | """ |
| 463 | import zmq |
| 464 | |
| 465 | from vllm.v1.engine import EngineCoreOutputs |
| 466 | from vllm.v1.engine.core import EngineCoreProc |
| 467 | |
| 468 | num_msgs = 100 |
| 469 | encoder = MsgpackEncoder() |
| 470 | decoder = MsgpackDecoder(EngineCoreOutputs) |
| 471 | # Enough requests that the payload frame is zero-copied rather than copied |
| 472 | # by pyzmq, which is what makes early reuse observable. |
| 473 | messages = [_logprobs_outputs(300, 24 + i % 8) for i in range(num_msgs)] |
| 474 | assert len(encoder.encode(messages[0])[0]) >= zmq.COPY_THRESHOLD |
| 475 | |
| 476 | reuse_buffers: list[bytearray] = [] |
| 477 | pending: list[tuple[zmq.MessageTracker, bytearray]] = [] |
| 478 | with zmq.Context() as ctx: |
| 479 | push = ctx.socket(zmq.PUSH) |
| 480 | push.bind("inproc://test-payload-reuse") |
| 481 | pull = ctx.socket(zmq.PULL) |
| 482 | pull.connect("inproc://test-payload-reuse") |
| 483 | |
| 484 | for outputs in messages: |
| 485 | while pending and pending[0][0].done: |
| 486 | reuse_buffers.append(pending.pop(0)[1]) |
| 487 | buffer = reuse_buffers.pop() if reuse_buffers else bytearray() |
| 488 | buffers = encoder.encode_into(outputs, buffer) |
| 489 | tracker = EngineCoreProc._send_msg_tracking_payload(push, buffers) |
| 490 | if tracker.done: |
| 491 | reuse_buffers.append(buffer) |
| 492 | else: |
| 493 | pending.append((tracker, buffer)) |
| 494 | |
| 495 | for i, sent in enumerate(messages): |
| 496 | received = decoder.decode(pull.recv_multipart(copy=False)) |
| 497 | assert len(received.outputs) == len(sent.outputs), f"message {i}" |
| 498 | for expected, actual in zip(sent.outputs, received.outputs): |
| 499 | sent_ids = expected.new_prompt_logprobs_tensors.logprob_token_ids |
| 500 | got_ids = actual.new_prompt_logprobs_tensors.logprob_token_ids |
| 501 | assert actual.request_id == expected.request_id, f"message {i}" |
| 502 | assert torch.equal(got_ids, sent_ids), ( |
| 503 | f"message {i} request {actual.request_id}: corrupted " |
| 504 | f"prompt logprobs, {got_ids.shape} vs {sent_ids.shape}" |
| 505 | ) |
| 506 | push.close(linger=0) |
| 507 | pull.close(linger=0) |
| 508 | |
| 509 | |
| 510 | def test_zero_copy_frames_survive_without_caller_side_references(): |
| 511 | """Callers don't need to retain the encoded object until zmq has sent it: |
| 512 | for a zero-copy frame, zmq holds its own reference to the backing buffer. |
| 513 | |
| 514 | The engine core clients rely on this when sending requests that carry |
| 515 | tensors (e.g. prompt embeds) without tracking the messages. |
| 516 | |
| 517 | What makes that safe is that `tensor_data()` hands zmq a memoryview which |
| 518 | transitively references the source tensor, so refcounting - not timing - |
| 519 | keeps the memory from being freed and reused underneath zmq. |
| 520 | """ |
| 521 | import gc |
| 522 | |
| 523 | import zmq |
| 524 | |
| 525 | from vllm.v1.utils import tensor_data |
| 526 | |
| 527 | num_elems = 100_000 # comfortably over zmq.COPY_THRESHOLD |
| 528 | expected = torch.arange(num_elems, dtype=torch.int64) |
| 529 | encoder = MsgpackEncoder() |
| 530 | decoder = MsgpackDecoder(RequestWithTensor) |
| 531 | |
| 532 | # The buffer handed to zmq must keep the tensor's storage alive by itself. |
| 533 | holder = tensor_data(expected).obj |
| 534 | while getattr(holder, "base", None) is not None: |
| 535 | holder = holder.base |
| 536 | assert isinstance(holder, torch.Tensor) |
| 537 | assert holder.data_ptr() == expected.data_ptr() |
| 538 | |
| 539 | with zmq.Context() as ctx: |
| 540 | push = ctx.socket(zmq.PUSH) |
| 541 | push.bind("inproc://test-zero-copy-lifetime") |
| 542 | pull = ctx.socket(zmq.PULL) |
| 543 | pull.connect("inproc://test-zero-copy-lifetime") |
| 544 | |
| 545 | request = RequestWithTensor(prompt_embeds=expected.clone(), data="req") |
| 546 | buffers = encoder.encode(request) |
| 547 | assert max(len(buf) for buf in buffers) >= zmq.COPY_THRESHOLD |
| 548 | push.send_multipart(buffers, copy=False) |
| 549 | |
| 550 | # Drop every reference the sender holds, then churn the allocator. |
| 551 | del request, buffers |
| 552 | gc.collect() |
| 553 | torch.arange(num_elems * 4, dtype=torch.int64) |
| 554 | |
| 555 | decoded = decoder.decode(pull.recv_multipart(copy=False)) |
| 556 | assert decoded.prompt_embeds is not None |
| 557 | assert torch.equal(decoded.prompt_embeds, expected) |
| 558 | push.close(linger=0) |
| 559 | pull.close(linger=0) |