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

[0oshowero0]

Motivation

A long-running RL job hit an occasional failure that took the whole run down:

Exception in thread TransferQueueControllerProcessRequestThread:
Traceback (most recent call last):
  ...
  File "transfer_queue/controller.py", line 1771, in _process_request
    request_msg = ZMQMessage.deserialize(serialized_msg)
  File "transfer_queue/utils/zmq_utils.py", line 182, in deserialize
    result = decode(frames)
  File "transfer_queue/utils/serial_utils.py", line 262, in decode
    result = self.decoder.decode(bufs[0])
msgspec.DecodeError: Input data was truncated

_process_request had no exception handling, so this one message killed
TransferQueueControllerProcessRequestThread permanently. The controller stayed alive as
an actor but stopped answering every request from that point on, and the job hung.

While investigating I found more defects on the same path.

What was wrong

1. One bad message killed the request loop. _process_request ran recv_multipart →
deserialize → dispatch → send_multipart with no guard anywhere. Note that
_wait_connection in the same file already logs and continues; the request loop did not.

2. Dropped requests hang the caller forever. Clients issue requests through
with_controller_socket with no RCVTIMEO and then block on await socket.recv_multipart(). Any controller path that fails to reply — an undecodable request,
an unhandled request type, a handler that raises — leaves that caller waiting for the rest
of the run.

3. decode() could never detect the pickle fallback marker. It located the marker with
frames[0] == _PICKLE_FALLBACK_SENTINEL. Every receiver calls recv_multipart(copy=False)
and therefore holds zmq.Frame objects, and zmq.Frame defines no __eq__ against
bytes, so the comparison was always False and the marker frame was handed to msgpack.

4. encode() almost never reached that fallback anyway. It caught only
TypeError/ValueError, but msgspec reports unrepresentable values with OverflowError
(an ArithmeticError), RecursionError (a RuntimeError) and its own MsgspecError —
none of which derive from ValueError. Those escaped to the caller instead of degrading
to pickle, even though pickle handles all of them (arbitrary-precision ints,
self-referential containers).

5. Unhandled request types replied with a stale response. Every branch in the dispatch
chain is an if/elif with no else, so an unknown request_type fell through to
send_multipart([identity, *response_msg.serialize()]) still holding the previous
iteration's response_msg, sending an unrelated requester someone else's answer.

What this changes

• _process_request is split into a thin supervisor plus _process_request_loop; anything
  that escapes the per-request handling restarts the loop, so the thread survives for the
  controller's lifetime. The ROUTER socket stays bound across a restart.
• The dispatch chain moves into _handle_request(), which returns the response or None.
  Every request now gets exactly one reply:
  • the handler's normal response;
  • a new generic ZMQRequestType.REQUEST_ERROR (built by _make_error_response()) when
    the handler raises — the reply body carries {request_type} failed: {error}, so the
    caller raises immediately instead of hanging;
  • REQUEST_ERROR ("no handler for request_type ...") when no branch matches, replacing
    the stale-response replay;
  • REQUEST_ERROR for undecodable requests too — the ROUTER identity frame is prepended
    by the transport and survives payload corruption, so the reply still reaches the
    requester.
    No client change is needed: every client call site already treats an unexpected response
    type as RuntimeError(body["message"]).
• ZMQMessage.deserialize self-checks the frame layout and raises ZMQMessageDecodeError
  carrying frame count and per-frame sizes. It lives in deserialize, so controller,
  storage, client and manager all get the diagnostics.
• decode() compares buffer contents via _is_pickle_fallback(). Size is tested first,
  so the msgpack path stays copy-free.
• encode() catches the errors msgspec actually raises, via a named
  _ENCODE_FALLBACK_ERRORS tuple documenting why none of them are ValueError. Falling
  back now logs at warning level — it is a performance degradation path and should be
  visible at default log levels.
• mypy now runs over the whole package in pre-commit (pass_filenames: false); with
  follow_imports = "skip", passing only changed files produced false errors.

Scope

This does not fix the root cause of the truncation. The evidence says the ZMQ multipart
frame boundaries shift, so frame 0 stops being the msgpack header — decoding an empty
buffer produces exactly Input data was truncated, and empty frames are routine here since
empty tensors serialize to zero-length frames. What this PR does is stop that from taking
the job down, and emit the frame layout needed to confirm it. The frame sizes make the
diagnosis immediate:

healthy message:            num_frames=3, frame_sizes=[184, 512, 512]
empty leading frame:        leading frame is empty; num_frames=4, frame_sizes=[0, 184, 512, 512]
boundary on a tensor frame: trailing characters (byte 1); num_frames=2, frame_sizes=[512, 512]
truncated header:           Input data was truncated; num_frames=3, frame_sizes=[12, 512, 512]

A healthy message is one small header frame followed by large buffers. A leading 0, a
missing small header, or an implausibly small header all point at shifted boundaries.

One caveat to note for rollout: the pickle fallback only works end to end when both peers
run this version — an old receiver (recv_multipart(copy=False) + == marker check)
cannot recognise fallback frames from a new sender, and new senders produce them in more
situations than before. Upgrade controller, storage units and clients together.

Tests

TestPickleFallback in tests/test_serial_utils_on_cpu.py (13 cases): oversized ints and
self-referential containers degrade instead of raising, round-trips through zmq.Frame,
marker detection across bytes/bytearray/memoryview/zmq.Frame, four near-miss
negatives, and an end-to-end round trip over a real ROUTER with recv_multipart(copy=False)
— the transport that hid the bug.

TestTransferQueueControllerBadRequests in tests/test_controller.py (3 cases):

• an empty leading frame (shifted multipart boundary) gets a REQUEST_ERROR reply and the
  controller keeps serving;
• an unhandled PUT_DATA gets a REQUEST_ERROR reply instead of replaying the previous
  iteration's stale response;
• a GET_META body missing its keys makes the handler raise KeyError, and the requester
  gets a REQUEST_ERROR ("GET_META failed: KeyError ...") while the loop keeps answering
  subsequent requests.

All three were confirmed to fail against the pre-fix code and pass after.

Results: 115 passed across the two serialization suites, 25 passed across
tests/test_controller.py (including the 3 new cases), and 76 passed across tests/e2e/.
pre-commit run --all-files is green (ruff, ruff-format, mypy).

Unrelated CI fix bundled here

The MooncakeStore e2e job started failing on libcudart.so.12: cannot open shared object file — upstream mooncake-transfer-engine-non-cuda 0.3.12 ships a mooncake_master
binary linked against the CUDA 12 runtime (verified by diffing the DT_NEEDED entries of
the 0.3.11.post1 and 0.3.12.post1 wheels; 0.3.11.post1 has no CUDA dependency). The two
workflows that install it now pin mooncake-transfer-engine-non-cuda<0.3.12 until upstream
fixes the wheel.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[ascend-robot]

CLA Signature Pass

0oshowero0, thanks for your pull request. All authors of the commits have signed the CLA. 👍

[ascend-robot]

CLA Signature Pass

0oshowero0, thanks for your pull request. All authors of the commits have signed the CLA. 👍

[ascend-robot]

CLA Signature Pass

0oshowero0, thanks for your pull request. All authors of the commits have signed the CLA. 👍