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

.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,89 @@ 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 or does not handle.
+
+    Either kind 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.
+    """
+
+    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()])
+
+            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_drops_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.
+        """
+        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()
+            )
+            with pytest.raises(zmq.error.Again):
+                sock.recv_multipart(copy=False)
+
+            self._assert_still_serving(sock)
+            print("✓ controller dropped 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()

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

transfer_queue/controller.py

@@ -22,7 +22,7 @@
 from itertools import groupby
 from operator import itemgetter
 from threading import Thread
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 from uuid import uuid4
 
 import numpy as np
@@ -1782,19 +1782,55 @@ def _start_process_request(self):
         self.process_request_thread.start()
 
     def _process_request(self):
-        """Main request processing loop - adapted for partition-based operations."""
+        """Keep the request loop running for the controller's lifetime.
+
+        Anything that escapes the per-request guards below would otherwise terminate this
+        thread and leave the controller permanently unable to answer requests. The ROUTER
+        socket stays bound across a restart, so the loop simply resumes with the next
+        queued request.
+        """
 
         logger.info(f"[{self.controller_id}]: start processing requests...")
 
         perf_monitor = IntervalPerfMonitor(caller_name=self.controller_id)
 
+        while True:
+            try:
+                self._process_request_loop(perf_monitor)
+            except zmq.ContextTerminated:
+                logger.info(f"[{self.controller_id}]: stopped processing requests (context terminated)")
+                return
+            except Exception as e:
+                if self.request_handle_socket.closed:
+                    logger.info(f"[{self.controller_id}]: stopped processing requests (socket closed)")
+                    return
+                logger.exception(f"[{self.controller_id}]: request loop raised {type(e).__name__}: {e}; restarting")
+                # Keep a persistent failure from turning into a busy log flood.
+                time.sleep(0.1)
+
+    def _process_request_loop(self, perf_monitor: IntervalPerfMonitor) -> None:
+        """Main request processing loop - adapted for partition-based operations."""
+
         while True:
             monitor = self._metrics if self._metrics is not None else perf_monitor
 
             messages = self.request_handle_socket.recv_multipart(copy=False)
             identity = messages.pop(0)
             serialized_msg = messages
-            request_msg = ZMQMessage.deserialize(serialized_msg)
+            try:
+                request_msg = ZMQMessage.deserialize(serialized_msg)
+            except Exception as e:
+                # An undecodable request says nothing about the controller's own state, so
+                # drop it and carry on. ZMQMessageDecodeError carries the frame layout,
+                # which is what tells a corrupt payload apart from shifted frame
+                # boundaries.
+                logger.error(
+                    f"[{self.controller_id}]: dropping undecodable request from "
+                    f"identity={bytes(identity)!r}: {type(e).__name__}: {e}"
+                )
+                continue
+
+            response_msg = None
 
             if request_msg.request_type == ZMQRequestType.GET_META:
                 with monitor.measure(op_type="GET_META"):
@@ -1824,7 +1860,7 @@ def _process_request(self):
 
                     # Update production status
                     success = self.update_production_status(
-                        partition_id=partition_id,
+                        partition_id=cast(str, partition_id),
                         global_indexes=global_indexes,
                         field_schema=message_data.get("field_schema", {}),
                         custom_backend_meta=message_data.get("custom_backend_meta", {}),
@@ -2077,6 +2113,15 @@ def _process_request(self):
                     body={"success": True},
                 )
 
+            if response_msg is None:
+                # No branch matched. Without this guard the reply below would send the
+                # previous iteration's response to an unrelated requester.
+                logger.error(
+                    f"[{self.controller_id}]: no handler for request_type={request_msg.request_type} "
+                    f"from sender={request_msg.sender_id}; dropping request {request_msg.request_id}"
+                )
+                continue
+
             self.request_handle_socket.send_multipart([identity, *response_msg.serialize()])
 
     def get_zmq_server_info(self) -> ZMQServerInfo: