diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py new file mode 100644 index 00000000..f43dcd96 --- /dev/null +++ b/tests/test_zmq_shared_context.py @@ -0,0 +1,587 @@ +# 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. + +Background: with_zmq_socket used to create a brand-new ``zmq.asyncio.Context()`` per RPC +call and ``context.term()`` it in the finally block. Under concurrency this churned +libzmq signaler file descriptors and crashed the process (``signaler.cpp`` Bad file +descriptor -> SIGABRT). The fix makes the decorator reuse the owner's long-lived context +(``get_context``) and only create/close the DEALER socket per call. + +These tests assert that concurrent decorated calls all reuse the SAME context object and +that the context is never terminated between calls, only when the client is closed. +""" + +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 KVStorageManager, StorageManager, StorageManagerFactory +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 + + +def test_factory_is_backend_agnostic(echo_controller): + """The client offers its context to every backend uniformly, naming none of them.""" + client = AsyncTransferQueueClient( + client_id="client_other_storage_context", + controller_info=echo_controller.zmq_server_info, + ) + config = {"client_name": "unused"} + + with patch("transfer_queue.client.StorageManagerFactory.create") as create_manager: + client.initialize_storage_manager("OtherStorage", config) + + create_manager.assert_called_once_with( + "OtherStorage", + controller_info=echo_controller.zmq_server_info, + config=config, + zmq_context=client.zmq_context, + ) + + client.close() + + +def test_kv_backends_keep_own_context(echo_controller): + """KV managers accept the shared context but deliberately keep an independent one. + + They move bulk data through their own SDKs and use ZMQ only for the controller + notify/handshake path, so they must not draw on the client's socket budget. + """ + client = AsyncTransferQueueClient( + client_id="client_kv_own_context", + controller_info=echo_controller.zmq_server_info, + ) + + with ( + patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"), + patch("transfer_queue.storage.managers.base.StorageClientFactory.create"), + ): + manager = KVStorageManager( + echo_controller.zmq_server_info, + {"client_name": "unused"}, + zmq_context=client.zmq_context, + ) + + assert manager.zmq_context is not client.zmq_context + assert manager._owns_zmq_context + + manager.close() + assert not client.zmq_context.closed + + client.close() + + +def test_factory_forwards_kwargs_to_registered_manager(echo_controller): + """The real factory forwards **kwargs verbatim, knowing no backend by name. + + The other factory test patches ``create`` out, so this one exercises the real + dispatch: a third-party manager registered from outside this package must receive + ``zmq_context`` without the factory special-casing its name. A regression to the old + ``if manager_type == "SimpleStorage"`` branch would silently drop the kwarg here, and + a manager whose signature drifts would raise TypeError -- neither of which mypy can + catch through ``**kwargs: Any``. + """ + received = {} + + @StorageManagerFactory.register("THIRD_PARTY_PROBE") + class ThirdPartyManager(StorageManager): + def __init__(self, controller_info, config, zmq_context=None): + received["zmq_context"] = zmq_context + received["config"] = config + super().__init__(controller_info, config, zmq_context=zmq_context) + + def _connect_to_controller(self): + pass + + def _do_handshake_with_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 + + async def notify_data_update(self, *args, **kwargs): + return None + + try: + client = AsyncTransferQueueClient( + client_id="client_third_party_factory", + controller_info=echo_controller.zmq_server_info, + ) + config = {"marker": "forwarded"} + + client.initialize_storage_manager("THIRD_PARTY_PROBE", config) + + # The kwarg survived dispatch through the unpatched factory... + assert received["zmq_context"] is client.zmq_context + assert received["config"] == config + # ...and a manager that opts in genuinely borrows rather than re-creating. + assert client.storage_manager.zmq_context is client.zmq_context + assert not client.storage_manager._owns_zmq_context + + # A borrower must not tear down a context it does not own. + client.storage_manager.close() + assert not client.zmq_context.closed + + client.close() + assert client.zmq_context.closed + finally: + StorageManagerFactory._registry.pop("THIRD_PARTY_PROBE", None) + + +@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_client_applies_max_sockets(echo_controller): + """The per-context socket ceiling is configurable, since it is now shared per client.""" + client = AsyncTransferQueueClient( + client_id="client_max_sockets", + controller_info=echo_controller.zmq_server_info, + zmq_max_sockets=2048, + ) + + assert client.zmq_context.get(zmq.MAX_SOCKETS) == 2048 + + client.close() + + +def test_client_rejects_max_sockets_above_build_limit(echo_controller): + """Values above this libzmq build's ZMQ_SOCKET_LIMIT are rejected up front.""" + probe = zmq.Context() + socket_limit = probe.get(zmq.SOCKET_LIMIT) + probe.term() + + with pytest.raises(ValueError, match="ZMQ_SOCKET_LIMIT"): + AsyncTransferQueueClient( + client_id="client_max_sockets_too_big", + controller_info=echo_controller.zmq_server_info, + zmq_max_sockets=socket_limit + 1, + ) + + +def test_max_sockets_from_env_var(echo_controller): + """TQ_CLIENT_ZMQ_MAX_SOCKETS configures the ceiling without touching call sites. + + The env var is the deployment-facing knob (the kwarg requires editing code), so it + needs its own coverage. Patched at the module constant because it is read at import. + """ + with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", "4096"): + client = AsyncTransferQueueClient( + client_id="client_max_sockets_env", + controller_info=echo_controller.zmq_server_info, + ) + + assert client.zmq_context.get(zmq.MAX_SOCKETS) == 4096 + client.close() + + +def test_explicit_max_sockets_overrides_env_var(echo_controller): + """An explicit kwarg wins over the env var, matching the io_threads precedence.""" + with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", "4096"): + client = AsyncTransferQueueClient( + client_id="client_max_sockets_precedence", + controller_info=echo_controller.zmq_server_info, + zmq_max_sockets=2048, + ) + + assert client.zmq_context.get(zmq.MAX_SOCKETS) == 2048 + client.close() + + +def test_unset_max_sockets_leaves_libzmq_default(echo_controller): + """Opt-in only: with nothing configured, libzmq's own default must be untouched.""" + probe = zmq.Context() + default = probe.get(zmq.MAX_SOCKETS) + probe.term() + + with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", None): + client = AsyncTransferQueueClient( + client_id="client_max_sockets_unset", + controller_info=echo_controller.zmq_server_info, + ) + + assert client.zmq_context.get(zmq.MAX_SOCKETS) == default + client.close() + + +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 + + def _do_handshake_with_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 + + async def notify_data_update(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 detects the failed shutdown but, because it does not own the context, has + no destroy() of its own to skip. If it stays silent the client proceeds to destroy a + context whose sockets that thread may still hold -- the documented non-thread-safe + Socket.close() hazard. The client must ask the manager 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 + + +def test_factory_tolerates_legacy_manager_signature(echo_controller): + """A manager on the old (controller_info, config) contract must still construct. + + Registration is an extension mechanism, so third-party managers are not required to + add ``zmq_context`` in lockstep. The factory drops keywords a constructor cannot + accept instead of raising TypeError. + """ + + @StorageManagerFactory.register("LEGACY_SIGNATURE_PROBE") + class LegacyManager(StorageManager): + def __init__(self, controller_info, config): # no zmq_context parameter + super().__init__(controller_info, config) + + def _connect_to_controller(self): + pass + + def _do_handshake_with_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 + + async def notify_data_update(self, *args, **kwargs): + return None + + try: + client = AsyncTransferQueueClient( + client_id="client_legacy_manager", + controller_info=echo_controller.zmq_server_info, + ) + + # Must not raise TypeError: unexpected keyword argument 'zmq_context'. + client.initialize_storage_manager("LEGACY_SIGNATURE_PROBE", {}) + + assert isinstance(client.storage_manager, LegacyManager) + # It never saw the context, so it owns the one it made and must not be vetoed on. + assert client.storage_manager.zmq_context is not client.zmq_context + assert client.storage_manager._owns_zmq_context + + client.storage_manager.close() + client.close() + finally: + StorageManagerFactory._registry.pop("LEGACY_SIGNATURE_PROBE", None) diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 4c0db125..ba45b60d 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -37,12 +37,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. The context serves every +# controller RPC (all backends) and, for SimpleStorage, storage-unit requests too, +# so this knob is client-scoped rather than backend-scoped. +TQ_CLIENT_ZMQ_IO_THREADS = int(os.environ.get("TQ_CLIENT_ZMQ_IO_THREADS", 8)) +# Per-context socket ceiling. Unset (or empty) means libzmq's default (1023). Because the +# context is now shared per client instead of created per call, all in-flight sockets draw +# on a single budget; raise this if a large num_data_storage_units fan-out exhausts it. +TQ_CLIENT_ZMQ_MAX_SOCKETS = os.environ.get("TQ_CLIENT_ZMQ_MAX_SOCKETS") or None # 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 +66,22 @@ 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 libzmq's own + default (1023) when that is unset. This ceiling is per-client: a single + ``put``/``get`` fans out one socket per storage unit, so deep fan-out + combined with high concurrency can approach it. Raising it also requires + enough file descriptors (``ulimit -n``). """ if controller_info is None: raise ValueError("controller_info cannot be None") @@ -70,6 +89,31 @@ 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. Its fixed native I/O-thread pool is shared + # by all controller RPCs and, for SimpleStorage only, storage-unit requests. + # Sockets remain 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 + 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 + if max_sockets is not None: + socket_limit = self.zmq_context.get(zmq.SOCKET_LIMIT) + 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}" + ) + self.zmq_context.set(zmq.MAX_SOCKETS, max_sockets) logger.info(f"[{self.client_id}]: Registered Controller server {controller_info.id} at {controller_info.ip}") def initialize_storage_manager( @@ -79,6 +123,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 +136,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 +1138,15 @@ 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 still in flight: this tears down the shared + ZMQ context via ``destroy()``, which calls ``Socket.close()`` internally and is + **not** thread-safe. Driving this client from your own event loop means closing + it only after outstanding tasks have been awaited or cancelled. Subclasses that + own a loop thread must join it before delegating here (see + :meth:`TransferQueueClient.close`). + """ try: if hasattr(self, "storage_manager") and self.storage_manager: if hasattr(self.storage_manager, "close"): @@ -1095,6 +1154,39 @@ def close(self) -> None: except Exception as e: logger.warning(f"Error closing storage manager: {e}") + # Tear down the shared context last. destroy(linger=0) force-closes any socket that + # leaked from an interrupted RPC then terminates, so shutdown cannot hang. + if not self._can_destroy_zmq_context(): + 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}") + + def _can_destroy_zmq_context(self) -> bool: + """Whether it is safe to call ``destroy()`` on the shared context. + + This class owns no background thread of its own, so the caller's quiescence + contract (documented on :meth:`close`) covers the client side. But 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 whether it finished + shutting down. Subclasses that run 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 +1322,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 libzmq's own + default (1023) 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 +1874,16 @@ 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. + + The ordering here is load-bearing and must not be rearranged: the background + loop thread is stopped and joined *before* delegating to + :meth:`AsyncTransferQueueClient.close`, which destroys the shared ZMQ context. + ``destroy()`` calls ``Socket.close()`` internally and is not thread-safe, so no + other thread may hold sockets on that context when it runs. If the join times + out, :meth:`_can_destroy_zmq_context` reports False and the context is leaked + instead of destroyed unsafely. + """ if hasattr(self, "_loop") and self._loop is not None: self._loop.call_soon_threadsafe(self._loop.stop) @@ -1789,3 +1899,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..aafa0691 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,12 @@ 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 whose fixed native I/O thread pool + # is shared by controller and storage-unit request sockets; SimpleStorage does. + # A manager that is handed nothing owns the context it creates, and 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 +395,59 @@ 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() + # Record the outcome even when the context is borrowed: the owner cannot see this + # thread, so a borrower that stayed silent here would let the owner destroy() a + # context whose sockets are still in use. See can_destroy_zmq_context(). + self._notify_thread_stopped = notify_thread_stopped + + if self._owns_zmq_context: + # Ordering below is load-bearing: destroy() calls Socket.close() internally, + # which is NOT thread-safe, so it must run only after the notify thread that + # owns sockets on this context is gone. If that thread outlived its join + # timeout, leak the context rather than risk a crash during shutdown -- the + # process is terminating anyway, so a leaked context is the cheaper outcome. + if notify_thread_stopped: + # destroy(linger=0) force-closes any socket still open (e.g. from an interrupted + # request or the notify path) then terminates, so shutdown 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()`` calls ``Socket.close()`` internally and is not thread-safe, so the + notify thread must be gone first. Only this manager can see that thread, so an + owner of a borrowed context must ask before tearing the context down. + + Checks the thread directly rather than trusting the flag recorded by ``close()``, + so this is also correct if 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 +471,57 @@ 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 keyword arguments are forwarded to the registered class. The factory + deliberately knows nothing about any individual backend: each manager decides for + itself what to do with what it receives (e.g. whether to borrow a caller-supplied + ``zmq_context`` or keep its own). + + Registration is an extension mechanism, so managers written against the older + ``(controller_info, config)`` contract must keep working without being updated in + lockstep. Any keyword the constructor does not accept is therefore dropped, with a + warning, rather than raising ``TypeError``. + """ 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 +530,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..aec260b6 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -49,6 +49,10 @@ "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 (base.py). Now shared by both the + # notify path (_notify_and_wait) and the per-call storage-unit request sockets below; + # this is 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 +73,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..09fc5fb5 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -356,6 +356,7 @@ 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, ): @@ -363,7 +364,16 @@ def with_zmq_socket( 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. + get owner's shared context -> create/connect socket -> inject socket -> close socket. + + The ZMQ context is owned by ``self`` (via ``get_context``) and is long-lived: it is + created once per owner and reused across all calls, then terminated once when the owner + is closed. Only the DEALER *socket* is created and closed per call. Do NOT create or + terminate a context here -- per-call context churn corrupts libzmq's signaler file + descriptors under concurrency (Bad file descriptor / SIGABRT) and can hang on term(). + Contexts are thread-safe and event-loop-agnostic, so a single shared context is safe + even when decorated methods run on different loops/threads; each socket is created and + fully used within one awaited call on one loop. Args: socket_name: Socket port key in ``ZMQServerInfo.ports``. @@ -373,6 +383,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 +410,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 +427,11 @@ 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 any unsent frames immediately (the reply is already + # received on the happy path) so close never blocks the event loop. + if sock is not None and not sock.closed: + sock.close(linger=0) return wrapper