diff --git a/deepinfer/distributed/shm_broadcast.py b/deepinfer/distributed/shm_broadcast.py index 4852391..75de71f 100644 --- a/deepinfer/distributed/shm_broadcast.py +++ b/deepinfer/distributed/shm_broadcast.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Derived from vLLM; modified by DeepSeek. +import io import os import pickle import time @@ -389,7 +390,12 @@ class MessageQueue: def enqueue(self, obj: Any, timeout: float | None = None) -> None: """Write to message queue with optional timeout (in seconds)""" assert self._is_writer, "Only writers can enqueue" - serialized_obj = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) + # Pickle is required here to transport arbitrary Python/torch objects + # (tensor metadata, dataclasses, etc.) between the trusted worker + # processes of the same distributed job over a local shared-memory + # ring buffer / loopback socket; the payload never comes from an + # external or untrusted source, so unpickling it is safe here. + serialized_obj = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) # nosec B301 if self.n_local_reader > 0: # first byte of the chunk is the overflow flag; the payload is stored # inline if it fits, otherwise it is sent through the local socket