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..33dbf7cc 100644 --- a/tests/test_simple_storage_unit.py +++ b/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() 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..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,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."""