diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py new file mode 100644 index 00000000..2bd9a6c7 --- /dev/null +++ b/tests/test_zmq_shared_context.py @@ -0,0 +1,338 @@ +# 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. + +"""Regression tests for the shared long-lived ZMQ context in with_zmq_socket. + +with_zmq_socket used to create and term() a context per RPC call, which churned libzmq +signaler file descriptors and crashed the process under concurrency (Bad file descriptor +-> SIGABRT). It now reuses the owner's context and only creates the socket per call, so +these tests assert every concurrent call sees the SAME context, alive until close(). +""" + +import asyncio +from threading import Thread +from unittest.mock import patch + +import pytest +import zmq + +import transfer_queue.utils.zmq_utils as zmq_utils +from transfer_queue.client import AsyncTransferQueueClient, TransferQueueClient +from transfer_queue.metadata import BatchMeta +from transfer_queue.storage.managers.base import StorageManager +from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager +from transfer_queue.utils.enum_utils import Role +from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo + + +class _EchoController: + """Minimal in-process ROUTER controller that answers GET_META requests.""" + + def __init__(self, controller_id="controller_0"): + self.controller_id = controller_id + self.context = zmq.Context() + self.request_socket = self.context.socket(zmq.ROUTER) + self.request_port = self.request_socket.bind_to_random_port("tcp://127.0.0.1") + self.zmq_server_info = ZMQServerInfo( + role=Role.CONTROLLER, + id=controller_id, + ip="127.0.0.1", + ports={"request_handle_socket": self.request_port}, + ) + self.running = True + self.request_thread = Thread(target=self._handle_requests, daemon=True) + self.request_thread.start() + + def _handle_requests(self): + poller = zmq.Poller() + poller.register(self.request_socket, zmq.POLLIN) + while self.running: + try: + socks = dict(poller.poll(100)) + if self.request_socket not in socks: + continue + messages = self.request_socket.recv_multipart(copy=False) + identity = messages.pop(0) + request_msg = ZMQMessage.deserialize(messages) + + batch_size = request_msg.body.get("batch_size", 1) + data_fields = request_msg.body.get("data_fields", []) + field_schema = { + name: {"dtype": None, "shape": None, "is_nested": False, "is_non_tensor": False} + for name in data_fields + } + metadata = BatchMeta( + global_indexes=list(range(batch_size)), + partition_ids=["0"] * batch_size, + field_schema=field_schema, + ) + response_msg = ZMQMessage.create( + request_type=ZMQRequestType.GET_META_RESPONSE, + sender_id=self.controller_id, + receiver_id=request_msg.sender_id, + body={"metadata": metadata}, + ) + self.request_socket.send_multipart([identity, *response_msg.serialize()]) + except zmq.Again: + continue + except Exception as e: # pragma: no cover - surfaced via test failure + print(f"_EchoController ERROR: {e}") + + def stop(self): + self.running = False + self.request_thread.join(timeout=2.0) + self.request_socket.close(linger=0) + self.context.term() + + +@pytest.fixture +def echo_controller(): + controller = _EchoController() + yield controller + controller.stop() + + +@pytest.mark.asyncio +async def test_shared_context_reused_across_concurrent_calls(echo_controller, monkeypatch): + """Many concurrent decorated RPCs must all reuse the client's single context. + + This is the core regression guard: pre-fix, each call created and term()ed its own + context, which is what corrupted libzmq's signaler FDs under concurrency. + """ + client = AsyncTransferQueueClient( + client_id="client_shared_ctx", + controller_info=echo_controller.zmq_server_info, + ) + + # Record the context object handed to every socket creation inside the decorator. + seen_contexts = [] + original_create = zmq_utils.create_zmq_socket + + def _spy_create(ctx, *args, **kwargs): + seen_contexts.append(ctx) + return original_create(ctx, *args, **kwargs) + + # The decorator resolves create_zmq_socket from the zmq_utils module globals. + monkeypatch.setattr(zmq_utils, "create_zmq_socket", _spy_create) + + num_calls = 200 + coros = [ + client.async_get_meta(data_fields=["tokens", "labels"], batch_size=2, partition_id="0") + for _ in range(num_calls) + ] + # wait_for guards against the hang failure mode. + results = await asyncio.wait_for(asyncio.gather(*coros), timeout=60) + + assert len(results) == num_calls + assert all(isinstance(meta, BatchMeta) for meta in results) + + # Every call must have used the SAME context, and it must be the client's context. + assert len(seen_contexts) == num_calls + assert all(ctx is client.zmq_context for ctx in seen_contexts) + # The shared context must NOT have been terminated by any call. + assert not client.zmq_context.closed + + client.close() + + +def test_client_context_has_fixed_io_thread_pool(echo_controller): + client = AsyncTransferQueueClient( + client_id="client_fixed_context_pool", + controller_info=echo_controller.zmq_server_info, + zmq_io_threads=4, + ) + + assert client.zmq_context.get(zmq.IO_THREADS) == 4 + + client.close() + + +def test_client_rejects_invalid_context_pool_size(echo_controller): + with pytest.raises(ValueError, match="at least 1"): + AsyncTransferQueueClient( + client_id="client_invalid_context_pool", + controller_info=echo_controller.zmq_server_info, + zmq_io_threads=0, + ) + + +def test_simple_storage_borrows_client_context(echo_controller): + client = AsyncTransferQueueClient( + client_id="client_simple_storage_context", + controller_info=echo_controller.zmq_server_info, + ) + config = {"zmq_info": {}} + + with patch("transfer_queue.client.StorageManagerFactory.create") as create_manager: + client.initialize_storage_manager("SimpleStorage", config) + + create_manager.assert_called_once_with( + "SimpleStorage", + controller_info=echo_controller.zmq_server_info, + config=config, + zmq_context=client.zmq_context, + ) + + client.close() + + +def test_simple_storage_does_not_destroy_borrowed_context(echo_controller): + client = AsyncTransferQueueClient( + client_id="client_borrowed_context_lifecycle", + controller_info=echo_controller.zmq_server_info, + ) + + with patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"): + manager = AsyncSimpleStorageManager( + echo_controller.zmq_server_info, + {"zmq_info": {"storage_0": echo_controller.zmq_server_info}}, + zmq_context=client.zmq_context, + ) + + assert manager.zmq_context is client.zmq_context + assert not manager._owns_zmq_context + + manager.close() + assert not client.zmq_context.closed + + client.close() + assert client.zmq_context.closed + + +@pytest.mark.asyncio +async def test_close_destroys_context(echo_controller): + """close() must terminate the shared context exactly once (no leak, no hang).""" + client = AsyncTransferQueueClient( + client_id="client_close_ctx", + controller_info=echo_controller.zmq_server_info, + ) + assert not client.zmq_context.closed + + # A normal call before shutdown leaves the context alive. + await asyncio.wait_for( + client.async_get_meta(data_fields=["tokens"], batch_size=1, partition_id="0"), + timeout=30, + ) + assert not client.zmq_context.closed + + client.close() + assert client.zmq_context.closed + + +def test_non_numeric_max_sockets_env_var_names_the_variable(echo_controller): + """A typo'd value must say which env var is wrong, not raise a bare int() error.""" + with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", "not-a-number"): + with pytest.raises(ValueError, match="TQ_CLIENT_ZMQ_MAX_SOCKETS must be an integer"): + AsyncTransferQueueClient( + client_id="client_max_sockets_garbage", + controller_info=echo_controller.zmq_server_info, + ) + + +def test_close_skips_destroy_while_loop_thread_alive(echo_controller): + """destroy() is not thread-safe, so a stuck loop thread must veto it. + + TransferQueueClient.close() joins its loop thread with a timeout that only warns on + expiry. Falling through to destroy() with that thread still holding sockets is the + documented hazard, so the context is leaked instead. + """ + client = TransferQueueClient( + client_id="client_stuck_thread", + controller_info=echo_controller.zmq_server_info, + ) + context = client.zmq_context + + # Simulate a loop thread that outlived its join timeout. + with patch.object(client._thread, "is_alive", return_value=True): + assert client._can_destroy_zmq_context() is False + client.close() + + assert not context.closed, "context must be leaked, not destroyed unsafely" + + # With the thread genuinely gone, the veto lifts. + assert client._can_destroy_zmq_context() is True + context.destroy(linger=0) + + +def _make_borrowing_manager(zmq_context): + """A minimal manager that borrows a caller's context, like SimpleStorage does.""" + + class Borrower(StorageManager): + def _connect_to_controller(self): + pass + + async def put_data(self, *args, **kwargs): + return None + + async def get_data(self, *args, **kwargs): + return None + + async def clear_data(self, *args, **kwargs): + return None + + return Borrower(None, {}, zmq_context=zmq_context) + + +def test_stuck_notify_thread_vetoes_destroy_of_borrowed_context(echo_controller): + """A borrowing manager's stuck notify thread must veto the owner's destroy(). + + The manager has no destroy() of its own to skip, so staying silent would let the client + destroy a context whose sockets that thread may still hold. It must be asked first. + """ + client = AsyncTransferQueueClient( + client_id="client_notify_thread_veto", + controller_info=echo_controller.zmq_server_info, + ) + client.storage_manager = _make_borrowing_manager(client.zmq_context) + context = client.zmq_context + + with patch.object(client.storage_manager._notify_thread, "is_alive", return_value=True): + assert client._can_destroy_zmq_context() is False + client.close() + assert not context.closed, "context must be leaked while the notify thread lives" + + # Veto lifts once the thread is genuinely gone. + assert client._can_destroy_zmq_context() is True + context.destroy(linger=0) + + +def test_healthy_notify_thread_does_not_block_destroy(echo_controller): + """The veto must not over-trigger: a clean manager shutdown still destroys.""" + client = AsyncTransferQueueClient( + client_id="client_notify_thread_clean", + controller_info=echo_controller.zmq_server_info, + ) + client.storage_manager = _make_borrowing_manager(client.zmq_context) + + client.close() + assert client.zmq_context.closed + + +def test_manager_with_own_context_does_not_veto(echo_controller): + """A manager holding its own context has no say in the client's teardown.""" + client = AsyncTransferQueueClient( + client_id="client_independent_manager", + controller_info=echo_controller.zmq_server_info, + ) + client.storage_manager = _make_borrowing_manager(None) # creates its own context + assert client.storage_manager.zmq_context is not client.zmq_context + + # Even a stuck notify thread on an unrelated context must not block the client. + with patch.object(client.storage_manager._notify_thread, "is_alive", return_value=True): + assert client._can_destroy_zmq_context() is True + + client.storage_manager.zmq_context.destroy(linger=0) + client.close() + assert client.zmq_context.closed diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 4c0db125..a0b82f92 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -16,6 +16,7 @@ import asyncio import os import threading +import weakref from typing import Any, Callable import torch @@ -37,12 +38,21 @@ logger = get_logger(__name__) TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8)) +# Size of the client context's native I/O-thread pool, shared by all controller RPCs and, +# for SimpleStorage, storage-unit requests -- so this knob is client- not backend-scoped. +TQ_CLIENT_ZMQ_IO_THREADS = int(os.environ.get("TQ_CLIENT_ZMQ_IO_THREADS", 8)) +# Per-context socket ceiling, shared by every in-flight socket on the client's context +# (including a borrowing storage manager's), so libzmq's own 1023 is reachable at scale. +# Raising it also needs enough file descriptors (``ulimit -n``). +TQ_CLIENT_ZMQ_MAX_SOCKETS = os.environ.get("TQ_CLIENT_ZMQ_MAX_SOCKETS") or None +DEFAULT_CLIENT_ZMQ_MAX_SOCKETS = 8192 # Pre-bound decorator for controller socket operations. with_controller_socket = with_zmq_socket( "request_handle_socket", get_identity=lambda self: self.client_id, get_peer=lambda self, target: self._controller, + get_context=lambda self: self.zmq_context, ) @@ -57,12 +67,23 @@ def __init__( self, client_id: str, controller_info: ZMQServerInfo, + zmq_io_threads: int | None = None, + zmq_max_sockets: int | None = None, ): """Initialize the asynchronous TransferQueue client. Args: client_id: Unique identifier for this client instance controller_info: Single controller ZMQ server information + zmq_io_threads: Fixed size of the client context's native I/O-thread pool. + Defaults to ``TQ_CLIENT_ZMQ_IO_THREADS`` (8). + zmq_max_sockets: Maximum number of sockets the client context may hold open + at once. Defaults to ``TQ_CLIENT_ZMQ_MAX_SOCKETS``, and to + ``DEFAULT_CLIENT_ZMQ_MAX_SOCKETS`` (8192) when that is unset -- well above + libzmq's own 1023, which large-scale training can exhaust. This ceiling is + per-client: a single ``put``/``get`` fans out one socket per storage unit, + and a borrowing storage manager shares the same budget. Raising it also + requires enough file descriptors (``ulimit -n``). """ if controller_info is None: raise ValueError("controller_info cannot be None") @@ -70,8 +91,65 @@ def __init__( raise TypeError(f"controller_info must be ZMQServerInfo, got {type(controller_info)}") self.client_id = client_id self._controller: ZMQServerInfo = controller_info + # One long-lived context per client; sockets stay per-request because ZMQ sockets + # are not thread-safe. + io_threads = TQ_CLIENT_ZMQ_IO_THREADS if zmq_io_threads is None else zmq_io_threads + if io_threads < 1: + raise ValueError(f"Client ZMQ I/O thread pool size must be at least 1, got {io_threads}") + self.zmq_context = zmq.asyncio.Context(io_threads=io_threads) + + max_sockets = zmq_max_sockets + explicitly_requested = max_sockets is not None + if max_sockets is None and TQ_CLIENT_ZMQ_MAX_SOCKETS is not None: + try: + max_sockets = int(TQ_CLIENT_ZMQ_MAX_SOCKETS) + except ValueError as e: + # Name the variable: a bare int() error gives no clue which knob is wrong. + raise ValueError( + f"TQ_CLIENT_ZMQ_MAX_SOCKETS must be an integer, got {TQ_CLIENT_ZMQ_MAX_SOCKETS!r}" + ) from e + explicitly_requested = True + if max_sockets is None: + max_sockets = DEFAULT_CLIENT_ZMQ_MAX_SOCKETS + + socket_limit = self.zmq_context.get(zmq.SOCKET_LIMIT) + if explicitly_requested: + # A value the caller asked for must not be silently reinterpreted. + if not 1 <= max_sockets <= socket_limit: + raise ValueError( + f"Client ZMQ max sockets must be between 1 and this build's " + f"ZMQ_SOCKET_LIMIT ({socket_limit}), got {max_sockets}" + ) + elif max_sockets > socket_limit: + # Nobody asked for the default, so clamp instead of failing to construct on a + # build whose ZMQ_SOCKET_LIMIT is below it. + logger.debug( + f"[{client_id}]: Clamping default ZMQ max sockets {max_sockets} to this " + f"build's ZMQ_SOCKET_LIMIT ({socket_limit})." + ) + max_sockets = socket_limit + self.zmq_context.set(zmq.MAX_SOCKETS, max_sockets) + + # Backstop for a client that is never closed, so the context and its I/O threads do + # not leak for the process lifetime. finalize() (not __del__) also runs at + # interpreter exit; close() detaches it. See _release_zmq_context(). + self._finalizer = weakref.finalize(self, self._release_zmq_context, self.zmq_context, client_id) logger.info(f"[{self.client_id}]: Registered Controller server {controller_info.id} at {controller_info.ip}") + @staticmethod + def _release_zmq_context(context: "zmq.asyncio.Context", client_id: str) -> None: + """Destroy *context* if it is still open. + + Must stay a staticmethod taking the context explicitly, or the bound method would + keep the client alive and never fire. TransferQueueClient's loop thread references + the client, deferring this to interpreter exit -- close() remains the supported path. + """ + try: + if not context.closed: + context.destroy(linger=0) + except Exception as e: + logger.warning(f"[{client_id}]: Error destroying zmq_context in finalizer: {e}") + def initialize_storage_manager( self, manager_type: str, @@ -79,6 +157,10 @@ def initialize_storage_manager( ): """Initialize the storage manager. + The client's long-lived ZMQ context is offered to every backend uniformly; each + registered manager decides whether to borrow it or keep its own, so the client + needs no knowledge of specific backend names. + Args: manager_type: Type of storage manager to create. Supported types include: AsyncSimpleStorageManager, KVStorageManager (under development), etc. @@ -88,7 +170,10 @@ def initialize_storage_manager( """ self.storage_manager = StorageManagerFactory.create( - manager_type, controller_info=self._controller, config=config + manager_type, + controller_info=self._controller, + config=config, + zmq_context=self.zmq_context, ) # ==================== Basic API ==================== @@ -1087,7 +1172,12 @@ async def async_kv_list( raise RuntimeError(f"[{self.client_id}]: Error in kv_list: {str(e)}") from e def close(self) -> None: - """Close the client and cleanup resources including storage manager.""" + """Close the client and cleanup resources including storage manager. + + The caller must ensure no RPCs are in flight: this destroys the shared context, and + ``destroy()`` calls ``Socket.close()`` internally, which is **not** thread-safe. + Subclasses owning a loop thread must join it first (see TransferQueueClient.close). + """ try: if hasattr(self, "storage_manager") and self.storage_manager: if hasattr(self.storage_manager, "close"): @@ -1095,6 +1185,47 @@ def close(self) -> None: except Exception as e: logger.warning(f"Error closing storage manager: {e}") + # Tear down the shared context last; linger=0 so shutdown cannot hang. + if not self._can_destroy_zmq_context(): + # Detach first, or the finalizer would later run the destroy() judged unsafe here. + self._detach_finalizer() + logger.warning( + f"[{self.client_id}]: Skipping zmq_context.destroy() because a thread owning " + f"sockets on it is still alive; destroy() is not thread-safe. Leaking the context." + ) + return + try: + if hasattr(self, "zmq_context") and self.zmq_context is not None: + self.zmq_context.destroy(linger=0) + except Exception as e: + logger.warning(f"[{self.client_id}]: Error terminating zmq_context: {e}") + finally: + # This close() superseded the finalizer; disarm it so it cannot run twice. + self._detach_finalizer() + + def _detach_finalizer(self) -> None: + """Disarm the context finalizer, if one was armed.""" + finalizer = getattr(self, "_finalizer", None) + if finalizer is not None: + finalizer.detach() + + def _can_destroy_zmq_context(self) -> bool: + """Whether it is safe to call ``destroy()`` on the shared context. + + A storage manager that *borrowed* this context runs a notify thread the client + cannot see, and that thread holds sockets on it, so ask the manager first. + Subclasses running their own loop thread extend this. + """ + manager = getattr(self, "storage_manager", None) + if manager is not None: + can_destroy = getattr(manager, "can_destroy_zmq_context", None) + # Only consult a manager that actually shares this context; one with its own + # context has no say in when the client's is destroyed. + if callable(can_destroy) and getattr(manager, "zmq_context", None) is self.zmq_context: + if not can_destroy(): + return False + return True + # ==================== Checkpoint API ==================== @with_controller_socket async def async_save_controller_checkpoint( @@ -1230,16 +1361,25 @@ def __init__( self, client_id: str, controller_info: ZMQServerInfo, + zmq_io_threads: int | None = None, + zmq_max_sockets: int | None = None, ): """Initialize the synchronous TransferQueue client. Args: client_id: Unique identifier for this client instance controller_info: Single controller ZMQ server information + zmq_io_threads: Fixed size of the client context's native I/O-thread pool. + Defaults to ``TQ_CLIENT_ZMQ_IO_THREADS`` (8). + zmq_max_sockets: Maximum number of sockets the client context may hold open + at once. Defaults to ``TQ_CLIENT_ZMQ_MAX_SOCKETS``, and to + ``DEFAULT_CLIENT_ZMQ_MAX_SOCKETS`` (8192) when that is unset. """ super().__init__( client_id, controller_info, + zmq_io_threads, + zmq_max_sockets, ) # create new event loop in a separate thread @@ -1773,7 +1913,12 @@ def load_storage_checkpoint(self, checkpoint_dir: str) -> None: return self._load_storage_checkpoint(checkpoint_dir) def close(self) -> None: - """Close the client and cleanup resources including event loop and thread.""" + """Close the client and cleanup resources including event loop and thread. + + Ordering is load-bearing: the loop thread is joined *before* the base close() + destroys the context, since ``destroy()`` is not thread-safe. If the join times + out, ``_can_destroy_zmq_context()`` returns False and the context is leaked instead. + """ if hasattr(self, "_loop") and self._loop is not None: self._loop.call_soon_threadsafe(self._loop.stop) @@ -1789,3 +1934,14 @@ def close(self) -> None: logger.warning(f"[{self.client_id}]: Error closing event loop: {e}") super().close() + + def _can_destroy_zmq_context(self) -> bool: + """False while the loop thread that owns sockets on the context is still alive. + + Also defers to the base check, which covers a borrowing storage manager's notify + thread -- both threads must be gone before ``destroy()`` is safe. + """ + thread = getattr(self, "_thread", None) + if thread is not None and thread.is_alive(): + return False + return super()._can_destroy_zmq_context() diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index e6b0faf4..9c32e213 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -14,6 +14,7 @@ # limitations under the License. import asyncio +import inspect import itertools import os import threading @@ -63,7 +64,12 @@ class StorageManager(ABC): """Base class for storage layer. It defines the interface for data operations and generally provides handshake & notification capabilities.""" - def __init__(self, controller_info: ZMQServerInfo, config: DictConfig): + def __init__( + self, + controller_info: ZMQServerInfo, + config: DictConfig, + zmq_context: zmq.asyncio.Context | None = None, + ): self.storage_manager_id = f"TQ_STORAGE_{uuid4().hex[:8]}" self.config = config self.controller_info = controller_info @@ -71,7 +77,10 @@ def __init__(self, controller_info: ZMQServerInfo, config: DictConfig): # Handshake socket is sync (used only during initialization) self.controller_handshake_socket: zmq.Socket | None = None - self.zmq_context = zmq.asyncio.Context() + # A manager may borrow a caller-owned context (SimpleStorage does) or own the one it + # creates when handed nothing. Only an owner tears its context down (see close()). + self._owns_zmq_context = zmq_context is None + self.zmq_context = zmq.asyncio.Context() if zmq_context is None else zmq_context self._connect_to_controller() # Dedicated asyncio loop for ZMQ notify traffic, isolated from the caller's loop @@ -384,21 +393,49 @@ def close(self) -> None: if hasattr(self, "_notify_loop") and self._notify_loop.is_running(): self._notify_loop.call_soon_threadsafe(self._notify_loop.stop) + notify_thread_stopped = True if hasattr(self, "_notify_thread") and self._notify_thread is not None: self._notify_thread.join(timeout=5.0) if self._notify_thread.is_alive(): + notify_thread_stopped = False logger.warning(f"[{self.storage_manager_id}]: Notify ZMQ thread did not stop within 5 second timeout.") else: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") - self.zmq_context.term() + if self._owns_zmq_context: + # destroy() calls Socket.close(), which is not thread-safe, so it must run only + # after the notify thread holding sockets is gone. If that thread outlived its + # join, leak the context rather than risk a crash on a terminating process. + if notify_thread_stopped: + # linger=0 force-closes sockets left by an interrupted request, so this + # cannot hang on term(). + self.zmq_context.destroy(linger=0) + else: + logger.warning( + f"[{self.storage_manager_id}]: Skipping zmq_context.destroy() because the notify " + f"thread is still alive; destroy() is not thread-safe while sockets are in use. " + f"The context will be leaked." + ) + + def can_destroy_zmq_context(self) -> bool: + """Whether an owner may safely ``destroy()`` a context this manager borrowed. + + ``destroy()`` is not thread-safe, so the notify thread must be gone first, and only + this manager can see it. Checks the thread directly so it stays correct when called + before ``close()`` or if the thread exited late. + """ + thread = getattr(self, "_notify_thread", None) + return thread is None or not thread.is_alive() def __del__(self): """Destructor to ensure resources are cleaned up.""" try: self.close() except Exception as e: - logger.error(f"[{self.storage_manager_id}]: Exception during __del__: {str(e)}") + # __init__ may have failed before storage_manager_id was set; reaching for it + # here would raise AttributeError and mask the exception we mean to report. + manager_id = getattr(self, "storage_manager_id", f"") + logger.error(f"[{manager_id}]: Exception during __del__: {str(e)}") class StorageManagerFactory: @@ -422,12 +459,52 @@ def decorator(manager_cls: type[StorageManager]): return decorator @classmethod - def create(cls, manager_type: str, controller_info: ZMQServerInfo, config: dict[str, Any]) -> StorageManager: - """Create and return a StorageManager instance.""" + def create( + cls, + manager_type: str, + controller_info: ZMQServerInfo, + config: dict[str, Any], + **kwargs: Any, + ) -> StorageManager: + """Create and return a StorageManager instance. + + Extra keywords are forwarded verbatim; the factory knows no backend by name, so each + manager decides what to do with what it receives. Keywords a constructor cannot + accept are dropped with a warning, keeping older ``(controller_info, config)`` + managers working without a lockstep update. + """ assert manager_type in cls._registry, ( f"Unknown manager_type: {manager_type}. Supported managers include: {list(cls._registry.keys())}" ) - return cls._registry[manager_type](controller_info, config) + manager_cls = cls._registry[manager_type] + accepted = cls._filter_supported_kwargs(manager_cls, kwargs, manager_type) + return manager_cls(controller_info, config, **accepted) + + @staticmethod + def _filter_supported_kwargs( + manager_cls: type[StorageManager], kwargs: dict[str, Any], manager_type: str + ) -> dict[str, Any]: + """Drop keywords ``manager_cls.__init__`` cannot accept, warning about each. + + A constructor taking ``**kwargs`` is assumed to accept everything. + """ + if not kwargs: + return kwargs + try: + params = inspect.signature(manager_cls.__init__).parameters + except (TypeError, ValueError): + # Un-introspectable constructor (e.g. a C extension): pass through unchanged + # rather than silently dropping arguments the class may well accept. + return kwargs + if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): + return kwargs + supported = {name: value for name, value in kwargs.items() if name in params} + for name in kwargs.keys() - supported.keys(): + logger.warning( + f"{manager_cls.__name__} (registered as '{manager_type}') does not accept " + f"'{name}'; ignoring it. Add '{name}' to its __init__ signature to use it." + ) + return supported class KVStorageManager(StorageManager): @@ -436,9 +513,22 @@ class KVStorageManager(StorageManager): It maps structured metadata (BatchMeta) to flat lists of keys and values for efficient KV operations. """ - def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): + def __init__( + self, + controller_info: ZMQServerInfo, + config: dict[str, Any], + zmq_context: zmq.asyncio.Context | None = None, + ): """ Initialize the KVStorageManager with configuration. + + Args: + controller_info: Controller ZMQ server information. + config: Backend configuration; must contain ``client_name``. + zmq_context: Accepted for interface uniformity but deliberately ignored. + KV backends move bulk data through their own SDKs and use ZMQ only for + the controller notify/handshake path, so they keep an independent + context rather than drawing on a caller's shared socket budget. """ client_name = config.get("client_name", None) if client_name is None: diff --git a/transfer_queue/storage/managers/mooncake_manager.py b/transfer_queue/storage/managers/mooncake_manager.py index c3e8f5ce..48fe6280 100644 --- a/transfer_queue/storage/managers/mooncake_manager.py +++ b/transfer_queue/storage/managers/mooncake_manager.py @@ -15,6 +15,8 @@ from typing import Any +import zmq + from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory from transfer_queue.utils.zmq_utils import ZMQServerInfo @@ -29,6 +31,11 @@ class MooncakeStorageManager(KVStorageManager): pybind bindings. """ - def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): + def __init__( + self, + controller_info: ZMQServerInfo, + config: dict[str, Any], + zmq_context: zmq.asyncio.Context | None = None, + ): config["client_name"] = "MooncakeStoreClient" - super().__init__(controller_info, config) + super().__init__(controller_info, config, zmq_context=zmq_context) diff --git a/transfer_queue/storage/managers/ray_storage_manager.py b/transfer_queue/storage/managers/ray_storage_manager.py index 0cc2a09c..a48176e4 100644 --- a/transfer_queue/storage/managers/ray_storage_manager.py +++ b/transfer_queue/storage/managers/ray_storage_manager.py @@ -15,6 +15,8 @@ from typing import Any +import zmq + from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory from transfer_queue.utils.zmq_utils import ZMQServerInfo @@ -23,8 +25,13 @@ class RayStorageManager(KVStorageManager): """Storage manager for Ray-RDT backend.""" - def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): + def __init__( + self, + controller_info: ZMQServerInfo, + config: dict[str, Any], + zmq_context: zmq.asyncio.Context | None = None, + ): config = (config or {}).copy() if config.get("client_name") not in (None, "RayStorageClient"): raise ValueError(f"RayStorageManager only supports 'RayStorageClient', got: {config.get('client_name')}") - super().__init__(controller_info, {**config, "client_name": "RayStorageClient"}) + super().__init__(controller_info, {**config, "client_name": "RayStorageClient"}, zmq_context=zmq_context) diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 0b6777fb..acc73a28 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -49,6 +49,9 @@ "put_get_socket", get_identity=lambda self: self.storage_manager_id, get_peer=lambda self, target: self.storage_unit_infos[target], + # Long-lived context from the base StorageManager, shared with the notify path. Safe + # because the context is loop-agnostic and each socket stays per-call. + get_context=lambda self: self.zmq_context, resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"), timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, ) @@ -69,8 +72,13 @@ class AsyncSimpleStorageManager(StorageManager): instances using ZMQ communication and dynamic socket management. """ - def __init__(self, controller_info: ZMQServerInfo, config: DictConfig): - super().__init__(controller_info, config) + def __init__( + self, + controller_info: ZMQServerInfo, + config: DictConfig, + zmq_context: zmq.asyncio.Context | None = None, + ): + super().__init__(controller_info, config, zmq_context=zmq_context) self.config = config server_infos: ZMQServerInfo | dict[str, ZMQServerInfo] | None = config.get("zmq_info", None) diff --git a/transfer_queue/storage/managers/yuanrong_manager.py b/transfer_queue/storage/managers/yuanrong_manager.py index f76b47b2..26c30c4b 100644 --- a/transfer_queue/storage/managers/yuanrong_manager.py +++ b/transfer_queue/storage/managers/yuanrong_manager.py @@ -15,6 +15,8 @@ from typing import Any +import zmq + from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ZMQServerInfo @@ -26,7 +28,12 @@ class YuanrongStorageManager(KVStorageManager): """Storage manager for Yuanrong backend.""" - def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): + def __init__( + self, + controller_info: ZMQServerInfo, + config: dict[str, Any], + zmq_context: zmq.asyncio.Context | None = None, + ): worker_port = config.get("worker_port", None) client_name = config.get("client_name", None) @@ -38,4 +45,4 @@ def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): config["client_name"] = "YuanrongStorageClient" elif client_name != "YuanrongStorageClient": raise ValueError(f"Invalid 'client_name': {client_name} in config. Expecting 'YuanrongStorageClient'") - super().__init__(controller_info, config) + super().__init__(controller_info, config, zmq_context=zmq_context) diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 49e8e674..2ca63d7f 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -356,14 +356,18 @@ def with_zmq_socket( *, get_identity: Callable[[Any], str], get_peer: Callable[[Any, str | None], ZMQServerInfo], + get_context: Callable[[Any], "zmq.asyncio.Context"], resolve_target: Callable[[tuple, dict], str | None] | None = None, timeout: int | None = None, ): """Create a reusable async decorator for request sockets. - This decorator encapsulates the common socket lifecycle used by both - client-side and storage-manager-side request paths: - create context/socket -> connect -> inject socket -> close/term. + Lifecycle: get owner's shared context -> create/connect socket -> inject -> close socket. + + The context comes from ``self`` via ``get_context`` and is long-lived; only the DEALER + socket is per-call. Do NOT create or terminate a context here -- per-call churn corrupts + libzmq's signaler file descriptors under concurrency (Bad file descriptor / SIGABRT) and + can hang on term(). Contexts are thread-safe and loop-agnostic, so sharing one is safe. Args: socket_name: Socket port key in ``ZMQServerInfo.ports``. @@ -373,6 +377,8 @@ def with_zmq_socket( For single-target scenarios, ignore the target parameter. Example: ``lambda self, target: self.server_info`` Example: ``lambda self, target: self.storage_unit_infos[target]`` + get_context: Callable that returns the owner's long-lived ``zmq.asyncio.Context``. + Example: ``lambda self: self.zmq_context`` resolve_target: Optional callable that extracts target identifier from function arguments. Receives (args, kwargs) and returns target name. Example: ``lambda args, kwargs: kwargs.get("target_storage_unit")`` @@ -398,7 +404,11 @@ async def wrapper(self, *args, **kwargs): if port is None: raise RuntimeError(f"Socket '{socket_name}' not configured for server '{server_info.id}'") - context = zmq.asyncio.Context() + # Reuse the owner's long-lived context; only the socket is per-call. + context = get_context(self) + if context is None: + raise RuntimeError("get_context returned None") + sock = None try: address = format_zmq_address(server_info.ip, port) @@ -411,14 +421,10 @@ async def wrapper(self, *args, **kwargs): kwargs["socket"] = sock return await func(self, *args, **kwargs) finally: - if sock is not None: - try: - if not sock.closed: - sock.close(linger=-1) - finally: - context.term() - else: - context.term() + # Close the per-call socket only; the context outlives the call. linger=0 + # drops unsent frames so close never blocks the event loop. + if sock is not None and not sock.closed: + sock.close(linger=0) return wrapper