perf: coalesce row payload frames with the existing packed codec

.github/workflows/run-tests.yml

@@ -8,10 +8,12 @@ on:
     branches:
       - main
       - v0.*
+      - release/v0.*
   pull_request:
     branches:
       - main
       - v0.*
+      - release/v0.*
 
 jobs:
   build:
@@ -48,4 +50,4 @@ jobs:
     - name: Run Yuanrong Backend Specific E2E Tests
       run: |
         TQ_TEST_BACKEND=Yuanrong pytest tests/e2e/test_e2e_lifecycle_consistency.py
-        TQ_TEST_BACKEND=Yuanrong pytest tests/e2e/test_kv_interface_e2e.py
+        TQ_TEST_BACKEND=Yuanrong pytest tests/e2e/test_kv_interface_e2e.py

README.md

@@ -98,6 +98,8 @@ This data structure design is motivated by the computational characteristics of
 
 The wire decoder rejects empty tensor or NumPy payload frames whose declared shape has nonzero elements. Legitimate zero-element shapes and NumPy scalar pickle serialization remain supported; a missing payload must never become uninitialized sample data.
 
+Large multipart messages coalesce small row frames into packed buffers of at most 64 MiB, using the same tensor serialization and decoded row objects. Large individual frames remain copy-free. Clients and services must use the same release when adopting this transport framing; the decoder continues to accept earlier multipart messages.
+
 ### User Interface: High-Level & Low-Level APIs
 
 | Level | Tier | Style | Fine-Grained Access | Streaming | Sampler | Multiple-Backends | 

tests/test_serial_utils_on_cpu.py

@@ -97,6 +97,65 @@ def test_zmq_msg_serialization():
         )
 
 
+@pytest.mark.parametrize(
+    "dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64, torch.int64, torch.bool]
+)
+def test_many_tensor_rows_cross_real_zmq_without_changing_bytes(dtype):
+    from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType
+
+    rows = [torch.arange((index % 6) * 4).reshape(index % 6, 4).t()[::2].to(dtype) for index in range(48)]
+    if dtype.is_floating_point:
+        rows[0] = torch.tensor([0.0, -0.0, float("nan"), float("inf"), -float("inf")], dtype=dtype)
+    message = ZMQMessage.create(ZMQRequestType.PUT_DATA, "sender", {"data": rows})
+    frames = message.serialize()
+    assert len(frames) < len(rows)
+    context = zmq.Context()
+    sender = context.socket(zmq.PAIR)
+    receiver = context.socket(zmq.PAIR)
+    try:
+        sender.bind("inproc://row-frames")
+        receiver.connect("inproc://row-frames")
+        sender.send_multipart(frames, copy=False)
+        decoded = ZMQMessage.deserialize(receiver.recv_multipart(copy=False))
+    finally:
+        sender.close()
+        receiver.close()
+        context.term()
+    assert isinstance(decoded.body["data"], list)
+    for expected, actual in zip(rows, decoded.body["data"], strict=True):
+        assert actual.shape == expected.shape
+        assert actual.dtype == expected.dtype
+        assert torch.equal(
+            actual.contiguous().reshape(-1).view(torch.uint8), expected.contiguous().reshape(-1).view(torch.uint8)
+        )
+
+    # Persisted or earlier-release multipart payloads keep the existing decoder.
+    original = encode({**vars(message), "request_type": message.request_type.value})
+    old_decoded = ZMQMessage.deserialize(original)
+    assert len(old_decoded.body["data"]) == len(rows)
+
+
+def test_coalesced_message_rejects_missing_and_invalid_layout():
+    from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQMessageDecodeError, ZMQRequestType
+
+    message = ZMQMessage.create(ZMQRequestType.PUT_DATA, "sender", {"data": [torch.ones(3) for _ in range(32)]})
+    frames = message.serialize()
+    for malformed in (frames[:2], [frames[0], b"\x02", *frames[2:]]):
+        with pytest.raises(ZMQMessageDecodeError, match="coalesced"):
+            ZMQMessage.deserialize(malformed)
+
+
+def test_coalesced_message_keeps_a_large_individual_frame_and_empty_rows():
+    from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType
+
+    large = torch.full((64 * 1024 * 1024 + 1,), 173, dtype=torch.uint8)
+    rows = [torch.empty(0, dtype=torch.int64) if index % 2 else torch.arange(index) for index in range(48)]
+    message = ZMQMessage.create(ZMQRequestType.GET_DATA_RESPONSE, "sender", {"large": large, "rows": rows})
+    decoded = ZMQMessage.deserialize(message.serialize())
+    assert torch.equal(decoded.body["large"], large)
+    assert all(torch.equal(actual, expected) for actual, expected in zip(decoded.body["rows"], rows, strict=True))
+
+
 @pytest.mark.parametrize(
     "make_view",
     [

transfer_queue/utils/zmq_utils.py

@@ -21,20 +21,77 @@
 from typing import Any, Callable, TypeAlias
 from uuid import uuid4
 
+import numpy as np
 import psutil
 import ray
 import zmq
 import zmq.asyncio
 
 from transfer_queue.utils.enum_utils import ExplicitEnum, Role
 from transfer_queue.utils.logging_utils import get_logger
-from transfer_queue.utils.serial_utils import decode, encode
+from transfer_queue.utils.serial_utils import calc_packed_size, decode, encode, pack_into, unpack_from
 
 logger = get_logger(__name__)
 
 
 bytestr: TypeAlias = bytes | bytearray | memoryview
 
+_COALESCED_FRAMES = b"\xc1TQPACK1"
+_PACKED_FRAME_BYTES = 64 * 1024 * 1024
+
+
+def _coalesce_frames(frames: list) -> list:
+    """Group small row payloads using the existing packed-buffer codec."""
+    if len(frames) <= 16:
+        return frames
+    output: list[bytestr] = []
+    flags: list[int] = []
+    pending: list[bytestr] = []
+    pending_bytes = 4
+
+    def flush() -> None:
+        nonlocal pending, pending_bytes
+        if len(pending) > 1:
+            buffer = np.empty(calc_packed_size(pending), dtype=np.uint8)
+            pack_into(buffer, pending)
+            output.append(memoryview(buffer))
+            flags.append(1)
+        elif pending:
+            output.append(pending[0])
+            flags.append(0)
+        pending, pending_bytes = [], 4
+
+    for frame in frames:
+        size = memoryview(frame).nbytes
+        if pending_bytes + 8 + size > _PACKED_FRAME_BYTES:
+            flush()
+        if size > _PACKED_FRAME_BYTES:
+            output.append(frame)
+            flags.append(0)
+        else:
+            pending.append(frame)
+            pending_bytes += 8 + size
+    flush()
+    if len(output) == len(frames):
+        return frames
+    return [_COALESCED_FRAMES, bytes(flags), *output]
+
+
+def _expand_frames(frames: list) -> list:
+    marker = memoryview(frames[0])
+    if marker.nbytes != len(_COALESCED_FRAMES) or marker.tobytes() != _COALESCED_FRAMES:
+        return frames
+    if len(frames) < 3:
+        raise ValueError("coalesced message is missing its layout or payload")
+    flags = memoryview(frames[1]).cast("B")
+    if len(flags) != len(frames) - 2 or any(flag not in (0, 1) for flag in flags):
+        raise ValueError("invalid coalesced frame layout")
+    return [
+        part
+        for flag, frame in zip(flags, frames[2:], strict=True)
+        for part in (unpack_from(frame) if flag else [frame])
+    ]
+
 
 class ZMQRequestType(ExplicitEnum):
     """
@@ -203,7 +260,7 @@ def create(
         )
 
     def serialize(self) -> list:
-        """Serialize using zero-copy msgpack; falls back to pickle for unsupported types."""
+        """Encode msgpack or pickle, then coalesce fragmented tensor payloads."""
         msg_dict = {
             "request_type": self.request_type.value,  # Enum -> str for msgpack
             "sender_id": self.sender_id,
@@ -212,7 +269,10 @@ def serialize(self) -> list:
             "timestamp": self.timestamp,
             "body": self.body,
         }
-        return encode(msg_dict)
+        # Per-row frames make large jagged fields expensive to send and leave
+        # thousands of receive allocations behind after the global field copy.
+        # Coalescing retains the existing encoded bytes and decoded objects.
+        return _coalesce_frames(encode(msg_dict))
 
     @classmethod
     def deserialize(cls, frames: list) -> "ZMQMessage":
@@ -228,6 +288,7 @@ def deserialize(cls, frames: list) -> "ZMQMessage":
             raise ZMQMessageDecodeError(f"leading frame is empty; {describe_frames(frames)}")
 
         try:
+            frames = _expand_frames(frames)
             result = decode(frames)
         except Exception as e:
             raise ZMQMessageDecodeError(f"{type(e).__name__}: {e}; {describe_frames(frames)}") from e