From 48c9d96054796adf6cf9ee0c731b6112d19449b7 Mon Sep 17 00:00:00 2001 From: Robert Washbourne Date: Wed, 21 Oct 2026 09:43:57 +0000 Subject: [PATCH 1/3] perf: coalesce row payload frames with the packed buffer codec Signed-off-by: Robert Washbourne --- README.md | 2 + tests/test_serial_utils_on_cpu.py | 59 +++++++++++++++++++++++++++++ transfer_queue/utils/zmq_utils.py | 63 ++++++++++++++++++++++++++++++- 3 files changed, 122 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 73f5ce48..4609840f 100644 --- a/README.md +++ b/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 | diff --git a/tests/test_serial_utils_on_cpu.py b/tests/test_serial_utils_on_cpu.py index 3e2f46cf..ceda1529 100644 --- a/tests/test_serial_utils_on_cpu.py +++ b/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", [ diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 2ca63d7f..4c6c3aab 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -21,6 +21,7 @@ from typing import Any, Callable, TypeAlias from uuid import uuid4 +import numpy as np import psutil import ray import zmq @@ -28,13 +29,67 @@ 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, flags, pending = [], [], [] + pending_bytes = 4 + + def flush(): + 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): """ @@ -212,7 +267,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 +286,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 From 3d73da1e38cd89e097e83df3896f03042b73eb31 Mon Sep 17 00:00:00 2001 From: Robert Washbourne Date: Wed, 21 Oct 2026 09:49:11 +0000 Subject: [PATCH 2/3] ci: run existing tests for release maintenance PRs Signed-off-by: Robert Washbourne --- .github/workflows/run-tests.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/run-tests.yml b/.github/workflows/run-tests.yml index 05bba3ae..96150cb9 100644 --- a/.github/workflows/run-tests.yml +++ b/.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 \ No newline at end of file + TQ_TEST_BACKEND=Yuanrong pytest tests/e2e/test_kv_interface_e2e.py From 51b3e058dca601842ccbd6887d008d1539c49330 Mon Sep 17 00:00:00 2001 From: Robert Washbourne Date: Wed, 21 Oct 2026 09:58:42 +0000 Subject: [PATCH 3/3] fix: type packed frame bookkeeping for static checks Signed-off-by: Robert Washbourne --- transfer_queue/utils/zmq_utils.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 4c6c3aab..aef641ee 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -44,10 +44,12 @@ def _coalesce_frames(frames: list) -> list: """Group small row payloads using the existing packed-buffer codec.""" if len(frames) <= 16: return frames - output, flags, pending = [], [], [] + output: list[bytestr] = [] + flags: list[int] = [] + pending: list[bytestr] = [] pending_bytes = 4 - def flush(): + def flush() -> None: nonlocal pending, pending_bytes if len(pending) > 1: buffer = np.empty(calc_packed_size(pending), dtype=np.uint8) @@ -258,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,