[fix] exit a storage unit whose worker dies instead of hanging callers

tests/test_storage_worker_death_exits.py

@@ -0,0 +1,134 @@
+# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""The unit process must exit when the thread that answers requests dies.
+
+Worker and proxy both run as daemon threads, so an exception in either only unwinds that
+thread: the process stays up with its ROUTER still bound, keeps accepting connections and
+never answers them, and every caller blocks until its own timeout. Exiting instead turns a
+silent hang into a visible actor failure.
+"""
+
+import threading
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from transfer_queue.storage import simple_storage
+from transfer_queue.storage.simple_storage import SimpleStorageUnit
+
+_StorageUnitClass = SimpleStorageUnit.__ray_metadata__.modified_class
+
+
+def _unit():
+    """A unit with only the attributes the abort path touches; no sockets, no Ray."""
+    unit = _StorageUnitClass.__new__(_StorageUnitClass)
+    unit.storage_unit_id = "TQ_STORAGE_UNIT_test0001"
+    unit._shutdown_event = threading.Event()
+    return unit
+
+
+def test_worker_death_exits_the_process():
+    unit = _unit()
+
+    with patch.object(simple_storage.os, "_exit") as exit_call:
+        unit._abort_process("the worker thread died")
+
+    exit_call.assert_called_once_with(simple_storage._WORKER_DEATH_EXIT_CODE)
+
+
+def test_abort_is_silent_during_shutdown():
+    """A shutdown races the poller, so an exception on the way out must not look like a crash."""
+    unit = _unit()
+    unit._shutdown_event.set()
+
+    with patch.object(simple_storage.os, "_exit") as exit_call:
+        unit._abort_process("the worker thread died")
+
+    exit_call.assert_not_called()
+
+
+def test_opt_out_keeps_the_process_alive_for_inspection():
+    unit = _unit()
+
+    with (
+        patch.object(simple_storage, "TQ_STORAGE_EXIT_ON_WORKER_DEATH", False),
+        patch.object(simple_storage.os, "_exit") as exit_call,
+    ):
+        unit._abort_process("the worker thread died")
+
+    exit_call.assert_not_called()
+
+
+def test_unhandled_worker_exception_triggers_the_exit():
+    """The whole point: a dead worker must take the process down, not linger unreachable."""
+    unit = _unit()
+    unit._metrics = None
+    unit._node_ip = "127.0.0.1"
+    unit._inproc_addr = "inproc://worker_death_test"
+    unit.zmq_context = MagicMock()
+
+    with (
+        patch.object(simple_storage, "create_zmq_socket", return_value=MagicMock()),
+        patch.object(simple_storage.zmq, "Poller", return_value=MagicMock()),
+        patch.object(_StorageUnitClass, "_worker_loop", side_effect=RuntimeError("boom")),
+        patch.object(simple_storage.os, "_exit") as exit_call,
+    ):
+        with pytest.raises(RuntimeError, match="boom"):
+            unit._worker_routine()
+
+    exit_call.assert_called_once_with(simple_storage._WORKER_DEATH_EXIT_CODE)
+
+
+def _failing_poller(shutdown_event: threading.Event | None = None) -> MagicMock:
+    """A proxy poller whose poll raises, optionally after a shutdown has begun."""
+
+    def poll(timeout):
+        if shutdown_event is not None:
+            shutdown_event.set()
+        raise RuntimeError("proxy boom")
+
+    poller = MagicMock()
+    poller.poll.side_effect = poll
+    return poller
+
+
+def test_proxy_death_exits_the_process():
+    """Without the proxy nothing reaches the worker, so the unit is equally unreachable."""
+    unit = _unit()
+    unit.put_get_socket = MagicMock()
+    unit.worker_socket = MagicMock()
+
+    with (
+        patch.object(simple_storage.zmq, "Poller", return_value=_failing_poller()),
+        patch.object(simple_storage.os, "_exit") as exit_call,
+    ):
+        unit._proxy_routine()
+
+    exit_call.assert_called_once_with(simple_storage._WORKER_DEATH_EXIT_CODE)
+
+
+def test_proxy_shutdown_does_not_exit():
+    unit = _unit()
+    unit.put_get_socket = MagicMock()
+    unit.worker_socket = MagicMock()
+
+    with (
+        patch.object(simple_storage.zmq, "Poller", return_value=_failing_poller(unit._shutdown_event)),
+        patch.object(simple_storage.os, "_exit") as exit_call,
+    ):
+        unit._proxy_routine()
+
+    exit_call.assert_not_called()

transfer_queue/storage/simple_storage.py

@@ -28,14 +28,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import math
 import os
 import pickle
 import shutil
+import sys
 import time
 import weakref
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import suppress
 from dataclasses import dataclass
 from pathlib import Path
 from threading import Event, Thread
@@ -49,7 +52,7 @@
 import zmq
 
 from transfer_queue.utils import compact_pickle
-from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads, log_heavy_operation
+from transfer_queue.utils.common import get_env_bool, limit_pytorch_auto_parallel_threads, log_heavy_operation
 from transfer_queue.utils.enum_utils import Role
 from transfer_queue.utils.logging_utils import get_logger
 from transfer_queue.utils.perf_utils import IntervalPerfMonitor
@@ -111,6 +114,31 @@ class _SSDValueRef:
 # Accept-queue sampling period in seconds; 0 disables the probe. Keep sub-second.
 TQ_ACCEPT_PROBE_INTERVAL = float(os.environ.get("TQ_ACCEPT_PROBE_INTERVAL", 0))
 
+# Exit the unit process when its worker or proxy thread dies, so the unit fails visibly. Both are
+# daemon threads, so without this the process keeps accepting requests it can never answer and
+# looks alive. Disable only to inspect a dead worker in place.
+TQ_STORAGE_EXIT_ON_WORKER_DEATH = get_env_bool("TQ_STORAGE_EXIT_ON_WORKER_DEATH", True)
+
+# Distinguishes a worker-death abort from a normal exit in Ray's actor-failure logs.
+_WORKER_DEATH_EXIT_CODE = 70
+
+
+def _flush_log_output() -> None:
+    """Flush every handler this module's records reach, and stdio, since ``os._exit`` skips both.
+
+    ``get_logger`` adds no handler when an ancestor already has one, so the records can live
+    entirely in a parent's handlers. A failing flush is ignored: the process exits regardless.
+    """
+    current: logging.Logger | None = logger
+    while current is not None:
+        for handler in current.handlers:
+            with suppress(Exception):
+                handler.flush()
+        current = current.parent if current.propagate else None
+    for stream in (sys.stdout, sys.stderr):
+        with suppress(Exception):
+            stream.flush()
+
 
 def _check_value_type(field: str, value: Any, is_tensor: bool, dtype: Any) -> None:
     """Reject a value that would change its field's tensor/non-tensor kind or tensor dtype."""
@@ -1068,7 +1096,10 @@ def _proxy_routine(self) -> None:
             if self._shutdown_event.is_set():
                 logger.info(f"[{self.storage_unit_id}]: ZMQ Proxy shutting down...")
             else:
-                logger.error(f"[{self.storage_unit_id}]: ZMQ Proxy unexpected error: {e}")
+                logger.exception(f"[{self.storage_unit_id}]: ZMQ Proxy unexpected error: {e}")
+                # Without the proxy nothing reaches the worker, so the unit is as unreachable
+                # as if the worker itself had died.
+                self._abort_process("the ZMQ proxy thread died")
 
     def _worker_routine(self) -> None:
         """Worker thread for processing requests."""
@@ -1082,6 +1113,46 @@ def _worker_routine(self) -> None:
         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)
+        except BaseException:
+            logger.critical(
+                f"[{self.storage_unit_id}]: worker thread is exiting on an unhandled exception.", exc_info=True
+            )
+            self._abort_process("the worker thread died")
+            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.
+
+        The worker runs as a daemon thread, so re-raising only unwinds that thread while the
+        process, its ROUTER and the proxy stay up and keep queueing requests nobody will ever
+        answer. ``os._exit`` skips atexit/flush handlers that could themselves block here.
+        """
+        if self._shutdown_event.is_set():
+            return
+
+        if not TQ_STORAGE_EXIT_ON_WORKER_DEATH:
+            logger.critical(
+                f"[{self.storage_unit_id}]: {reason}; staying up because "
+                f"TQ_STORAGE_EXIT_ON_WORKER_DEATH is disabled. Requests to this unit will hang."
+            )
+            return
+
+        logger.critical(
+            f"[{self.storage_unit_id}]: {reason}; exiting with code {_WORKER_DEATH_EXIT_CODE} so "
+            f"the unit fails visibly rather than lingering as one that cannot reply. "
+            f"Set TQ_STORAGE_EXIT_ON_WORKER_DEATH=0 to keep the process alive for inspection."
+        )
+        _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:
@@ -1103,10 +1174,6 @@ def _worker_routine(self) -> None:
                 # released before the worker waits for the next request.
                 self._process_one_worker_request(worker_socket, monitor)
 
-        logger.info(f"[{self.storage_unit_id}]: worker stopped.")
-        poller.unregister(worker_socket)
-        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."""
         # Messages received from proxy: [identity, serialized_msg_frame1, ...]