[Core] Preserve native scalar encoding in stage tensor frames

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

tests/distributed/omni_connectors/test_shm_ring.py

@@ -10,6 +10,7 @@ import uuid
 from dataclasses import asdict
 from multiprocessing import shared_memory
 
+import msgspec
 import pytest
 import torch
 from vllm.outputs import CompletionOutput, RequestOutput
@@ -200,11 +201,34 @@ def test_disabled_ring_and_unsupported_tree_keep_legacy_serializer(edge):
         _, _, metadata = legacy.put("0", "1", "legacy", {"codes": torch.arange(3)})
         assert "shm" in metadata
         assert receiver.get("0", "1", "legacy", metadata)[0]["codes"].tolist() == [0, 1, 2]
-    assert prepare_tensor_frame({"codes": torch.arange(3), "opaque": {1, 2}}) is None
+    assert prepare_tensor_frame({"codes": torch.arange(3), "opaque": object()}) is None
     assert sender.put("0", "1", "falsey", False)[0]
     assert receiver.get("0", "1", "falsey")[0] is False
 
 
+def test_tensor_free_lists_reuse_native_serialization():
+    payload = {"prompt_token_ids": list(range(8192)), "condition": [0.25] * 8192}
+    with SharedMemoryConnector({}) as connector:
+        frame = prepare_tensor_frame(payload)
+        assert isinstance(frame, bytes)
+        assert frame == connector.serialize_obj(payload)
+        assert connector.put("0", "1", "native-lists", payload)[0]
+        assert connector.get("0", "1", "native-lists")[0] == payload
+
+
+def test_scalar_lists_and_literal_extensions_preserve_wire_semantics():
+    literal = msgspec.msgpack.Ext(42, bytes(20))
+    payload = {"condition": [0.25] * 8192, "codes": torch.arange(28), "extension": literal}
+    with SharedMemoryConnector({}) as connector:
+        _, _, metadata = connector.put("0", "1", "mixed-lists", payload)
+        assert "host_ring" in metadata
+        result = connector.get("0", "1", "mixed-lists")[0]
+        assert result["condition"] == payload["condition"]
+        assert torch.equal(result["codes"], payload["codes"])
+        assert result["extension"].code == literal.code
+        assert result["extension"].data == literal.data
+
+
 def _locked_writer(name, ready, release):
     fd = os.open(f"/dev/shm/{name}", os.O_RDWR)
     try:

vllm_omni/distributed/omni_connectors/utils/tensor_frame.py

@@ -2,7 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Owned CPU tensor trees with a msgpack header and raw tensor regions.
 
-The format uses tagged nodes rather than reserved dictionary keys. It never
+Tensor references use message-scoped msgpack extensions, preserving literal
+extensions and native scalar containers without reserved dictionary keys. It never
 exposes a view of transport storage to a caller. Unsupported objects retain
 the ordinary OmniSerializer path.
 """
@@ -12,6 +13,7 @@ from __future__ import annotations
 import ctypes
 import math
 import struct
+import uuid
 from dataclasses import dataclass
 
 import msgspec
@@ -19,8 +21,9 @@ import torch
 
 _LENGTH = struct.Struct("<I")
 _ENCODER = msgspec.msgpack.Encoder()
+_REFERENCE = struct.Struct("<16sI")
+_REFERENCE_CODE = 42
 _DECODER = msgspec.msgpack.Decoder()
-_SCALAR, _DICT, _LIST, _TENSOR = range(4)
 
 
 def _align(size: int) -> int:
@@ -51,61 +54,47 @@ class TensorFrame:
                 ctypes.memmove(base + offset, tensor.data_ptr(), tensor.nbytes)
 
 
-def prepare_tensor_frame(payload: object, *, max_bytes: int | None = None) -> TensorFrame | None:
-    """Prepare a plain tree containing dense CPU tensors; otherwise return None.
+def prepare_tensor_frame(payload: object, *, max_bytes: int | None = None) -> TensorFrame | bytes | None:
+    """Encode native containers once, or prepare their dense CPU tensor regions.
 
-    Retaining tensors is enough because the caller writes before put() returns.
-    This is not an asynchronous snapshot or a CUDA source ownership protocol.
+    Scalar/token lists stay in msgpack's C encoder. Tensor-free messages reuse
+    those encoded bytes, rather than constructing and discarding a tagged tree.
+    Sources remain caller-owned until the synchronous put() returns.
     """
     tensors: list[torch.Tensor] = []
     offsets: list[int] = []
     descriptors: list[tuple[str, tuple[int, ...], int]] = []
     data_size = 0
+    nonce = uuid.uuid4().bytes
 
-    def visit(value: object, depth: int) -> list:
+    def encode_tensor(value: object) -> msgspec.msgpack.Ext:
         nonlocal data_size
-        if depth > 64:
+        if type(value) is not torch.Tensor:
             raise _UnsupportedTreeError
-        if type(value) in (type(None), bool, int, float, str, bytes):
-            return [_SCALAR, value]
-        if type(value) is torch.Tensor:
-            if value.device.type != "cpu" or value.layout != torch.strided or value.is_quantized:
-                raise _UnsupportedTreeError
-            tensor = value
-            offset = _align(data_size)
-            index = len(tensors)
-            tensors.append(tensor)
-            offsets.append(offset)
-            descriptors.append((str(tensor.dtype).removeprefix("torch."), tuple(tensor.shape), offset))
-            data_size = offset + tensor.nbytes
-            if max_bytes is not None and data_size > max_bytes:
-                raise _UnsupportedTreeError
-            return [_TENSOR, index]
-        if type(value) is dict and all(type(key) is str for key in value):
-            return [_DICT, {key: visit(item, depth + 1) for key, item in value.items()}]
-        if isinstance(value, (list, tuple)) and type(value) in (list, tuple):
-            # msgpack represents both as arrays; match OmniSerializer's wire type.
-            return [_LIST, [visit(item, depth + 1) for item in value]]
-        if isinstance(value, msgspec.Struct):
-            try:
-                fields = msgspec.to_builtins(value, builtin_types=(torch.Tensor, bytes))
-            except TypeError:
-                raise _UnsupportedTreeError from None
-            return visit(fields, depth + 1)
-        raise _UnsupportedTreeError
+        if value.device.type != "cpu" or value.layout != torch.strided or value.is_quantized:
+            raise _UnsupportedTreeError
+        offset = _align(data_size)
+        index = len(tensors)
+        tensors.append(value)
+        offsets.append(offset)
+        descriptors.append((str(value.dtype).removeprefix("torch."), tuple(value.shape), offset))
+        data_size = offset + value.nbytes
+        if max_bytes is not None and data_size > max_bytes:
+            raise _UnsupportedTreeError
+        return msgspec.msgpack.Ext(_REFERENCE_CODE, _REFERENCE.pack(nonce, index))
 
     try:
-        tree = visit(payload, 0)
+        tree = msgspec.msgpack.Encoder(enc_hook=encode_tensor).encode(payload)
         if not tensors:
-            return None
-        header = _ENCODER.encode([descriptors, tree])
-    except (_UnsupportedTreeError, OverflowError):
+            return tree
+        header = _ENCODER.encode([descriptors, tree, nonce])
+    except (_UnsupportedTreeError, TypeError, OverflowError, RecursionError):
         return None
     size = _align(4 + len(header)) + data_size
     if max_bytes is not None and size > max_bytes:
         return None
-    # Only materialize non-contiguous/flagged storage after the entire tree
-    # is known to fit. Large and unsupported payloads avoid an unused copy.
+    # Only materialize non-contiguous/flagged storage after the entire message
+    # is known to fit. Unsupported and large frames avoid an unused copy.
     for index, tensor in enumerate(tensors):
         if tensor.is_conj():
             tensor = tensor.resolve_conj()
@@ -127,7 +116,7 @@ def read_tensor_frame(buffer: memoryview) -> object:
         raise ValueError("tensor frame header exceeds payload")
     header_view = buffer[4 : 4 + header_size]
     try:
-        descriptors, tree = _DECODER.decode(header_view)
+        descriptors, tree, nonce = _DECODER.decode(header_view)
     finally:
         header_view.release()
     # All leaves share one independently owned allocation. This amortizes CPU
@@ -159,16 +148,14 @@ def read_tensor_frame(buffer: memoryview) -> object:
                 ctypes.memmove(tensor.data_ptr(), ctypes.addressof(anchor) + start, size)
         tensors.append(tensor)
 
-    def restore(node: list) -> object:
-        tag, value = node
-        if tag == _SCALAR:
-            return value
-        if tag == _TENSOR:
-            return tensors[value]
-        if tag == _DICT:
-            return {key: restore(item) for key, item in value.items()}
-        if tag == _LIST:
-            return [restore(item) for item in value]
-        raise ValueError("unknown tensor frame node")
+    def restore_reference(code: int, data: memoryview) -> object:
+        if code == _REFERENCE_CODE and len(data) == _REFERENCE.size:
+            reference_nonce, index = _REFERENCE.unpack(data)
+            if reference_nonce == nonce:
+                if index >= len(tensors):
+                    raise ValueError("invalid tensor frame reference")
+                return tensors[index]
+        # Literal user extensions keep their existing msgpack wire semantics.
+        return msgspec.msgpack.Ext(code, bytes(data))
 
-    return restore(tree)
+    return msgspec.msgpack.decode(tree, ext_hook=restore_reference)