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..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 @@ -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, ...]