From bf12613bd150fdbaa516764fa3c7f5797213949e Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 9 Sep 2026 20:57:08 +0800 Subject: [PATCH 1/4] [storage] fix: exit a storage unit whose worker dies instead of hanging callers The worker and proxy threads are daemons, so an unhandled exception in either only unwinds that thread. The unit process stays up with its ROUTER still bound, the kernel keeps completing handshakes and queueing requests (RCVHWM/SNDHWM are 0, so nothing rejects them), and no one is left to answer. Callers cannot tell this from a slow unit, so each one blocks until its own recv timeout and raising that timeout only makes the stall longer. Exit the process instead. A unit that cannot answer is worth less than no unit at all, because its silence is indistinguishable from work in progress; exiting turns it into a visible Ray actor failure at the moment it breaks. Exit code 70 marks the cause, and os._exit avoids the atexit and flush handlers that could block on the same ZMQ state that just failed. The proxy gets the same treatment: without it nothing reaches the worker, so the unit is equally unreachable. Aborts are suppressed once _shutdown_event is set, since an exception racing a normal teardown is expected rather than a crash. TQ_STORAGE_EXIT_ON_WORKER_DEATH=0 restores the old behavior for anyone who needs to attach a profiler to the dead process, which is what TQ_HANG_HOLD_SECONDS already does on the manager side. Signed-off-by: OutstanderWang --- tests/test_storage_worker_death_exits.py | 134 +++++++++++++++++++++++ transfer_queue/storage/simple_storage.py | 63 ++++++++++- 2 files changed, 191 insertions(+), 6 deletions(-) create mode 100644 tests/test_storage_worker_death_exits.py diff --git a/tests/test_storage_worker_death_exits.py b/tests/test_storage_worker_death_exits.py new file mode 100644 index 00000000..863b6189 --- /dev/null +++ b/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() diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 60c2eda7..82039094 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -49,7 +49,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 +111,14 @@ 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 thread dies. On by default: the worker is a daemon +# thread, so without this the process keeps accepting requests it can never answer and every +# caller blocks until its own timeout. 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 _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 +1076,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 +1093,50 @@ 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. " + f"This storage unit can no longer answer any request even though its process and " + f"ROUTER socket stay up; every client will time out until the job is restarted.", + 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 callers fail fast instead of blocking on 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"callers fail fast rather than blocking on a unit that cannot reply. " + f"Set TQ_STORAGE_EXIT_ON_WORKER_DEATH=0 to keep the process alive for inspection." + ) + for handler in logger.handlers: + handler.flush() + 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 +1158,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, ...] From fb9e708eeddb54bf3cf5c62cbf7881d4c7f4e427 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Thu, 22 Oct 2026 20:24:09 +0800 Subject: [PATCH 2/4] [storage] doc: say a dead storage unit fails visibly, not that callers fail fast Exiting the unit does not shorten a caller's wait: a DEALER keeps queueing to a dead peer, so each request still runs to its send/recv timeout and retries. What the exit changes is that the unit shows up as a dead Ray actor with exit code 70, and the post-timeout probe reports tcp=down. Signed-off-by: OutstanderWang --- transfer_queue/storage/simple_storage.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index 82039094..e2700e62 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/transfer_queue/storage/simple_storage.py @@ -111,9 +111,9 @@ 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 thread dies. On by default: the worker is a daemon -# thread, so without this the process keeps accepting requests it can never answer and every -# caller blocks until its own timeout. Disable only to inspect a dead worker in place. +# 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. @@ -1110,7 +1110,7 @@ def _worker_routine(self) -> None: worker_socket.close(linger=0) def _abort_process(self, reason: str) -> None: - """Kill the whole unit process so callers fail fast instead of blocking on a dead unit. + """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 @@ -1128,7 +1128,7 @@ def _abort_process(self, reason: str) -> None: logger.critical( f"[{self.storage_unit_id}]: {reason}; exiting with code {_WORKER_DEATH_EXIT_CODE} so " - f"callers fail fast rather than blocking on a unit that cannot reply. " + 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." ) for handler in logger.handlers: From 137fbb694547a9d25427778f661f29f4823ec83f Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Thu, 22 Oct 2026 20:26:56 +0800 Subject: [PATCH 3/4] [storage] fix: flush inherited log handlers before exiting, and log a dead worker once get_logger adds no handler when an ancestor already has one, so flushing only the module logger's own handlers could flush nothing, and os._exit then drops the abort log that explains the exit. Flush every handler up the propagation chain, and stdout and stderr. The worker's CRITICAL line also said the process and ROUTER stay up, right before the exit; it now only reports the death with its traceback. The stay-up warning remains in the TQ_STORAGE_EXIT_ON_WORKER_DEATH=0 branch, where it holds. Signed-off-by: OutstanderWang --- tests/test_storage_worker_death_exits.py | 21 ++++++++++++++++++ transfer_queue/storage/simple_storage.py | 28 +++++++++++++++++++----- 2 files changed, 43 insertions(+), 6 deletions(-) diff --git a/tests/test_storage_worker_death_exits.py b/tests/test_storage_worker_death_exits.py index 863b6189..df88654b 100644 --- a/tests/test_storage_worker_death_exits.py +++ b/tests/test_storage_worker_death_exits.py @@ -21,6 +21,7 @@ silent hang into a visible actor failure. """ +import logging import threading from unittest.mock import MagicMock, patch @@ -49,6 +50,26 @@ def test_worker_death_exits_the_process(): exit_call.assert_called_once_with(simple_storage._WORKER_DEATH_EXIT_CODE) +def test_exit_flushes_handlers_inherited_from_a_parent_logger(): + """os._exit skips flushing, so the abort log must be flushed even when only an ancestor holds the handler.""" + unit = _unit() + parent = logging.getLogger("tq_worker_death_test_parent") + child = logging.getLogger("tq_worker_death_test_parent.child") + handler = MagicMock(spec=logging.Handler, level=logging.NOTSET) + parent.addHandler(handler) + try: + with ( + patch.object(simple_storage, "logger", child), + patch.object(simple_storage.os, "_exit") as exit_call, + ): + unit._abort_process("the worker thread died") + finally: + parent.removeHandler(handler) + + handler.flush.assert_called() + 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() diff --git a/transfer_queue/storage/simple_storage.py b/transfer_queue/storage/simple_storage.py index e2700e62..8b904bf4 100644 --- a/transfer_queue/storage/simple_storage.py +++ b/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 @@ -120,6 +123,23 @@ class _SSDValueRef: _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.""" if not is_tensor: @@ -1097,10 +1117,7 @@ def _worker_routine(self) -> None: 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. " - f"This storage unit can no longer answer any request even though its process and " - f"ROUTER socket stay up; every client will time out until the job is restarted.", - exc_info=True, + f"[{self.storage_unit_id}]: worker thread is exiting on an unhandled exception.", exc_info=True ) self._abort_process("the worker thread died") raise @@ -1131,8 +1148,7 @@ def _abort_process(self, reason: str) -> None: 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." ) - for handler in logger.handlers: - handler.flush() + _flush_log_output() os._exit(_WORKER_DEATH_EXIT_CODE) def _worker_loop(self, worker_socket: zmq.Socket, poller: zmq.Poller, perf_monitor: Any) -> None: From 4824c3cedb3ebc801d2f4e83edcdbebc2ec29a3e Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Thu, 22 Oct 2026 20:29:22 +0800 Subject: [PATCH 4/4] [test] drop the parent-logger flush test Signed-off-by: OutstanderWang --- tests/test_storage_worker_death_exits.py | 21 --------------------- 1 file changed, 21 deletions(-) diff --git a/tests/test_storage_worker_death_exits.py b/tests/test_storage_worker_death_exits.py index df88654b..863b6189 100644 --- a/tests/test_storage_worker_death_exits.py +++ b/tests/test_storage_worker_death_exits.py @@ -21,7 +21,6 @@ silent hang into a visible actor failure. """ -import logging import threading from unittest.mock import MagicMock, patch @@ -50,26 +49,6 @@ def test_worker_death_exits_the_process(): exit_call.assert_called_once_with(simple_storage._WORKER_DEATH_EXIT_CODE) -def test_exit_flushes_handlers_inherited_from_a_parent_logger(): - """os._exit skips flushing, so the abort log must be flushed even when only an ancestor holds the handler.""" - unit = _unit() - parent = logging.getLogger("tq_worker_death_test_parent") - child = logging.getLogger("tq_worker_death_test_parent.child") - handler = MagicMock(spec=logging.Handler, level=logging.NOTSET) - parent.addHandler(handler) - try: - with ( - patch.object(simple_storage, "logger", child), - patch.object(simple_storage.os, "_exit") as exit_call, - ): - unit._abort_process("the worker thread died") - finally: - parent.removeHandler(handler) - - handler.flush.assert_called() - 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()