diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py new file mode 100644 index 00000000..a4dde94f --- /dev/null +++ b/tests/test_zmq_shared_context.py @@ -0,0 +1,253 @@ +# 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 +from transfer_queue.metadata import BatchMeta +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, + simple_storage_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, + simple_storage_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_other_backends_do_not_borrow_client_context(echo_controller): + 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, + ) + + client.close() + + +@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 diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 4c0db125..656331cc 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -37,12 +37,14 @@ logger = get_logger(__name__) TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8)) +TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS = int(os.environ.get("TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS", 8)) # 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 +59,16 @@ def __init__( self, client_id: str, controller_info: ZMQServerInfo, + simple_storage_zmq_io_threads: int | None = None, ): """Initialize the asynchronous TransferQueue client. Args: client_id: Unique identifier for this client instance controller_info: Single controller ZMQ server information + simple_storage_zmq_io_threads: Fixed size of the client context's + native I/O-thread pool. Defaults to + ``TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS`` (8). """ if controller_info is None: raise ValueError("controller_info cannot be None") @@ -70,6 +76,20 @@ 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_SIMPLE_STORAGE_ZMQ_IO_THREADS + if simple_storage_zmq_io_threads is None + else simple_storage_zmq_io_threads + ) + if io_threads < 1: + raise ValueError( + "SimpleStorage ZMQ I/O thread pool size must be at least 1, " + f"got {io_threads}" + ) + self.zmq_context = zmq.asyncio.Context(io_threads=io_threads) logger.info(f"[{self.client_id}]: Registered Controller server {controller_info.id} at {controller_info.ip}") def initialize_storage_manager( @@ -87,8 +107,14 @@ def initialize_storage_manager( - zmq_info: ZMQ server information about the storage units """ + create_kwargs = {} + if manager_type == "SimpleStorage": + create_kwargs["zmq_context"] = self.zmq_context self.storage_manager = StorageManagerFactory.create( - manager_type, controller_info=self._controller, config=config + manager_type, + controller_info=self._controller, + config=config, + **create_kwargs, ) # ==================== Basic API ==================== @@ -1095,6 +1121,14 @@ 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. + 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}") + # ==================== Checkpoint API ==================== @with_controller_socket async def async_save_controller_checkpoint( @@ -1230,16 +1264,21 @@ def __init__( self, client_id: str, controller_info: ZMQServerInfo, + simple_storage_zmq_io_threads: int | None = None, ): """Initialize the synchronous TransferQueue client. Args: client_id: Unique identifier for this client instance controller_info: Single controller ZMQ server information + simple_storage_zmq_io_threads: Fixed size of the client context's + native I/O-thread pool. Defaults to + ``TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS`` (8). """ super().__init__( client_id, controller_info, + simple_storage_zmq_io_threads, ) # create new event loop in a separate thread diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index e6b0faf4..f435fd1e 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -63,7 +63,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 +76,11 @@ 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() + # SimpleStorage may borrow a client-owned context whose fixed native I/O + # thread pool is shared by controller and storage-unit request sockets. + # Other backends and standalone managers retain their own context. + self._owns_zmq_context = zmq_context is None + self.zmq_context = zmq_context or zmq.asyncio.Context() self._connect_to_controller() # Dedicated asyncio loop for ZMQ notify traffic, isolated from the caller's loop @@ -391,7 +400,10 @@ def close(self) -> None: else: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") - self.zmq_context.term() + if self._owns_zmq_context: + # 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) def __del__(self): """Destructor to ensure resources are cleaned up.""" @@ -422,12 +434,21 @@ def decorator(manager_cls: type[StorageManager]): return decorator @classmethod - def create(cls, manager_type: str, controller_info: ZMQServerInfo, config: dict[str, Any]) -> StorageManager: + def create( + cls, + manager_type: str, + controller_info: ZMQServerInfo, + config: dict[str, Any], + zmq_context: zmq.asyncio.Context | None = None, + ) -> StorageManager: """Create and return a StorageManager instance.""" 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] + if manager_type == "SimpleStorage" and zmq_context is not None: + return manager_cls(controller_info, config, zmq_context=zmq_context) + return manager_cls(controller_info, config) class KVStorageManager(StorageManager): 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/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 93656542..147340e4 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -316,6 +316,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, ): @@ -323,7 +324,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``. @@ -333,6 +343,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")`` @@ -358,7 +370,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) @@ -371,14 +387,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