From 227434ef4b909019c1c61aefa23d94413b80fbfe Mon Sep 17 00:00:00 2001 From: huniu20 Date: Fri, 23 Oct 2026 16:55:45 +0800 Subject: [PATCH 1/2] [fix] Keep the storage worker alive when a request fails outside its handlers A storage unit has a single worker thread. `recv_multipart` and the final `send_multipart` in `_process_one_worker_request` were outside any try block, so one failure there ended the thread while the proxy kept queueing requests nobody would answer, and `shutdown()` then blocked in `zmq_context.term()` on the dead worker's socket. Catch and log per request so only that caller times out. Also warn (once per hang, with ip:port) when a metrics probe times out: it is served by the same worker, so it is the earliest sign a unit stopped answering. Signed-off-by: huniu20 --- tests/test_metrics.py | 22 ++++++++++++++++++++ tests/test_simple_storage_unit.py | 26 ++++++++++++++++++++++++ transfer_queue/metrics.py | 13 +++++++++++- transfer_queue/storage/simple_storage.py | 7 ++++++- 4 files changed, 66 insertions(+), 2 deletions(-) diff --git a/tests/test_metrics.py b/tests/test_metrics.py index f3e338b0..1ebc2155 100644 --- a/tests/test_metrics.py +++ b/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): diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index 001c924e..12210cff 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -1228,3 +1228,29 @@ 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_failed_request(): + """A failure outside the per-operation handlers, e.g. in send, must not end the only worker.""" + unit = SimpleStorageUnit.__ray_metadata__.modified_class({}) + process_request = unit._process_one_worker_request + failed = [] + + def fail_first_request(worker_socket, monitor): + if not failed: + failed.append(True) + worker_socket.recv_multipart() + raise zmq.ZMQError(errno.EAGAIN) + process_request(worker_socket, monitor) + + unit._process_one_worker_request = fail_first_request + client = MockStorageClient(unit.zmq_server_info.to_addr("put_get_socket"), unit.zmq_server_info.ip) + client.socket.setsockopt(zmq.RCVTIMEO, 1000) + try: + with pytest.raises(zmq.Again): + client.send_put(0, [0], {"val": [torch.tensor([1.0])]}) + response = client.send_put(0, [0], {"val": [torch.tensor([1.0])]}) + assert response.request_type == ZMQRequestType.PUT_DATA_RESPONSE + finally: + client.close() + unit.shutdown() diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index 15476d3a..54601b3f 100644 --- a/transfer_queue/metrics.py +++ b/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}") diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 8b904bf4..19faa6b9 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -1172,7 +1172,12 @@ def _worker_loop(self, worker_socket: zmq.Socket, poller: zmq.Poller, perf_monit 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) + try: + self._process_one_worker_request(worker_socket, monitor) + except Exception as e: + # This is the unit's only worker: if a failed recv or send ended it, the + # proxy would keep queueing requests that nobody answers. Strand one caller. + logger.exception(f"[{self.storage_unit_id}]: worker failed a request: {type(e).__name__}: {e}") def _process_one_worker_request(self, worker_socket: zmq.Socket, monitor: Any) -> None: """Process one storage request and send its response.""" From 4ac95abefcc750f485e453148d773357b140317e Mon Sep 17 00:00:00 2001 From: huniu20 Date: Tue, 27 Oct 2026 10:45:15 +0800 Subject: [PATCH 2/2] [fix] Recreate the worker socket after a failed request Multipart I/O can fail between frames and leave the DEALER mid-message, so the next reply would be appended to the stranded caller's route. Close the socket, which drops the unfinished message, and connect a fresh one. The regression test fails after the identity frame and checks the next caller gets its own reply. Signed-off-by: huniu20 --- tests/test_simple_storage_unit.py | 26 +++++---- transfer_queue/storage/simple_storage.py | 72 +++++++++++++----------- 2 files changed, 53 insertions(+), 45 deletions(-) diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index 12210cff..33dbf7cc 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -1230,27 +1230,31 @@ def test_undecodable_request_answered_and_worker_survives(storage_setup): client.close() -def test_worker_survives_a_failed_request(): - """A failure outside the per-operation handlers, e.g. in send, must not end the only worker.""" +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_first_request(worker_socket, monitor): + def fail_after_identity_frame(worker_socket, monitor): if not failed: failed.append(True) - worker_socket.recv_multipart() - raise zmq.ZMQError(errno.EAGAIN) + 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_first_request - client = MockStorageClient(unit.zmq_server_info.to_addr("put_get_socket"), unit.zmq_server_info.ip) - client.socket.setsockopt(zmq.RCVTIMEO, 1000) + 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): - client.send_put(0, [0], {"val": [torch.tensor([1.0])]}) - response = client.send_put(0, [0], {"val": [torch.tensor([1.0])]}) + 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: - client.close() + stranded.close() + next_client.close() unit.shutdown() diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 19faa6b9..8dfa8242 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/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,33 +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. + 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: - self._process_one_worker_request(worker_socket, monitor) + # 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: - # This is the unit's only worker: if a failed recv or send ended it, the - # proxy would keep queueing requests that nobody answers. Strand one caller. - logger.exception(f"[{self.storage_unit_id}]: worker failed a request: {type(e).__name__}: {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."""