[fix] Make the controller survive bad ZMQ messages and repair the pickle fallback

.github/workflows/python-package.yml

@@ -33,7 +33,9 @@ jobs:
         python -m pip install --upgrade pip
         pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
         pip install -e ".[test,build,yuanrong]"
-        pip install mooncake-transfer-engine-non-cuda
+        # Pin below 0.3.12: the 0.3.12 non-cuda wheel ships a mooncake_master binary
+        # linked against libcudart.so.12, which fails to start on runners without CUDA.
+        pip install "mooncake-transfer-engine-non-cuda<0.3.12"
     - name: Lint with flake8
       run: |
         # stop the build if there are Python syntax errors or undefined names

.github/workflows/run-tests.yml

@@ -33,7 +33,9 @@ jobs:
         python -m pip install --upgrade pip
         pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
         pip install -e ".[test,build,yuanrong]"
-        pip install mooncake-transfer-engine-non-cuda
+        # Pin below 0.3.12: the 0.3.12 non-cuda wheel ships a mooncake_master binary
+        # linked against libcudart.so.12, which fails to start on runners without CUDA.
+        pip install "mooncake-transfer-engine-non-cuda<0.3.12"
     - name: Run Tests
       run: |
         pytest tests

.pre-commit-config.yaml

@@ -10,3 +10,8 @@ repos:
     rev: 'v1.17.0'
     hooks:
       - id: mypy
+        # Check the whole package: with `follow_imports = "skip"` in pyproject.toml,
+        # passing only changed files makes types from other modules (e.g. ExplicitEnum)
+        # unresolved and produces false errors.
+        pass_filenames: false
+        args: [transfer_queue]

tests/test_controller.py

@@ -14,12 +14,15 @@
 # limitations under the License.
 
 import logging
+from uuid import uuid4
 
 import pytest
 import ray
 import torch
+import zmq
 
 from transfer_queue.controller import TransferQueueController
+from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -1263,3 +1266,126 @@ def test_controller_checkpoint_load_nonexistent_file(self, ray_setup, tmp_path):
             ray.get(tq_controller.load_checkpoint.remote(missing))
 
         print("✓ load_checkpoint raises on missing file")
+
+
+class TestTransferQueueControllerBadRequests:
+    """The request loop must survive requests it cannot decode, does not handle, or fails on.
+
+    Any of these used to propagate out of ``_process_request``, killing
+    ``TransferQueueControllerProcessRequestThread``. The controller then stopped answering
+    anything at all for the rest of the run, so one malformed message took down the job.
+    Every bad request must now get a REQUEST_ERROR reply: requesters block on recv until a
+    reply arrives, so silently dropping a request would hang the caller instead.
+    """
+
+    RECV_TIMEOUT_MS = 5000
+
+    @staticmethod
+    def _connect(ctx, tq_controller):
+        info = ray.get(tq_controller.get_zmq_server_info.remote())
+        sock = create_zmq_socket(ctx, zmq.DEALER, info.ip, identity=f"probe-{uuid4().hex[:8]}".encode())
+        sock.connect(info.to_addr("request_handle_socket"))
+        sock.setsockopt(zmq.RCVTIMEO, TestTransferQueueControllerBadRequests.RECV_TIMEOUT_MS)
+        sock.setsockopt(zmq.SNDTIMEO, TestTransferQueueControllerBadRequests.RECV_TIMEOUT_MS)
+        return sock
+
+    @staticmethod
+    def _list_partitions_request():
+        return ZMQMessage.create(
+            request_type=ZMQRequestType.GET_LIST_PARTITIONS,
+            sender_id="probe",
+            body={},
+        ).serialize()
+
+    @classmethod
+    def _assert_still_serving(cls, sock):
+        """A well-formed request must still get its own correct reply."""
+        sock.send_multipart(cls._list_partitions_request())
+        reply = ZMQMessage.deserialize(sock.recv_multipart(copy=False))
+        assert reply.request_type == ZMQRequestType.LIST_PARTITIONS_RESPONSE
+
+    def test_controller_survives_undecodable_request(self, ray_setup):
+        """An empty leading frame is what a shifted multipart boundary looks like on the wire."""
+        tq_controller = TransferQueueController.remote()
+        ctx = zmq.Context()
+        sock = None
+        try:
+            sock = self._connect(ctx, tq_controller)
+            self._assert_still_serving(sock)
+
+            # Frame 0 must be the msgpack header; prepending an empty frame shifts every
+            # boundary by one, which is exactly the corruption seen in production.
+            sock.send_multipart([b"", *self._list_partitions_request()])
+            reply = ZMQMessage.deserialize(sock.recv_multipart(copy=False))
+            assert reply.request_type == ZMQRequestType.REQUEST_ERROR
+            assert "undecodable" in reply.body["message"]
+
+            self._assert_still_serving(sock)
+            print("✓ controller still serving after an undecodable request")
+        finally:
+            if sock is not None and not sock.closed:
+                sock.close(linger=0)
+            ctx.term()
+
+    def test_controller_rejects_unhandled_request_type(self, ray_setup):
+        """PUT_DATA is a storage-unit request; the controller has no branch for it.
+
+        Without the guard, the reply carried whatever response the previous loop iteration
+        had left in ``response_msg``, so this probe received a stale LIST_PARTITIONS_RESPONSE.
+        Now the requester gets a REQUEST_ERROR it can react to instead of a stale reply.
+        """
+        tq_controller = TransferQueueController.remote()
+        ctx = zmq.Context()
+        sock = None
+        try:
+            sock = self._connect(ctx, tq_controller)
+            # Leaves a LIST_PARTITIONS_RESPONSE behind as the previous iteration's response.
+            self._assert_still_serving(sock)
+
+            sock.send_multipart(
+                ZMQMessage.create(
+                    request_type=ZMQRequestType.PUT_DATA,
+                    sender_id="probe",
+                    body={},
+                ).serialize()
+            )
+            reply = ZMQMessage.deserialize(sock.recv_multipart(copy=False))
+            assert reply.request_type == ZMQRequestType.REQUEST_ERROR
+            assert "no handler" in reply.body["message"]
+
+            self._assert_still_serving(sock)
+            print("✓ controller rejected an unhandled request type without replaying a stale response")
+        finally:
+            if sock is not None and not sock.closed:
+                sock.close(linger=0)
+            ctx.term()
+
+    def test_controller_replies_error_when_handler_raises(self, ray_setup):
+        """A GET_META body missing its keys makes the handler raise KeyError.
+
+        The exception used to kill the request thread, hanging every client. Now the
+        requester gets a REQUEST_ERROR and the loop keeps serving subsequent requests.
+        """
+        tq_controller = TransferQueueController.remote()
+        ctx = zmq.Context()
+        sock = None
+        try:
+            sock = self._connect(ctx, tq_controller)
+
+            sock.send_multipart(
+                ZMQMessage.create(
+                    request_type=ZMQRequestType.GET_META,
+                    sender_id="probe",
+                    body={},  # no data_fields/batch_size/partition_id -> KeyError in the handler
+                ).serialize()
+            )
+            reply = ZMQMessage.deserialize(sock.recv_multipart(copy=False))
+            assert reply.request_type == ZMQRequestType.REQUEST_ERROR
+            assert "GET_META failed" in reply.body["message"]
+
+            self._assert_still_serving(sock)
+            print("✓ controller replied with an error when the handler raised")
+        finally:
+            if sock is not None and not sock.closed:
+                sock.close(linger=0)
+            ctx.term()

tests/test_serial_utils_on_cpu.py

@@ -17,9 +17,17 @@
 import numpy as np
 import pytest
 import torch
+import zmq
 from tensordict import TensorDict
 
-from transfer_queue.utils.serial_utils import MsgpackDecoder, MsgpackEncoder
+from transfer_queue.utils.serial_utils import (
+    _PICKLE_FALLBACK_SENTINEL,
+    MsgpackDecoder,
+    MsgpackEncoder,
+    _is_pickle_fallback,
+    decode,
+    encode,
+)
 
 
 @pytest.mark.parametrize(
@@ -1087,3 +1095,138 @@ def test_numpy_object_array_still_uses_pickle(self):
         deserialized = decoder.decode(serialized)
         assert isinstance(deserialized, np.ndarray)
         assert np.array_equal(deserialized, arr)
+
+
+# ============================================================================
+# Whole-Message Pickle Fallback Tests
+# ============================================================================
+class TestPickleFallback:
+    """Tests for the whole-message pickle fallback shared by ``encode`` and ``decode``.
+
+    Two independent defects made this path unusable end to end:
+
+    1. ``encode`` only caught ``TypeError``/``ValueError``, but msgspec reports
+       unrepresentable values with ``OverflowError``, ``RecursionError`` and its own
+       ``MsgspecError`` — none of which derive from ``ValueError`` — so the fallback was
+       skipped and the exception escaped to the caller.
+    2. ``decode`` located the marker frame with ``frames[0] == _PICKLE_FALLBACK_SENTINEL``.
+       Every receiver calls ``recv_multipart(copy=False)`` and therefore holds
+       ``zmq.Frame`` objects, which define no ``__eq__`` against ``bytes``, so the
+       comparison was always False and the marker was handed to msgpack instead.
+    """
+
+    @staticmethod
+    def _as_received(frames):
+        """Rebuild frames the way ``recv_multipart(copy=False)`` delivers them."""
+        return [zmq.Frame(memoryview(frame).tobytes()) for frame in frames]
+
+    @pytest.mark.parametrize(
+        "value",
+        [
+            pytest.param(2**70, id="above_uint64_max"),
+            pytest.param(-(2**70), id="below_int64_min"),
+        ],
+    )
+    def test_oversized_int_falls_back_instead_of_raising(self, value):
+        """msgspec raises OverflowError here, which is an ArithmeticError, not a ValueError."""
+        obj = {"global_step": 7, "offset": value}
+
+        frames = encode(obj)
+
+        assert len(frames) == 2, "expected the two-frame pickle fallback layout"
+        assert bytes(frames[0]) == _PICKLE_FALLBACK_SENTINEL
+        assert decode(frames) == obj
+
+    def test_self_referential_container_falls_back(self):
+        """msgspec raises RecursionError on cycles; pickle memoizes them."""
+        obj = {"name": "cyclic"}
+        obj["self"] = obj
+
+        frames = encode(obj)
+
+        assert len(frames) == 2
+        decoded = decode(frames)
+        assert decoded["name"] == "cyclic"
+        assert decoded["self"] is decoded, "pickle should restore the cycle"
+
+    def test_fallback_round_trips_through_zmq_frames(self):
+        """The core regression: the marker must still be recognised inside a zmq.Frame."""
+        obj = {"offset": 2**70, "indexes": torch.arange(8)}
+
+        frames = self._as_received(encode(obj))
+        decoded = decode(frames)
+
+        assert decoded["offset"] == 2**70
+        assert torch.equal(decoded["indexes"], obj["indexes"])
+
+    @pytest.mark.parametrize(
+        "wrap",
+        [
+            pytest.param(bytes, id="bytes"),
+            pytest.param(bytearray, id="bytearray"),
+            pytest.param(memoryview, id="memoryview"),
+            pytest.param(zmq.Frame, id="zmq_frame"),
+        ],
+    )
+    def test_marker_detected_for_every_buffer_type(self, wrap):
+        """Senders emit bytes; receivers hold zmq.Frame. Both must be recognised."""
+        assert _is_pickle_fallback(wrap(_PICKLE_FALLBACK_SENTINEL))
+
+    @pytest.mark.parametrize(
+        "payload",
+        [
+            pytest.param(b"", id="empty"),
+            pytest.param(b"\xa2ab", id="same_size_msgpack_str"),
+            pytest.param(b"\xc1\xfe", id="truncated_marker"),
+            pytest.param(b"\xc1\xfe\xed\x00", id="marker_plus_trailing_byte"),
+        ],
+    )
+    def test_non_marker_frames_are_not_mistaken_for_fallback(self, payload):
+        """A msgpack header frame that merely resembles the marker must not divert decode."""
+        assert not _is_pickle_fallback(payload)
+        assert not _is_pickle_fallback(zmq.Frame(payload))
+
+    def test_zmq_message_fallback_survives_a_real_socket(self):
+        """End-to-end over the transport that hid the bug: ROUTER + recv_multipart(copy=False)."""
+        from transfer_queue.utils.zmq_utils import (
+            ZMQMessage,
+            ZMQRequestType,
+            create_zmq_socket,
+            format_zmq_address,
+            get_free_port,
+        )
+
+        ip = "127.0.0.1"
+        ctx = zmq.Context()
+        router = dealer = None
+        try:
+            port = get_free_port(ip)
+            router = create_zmq_socket(ctx, zmq.ROUTER, ip)
+            router.bind(format_zmq_address(ip, port))
+            dealer = create_zmq_socket(ctx, zmq.DEALER, ip, identity=b"fallback-probe")
+            dealer.connect(format_zmq_address(ip, port))
+
+            # 2**70 has no msgpack representation, so this message takes the pickle path.
+            msg = ZMQMessage.create(
+                request_type=ZMQRequestType.NOTIFY_DATA_UPDATE,
+                sender_id="storage-1",
+                body={"partition_id": "p0", "offset": 2**70},
+            )
+            sent = msg.serialize()
+            assert len(sent) == 2, "expected the pickle fallback layout"
+            dealer.send_multipart(sent)
+
+            assert router.poll(10_000), "message never arrived"
+            frames = router.recv_multipart(copy=False)
+            frames.pop(0)  # ROUTER identity prefix, as the controller does
+
+            received = ZMQMessage.deserialize(frames)
+            assert received.sender_id == "storage-1"
+            assert received.request_type == ZMQRequestType.NOTIFY_DATA_UPDATE
+            assert received.body["offset"] == 2**70
+            assert received.body["partition_id"] == "p0"
+        finally:
+            for sock in (dealer, router):
+                if sock is not None and not sock.closed:
+                    sock.close(linger=0)
+            ctx.term()