| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pickle |
| 4 | from pathlib import Path |
| 5 | |
| 6 | import pytest |
| 7 | from PIL import Image |
| 8 | |
| 9 | from vllm.multimodal.media import MediaWithBytes |
| 10 | |
| 11 | pytestmark = pytest.mark.cpu_test |
| 12 | |
| 13 | ASSETS_DIR = Path(__file__).parent.parent / "assets" |
| 14 | assert ASSETS_DIR.exists() |
| 15 | |
| 16 | |
| 17 | def test_media_with_bytes_pickle_roundtrip(): |
| 18 | """Regression test for pickle/unpickle of MediaWithBytes. |
| 19 | |
| 20 | Verifies that MediaWithBytes can be pickled and unpickled without |
| 21 | RecursionError. See: https://github.com/vllm-project/vllm/issues/30818 |
| 22 | """ |
| 23 | original_image = Image.open(ASSETS_DIR / "image1.png").convert("RGB") |
| 24 | original_bytes = b"test_bytes_data" |
| 25 | |
| 26 | wrapper = MediaWithBytes(media=original_image, original_bytes=original_bytes) |
| 27 | |
| 28 | # Verify attribute delegation works before pickling |
| 29 | assert wrapper.width == original_image.width |
| 30 | assert wrapper.height == original_image.height |
| 31 | assert wrapper.mode == original_image.mode |
| 32 | |
| 33 | # Pickle and unpickle (this would cause RecursionError before the fix) |
| 34 | pickled = pickle.dumps(wrapper) |
| 35 | unpickled = pickle.loads(pickled) |
| 36 | |
| 37 | # Verify the unpickled object works correctly |
| 38 | assert unpickled.original_bytes == original_bytes |
| 39 | assert unpickled.media.width == original_image.width |
| 40 | assert unpickled.media.height == original_image.height |
| 41 | |
| 42 | # Verify attribute delegation works after unpickling |
| 43 | assert unpickled.width == original_image.width |
| 44 | assert unpickled.height == original_image.height |
| 45 | assert unpickled.mode == original_image.mode |