Fix RecursionError in MediaWithBytes unpickling (#31191)

Signed-off-by: Nikhil Ghosh <nikhil@anyscale.com>

tests/multimodal/test_image.py

@@ -1,11 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pickle
 from pathlib import Path
 
 import numpy as np
 import pytest
 from PIL import Image, ImageChops
 
+from vllm.multimodal.base import MediaWithBytes
 from vllm.multimodal.image import ImageMediaIO, convert_image_mode
 
 pytestmark = pytest.mark.cpu_test
@@ -157,3 +159,34 @@ def test_rgba_background_color_validation():
     ImageMediaIO(rgba_background_color=(0, 0, 0))  # Should not raise
     ImageMediaIO(rgba_background_color=[255, 255, 255])  # Should not raise
     ImageMediaIO(rgba_background_color=(128, 128, 128))  # Should not raise
+
+
+def test_media_with_bytes_pickle_roundtrip():
+    """Regression test for pickle/unpickle of MediaWithBytes.
+
+    Verifies that MediaWithBytes can be pickled and unpickled without
+    RecursionError. See: https://github.com/vllm-project/vllm/issues/30818
+    """
+    original_image = Image.open(ASSETS_DIR / "image1.png").convert("RGB")
+    original_bytes = b"test_bytes_data"
+
+    wrapper = MediaWithBytes(media=original_image, original_bytes=original_bytes)
+
+    # Verify attribute delegation works before pickling
+    assert wrapper.width == original_image.width
+    assert wrapper.height == original_image.height
+    assert wrapper.mode == original_image.mode
+
+    # Pickle and unpickle (this would cause RecursionError before the fix)
+    pickled = pickle.dumps(wrapper)
+    unpickled = pickle.loads(pickled)
+
+    # Verify the unpickled object works correctly
+    assert unpickled.original_bytes == original_bytes
+    assert unpickled.media.width == original_image.width
+    assert unpickled.media.height == original_image.height
+
+    # Verify attribute delegation works after unpickling
+    assert unpickled.width == original_image.width
+    assert unpickled.height == original_image.height
+    assert unpickled.mode == original_image.mode

tools/pre_commit/check_pickle_imports.py

@@ -27,6 +27,7 @@ ALLOWED_FILES = {
     "vllm/distributed/device_communicators/shm_broadcast.py",
     "vllm/distributed/device_communicators/shm_object_storage.py",
     "vllm/utils/hashing.py",
+    "tests/multimodal/test_image.py",
     "tests/tokenizers_/test_hf.py",
     "tests/utils_/test_hashing.py",
     "benchmarks/kernels/graph_machete_bench.py",

vllm/multimodal/base.py

@@ -34,7 +34,11 @@ class MediaWithBytes(Generic[_T]):
 
     def __getattr__(self, name: str):
         """Delegate attribute access to the underlying media object."""
-        # This is only called when the attribute is not found on self
+        # Guard against recursion during unpickling when media isn't set yet.
+        # pickle creates objects without calling __init__, so self.media may
+        # not exist when __getattr__ is called for methods like __setstate__.
+        if "media" not in self.__dict__:
+            raise AttributeError(name)
         return getattr(self.media, name)