[fix] Keep the storage worker alive when a request fails outside its handlers

tests/test_metrics.py

@@ -309,6 +309,28 @@ def serve():
             router.close(linger=0)
             ctx_peer.term()
 
+    def test_unanswered_probe_warns_once_per_hang(self):
+        """A unit that stops answering must be visible, without a warning every collection cycle."""
+        ctx_peer = zmq.Context()
+        silent_unit = ctx_peer.socket(zmq.ROUTER)
+        port = silent_unit.bind_to_random_port("tcp://127.0.0.1")
+        su_info = ZMQServerInfo(role=Role.STORAGE, id="storage_0", ip="127.0.0.1", ports={"put_get_socket": port})
+        ctx = zmq.Context()
+        try:
+            with (
+                patch("transfer_queue.metrics.TQ_METRICS_STORAGE_TIMEOUT", 1),
+                patch("transfer_queue.metrics.logger.warning") as warning,
+            ):
+                exporter = TQMetricsExporter(zmq_context=ctx)
+                for _ in range(2):
+                    assert exporter._query_storage_unit(su_info, "storage_0") is None
+            warning.assert_called_once()
+            assert f"127.0.0.1:{port}" in warning.call_args.args[0]
+        finally:
+            ctx.destroy(linger=0)
+            silent_unit.close(linger=0)
+            ctx_peer.term()
+
 
 class TestStorageMetricsCollection:
     def test_collect_with_no_storage_units(self):

tests/test_simple_storage_unit.py

@@ -1228,3 +1228,33 @@ def test_undecodable_request_answered_and_worker_survives(storage_setup):
     response = client.send_put(0, [0], {"val": [torch.tensor([1.0])]})
     assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE
     client.close()
+
+
+def test_worker_survives_a_reply_failing_between_frames():
+    """A reply that fails mid-message must neither end the only worker nor misroute the next reply."""
+    unit = SimpleStorageUnit.__ray_metadata__.modified_class({})
+    process_request = unit._process_one_worker_request
+    failed = []
+
+    def fail_after_identity_frame(worker_socket, monitor):
+        if not failed:
+            failed.append(True)
+            identity = worker_socket.recv_multipart()[0]
+            worker_socket.send(identity, zmq.SNDMORE)
+            raise MemoryError
+        process_request(worker_socket, monitor)
+
+    unit._process_one_worker_request = fail_after_identity_frame
+    address, ip = unit.zmq_server_info.to_addr("put_get_socket"), unit.zmq_server_info.ip
+    stranded, next_client = MockStorageClient(address, ip), MockStorageClient(address, ip)
+    stranded.socket.setsockopt(zmq.RCVTIMEO, 1000)
+    try:
+        with pytest.raises(zmq.Again):
+            stranded.send_put(0, [0], {"val": [torch.tensor([1.0])]})
+        response = next_client.send_put(1, [1], {"val": [torch.tensor([2.0])]})
+        assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE
+        assert not stranded.socket.poll(500), "the next reply was routed to the stranded caller"
+    finally:
+        stranded.close()
+        next_client.close()
+        unit.shutdown()

transfer_queue/metrics.py

@@ -97,6 +97,8 @@ def __init__(self, role: str = "controller", zmq_context: zmq.Context | None = N
         self._storage_unit_infos: dict[str, ZMQServerInfo] = {}
         self._zmq_ctx = zmq_context
         self._zmq_socket_pool: ZMQSocketPool | None = None
+        # Units whose last probe timed out, so a hung unit warns once rather than every cycle.
+        self._unresponsive_storage_units: set[str] = set()
         self._known_partition_ids: set[str] = set()
         self._known_production_labels: set[tuple[str, str]] = set()
         self._known_consumption_labels: set[tuple[str, str]] = set()
@@ -541,12 +543,21 @@ def _query_storage_unit(self, su_info: ZMQServerInfo, su_id: str) -> dict[str, A
                 # a kept socket holds budget for the whole walk to save one handshake.
                 sock.close(linger=0)
                 if response_msg.request_type == ZMQRequestType.METRICS_RESPONSE:
+                    if su_id in self._unresponsive_storage_units:
+                        self._unresponsive_storage_units.discard(su_id)
+                        logger.warning(f"Storage unit {su_id} is answering metrics probes again.")
                     return response_msg.body
                 return None
         except zmq.error.Again:
             # The pool discarded the socket, so the next cycle starts clean. Reusing it
             # would read this reply, if it lands late, as the answer to that cycle's query.
-            logger.debug(f"Timeout querying metrics from {su_id}")
+            # The probe shares the unit's only worker with put/get, so this is the earliest sign it stopped.
+            if su_id not in self._unresponsive_storage_units:
+                self._unresponsive_storage_units.add(su_id)
+                logger.warning(
+                    f"Timeout ({TQ_METRICS_STORAGE_TIMEOUT}s) querying metrics from storage unit {su_id} at "
+                    f"{su_info.ip}:{su_info.ports.get('put_get_socket')}; put/get on it may stall too."
+                )
             return None
         except Exception as e:
             logger.warning(f"Error querying metrics from {su_id}: {e}")

transfer_queue/storage/simple_storage.py

@@ -1103,18 +1103,11 @@ def _proxy_routine(self) -> None:
 
     def _worker_routine(self) -> None:
         """Worker thread for processing requests."""
-
-        worker_socket = create_zmq_socket(self.zmq_context, zmq.DEALER, self._node_ip)
-        worker_socket.connect(self._inproc_addr)
-
-        poller = zmq.Poller()
-        poller.register(worker_socket, zmq.POLLIN)
-
         logger.info(f"[{self.storage_unit_id}]: worker thread started...")
         perf_monitor = IntervalPerfMonitor(caller_name=f"{self.storage_unit_id}")
 
         try:
-            self._worker_loop(worker_socket, poller, perf_monitor)
+            self._worker_loop(perf_monitor)
         except BaseException:
             logger.critical(
                 f"[{self.storage_unit_id}]: worker thread is exiting on an unhandled exception.", exc_info=True
@@ -1123,8 +1116,6 @@ def _worker_routine(self) -> None:
             raise
         finally:
             logger.info(f"[{self.storage_unit_id}]: worker stopped.")
-            poller.unregister(worker_socket)
-            worker_socket.close(linger=0)
 
     def _abort_process(self, reason: str) -> None:
         """Kill the whole unit process so it fails visibly instead of lingering as a dead unit.
@@ -1151,28 +1142,46 @@ def _abort_process(self, reason: str) -> None:
         _flush_log_output()
         os._exit(_WORKER_DEATH_EXIT_CODE)
 
-    def _worker_loop(self, worker_socket: zmq.Socket, poller: zmq.Poller, perf_monitor: Any) -> None:
-        """Poll for requests and answer them until shutdown is signalled."""
-        while not self._shutdown_event.is_set():
-            monitor = self._metrics if self._metrics is not None else perf_monitor
-            try:
-                # The event cannot wake a ZMQ poll, so bound idle shutdown latency.
-                socks = dict(poller.poll(min(TQ_STORAGE_POLLER_TIMEOUT, 1) * 1000))
-            except zmq.error.ContextTerminated:
-                # ZMQ context was terminated, exit gracefully
-                logger.info(f"[{self.storage_unit_id}]: worker stopped gracefully (Context Terminated)")
-                break
-            except Exception as e:
-                logger.warning(f"[{self.storage_unit_id}]: worker poll error: {e}")
-                continue
-
-            if self._shutdown_event.is_set():
-                break
+    def _connect_worker_socket(self, poller: zmq.Poller) -> zmq.Socket:
+        worker_socket = create_zmq_socket(self.zmq_context, zmq.DEALER, self._node_ip)
+        worker_socket.connect(self._inproc_addr)
+        poller.register(worker_socket, zmq.POLLIN)
+        return worker_socket
 
-            if worker_socket in socks:
-                # Keep request handling in a separate scope so payload references are
-                # released before the worker waits for the next request.
-                self._process_one_worker_request(worker_socket, monitor)
+    def _worker_loop(self, perf_monitor: Any) -> None:
+        """Poll for requests and answer them until shutdown is signalled."""
+        poller = zmq.Poller()
+        worker_socket = self._connect_worker_socket(poller)
+        try:
+            while not self._shutdown_event.is_set():
+                monitor = self._metrics if self._metrics is not None else perf_monitor
+                try:
+                    # The event cannot wake a ZMQ poll, so bound idle shutdown latency.
+                    socks = dict(poller.poll(min(TQ_STORAGE_POLLER_TIMEOUT, 1) * 1000))
+                except zmq.error.ContextTerminated:
+                    # ZMQ context was terminated, exit gracefully
+                    logger.info(f"[{self.storage_unit_id}]: worker stopped gracefully (Context Terminated)")
+                    break
+                except Exception as e:
+                    logger.warning(f"[{self.storage_unit_id}]: worker poll error: {e}")
+                    continue
+
+                if self._shutdown_event.is_set():
+                    break
+
+                if worker_socket in socks:
+                    # Keep request handling in a separate scope so payload references are
+                    # released before the worker waits for the next request.
+                    try:
+                        self._process_one_worker_request(worker_socket, monitor)
+                    except Exception as e:
+                        logger.exception(f"[{self.storage_unit_id}]: worker failed a request: {type(e).__name__}: {e}")
+                        # The socket may be left mid-message; a fresh one keeps the next reply routed correctly.
+                        poller.unregister(worker_socket)
+                        worker_socket.close(linger=0)
+                        worker_socket = self._connect_worker_socket(poller)
+        finally:
+            worker_socket.close(linger=0)
 
     def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) -> None:
         """Process one storage request and send its response."""