[Core] Optimize mixed stage payloads and scalar restoration

Signed-off-by: Sy03 <1370724210@qq.com>

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))

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)

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
 

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("<I")
 _ENCODER = msgspec.msgpack.Encoder()
@@ -65,12 +69,20 @@ def prepare_tensor_frame(payload: object, *, max_bytes: int | None = None) -> 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: