From 97a58606816576a46158cd8c4d3442ffb2944f97 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Mon, 22 Dec 2025 08:21:16 +0800 Subject: [PATCH] [Security][Feature] Use Msgpack in OmniSerializer instead of pickle (#310) Signed-off-by: gcanlin --- .pre-commit-config.yaml | 7 + .../omni_connectors/test_basic_connectors.py | 30 +- tests/entrypoints/test_omni_llm.py | 9 +- tools/pre_commit/check_pickle_imports.py | 84 +++++ .../omni_connectors/utils/serialization.py | 314 ++++++++++++++---- 5 files changed, 380 insertions(+), 64 deletions(-) create mode 100644 tools/pre_commit/check_pickle_imports.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 643cf7a46..87117f1da 100644 --- a/.pre-commit-config.yaml +++ b/.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] diff --git a/tests/distributed/omni_connectors/test_basic_connectors.py b/tests/distributed/omni_connectors/test_basic_connectors.py index 18a3262af..a7b4b2d01 100644 --- a/tests/distributed/omni_connectors/test_basic_connectors.py +++ b/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}) diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py index a3cb236de..1b36ed509 100644 --- a/tests/entrypoints/test_omni_llm.py +++ b/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) diff --git a/tools/pre_commit/check_pickle_imports.py b/tools/pre_commit/check_pickle_imports.py new file mode 100644 index 000000000..ab7162400 --- /dev/null +++ b/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()) diff --git a/vllm_omni/distributed/omni_connectors/utils/serialization.py b/vllm_omni/distributed/omni_connectors/utils/serialization.py index a30b20508..c1055e5eb 100644 --- a/vllm_omni/distributed/omni_connectors/utils/serialization.py +++ b/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()