[Security][Feature] Use Msgpack in OmniSerializer instead of pickle (#310)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

.pre-commit-config.yaml

@@ -58,3 +58,10 @@ repos:
       verbose: true
       pass_filenames: false
     # Insert new entries above the `suggestion` entry
+
+    - id: check-pickle-imports
+      name: Prevent new pickle/cloudpickle imports
+      entry: python tools/pre_commit/check_pickle_imports.py
+      language: python
+      types: [python]
+      additional_dependencies: [regex]

tests/distributed/omni_connectors/test_basic_connectors.py

@@ -10,16 +10,38 @@ from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
 
 
-def test_pickle_serialization():
-    """Test basic pickle serialization."""
+def test_basic_serialization():
+    """Test basic msgpack serialization."""
     data = {"key": "value", "list": [1, 2, 3]}
-    serialized = OmniSerializer.serialize(data, method="cloudpickle")
+    serialized = OmniSerializer.serialize(data)
     assert isinstance(serialized, bytes)
 
-    deserialized = OmniSerializer.deserialize(serialized, method="cloudpickle")
+    deserialized = OmniSerializer.deserialize(serialized)
     assert data == deserialized
 
 
+def test_tensor_serialization():
+    """Test torch.Tensor serialization."""
+    import torch
+
+    tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+    serialized = OmniSerializer.serialize(tensor)
+    deserialized = OmniSerializer.deserialize(serialized)
+
+    assert torch.equal(tensor, deserialized)
+
+
+def test_ndarray_serialization():
+    """Test numpy.ndarray serialization."""
+    import numpy as np
+
+    arr = np.array([[1, 2, 3], [4, 5, 6]])
+    serialized = OmniSerializer.serialize(arr)
+    deserialized = OmniSerializer.deserialize(serialized)
+
+    assert np.array_equal(arr, deserialized)
+
+
 def test_create_shm_connector():
     """Test creating SharedMemoryConnector via Factory."""
     spec = ConnectorSpec(name="SharedMemoryConnector", extra={"shm_threshold_bytes": 1024})

tests/entrypoints/test_omni_llm.py

@@ -676,7 +676,8 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
         }
     )
 
-    sampling_params_list = [object(), object()]
+    # Use dicts instead of object() for serializable sampling params
+    sampling_params_list = [{"temperature": 0.7}, {"temperature": 0.8}]
     prompts = ["hi"]
     outputs = llm.generate(prompts=prompts, sampling_params_list=sampling_params_list)
 
@@ -771,7 +772,8 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
         }
     )
 
-    outputs = llm.generate(prompts=["p"], sampling_params_list=[object(), object()])
+    # Use dicts instead of object() for serializable sampling params
+    outputs = llm.generate(prompts=["p"], sampling_params_list=[{"temperature": 0.7}, {"temperature": 0.8}])
     assert outputs == []
 
 
@@ -970,7 +972,8 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     )
 
     # Generate should handle error gracefully (log but continue)
-    sampling_params_list = [object()]
+    # Use dict instead of object() for serializable sampling params
+    sampling_params_list = [{"temperature": 0.7}]
     outputs = llm.generate(prompts=["hi"], sampling_params_list=sampling_params_list)
     # Should return final output (error was logged but didn't stop processing)
     assert isinstance(outputs, list)

tools/pre_commit/check_pickle_imports.py

@@ -0,0 +1,84 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import sys
+
+import regex as re
+
+# List of files (relative to repo root) that are allowed to import pickle or
+# cloudpickle
+#
+# STOP AND READ BEFORE YOU ADD ANYTHING ELSE TO THIS LIST:
+#  The pickle and cloudpickle modules are known to be unsafe when deserializing
+#  data from potentially untrusted parties. They have resulted in multiple CVEs
+#  for vLLM and numerous vulnerabilities in the Python ecosystem more broadly.
+#  Before adding new uses of pickle/cloudpickle, please consider safer
+#  alternatives like msgpack or pydantic that are already in use in vLLM. Only
+#  add to this list if absolutely necessary and after careful security review.
+ALLOWED_FILES = {
+    "vllm_omni/entrypoints/omni_llm.py",
+    "tests/e2e/offline_inference/utils.py",
+    "vllm_omni/diffusion/distributed/group_coordinator.py",
+    "tests/diffusion/attention/test_ulysses_sequence_parallel.py",
+}
+
+PICKLE_RE = re.compile(
+    r"^\s*(import\s+(pickle|cloudpickle)(\s|$|\sas)"
+    r"|from\s+(pickle|cloudpickle)\s+import\b)"
+)
+
+
+def scan_file(path: str) -> int:
+    with open(path, encoding="utf-8") as f:
+        for i, line in enumerate(f, 1):
+            if PICKLE_RE.match(line):
+                print(
+                    f"{path}:{i}: "
+                    "\033[91merror:\033[0m "  # red color
+                    "Found pickle/cloudpickle import"
+                )
+                return 1
+    return 0
+
+
+def main():
+    returncode = 0
+    for filename in sys.argv[1:]:
+        if filename in ALLOWED_FILES:
+            continue
+        returncode |= scan_file(filename)
+    return returncode
+
+
+def test_regex():
+    test_cases = [
+        # Should match
+        ("import pickle", True),
+        ("import cloudpickle", True),
+        ("import pickle as pkl", True),
+        ("import cloudpickle as cpkl", True),
+        ("from pickle import *", True),
+        ("from cloudpickle import dumps", True),
+        ("from pickle import dumps, loads", True),
+        ("from cloudpickle import (dumps, loads)", True),
+        ("    import pickle", True),
+        ("\timport cloudpickle", True),
+        ("from   pickle   import   loads", True),
+        # Should not match
+        ("import somethingelse", False),
+        ("from somethingelse import pickle", False),
+        ("# import pickle", False),
+        ("print('import pickle')", False),
+        ("import pickleas as asdf", False),
+    ]
+    for i, (line, should_match) in enumerate(test_cases):
+        result = bool(PICKLE_RE.match(line))
+        assert result == should_match, f"Test case {i} failed: '{line}' (expected {should_match}, got {result})"
+    print("All regex tests passed.")
+
+
+if __name__ == "__main__":
+    if "--test-regex" in sys.argv:
+        test_regex()
+    else:
+        sys.exit(main())

vllm_omni/distributed/omni_connectors/utils/serialization.py

@@ -1,79 +1,279 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import io
+from dataclasses import asdict, is_dataclass
 from typing import Any
 
-try:
-    import torch
+import msgspec
+import numpy as np
+import torch
+from msgspec import msgpack
+from PIL import Image
+from vllm.outputs import CompletionOutput, RequestOutput
 
-    _has_torch = True
-except ImportError:
-    _has_torch = False
+# Type markers for custom serialization
+_TENSOR_MARKER = "__tensor__"
+_NDARRAY_MARKER = "__ndarray__"
+_PIL_IMAGE_MARKER = "__pil_image__"
 
-try:
-    import cloudpickle
+# Keys that identify a RequestOutput dict (for reconstruction)
+_REQUEST_OUTPUT_KEYS = frozenset({"request_id", "prompt", "prompt_token_ids", "outputs", "finished"})
 
-    _has_cloudpickle = True
-except ImportError:
-    _has_cloudpickle = False
-
-from .logging import get_connector_logger
-
-logger = get_connector_logger(__name__)
+# Keys that identify a CompletionOutput dict (for reconstruction)
+_COMPLETION_OUTPUT_KEYS = frozenset({"index", "text", "token_ids", "finish_reason"})
 
 
-class OmniSerializer:
+class OmniMsgpackEncoder:
     """
-    Centralized serialization handler for OmniConnectors.
-    Supports multiple backends (torch, cloudpickle) to ensure consistency across connectors.
+    This implementation is adapted from vLLM’s MsgpackEncoder.
+    However, zero-copy support has not been implemented yet.
+    Handles torch.Tensor, numpy.ndarray, PIL.Image, RequestOutput and
+    CompletionOutput by converting them to serializable dict representations.
+    TODO: Enable zero-copy support.
     """
 
-    @staticmethod
-    def serialize(obj: Any, method: str = "cloudpickle") -> bytes:
+    def __init__(self):
+        self.encoder = msgpack.Encoder(enc_hook=self._enc_hook)
+
+    def encode(self, obj: Any) -> bytes:
+        """Encode an object to bytes."""
+        return self.encoder.encode(obj)
+
+    def _enc_hook(self, obj: Any) -> Any:
+        """Custom encoding hook for non-standard types."""
+        # torch.Tensor
+        if isinstance(obj, torch.Tensor):
+            return self._encode_tensor(obj)
+
+        # numpy.ndarray (exclude object/void dtypes)
+        if isinstance(obj, np.ndarray) and obj.dtype.kind not in ("O", "V"):
+            return self._encode_ndarray(obj)
+
+        # PIL.Image
+        if isinstance(obj, Image.Image):
+            return self._encode_pil_image(obj)
+
+        # RequestOutput (not a dataclass, needs special handling)
+        if isinstance(obj, RequestOutput):
+            return self._encode_request_output(obj)
+
+        # CompletionOutput (dataclass)
+        if isinstance(obj, CompletionOutput):
+            return self._encode_completion_output(obj)
+
+        # Other dataclasses
+        if is_dataclass(obj) and not isinstance(obj, type):
+            return asdict(obj)
+
+        # slice
+        if isinstance(obj, slice):
+            return (obj.start, obj.stop, obj.step)
+
+        raise TypeError(
+            f"Object of type {type(obj).__name__} is not serializable. "
+            "Supported types: torch.Tensor, np.ndarray, PIL.Image, dataclass, "
+            "RequestOutput, and standard Python types (dict, list, str, int, float, bool, None, bytes)."
+        )
+
+    def _encode_tensor(self, tensor: torch.Tensor) -> dict[str, Any]:
+        """Encode torch.Tensor to dict."""
+        t = tensor.detach().contiguous().cpu().view(torch.uint8)
+        return {
+            _TENSOR_MARKER: True,
+            "dtype": str(tensor.dtype).removeprefix("torch."),
+            "shape": list(tensor.shape),
+            "data": t.numpy().tobytes(),
+        }
+
+    def _encode_ndarray(self, arr: np.ndarray) -> dict[str, Any]:
+        """Encode numpy.ndarray to dict."""
+        if not arr.flags.c_contiguous:
+            arr = np.ascontiguousarray(arr)
+        return {
+            _NDARRAY_MARKER: True,
+            "dtype": arr.dtype.str,
+            "shape": list(arr.shape),
+            "data": arr.tobytes(),
+        }
+
+    def _encode_pil_image(self, img: Image.Image) -> dict[str, Any]:
+        """Encode PIL.Image to dict."""
+        arr = np.asarray(img, dtype=np.uint8)
+        if not arr.flags.c_contiguous:
+            arr = np.ascontiguousarray(arr)
+        return {
+            _PIL_IMAGE_MARKER: True,
+            "mode": img.mode,
+            "shape": list(arr.shape),
+            "data": arr.tobytes(),
+        }
+
+    def _encode_request_output(self, obj: RequestOutput) -> dict[str, Any]:
+        """Encode RequestOutput to dict.
+
+        RequestOutput is not a dataclass, so we manually extract its attributes.
+        Also handles dynamically added 'multimodal_output' attribute.
         """
-        Serialize an object to bytes.
+        # msgspec can serialize CompletionOutput dataclasses directly, but it
+        # drops dynamic fields such as multimodal_output. Encode them manually
+        # to preserve multimodal payloads across IPC.
+        encoded_outputs = []
+        for o in obj.outputs:
+            if isinstance(o, CompletionOutput):
+                encoded_outputs.append(self._encode_completion_output(o))
+            else:
+                encoded_outputs.append(o)
 
-        Args:
-            obj: The object to serialize.
-            method: Serialization method ("cloudpickle" or "torch").
-                    Defaults to "cloudpickle" (consistent with legacy stage_utils).
+        result = {
+            "request_id": obj.request_id,
+            "prompt": obj.prompt,
+            "prompt_token_ids": obj.prompt_token_ids,
+            "prompt_logprobs": obj.prompt_logprobs,
+            "outputs": encoded_outputs,
+            "finished": obj.finished,
+            "metrics": obj.metrics,
+            "lora_request": obj.lora_request,
+            "encoder_prompt": obj.encoder_prompt,
+            "encoder_prompt_token_ids": obj.encoder_prompt_token_ids,
+            "num_cached_tokens": obj.num_cached_tokens,
+            "multi_modal_placeholders": obj.multi_modal_placeholders,
+            "kv_transfer_params": obj.kv_transfer_params,
+        }
+        # Handle dynamically added multimodal_output attribute
+        mm_output = getattr(obj, "multimodal_output", None)
+        if mm_output is not None:
+            result["multimodal_output"] = mm_output
+        return result
+
+    def _encode_completion_output(self, obj: CompletionOutput) -> dict[str, Any]:
+        """Encode CompletionOutput to dict, preserving multimodal payloads."""
+        result = asdict(obj)
+        mm_output = getattr(obj, "multimodal_output", None)
+        if mm_output is not None:
+            result["multimodal_output"] = mm_output
+        return result
+
+
+class OmniMsgpackDecoder:
+    """
+    This implementation is adapted from vLLM’s MsgpackDecoder.
+    However, zero-copy support has not been implemented yet.
+
+    Automatically reconstructs torch.Tensor, numpy.ndarray, PIL.Image,
+    RequestOutput and CompletionOutput from their dict representations.
+    TODO: Enable zero-copy support.
+    """
+
+    def __init__(self):
+        self.decoder = msgpack.Decoder()
+
+    def decode(self, data: bytes | bytearray | memoryview) -> Any:
+        """Decode bytes to object."""
+        result = self.decoder.decode(data)
+        return self._post_process(result)
+
+    def _post_process(self, obj: Any) -> Any:
+        """Recursively restore tensor/ndarray/image/RequestOutput from their dict representations."""
+        if isinstance(obj, dict):
+            # Check for type markers first
+            if obj.get(_TENSOR_MARKER):
+                return self._decode_tensor(obj)
+            if obj.get(_NDARRAY_MARKER):
+                return self._decode_ndarray(obj)
+            if obj.get(_PIL_IMAGE_MARKER):
+                return self._decode_pil_image(obj)
+
+            # Process values recursively first
+            processed = {k: self._post_process(v) for k, v in obj.items()}
+
+            # Check if this looks like a RequestOutput
+            if _REQUEST_OUTPUT_KEYS.issubset(processed.keys()):
+                return self._decode_request_output(processed)
+
+            # Check if this looks like a CompletionOutput
+            if _COMPLETION_OUTPUT_KEYS.issubset(processed.keys()):
+                return self._decode_completion_output(processed)
+
+            return processed
+
+        if isinstance(obj, list):
+            return [self._post_process(item) for item in obj]
+
+        if isinstance(obj, tuple):
+            return tuple(self._post_process(item) for item in obj)
+
+        return obj
+
+    def _decode_tensor(self, obj: dict[str, Any]) -> torch.Tensor:
+        """Decode dict to torch.Tensor."""
+        dtype_str = obj["dtype"]
+        shape = obj["shape"]
+        data = obj["data"]
+
+        torch_dtype = getattr(torch, dtype_str)
+        if not data:
+            return torch.empty(shape, dtype=torch_dtype)
+
+        buffer = bytearray(data) if isinstance(data, (bytes, memoryview)) else data
+        arr = torch.frombuffer(buffer, dtype=torch.uint8)
+        return arr.view(torch_dtype).reshape(shape)
+
+    def _decode_ndarray(self, obj: dict[str, Any]) -> np.ndarray:
+        """Decode dict to numpy.ndarray."""
+        dtype = obj["dtype"]
+        shape = obj["shape"]
+        data = obj["data"]
+        return np.frombuffer(data, dtype=dtype).reshape(shape)
+
+    def _decode_pil_image(self, obj: dict[str, Any]) -> Image.Image:
+        """Decode dict to PIL.Image."""
+        mode = obj["mode"]
+        shape = obj["shape"]
+        data = obj["data"]
+        arr = np.frombuffer(data, dtype=np.uint8).reshape(shape)
+        return Image.fromarray(arr, mode=mode)
+
+    def _decode_completion_output(self, obj: dict[str, Any]) -> CompletionOutput:
+        """Decode dict to CompletionOutput using msgspec.convert."""
+        mm_output = obj.pop("multimodal_output", None)
+        co = msgspec.convert(obj, CompletionOutput)
+        if mm_output is not None:
+            setattr(co, "multimodal_output", mm_output)
+        return co
+
+    def _decode_request_output(self, obj: dict[str, Any]) -> RequestOutput:
+        """Decode dict to RequestOutput.
+
+        RequestOutput is not a dataclass, so msgspec.convert doesn't work.
+        We construct it manually, passing all known fields via **kwargs.
         """
-        if method == "torch" and _has_torch:
-            buf = io.BytesIO()
-            torch.save(obj, buf)
-            return buf.getvalue()
-        elif method == "cloudpickle" or (method == "torch" and not _has_torch):
-            if not _has_cloudpickle:
-                # Fallback to standard pickle if cloudpickle is missing
-                import pickle
+        # Extract multimodal_output before constructing (it's dynamically added)
+        mm_output = obj.pop("multimodal_output", None)
 
-                return pickle.dumps(obj)
-            return cloudpickle.dumps(obj)
-        else:
-            # Default fallback
-            import pickle
+        # RequestOutput.__init__ accepts **kwargs for forward compatibility
+        ro = RequestOutput(**obj)
 
-            return pickle.dumps(obj)
+        # Restore dynamically added multimodal_output attribute
+        if mm_output is not None:
+            setattr(ro, "multimodal_output", mm_output)
+        return ro
 
-    @staticmethod
-    def deserialize(data: bytes, method: str = "cloudpickle") -> Any:
-        """
-        Deserialize bytes to an object.
 
-        Args:
-            data: The bytes to deserialize.
-            method: Serialization method used ("cloudpickle" or "torch").
-        """
-        if method == "torch" and _has_torch:
-            return torch.load(io.BytesIO(data), map_location="cpu", weights_only=False)
-        elif method == "cloudpickle" or (method == "torch" and not _has_torch):
-            if not _has_cloudpickle:
-                import pickle
+class OmniSerde:
+    """Serialization/deserialization handler for Omni IPC."""
 
-                return pickle.loads(data)
-            return cloudpickle.loads(data)
-        else:
-            import pickle
+    def __init__(self):
+        self.encoder = OmniMsgpackEncoder()
+        self.decoder = OmniMsgpackDecoder()
 
-            return pickle.loads(data)
+    def serialize(self, obj: Any) -> bytes:
+        """Serialize an object to bytes."""
+        return self.encoder.encode(obj)
+
+    def deserialize(self, data: bytes | bytearray | memoryview) -> Any:
+        """Deserialize bytes to an object."""
+        return self.decoder.decode(data)
+
+
+# Global instance for simple interface
+OmniSerializer = OmniSerde()