diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index ee2ec053..b734f15d 100644 --- a/.github/workflows/python-package.yml +++ b/.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 diff --git a/.github/workflows/run-tests.yml b/.github/workflows/run-tests.yml index 2c4429ae..05bba3ae 100644 --- a/.github/workflows/run-tests.yml +++ b/.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 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5b056757..7ebeec66 100644 --- a/.pre-commit-config.yaml +++ b/.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] diff --git a/tests/test_controller.py b/tests/test_controller.py index b7638cc5..604fd566 100644 --- a/tests/test_controller.py +++ b/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() diff --git a/tests/test_serial_utils_on_cpu.py b/tests/test_serial_utils_on_cpu.py index ffa08d75..29ad9c73 100644 --- a/tests/test_serial_utils_on_cpu.py +++ b/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() diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index db9d6a64..719b4d71 100644 --- a/transfer_queue/controller.py +++ b/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 @@ -1817,314 +1817,403 @@ 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. + + Every decoded request gets exactly one reply: either its handler's response or a + REQUEST_ERROR explaining why it could not be served. Requesters block on recv + until a reply arrives, so silently dropping a request would hang the caller. + """ + 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: + # ZMQMessageDecodeError carries the frame layout, which is what tells a + # corrupt payload apart from shifted frame boundaries. The ROUTER identity + # frame is prepended by the transport and survives such corruption, so the + # requester can still get an error reply instead of hanging on recv. + logger.error( + f"[{self.controller_id}]: undecodable request from " + f"identity={bytes(identity)!r}: {type(e).__name__}: {e}" + ) + error_msg = self._make_error_response( + receiver_id=None, + message=f"undecodable request: {type(e).__name__}: {e}", + ) + self.request_handle_socket.send_multipart([identity, *error_msg.serialize()]) + continue - if request_msg.request_type == ZMQRequestType.GET_META: - with monitor.measure(op_type="GET_META"): - params = request_msg.body + try: + response_msg: ZMQMessage | None = self._handle_request(request_msg, monitor) + except Exception as e: + # A failing handler must not kill the loop (it used to terminate this + # thread) nor hang the requester, so report the failure back. + logger.exception( + f"[{self.controller_id}]: handler for request_type={request_msg.request_type} " + f"from sender={request_msg.sender_id} raised {type(e).__name__}: {e}" + ) + response_msg = self._make_error_response( + receiver_id=request_msg.sender_id, + message=f"{request_msg.request_type} failed: {type(e).__name__}: {e}", + ) - metadata = self.get_metadata( - data_fields=params["data_fields"], - batch_size=params["batch_size"], - partition_id=params["partition_id"], - mode=params.get("mode", "fetch"), - task_name=params.get("task_name"), - sampling_config=params.get("sampling_config", {}), - ) + 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}; request {request_msg.request_id}" + ) + response_msg = self._make_error_response( + receiver_id=request_msg.sender_id, + message=f"no handler for request_type {request_msg.request_type}", + ) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_META_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"metadata": metadata}, - ) + self.request_handle_socket.send_multipart([identity, *response_msg.serialize()]) - elif request_msg.request_type == ZMQRequestType.NOTIFY_DATA_UPDATE: - with monitor.measure(op_type="NOTIFY_DATA_UPDATE"): - message_data = request_msg.body - partition_id = message_data.get("partition_id") - global_indexes = message_data.get("global_indexes", []) + def _make_error_response(self, receiver_id: str | None, message: str) -> ZMQMessage: + """Build a generic error reply so the requester raises instead of hanging on recv.""" + return ZMQMessage.create( + request_type=ZMQRequestType.REQUEST_ERROR, + sender_id=self.controller_id, + receiver_id=receiver_id, + body={"success": False, "message": message}, + ) - # Update production status - success = self.update_production_status( - partition_id=partition_id, - global_indexes=global_indexes, - field_schema=message_data.get("field_schema", {}), - custom_backend_meta=message_data.get("custom_backend_meta", {}), - ) - if success: - if self._metrics is not None: - self._metrics.record_samples("NOTIFY_DATA_UPDATE", len(global_indexes)) - logger.debug(f"[{self.controller_id}]: Updated production status for partition {partition_id}") + def _handle_request(self, request_msg: ZMQMessage, monitor: Any) -> ZMQMessage | None: + """Build the response for a decoded request, or None if no handler matches. - # Send acknowledgment - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.NOTIFY_DATA_UPDATE_ACK, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={ - "controller_id": self.controller_id, - "partition_id": partition_id, - "success": success, - }, - ) + Whatever the matching handler raises propagates to the caller, which turns it + into an error response. + """ + response_msg = None - elif request_msg.request_type == ZMQRequestType.GET_PARTITION_META: - with monitor.measure(op_type="GET_PARTITION_META"): - params = request_msg.body - partition_id = params["partition_id"] - partition = self._get_partition(partition_id) - if partition is not None: - partition_data_fields = list(partition.field_name_mapping.keys()) - - metadata = self.get_metadata( - data_fields=partition_data_fields, - partition_id=partition_id, - mode="force_fetch", - ) - else: - metadata = None + if request_msg.request_type == ZMQRequestType.GET_META: + with monitor.measure(op_type="GET_META"): + params = request_msg.body - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_PARTITION_META_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"metadata": metadata}, - ) - elif request_msg.request_type == ZMQRequestType.SET_CUSTOM_META: - with monitor.measure(op_type="SET_CUSTOM_META"): - params = request_msg.body - partition_custom_meta = params["partition_custom_meta"] + metadata = self.get_metadata( + data_fields=params["data_fields"], + batch_size=params["batch_size"], + partition_id=params["partition_id"], + mode=params.get("mode", "fetch"), + task_name=params.get("task_name"), + sampling_config=params.get("sampling_config", {}), + ) - self.set_custom_meta(partition_custom_meta=partition_custom_meta) + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.GET_META_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"metadata": metadata}, + ) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.SET_CUSTOM_META_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"message": "Successfully set custom_meta"}, - ) + elif request_msg.request_type == ZMQRequestType.NOTIFY_DATA_UPDATE: + with monitor.measure(op_type="NOTIFY_DATA_UPDATE"): + message_data = request_msg.body + partition_id = message_data.get("partition_id") + global_indexes = message_data.get("global_indexes", []) + + # Update production status + success = self.update_production_status( + 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", {}), + ) + if success: + if self._metrics is not None: + self._metrics.record_samples("NOTIFY_DATA_UPDATE", len(global_indexes)) + logger.debug(f"[{self.controller_id}]: Updated production status for partition {partition_id}") - elif request_msg.request_type == ZMQRequestType.MARK_CLEARING: - with monitor.measure(op_type="MARK_CLEARING"): - params = request_msg.body - self.mark_clearing(params["global_indexes"], params["partition_ids"]) + # Send acknowledgment + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.NOTIFY_DATA_UPDATE_ACK, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={ + "controller_id": self.controller_id, + "partition_id": partition_id, + "success": success, + }, + ) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.MARK_CLEARING_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"message": "Mark clearing completed"}, + elif request_msg.request_type == ZMQRequestType.GET_PARTITION_META: + with monitor.measure(op_type="GET_PARTITION_META"): + params = request_msg.body + partition_id = params["partition_id"] + partition = self._get_partition(partition_id) + if partition is not None: + partition_data_fields = list(partition.field_name_mapping.keys()) + + metadata = self.get_metadata( + data_fields=partition_data_fields, + partition_id=partition_id, + mode="force_fetch", ) + else: + metadata = None - elif request_msg.request_type == ZMQRequestType.CLEAR_META: - with monitor.measure(op_type="CLEAR_META"): - params = request_msg.body - global_indexes = params["global_indexes"] - partition_ids = params["partition_ids"] + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.GET_PARTITION_META_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"metadata": metadata}, + ) + elif request_msg.request_type == ZMQRequestType.SET_CUSTOM_META: + with monitor.measure(op_type="SET_CUSTOM_META"): + params = request_msg.body + partition_custom_meta = params["partition_custom_meta"] - self.clear_meta(global_indexes, partition_ids) - if self._metrics is not None: - self._metrics.record_samples("CLEAR_META", len(global_indexes)) + self.set_custom_meta(partition_custom_meta=partition_custom_meta) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.CLEAR_META_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"message": f"Clear samples operation completed by controller {self.controller_id}"}, - ) + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.SET_CUSTOM_META_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"message": "Successfully set custom_meta"}, + ) - elif request_msg.request_type == ZMQRequestType.CLEAR_PARTITION: - with monitor.measure(op_type="CLEAR_PARTITION"): - params = request_msg.body - partition_id = params["partition_id"] + elif request_msg.request_type == ZMQRequestType.MARK_CLEARING: + with monitor.measure(op_type="MARK_CLEARING"): + params = request_msg.body + self.mark_clearing(params["global_indexes"], params["partition_ids"]) - self.clear_partition(partition_id) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.CLEAR_PARTITION_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"message": f"Clear partition operation completed by controller {self.controller_id}"}, - ) + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.MARK_CLEARING_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"message": "Mark clearing completed"}, + ) - elif request_msg.request_type == ZMQRequestType.GET_CONSUMPTION: - with monitor.measure(op_type="GET_CONSUMPTION"): - # Handle consumption status checks - params = request_msg.body + elif request_msg.request_type == ZMQRequestType.CLEAR_META: + with monitor.measure(op_type="CLEAR_META"): + params = request_msg.body + global_indexes = params["global_indexes"] + partition_ids = params["partition_ids"] - global_index, consumption_status = self.get_consumption_status( - params["partition_id"], params["task_name"] - ) - sample_filter = params.get("sample_filter") # TODO: DEPRECATED in future + self.clear_meta(global_indexes, partition_ids) + if self._metrics is not None: + self._metrics.record_samples("CLEAR_META", len(global_indexes)) - if sample_filter and consumption_status is not None: - # TODO: DEPRECATED in future - consumption_status = consumption_status[sample_filter] + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.CLEAR_META_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"message": f"Clear samples operation completed by controller {self.controller_id}"}, + ) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.CONSUMPTION_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={ - "partition_id": params["partition_id"], - "global_index": global_index, - "consumption_status": consumption_status, - }, - ) + elif request_msg.request_type == ZMQRequestType.CLEAR_PARTITION: + with monitor.measure(op_type="CLEAR_PARTITION"): + params = request_msg.body + partition_id = params["partition_id"] - elif request_msg.request_type == ZMQRequestType.RESET_CONSUMPTION: - with monitor.measure(op_type="RESET_CONSUMPTION"): - # Handle reset consumption status request - params = request_msg.body - partition_id = params["partition_id"] - task_name = params.get("task_name") # Optional - try: - self.reset_consumption(partition_id, task_name) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.RESET_CONSUMPTION_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={ - "partition_id": partition_id, - "success": True, - "message": f"Consumption reset for partition {partition_id}", - }, - ) - except Exception as e: - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.RESET_CONSUMPTION_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={ - "partition_id": partition_id, - "success": False, - "message": str(e), - }, - ) + self.clear_partition(partition_id) + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.CLEAR_PARTITION_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"message": f"Clear partition operation completed by controller {self.controller_id}"}, + ) - elif request_msg.request_type == ZMQRequestType.GET_PRODUCTION: - with monitor.measure(op_type="GET_PRODUCTION"): - # Handle production status checks - params = request_msg.body + elif request_msg.request_type == ZMQRequestType.GET_CONSUMPTION: + with monitor.measure(op_type="GET_CONSUMPTION"): + # Handle consumption status checks + params = request_msg.body - global_index, production_status = self.get_production_status( - params["partition_id"], params["data_fields"] - ) + global_index, consumption_status = self.get_consumption_status( + params["partition_id"], params["task_name"] + ) + sample_filter = params.get("sample_filter") # TODO: DEPRECATED in future + + if sample_filter and consumption_status is not None: + # TODO: DEPRECATED in future + consumption_status = consumption_status[sample_filter] + + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.CONSUMPTION_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={ + "partition_id": params["partition_id"], + "global_index": global_index, + "consumption_status": consumption_status, + }, + ) + elif request_msg.request_type == ZMQRequestType.RESET_CONSUMPTION: + with monitor.measure(op_type="RESET_CONSUMPTION"): + # Handle reset consumption status request + params = request_msg.body + partition_id = params["partition_id"] + task_name = params.get("task_name") # Optional + try: + self.reset_consumption(partition_id, task_name) response_msg = ZMQMessage.create( - request_type=ZMQRequestType.PRODUCTION_RESPONSE, + request_type=ZMQRequestType.RESET_CONSUMPTION_RESPONSE, sender_id=self.controller_id, receiver_id=request_msg.sender_id, body={ - "partition_id": params["partition_id"], - "global_index": global_index, - "production_status": production_status, + "partition_id": partition_id, + "success": True, + "message": f"Consumption reset for partition {partition_id}", }, ) - - elif request_msg.request_type == ZMQRequestType.GET_LIST_PARTITIONS: - with monitor.measure(op_type="GET_LIST_PARTITIONS"): - # Handle list partitions request - partition_ids = self.list_partitions() + except Exception as e: response_msg = ZMQMessage.create( - request_type=ZMQRequestType.LIST_PARTITIONS_RESPONSE, + request_type=ZMQRequestType.RESET_CONSUMPTION_RESPONSE, sender_id=self.controller_id, receiver_id=request_msg.sender_id, - body={"partition_ids": partition_ids}, + body={ + "partition_id": partition_id, + "success": False, + "message": str(e), + }, ) - elif request_msg.request_type == ZMQRequestType.KV_RETRIEVE_META: - with monitor.measure(op_type="KV_RETRIEVE_META"): - params = request_msg.body - keys = params["keys"] - partition_id = params["partition_id"] - create = params["create"] + elif request_msg.request_type == ZMQRequestType.GET_PRODUCTION: + with monitor.measure(op_type="GET_PRODUCTION"): + # Handle production status checks + params = request_msg.body - metadata = self.kv_retrieve_meta(keys=keys, partition_id=partition_id, create=create) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.KV_RETRIEVE_META_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"metadata": metadata}, - ) + global_index, production_status = self.get_production_status( + params["partition_id"], params["data_fields"] + ) - elif request_msg.request_type == ZMQRequestType.KV_RETRIEVE_KEYS: - with monitor.measure(op_type="KV_RETRIEVE_KEYS"): - params = request_msg.body - global_indexes = params["global_indexes"] - partition_id = params["partition_id"] + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.PRODUCTION_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={ + "partition_id": params["partition_id"], + "global_index": global_index, + "production_status": production_status, + }, + ) - keys = self.kv_retrieve_keys(global_indexes=global_indexes, partition_id=partition_id) - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.KV_RETRIEVE_KEYS_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"keys": keys}, - ) + elif request_msg.request_type == ZMQRequestType.GET_LIST_PARTITIONS: + with monitor.measure(op_type="GET_LIST_PARTITIONS"): + # Handle list partitions request + partition_ids = self.list_partitions() + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.LIST_PARTITIONS_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"partition_ids": partition_ids}, + ) - elif request_msg.request_type == ZMQRequestType.KV_LIST: - with monitor.measure(op_type="KV_LIST"): - params = request_msg.body - partition_id = params["partition_id"] - if partition_id is None: - partition_id = list(self.partitions.keys()) - else: - partition_id = [partition_id] - - message = "success" - partition_info = {} - for pid in partition_id: - partition = self._get_partition(pid) - if partition: - keys = list(partition.keys_mapping.keys()) - single_partition_info = { - k: partition.custom_meta.get(partition.keys_mapping[k], {}) for k in keys - } - partition_info[pid] = single_partition_info - else: - # this only happens when params["partition_id"] is not None - message = f"partition {pid} does not exist" + elif request_msg.request_type == ZMQRequestType.KV_RETRIEVE_META: + with monitor.measure(op_type="KV_RETRIEVE_META"): + params = request_msg.body + keys = params["keys"] + partition_id = params["partition_id"] + create = params["create"] - response_msg = ZMQMessage.create( - request_type=ZMQRequestType.KV_LIST_RESPONSE, - sender_id=self.controller_id, - receiver_id=request_msg.sender_id, - body={"partition_info": partition_info, "message": message}, - ) + metadata = self.kv_retrieve_meta(keys=keys, partition_id=partition_id, create=create) + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.KV_RETRIEVE_META_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"metadata": metadata}, + ) - elif request_msg.request_type == ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT: - path = request_msg.body["path"] - self.save_checkpoint(path) + elif request_msg.request_type == ZMQRequestType.KV_RETRIEVE_KEYS: + with monitor.measure(op_type="KV_RETRIEVE_KEYS"): + params = request_msg.body + global_indexes = params["global_indexes"] + partition_id = params["partition_id"] + + keys = self.kv_retrieve_keys(global_indexes=global_indexes, partition_id=partition_id) response_msg = ZMQMessage.create( - request_type=ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT_RESPONSE, + request_type=ZMQRequestType.KV_RETRIEVE_KEYS_RESPONSE, sender_id=self.controller_id, receiver_id=request_msg.sender_id, - body={"success": True}, + body={"keys": keys}, ) - elif request_msg.request_type == ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT: - path = request_msg.body["path"] - self.load_checkpoint(path) + elif request_msg.request_type == ZMQRequestType.KV_LIST: + with monitor.measure(op_type="KV_LIST"): + params = request_msg.body + partition_id = params["partition_id"] + if partition_id is None: + partition_id = list(self.partitions.keys()) + else: + partition_id = [partition_id] + + message = "success" + partition_info = {} + for pid in partition_id: + partition = self._get_partition(pid) + if partition: + keys = list(partition.keys_mapping.keys()) + single_partition_info = { + k: partition.custom_meta.get(partition.keys_mapping[k], {}) for k in keys + } + partition_info[pid] = single_partition_info + else: + # this only happens when params["partition_id"] is not None + message = f"partition {pid} does not exist" + response_msg = ZMQMessage.create( - request_type=ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT_RESPONSE, + request_type=ZMQRequestType.KV_LIST_RESPONSE, sender_id=self.controller_id, receiver_id=request_msg.sender_id, - body={"success": True}, + body={"partition_info": partition_info, "message": message}, ) - self.request_handle_socket.send_multipart([identity, *response_msg.serialize()]) + elif request_msg.request_type == ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT: + path = request_msg.body["path"] + self.save_checkpoint(path) + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.SAVE_CONTROLLER_CHECKPOINT_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"success": True}, + ) + + elif request_msg.request_type == ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT: + path = request_msg.body["path"] + self.load_checkpoint(path) + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.LOAD_CONTROLLER_CHECKPOINT_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"success": True}, + ) + + return response_msg def get_zmq_server_info(self) -> ZMQServerInfo: """Get ZMQ server connection information.""" diff --git a/transfer_queue/utils/serial_utils.py b/transfer_queue/utils/serial_utils.py index 170e2e88..e8ab267a 100644 --- a/transfer_queue/utils/serial_utils.py +++ b/transfer_queue/utils/serial_utils.py @@ -26,6 +26,7 @@ from typing import Any, TypeAlias import cloudpickle +import msgspec import numpy as np import torch import zmq @@ -42,6 +43,7 @@ # 0xC1 is permanently reserved (invalid) in msgpack spec — safe to use as pickle fallback sentinel. _PICKLE_FALLBACK_SENTINEL = b"\xc1\xfe\xed" +_PICKLE_FALLBACK_SENTINEL_SIZE = len(_PICKLE_FALLBACK_SENTINEL) bytestr: TypeAlias = bytes | bytearray | memoryview | zmq.Frame @@ -364,6 +366,34 @@ def ext_hook(self, code: int, data: memoryview) -> Any: _decoder = MsgpackDecoder() +# Values msgpack cannot represent. None of these derive from ValueError: OverflowError is +# an ArithmeticError, RecursionError a RuntimeError, and msgspec's own errors subclass +# Exception directly. pickle handles all of them (arbitrary-precision ints, +# self-referential containers, ...), so they are degradation paths rather than failures. +_ENCODE_FALLBACK_ERRORS = ( + TypeError, + ValueError, + OverflowError, + RecursionError, + msgspec.MsgspecError, +) + + +def _is_pickle_fallback(frame: bytestr) -> bool: + """Whether ``frame`` is the pickle fallback marker. + + Compares buffer contents rather than using ``==``: every receiver calls + ``recv_multipart(copy=False)`` and so holds ``zmq.Frame`` objects, which do not + implement ``__eq__`` against ``bytes``. Testing the size first keeps the msgpack + path copy-free. + """ + try: + view = memoryview(frame) + except TypeError: + return False + return view.nbytes == _PICKLE_FALLBACK_SENTINEL_SIZE and view.tobytes() == _PICKLE_FALLBACK_SENTINEL + + def encode(obj: Any) -> list[bytestr]: """Encode an object via msgpack zero-copy; falls back to pickle on failure. @@ -372,8 +402,8 @@ def encode(obj: Any) -> list[bytestr]: """ try: return list(_encoder.encode(obj)) - except (TypeError, ValueError) as e: - logger.debug( + except _ENCODE_FALLBACK_ERRORS as e: + logger.warning( "encode: msgpack failed (%s), falling back to pickle.", type(e).__name__, ) @@ -386,7 +416,7 @@ def decode(frames: list) -> Any: Transparently handles both the msgpack zero-copy path and the pickle fallback path based on the leading sentinel frame. """ - if len(frames) >= 2 and frames[0] == _PICKLE_FALLBACK_SENTINEL: + if len(frames) >= 2 and _is_pickle_fallback(frames[0]): return pickle.loads(frames[1]) return _decoder.decode(frames) diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 93656542..49e8e674 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -15,6 +15,7 @@ import socket import time +from collections.abc import Sequence from dataclasses import dataclass from functools import wraps from typing import Any, Callable, TypeAlias @@ -44,6 +45,9 @@ class ZMQRequestType(ExplicitEnum): HANDSHAKE = "HANDSHAKE" # TransferQueueStorageUnit -> TransferQueueController HANDSHAKE_ACK = "HANDSHAKE_ACK" # TransferQueueController -> TransferQueueStorageUnit + # GENERIC + REQUEST_ERROR = "REQUEST_ERROR" # TransferQueueController -> requester, when a request cannot be served + # DATA_OPERATION GET_DATA = "GET" PUT_DATA = "PUT" @@ -142,6 +146,31 @@ def __str__(self) -> str: return f"ZMQSocketInfo(role={self.role}, id={self.id}, ip={self.ip}, ports={self.ports})" +class ZMQMessageDecodeError(ValueError): + """Raised when a received multipart message cannot be decoded into a ZMQMessage.""" + + +def frame_nbytes(frame: Any) -> int | None: + """Byte length of a single ZMQ frame, or None if the object exposes no buffer.""" + try: + return memoryview(frame).nbytes + except TypeError: + return None + + +def describe_frames(frames: Sequence[Any], max_reported: int = 32) -> str: + """Summarize a multipart message's frame layout: frame count and per-frame byte sizes. + + ``encode()`` emits a fixed layout (msgpack header in frame 0, one buffer per + tensor/ndarray after it), so the frame count and sizes are what distinguish a + genuinely corrupt payload from shifted multipart boundaries. + """ + sizes = [frame_nbytes(frame) for frame in frames[:max_reported]] + shown = ", ".join("?" if size is None else str(size) for size in sizes) + ellipsis = ", ..." if len(frames) > max_reported else "" + return f"num_frames={len(frames)}, frame_sizes=[{shown}{ellipsis}]" + + @dataclass class ZMQMessage: """ @@ -191,7 +220,18 @@ def deserialize(cls, frames: list) -> "ZMQMessage": if not frames: raise ValueError("Empty frames received") - result = decode(frames) + # Frame 0 is always the msgpack header; buffers for tensors/ndarrays follow it. A + # zero-length frame 0 therefore means the multipart boundaries have shifted rather + # than that the payload itself is bad, and decoding it would only report a + # misleading msgpack error. Report the frame layout instead. + if frame_nbytes(frames[0]) == 0: + raise ZMQMessageDecodeError(f"leading frame is empty; {describe_frames(frames)}") + + try: + result = decode(frames) + except Exception as e: + raise ZMQMessageDecodeError(f"{type(e).__name__}: {e}; {describe_frames(frames)}") from e + return cls( request_type=ZMQRequestType(result["request_type"]), sender_id=result["sender_id"],