# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import UserDict from dataclasses import dataclass from typing import Optional import msgspec import numpy as np import pytest import torch from vllm.multimodal.inputs import (MultiModalBatchedField, MultiModalFieldElem, MultiModalFlatField, MultiModalKwargs, MultiModalKwargsItem, MultiModalSharedField, NestedTensors) from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder class UnrecognizedType(UserDict): def __init__(self, an_int: int): super().__init__() self.an_int = an_int @dataclass class MyType: tensor1: torch.Tensor a_string: str list_of_tensors: list[torch.Tensor] numpy_array: np.ndarray unrecognized: UnrecognizedType small_f_contig_tensor: torch.Tensor large_f_contig_tensor: torch.Tensor small_non_contig_tensor: torch.Tensor large_non_contig_tensor: torch.Tensor empty_tensor: torch.Tensor def test_encode_decode(monkeypatch: pytest.MonkeyPatch): """Test encode/decode loop with zero-copy tensors.""" with monkeypatch.context() as m: m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") obj = MyType( tensor1=torch.randint(low=0, high=100, size=(1024, ), dtype=torch.int32), a_string="hello", list_of_tensors=[ torch.rand((1, 10), dtype=torch.float32), torch.rand((3, 5, 4000), dtype=torch.float64), torch.tensor(1984), # test scalar too # Make sure to test bf16 which numpy doesn't support. torch.rand((3, 5, 1000), dtype=torch.bfloat16), torch.tensor([float("-inf"), float("inf")] * 1024, dtype=torch.bfloat16), ], numpy_array=np.arange(512), unrecognized=UnrecognizedType(33), small_f_contig_tensor=torch.rand(5, 4).t(), large_f_contig_tensor=torch.rand(1024, 4).t(), small_non_contig_tensor=torch.rand(2, 4)[:, 1:3], large_non_contig_tensor=torch.rand(1024, 512)[:, 10:20], empty_tensor=torch.empty(0), ) encoder = MsgpackEncoder(size_threshold=256) decoder = MsgpackDecoder(MyType) encoded = encoder.encode(obj) # There should be the main buffer + 4 large tensor buffers # + 1 large numpy array. "large" is <= 512 bytes. # The two small tensors are encoded inline. assert len(encoded) == 8 decoded: MyType = decoder.decode(encoded) assert_equal(decoded, obj) # Test encode_into case preallocated = bytearray() encoded2 = encoder.encode_into(obj, preallocated) assert len(encoded2) == 8 assert encoded2[0] is preallocated decoded2: MyType = decoder.decode(encoded2) assert_equal(decoded2, obj) class MyRequest(msgspec.Struct): mm: Optional[list[MultiModalKwargs]] def test_multimodal_kwargs(): e1 = MultiModalFieldElem("audio", "a0", torch.zeros(1000, dtype=torch.bfloat16), MultiModalBatchedField()) e2 = MultiModalFieldElem( "video", "v0", [torch.zeros(1000, dtype=torch.int8) for _ in range(4)], MultiModalFlatField( [[slice(1, 2, 3), slice(4, 5, 6)], [slice(None, 2)]], 0), ) e3 = MultiModalFieldElem("image", "i0", torch.zeros(1000, dtype=torch.int32), MultiModalSharedField(4)) e4 = MultiModalFieldElem( "image", "i1", torch.zeros(1000, dtype=torch.int32), MultiModalFlatField([slice(1, 2, 3), slice(4, 5, 6)], 2)) audio = MultiModalKwargsItem.from_elems([e1]) video = MultiModalKwargsItem.from_elems([e2]) image = MultiModalKwargsItem.from_elems([e3, e4]) mm = MultiModalKwargs([audio, video, image]) # pack mm kwargs into a mock request so that it can be decoded properly req = MyRequest([mm]) encoder = MsgpackEncoder() decoder = MsgpackDecoder(MyRequest) encoded = encoder.encode(req) assert len(encoded) == 8 total_len = sum(memoryview(x).cast("B").nbytes for x in encoded) # expected total encoding length, should be 14255, +-20 for minor changes assert 14250 <= total_len <= 14300 decoded: MultiModalKwargs = decoder.decode(encoded).mm[0] # check all modalities were recovered and do some basic sanity checks assert len(decoded.modalities) == 3 images = decoded.get_items("image") assert len(images) == 1 assert len(images[0].items()) == 2 assert list(images[0].keys()) == ["i0", "i1"] # check the tensor contents and layout in the main dict assert all(nested_equal(mm[k], decoded[k]) for k in mm) def nested_equal(a: NestedTensors, b: NestedTensors): if isinstance(a, torch.Tensor): return torch.equal(a, b) return all(nested_equal(x, y) for x, y in zip(a, b)) def assert_equal(obj1: MyType, obj2: MyType): assert torch.equal(obj1.tensor1, obj2.tensor1) assert obj1.a_string == obj2.a_string assert all( torch.equal(a, b) for a, b in zip(obj1.list_of_tensors, obj2.list_of_tensors)) assert np.array_equal(obj1.numpy_array, obj2.numpy_array) assert obj1.unrecognized.an_int == obj2.unrecognized.an_int assert torch.equal(obj1.small_f_contig_tensor, obj2.small_f_contig_tensor) assert torch.equal(obj1.large_f_contig_tensor, obj2.large_f_contig_tensor) assert torch.equal(obj1.small_non_contig_tensor, obj2.small_non_contig_tensor) assert torch.equal(obj1.large_non_contig_tensor, obj2.large_non_contig_tensor) assert torch.equal(obj1.empty_tensor, obj2.empty_tensor) def test_dict_serialization(): """Test encoding and decoding of a generic Python object using pickle.""" encoder = MsgpackEncoder() decoder = MsgpackDecoder() # Create a sample Python object obj = {"key": "value", "number": 42} # Encode the object encoded = encoder.encode(obj) # Decode the object decoded = decoder.decode(encoded) # Verify the decoded object matches the original assert obj == decoded, "Decoded object does not match the original object." def test_tensor_serialization(): """Test encoding and decoding of a torch.Tensor.""" encoder = MsgpackEncoder() decoder = MsgpackDecoder(torch.Tensor) # Create a sample tensor tensor = torch.rand(10, 10) # Encode the tensor encoded = encoder.encode(tensor) # Decode the tensor decoded = decoder.decode(encoded) # Verify the decoded tensor matches the original assert torch.allclose( tensor, decoded), "Decoded tensor does not match the original tensor." def test_numpy_array_serialization(): """Test encoding and decoding of a numpy array.""" encoder = MsgpackEncoder() decoder = MsgpackDecoder(np.ndarray) # Create a sample numpy array array = np.random.rand(10, 10) # Encode the numpy array encoded = encoder.encode(array) # Decode the numpy array decoded = decoder.decode(encoded) # Verify the decoded array matches the original assert np.allclose( array, decoded), "Decoded numpy array does not match the original array." class CustomClass: def __init__(self, value): self.value = value def __eq__(self, other): return isinstance(other, CustomClass) and self.value == other.value def test_custom_class_serialization_allowed_with_pickle( monkeypatch: pytest.MonkeyPatch): """Test that serializing a custom class succeeds when allow_pickle=True.""" with monkeypatch.context() as m: m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") encoder = MsgpackEncoder() decoder = MsgpackDecoder(CustomClass) obj = CustomClass("test_value") # Encode the custom class encoded = encoder.encode(obj) # Decode the custom class decoded = decoder.decode(encoded) # Verify the decoded object matches the original assert obj == decoded, ( "Decoded object does not match the original object.") def test_custom_class_serialization_disallowed_without_pickle(): """Test that serializing a custom class fails when allow_pickle=False.""" encoder = MsgpackEncoder() obj = CustomClass("test_value") with pytest.raises(TypeError): # Attempt to encode the custom class encoder.encode(obj)