diff --git a/tests/distributed/omni_connectors/test_shm_ring.py b/tests/distributed/omni_connectors/test_shm_ring.py index 2a7fc58a3..c2ecc3dec 100644 --- a/tests/distributed/omni_connectors/test_shm_ring.py +++ b/tests/distributed/omni_connectors/test_shm_ring.py @@ -11,12 +11,15 @@ from dataclasses import asdict from multiprocessing import shared_memory import msgspec +import numpy as np import pytest import torch +from PIL import Image from vllm.outputs import CompletionOutput, RequestOutput from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector +from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer from vllm_omni.distributed.omni_connectors.utils.tensor_frame import prepare_tensor_frame pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -107,6 +110,49 @@ def test_source_can_be_overwritten_as_soon_as_put_returns(edge): assert receiver.get("0", "1", "source-reuse")[0].tolist() == list(range(56)) +@pytest.mark.parametrize("dtype", [np.bool_, np.int32, np.float16, np.float64, np.complex64]) +def test_mixed_numpy_image_and_tensor_keep_shared_wire_semantics(edge, dtype): + sender, receiver = edge + array = np.arange(24).astype(dtype).reshape(4, 6)[::-1, ::2] + image = Image.fromarray(np.arange(18, dtype=np.uint8).reshape(2, 3, 3)) + payload = {"codes": torch.arange(28), "array": array, "image": image, "slice": slice(1, 7, 2)} + _, _, metadata = sender.put("0", "1", "mixed-native", payload) + assert "host_ring" in metadata + output = receiver.get("0", "1", "mixed-native", metadata)[0] + assert np.array_equal(output["array"], array) + assert output["array"].dtype == array.dtype + assert not output["array"].flags.writeable + assert output["image"].mode == image.mode and output["image"].tobytes() == image.tobytes() + assert output["slice"] == [1, 7, 2] + assert torch.equal(output["codes"], payload["codes"]) + array.fill(0) + image.paste(0, (0, 0, 3, 2)) + assert np.count_nonzero(output["array"]) and any(output["image"].tobytes()) + + +def test_large_array_rejects_unused_framing_snapshot(monkeypatch): + def unexpected_snapshot(value): + pytest.fail("a too-large array must not be snapshotted while testing ring admission") + + monkeypatch.setattr(OmniSerializer.encoder, "_enc_hook", unexpected_snapshot) + assert prepare_tensor_frame({"array": np.zeros(1024)}, max_bytes=1024) is None + + +def test_native_scalar_restore_preserves_nested_markers_and_input_containers(): + array = np.arange(3) + image = Image.fromarray(np.array([[1, 2]], dtype=np.uint8)) + tensor = torch.arange(2) + payload = {"ids": list(range(8192)), "mixed": [1.5, None, (array, {"image": image, "tensor": tensor})]} + decoded = msgspec.msgpack.decode(OmniSerializer.serialize(payload)) + output = OmniSerializer.restore(decoded) + assert output is not decoded and output["ids"] is not decoded["ids"] + assert output["ids"] == payload["ids"] + assert np.array_equal(output["mixed"][2][0], array) + assert output["mixed"][2][1]["image"].tobytes() == image.tobytes() + assert torch.equal(output["mixed"][2][1]["tensor"], tensor) + assert isinstance(decoded["mixed"][2][0], dict) + + def test_tensor_frames_preserve_shared_output_reconstruction(edge): sender, receiver = edge completion = asdict(CompletionOutput(index=0, text="", token_ids=[1], cumulative_logprob=0.0, logprobs=None)) diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 769643070..9fc3cc42d 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -210,7 +210,9 @@ class SharedMemoryConnector(OmniConnectorBase): data: Any, ) -> tuple[bool, int, dict[str, Any] | None]: try: - self.reap_consumed() + # Ring put() reclaims its own edge under the publication lock. + # Keep the per-key sweep here without taking every ring lock twice. + self._reap_segments() use_ring = self._host_ring.capacity and str(from_stage).isdigit() and str(to_stage).isdigit() frame = ( prepare_tensor_frame(data, max_bytes=self._host_ring.capacity // 4 - 16 - len(put_key.encode())) @@ -414,6 +416,9 @@ class SharedMemoryConnector(OmniConnectorBase): def reap_consumed(self) -> None: """Bounded round-robin sweep; receivers unlink SHM in another process.""" self._host_ring.reap() + self._reap_segments() + + def _reap_segments(self) -> None: with self._pending_keys_lock: for _ in range(min(64, len(self._pending_keys))): key, _ = self._pending_keys.popitem(last=False) diff --git a/vllm_omni/distributed/omni_connectors/utils/serialization.py b/vllm_omni/distributed/omni_connectors/utils/serialization.py index dd5c0e5c9..4b06fa1f3 100644 --- a/vllm_omni/distributed/omni_connectors/utils/serialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/serialization.py @@ -205,7 +205,7 @@ class OmniMsgpackDecoder: return self._decode_pil_image(obj) # Process values recursively first - processed = {k: self._post_process(v) for k, v in obj.items()} + processed = {k: self._post_process(v) if isinstance(v, (dict, list, tuple)) else v for k, v in obj.items()} # Check if this looks like an OmniRequestOutput (check before RequestOutput # since OmniRequestOutput may also have some RequestOutput-like fields) @@ -223,10 +223,13 @@ class OmniMsgpackDecoder: return processed if isinstance(obj, list): - return [self._post_process(item) for item in obj] + # Native msgpack already owns scalar values. Only containers can + # contain wire markers; avoid a recursive call for every token or + # float in long model payloads, while preserving new-list semantics. + return [self._post_process(item) if isinstance(item, (dict, list, tuple)) else item for item in obj] if isinstance(obj, tuple): - return tuple(self._post_process(item) for item in obj) + return tuple(self._post_process(item) if isinstance(item, (dict, list, tuple)) else item for item in obj) return obj diff --git a/vllm_omni/distributed/omni_connectors/utils/tensor_frame.py b/vllm_omni/distributed/omni_connectors/utils/tensor_frame.py index 42868907b..553df9ba9 100644 --- a/vllm_omni/distributed/omni_connectors/utils/tensor_frame.py +++ b/vllm_omni/distributed/omni_connectors/utils/tensor_frame.py @@ -17,7 +17,11 @@ import uuid from dataclasses import dataclass import msgspec +import numpy as np import torch +from PIL import Image + +from .serialization import OmniSerializer _LENGTH = struct.Struct(" Te offsets: list[int] = [] descriptors: list[tuple[str, tuple[int, ...], int]] = [] data_size = 0 - nonce = uuid.uuid4().bytes + nonce: bytes | None = None - def encode_tensor(value: object) -> msgspec.msgpack.Ext: - nonlocal data_size + def encode_tensor(value: object) -> object: + nonlocal data_size, nonce if type(value) is not torch.Tensor: - raise _UnsupportedTreeError + # Reuse the common wire representations for mixed model payloads. + # Large arrays/images keep the ordinary path without first making + # an extra byte snapshot just to discover that the frame cannot fit. + if max_bytes is not None: + if isinstance(value, np.ndarray) and value.nbytes > max_bytes: + raise _UnsupportedTreeError + if isinstance(value, Image.Image) and value.width * value.height * len(value.getbands()) > max_bytes: + raise _UnsupportedTreeError + return OmniSerializer.encoder._enc_hook(value) if value.device.type != "cpu" or value.layout != torch.strided or value.is_quantized: raise _UnsupportedTreeError offset = _align(data_size) @@ -81,6 +93,8 @@ def prepare_tensor_frame(payload: object, *, max_bytes: int | None = None) -> Te data_size = offset + value.nbytes if max_bytes is not None and data_size > max_bytes: raise _UnsupportedTreeError + if nonce is None: + nonce = uuid.uuid4().bytes return msgspec.msgpack.Ext(_REFERENCE_CODE, _REFERENCE.pack(nonce, index)) try: