From f6b3b64c5b8ff1655a8770a8a0fe17650d3566f4 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Thu, 23 Jul 2026 17:04:16 +0800 Subject: [PATCH 01/10] feat: reuse long-lived ZMQ context in with_zmq_socket instead of per-call The with_zmq_socket decorator created a brand-new zmq.asyncio.Context() per RPC call and context.term()'d it in the finally block. Under the high-concurrency agent-loop store path this churned libzmq's signaler file descriptors and crashed the worker (signaler.cpp Bad file descriptor -> SIGABRT), and could also hang on the blocking term() (aggravated by sock.close(linger=-1)). Fix: the decorator now reuses the owner's long-lived context via a required get_context callable, and only creates/closes the DEALER socket per call. The context is created once per owner and terminated once at close(). Contexts are thread-safe and event-loop-agnostic, so a single shared context is safe across loops/threads; each socket stays per-call on one loop. - zmq_utils.with_zmq_socket: add required get_context; drop per-call Context()/term(); change sock.close(linger=-1) -> linger=0. - client.AsyncTransferQueueClient: own a shared self.zmq_context; destroy(linger=0) in close(). - simple_storage_manager: feed the base StorageManager's self.zmq_context via get_context. - base.StorageManager.close(): term() -> destroy(linger=0) so a leaked socket cannot hang shutdown. - tests: add test_zmq_shared_context.py asserting concurrent RPCs reuse one context and it is closed exactly once. Microbenchmark: ~7.6x faster socket setup/teardown (~210us saved per call). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 168 ++++++++++++++++++ transfer_queue/client.py | 13 ++ transfer_queue/storage/managers/base.py | 4 +- .../managers/simple_storage_manager.py | 4 + transfer_queue/utils/zmq_utils.py | 33 ++-- 5 files changed, 211 insertions(+), 11 deletions(-) create mode 100644 tests/test_zmq_shared_context.py diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py new file mode 100644 index 00000000..37700189 --- /dev/null +++ b/tests/test_zmq_shared_context.py @@ -0,0 +1,168 @@ +# 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 + +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.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() + + +@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..6b088785 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -43,6 +43,7 @@ "request_handle_socket", get_identity=lambda self: self.client_id, get_peer=lambda self, target: self._controller, + get_context=lambda self: self.zmq_context, ) @@ -70,6 +71,10 @@ def __init__( raise TypeError(f"controller_info must be ZMQServerInfo, got {type(controller_info)}") self.client_id = client_id self._controller: ZMQServerInfo = controller_info + # Long-lived ZMQ context shared by all controller RPCs on this client. Contexts are + # thread-safe and event-loop-agnostic; each RPC creates and closes its own DEALER + # socket from this context (see with_controller_socket). Terminated once in close(). + self.zmq_context = zmq.asyncio.Context() logger.info(f"[{self.client_id}]: Registered Controller server {controller_info.id} at {controller_info.ip}") def initialize_storage_manager( @@ -1095,6 +1100,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( diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index e6b0faf4..c6c959c4 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -391,7 +391,9 @@ def close(self) -> None: else: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") - self.zmq_context.term() + # 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.""" diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 0b6777fb..c6b20eff 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, ) 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 From 04224ba35696eeff5a4a5f65764574694beca538 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 29 Jul 2026 22:10:36 +0800 Subject: [PATCH 02/10] feat: share fixed ZMQ context pool across client requests Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 49 +++++++++++++++++- transfer_queue/client.py | 22 ++++++-- transfer_queue/storage/managers/base.py | 51 +++++++++++++++---- .../storage/managers/mooncake_manager.py | 11 +++- .../storage/managers/ray_storage_manager.py | 15 +++++- .../managers/simple_storage_manager.py | 9 +++- .../storage/managers/yuanrong_manager.py | 11 +++- transfer_queue/utils/zmq_utils.py | 22 +++++++- 8 files changed, 163 insertions(+), 27 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 37700189..808df2a3 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -21,12 +21,13 @@ 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. +These tests assert that concurrent decorated calls all reuse the SAME fixed-size context +pool 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 @@ -148,6 +149,50 @@ def _spy_create(ctx, *args, **kwargs): client.close() +def test_context_uses_configured_fixed_io_thread_pool(echo_controller): + """All sockets from a client share the configured native ZMQ I/O-thread pool.""" + client = AsyncTransferQueueClient( + client_id="client_fixed_ctx_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_context_rejects_invalid_io_thread_pool_size(echo_controller): + with pytest.raises(ValueError, match="at least 1"): + AsyncTransferQueueClient( + client_id="client_invalid_ctx_pool", + controller_info=echo_controller.zmq_server_info, + zmq_io_threads=0, + ) + + +def test_client_shares_context_pool_with_storage_manager(echo_controller): + client = AsyncTransferQueueClient( + client_id="client_storage_ctx_pool", + controller_info=echo_controller.zmq_server_info, + zmq_io_threads=4, + ) + config = {"client_name": "unused"} + + with patch("transfer_queue.client.StorageManagerFactory.create") as create_manager: + client.initialize_storage_manager("unused", config) + + create_manager.assert_called_once_with( + "unused", + controller_info=echo_controller.zmq_server_info, + config=config, + zmq_context=client.zmq_context, + ) + assert config == {"client_name": "unused"} + + 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).""" diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 6b088785..fda39a0b 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -31,6 +31,7 @@ ZMQMessage, ZMQRequestType, ZMQServerInfo, + create_zmq_context, with_zmq_socket, ) @@ -58,12 +59,15 @@ def __init__( self, client_id: str, controller_info: ZMQServerInfo, + 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 + zmq_io_threads: Size of the long-lived ZMQ context's native I/O-thread + pool. Defaults to ``TQ_ZMQ_IO_THREADS`` (8). """ if controller_info is None: raise ValueError("controller_info cannot be None") @@ -71,10 +75,11 @@ def __init__( raise TypeError(f"controller_info must be ZMQServerInfo, got {type(controller_info)}") self.client_id = client_id self._controller: ZMQServerInfo = controller_info - # Long-lived ZMQ context shared by all controller RPCs on this client. Contexts are - # thread-safe and event-loop-agnostic; each RPC creates and closes its own DEALER - # socket from this context (see with_controller_socket). Terminated once in close(). - self.zmq_context = zmq.asyncio.Context() + self._zmq_io_threads = zmq_io_threads + # One long-lived ZMQ context per client. Its fixed native I/O-thread pool is shared + # by all concurrent RPC sockets; sockets remain per-request because ZMQ sockets are + # not thread-safe. The context is terminated once in close(). + self.zmq_context = create_zmq_context(zmq_io_threads) logger.info(f"[{self.client_id}]: Registered Controller server {controller_info.id} at {controller_info.ip}") def initialize_storage_manager( @@ -93,7 +98,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 ==================== @@ -1243,16 +1251,20 @@ def __init__( self, client_id: str, controller_info: ZMQServerInfo, + 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 + zmq_io_threads: Size of the long-lived ZMQ context's native I/O-thread + pool. Defaults to ``TQ_ZMQ_IO_THREADS`` (8). """ super().__init__( client_id, controller_info, + 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 c6c959c4..a0effa6f 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 @@ -35,7 +36,13 @@ from transfer_queue.metadata import BatchMeta, extract_field_schema from transfer_queue.storage.clients.base import StorageClientFactory from transfer_queue.utils.logging_utils import get_logger -from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo, create_zmq_socket +from transfer_queue.utils.zmq_utils import ( + ZMQMessage, + ZMQRequestType, + ZMQServerInfo, + create_zmq_context, + create_zmq_socket, +) logger = get_logger(__name__) @@ -63,7 +70,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 +83,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() + # A manager created by TransferQueueClient borrows the client's context so + # controller and storage requests share one fixed native I/O-thread pool. + # Standalone managers create and own an equivalent long-lived context. + self._owns_zmq_context = zmq_context is None + self.zmq_context = zmq_context or create_zmq_context(config.get("zmq_io_threads", None)) self._connect_to_controller() # Dedicated asyncio loop for ZMQ notify traffic, isolated from the caller's loop @@ -391,9 +407,10 @@ def close(self) -> None: else: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") - # 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) + 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.""" @@ -424,12 +441,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 zmq_context is not None and "zmq_context" in inspect.signature(manager_cls).parameters: + return manager_cls(controller_info, config, zmq_context=zmq_context) + return manager_cls(controller_info, config) class KVStorageManager(StorageManager): @@ -438,14 +464,19 @@ 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. """ client_name = config.get("client_name", None) if client_name is None: raise ValueError("Missing client_name in config") - super().__init__(controller_info, config) + super().__init__(controller_info, config, zmq_context=zmq_context) self.storage_client = StorageClientFactory.create(client_name, config) self._multi_threads_executor: ThreadPoolExecutor | None = None self._executor_finalizer = weakref.finalize(self, self._shutdown_executor, self._multi_threads_executor) diff --git a/transfer_queue/storage/managers/mooncake_manager.py b/transfer_queue/storage/managers/mooncake_manager.py index c3e8f5ce..a929d6b7 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.asyncio + 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..f91f008a 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.asyncio + from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory from transfer_queue.utils.zmq_utils import ZMQServerInfo @@ -23,8 +25,17 @@ 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 c6b20eff..aec260b6 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -73,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..a409cb66 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.asyncio + 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 09fc5fb5..713e6396 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os import socket import time from collections.abc import Sequence @@ -305,6 +306,22 @@ def get_free_port(ip: str) -> int: return sock.getsockname()[1] +TQ_ZMQ_IO_THREADS = int(os.environ.get("TQ_ZMQ_IO_THREADS", 8)) + + +def create_zmq_context(io_threads: int | None = None) -> "zmq.asyncio.Context": + """Create a long-lived async ZMQ context with a fixed I/O-thread pool. + + A ZMQ context owns the native I/O-thread pool used by all sockets created from + that context. Keeping one context per owner lets concurrent request sockets share + the whole pool without creating or terminating contexts per request. + """ + pool_size = TQ_ZMQ_IO_THREADS if io_threads is None else io_threads + if pool_size < 1: + raise ValueError(f"ZMQ I/O thread pool size must be at least 1, got {pool_size}") + return zmq.asyncio.Context(io_threads=pool_size) + + def create_zmq_socket( ctx: zmq.Context, socket_type: Any, @@ -372,8 +389,9 @@ def with_zmq_socket( 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. + even when decorated methods run on different loops/threads. The context's fixed native + I/O-thread pool is shared by all request sockets; each socket is created and fully used + within one awaited call on one loop. Args: socket_name: Socket port key in ``ZMQServerInfo.ports``. From 8c405815bb359d9e2171fd25264932b415418253 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 29 Jul 2026 22:16:43 +0800 Subject: [PATCH 03/10] fix: limit shared ZMQ context to SimpleStorage Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 26 ++++++++++++++++--- transfer_queue/client.py | 5 +++- transfer_queue/storage/managers/base.py | 20 +++++--------- .../storage/managers/mooncake_manager.py | 11 ++------ .../storage/managers/ray_storage_manager.py | 15 ++--------- .../storage/managers/yuanrong_manager.py | 11 ++------ 6 files changed, 39 insertions(+), 49 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 808df2a3..ce0f5b62 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -171,7 +171,7 @@ def test_context_rejects_invalid_io_thread_pool_size(echo_controller): ) -def test_client_shares_context_pool_with_storage_manager(echo_controller): +def test_client_shares_context_pool_with_simple_storage_manager(echo_controller): client = AsyncTransferQueueClient( client_id="client_storage_ctx_pool", controller_info=echo_controller.zmq_server_info, @@ -180,10 +180,10 @@ def test_client_shares_context_pool_with_storage_manager(echo_controller): config = {"client_name": "unused"} with patch("transfer_queue.client.StorageManagerFactory.create") as create_manager: - client.initialize_storage_manager("unused", config) + client.initialize_storage_manager("SimpleStorage", config) create_manager.assert_called_once_with( - "unused", + "SimpleStorage", controller_info=echo_controller.zmq_server_info, config=config, zmq_context=client.zmq_context, @@ -193,6 +193,26 @@ def test_client_shares_context_pool_with_storage_manager(echo_controller): client.close() +def test_client_does_not_pass_context_to_other_storage_backends(echo_controller): + client = AsyncTransferQueueClient( + client_id="client_other_storage_ctx", + controller_info=echo_controller.zmq_server_info, + zmq_io_threads=4, + ) + 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).""" diff --git a/transfer_queue/client.py b/transfer_queue/client.py index fda39a0b..26ccd0b5 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -97,11 +97,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, - zmq_context=self.zmq_context, + **create_kwargs, ) # ==================== Basic API ==================== diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index a0effa6f..6c683670 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -14,7 +14,6 @@ # limitations under the License. import asyncio -import inspect import itertools import os import threading @@ -40,7 +39,6 @@ ZMQMessage, ZMQRequestType, ZMQServerInfo, - create_zmq_context, create_zmq_socket, ) @@ -83,11 +81,10 @@ def __init__( # Handshake socket is sync (used only during initialization) self.controller_handshake_socket: zmq.Socket | None = None - # A manager created by TransferQueueClient borrows the client's context so - # controller and storage requests share one fixed native I/O-thread pool. - # Standalone managers create and own an equivalent long-lived context. + # SimpleStorage can borrow the client's long-lived context. Other storage + # backends retain the original behavior and own their default ZMQ context. self._owns_zmq_context = zmq_context is None - self.zmq_context = zmq_context or create_zmq_context(config.get("zmq_io_threads", 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 @@ -453,7 +450,7 @@ def create( f"Unknown manager_type: {manager_type}. Supported managers include: {list(cls._registry.keys())}" ) manager_cls = cls._registry[manager_type] - if zmq_context is not None and "zmq_context" in inspect.signature(manager_cls).parameters: + 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) @@ -464,19 +461,14 @@ 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], - zmq_context: zmq.asyncio.Context | None = None, - ): + def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): """ Initialize the KVStorageManager with configuration. """ client_name = config.get("client_name", None) if client_name is None: raise ValueError("Missing client_name in config") - super().__init__(controller_info, config, zmq_context=zmq_context) + super().__init__(controller_info, config) self.storage_client = StorageClientFactory.create(client_name, config) self._multi_threads_executor: ThreadPoolExecutor | None = None self._executor_finalizer = weakref.finalize(self, self._shutdown_executor, self._multi_threads_executor) diff --git a/transfer_queue/storage/managers/mooncake_manager.py b/transfer_queue/storage/managers/mooncake_manager.py index a929d6b7..c3e8f5ce 100644 --- a/transfer_queue/storage/managers/mooncake_manager.py +++ b/transfer_queue/storage/managers/mooncake_manager.py @@ -15,8 +15,6 @@ from typing import Any -import zmq.asyncio - from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory from transfer_queue.utils.zmq_utils import ZMQServerInfo @@ -31,11 +29,6 @@ class MooncakeStorageManager(KVStorageManager): pybind bindings. """ - def __init__( - self, - controller_info: ZMQServerInfo, - config: dict[str, Any], - zmq_context: zmq.asyncio.Context | None = None, - ): + def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): config["client_name"] = "MooncakeStoreClient" - super().__init__(controller_info, config, zmq_context=zmq_context) + super().__init__(controller_info, config) diff --git a/transfer_queue/storage/managers/ray_storage_manager.py b/transfer_queue/storage/managers/ray_storage_manager.py index f91f008a..0cc2a09c 100644 --- a/transfer_queue/storage/managers/ray_storage_manager.py +++ b/transfer_queue/storage/managers/ray_storage_manager.py @@ -15,8 +15,6 @@ from typing import Any -import zmq.asyncio - from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory from transfer_queue.utils.zmq_utils import ZMQServerInfo @@ -25,17 +23,8 @@ class RayStorageManager(KVStorageManager): """Storage manager for Ray-RDT backend.""" - def __init__( - self, - controller_info: ZMQServerInfo, - config: dict[str, Any], - zmq_context: zmq.asyncio.Context | None = None, - ): + def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): 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"}, - zmq_context=zmq_context, - ) + super().__init__(controller_info, {**config, "client_name": "RayStorageClient"}) diff --git a/transfer_queue/storage/managers/yuanrong_manager.py b/transfer_queue/storage/managers/yuanrong_manager.py index a409cb66..f76b47b2 100644 --- a/transfer_queue/storage/managers/yuanrong_manager.py +++ b/transfer_queue/storage/managers/yuanrong_manager.py @@ -15,8 +15,6 @@ from typing import Any -import zmq.asyncio - 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 @@ -28,12 +26,7 @@ class YuanrongStorageManager(KVStorageManager): """Storage manager for Yuanrong backend.""" - def __init__( - self, - controller_info: ZMQServerInfo, - config: dict[str, Any], - zmq_context: zmq.asyncio.Context | None = None, - ): + def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): worker_port = config.get("worker_port", None) client_name = config.get("client_name", None) @@ -45,4 +38,4 @@ def __init__( 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, zmq_context=zmq_context) + super().__init__(controller_info, config) From 762728b07c743ef8dc6bca750e28546a640d5ca4 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 29 Jul 2026 22:28:45 +0800 Subject: [PATCH 04/10] Revert "fix: limit shared ZMQ context to SimpleStorage" This reverts commit b54c4093d71348b3c82572ff18c3a35a0bb9dd2b. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 26 +++---------------- transfer_queue/client.py | 5 +--- transfer_queue/storage/managers/base.py | 20 +++++++++----- .../storage/managers/mooncake_manager.py | 11 ++++++-- .../storage/managers/ray_storage_manager.py | 15 +++++++++-- .../storage/managers/yuanrong_manager.py | 11 ++++++-- 6 files changed, 49 insertions(+), 39 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index ce0f5b62..808df2a3 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -171,7 +171,7 @@ def test_context_rejects_invalid_io_thread_pool_size(echo_controller): ) -def test_client_shares_context_pool_with_simple_storage_manager(echo_controller): +def test_client_shares_context_pool_with_storage_manager(echo_controller): client = AsyncTransferQueueClient( client_id="client_storage_ctx_pool", controller_info=echo_controller.zmq_server_info, @@ -180,10 +180,10 @@ def test_client_shares_context_pool_with_simple_storage_manager(echo_controller) config = {"client_name": "unused"} with patch("transfer_queue.client.StorageManagerFactory.create") as create_manager: - client.initialize_storage_manager("SimpleStorage", config) + client.initialize_storage_manager("unused", config) create_manager.assert_called_once_with( - "SimpleStorage", + "unused", controller_info=echo_controller.zmq_server_info, config=config, zmq_context=client.zmq_context, @@ -193,26 +193,6 @@ def test_client_shares_context_pool_with_simple_storage_manager(echo_controller) client.close() -def test_client_does_not_pass_context_to_other_storage_backends(echo_controller): - client = AsyncTransferQueueClient( - client_id="client_other_storage_ctx", - controller_info=echo_controller.zmq_server_info, - zmq_io_threads=4, - ) - 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).""" diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 26ccd0b5..fda39a0b 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -97,14 +97,11 @@ 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, - **create_kwargs, + zmq_context=self.zmq_context, ) # ==================== Basic API ==================== diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 6c683670..a0effa6f 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 @@ -39,6 +40,7 @@ ZMQMessage, ZMQRequestType, ZMQServerInfo, + create_zmq_context, create_zmq_socket, ) @@ -81,10 +83,11 @@ def __init__( # Handshake socket is sync (used only during initialization) self.controller_handshake_socket: zmq.Socket | None = None - # SimpleStorage can borrow the client's long-lived context. Other storage - # backends retain the original behavior and own their default ZMQ context. + # A manager created by TransferQueueClient borrows the client's context so + # controller and storage requests share one fixed native I/O-thread pool. + # Standalone managers create and own an equivalent long-lived context. self._owns_zmq_context = zmq_context is None - self.zmq_context = zmq_context or zmq.asyncio.Context() + self.zmq_context = zmq_context or create_zmq_context(config.get("zmq_io_threads", None)) self._connect_to_controller() # Dedicated asyncio loop for ZMQ notify traffic, isolated from the caller's loop @@ -450,7 +453,7 @@ def create( f"Unknown manager_type: {manager_type}. Supported managers include: {list(cls._registry.keys())}" ) manager_cls = cls._registry[manager_type] - if manager_type == "SimpleStorage" and zmq_context is not None: + if zmq_context is not None and "zmq_context" in inspect.signature(manager_cls).parameters: return manager_cls(controller_info, config, zmq_context=zmq_context) return manager_cls(controller_info, config) @@ -461,14 +464,19 @@ 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. """ client_name = config.get("client_name", None) if client_name is None: raise ValueError("Missing client_name in config") - super().__init__(controller_info, config) + super().__init__(controller_info, config, zmq_context=zmq_context) self.storage_client = StorageClientFactory.create(client_name, config) self._multi_threads_executor: ThreadPoolExecutor | None = None self._executor_finalizer = weakref.finalize(self, self._shutdown_executor, self._multi_threads_executor) diff --git a/transfer_queue/storage/managers/mooncake_manager.py b/transfer_queue/storage/managers/mooncake_manager.py index c3e8f5ce..a929d6b7 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.asyncio + 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..f91f008a 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.asyncio + from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory from transfer_queue.utils.zmq_utils import ZMQServerInfo @@ -23,8 +25,17 @@ 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/yuanrong_manager.py b/transfer_queue/storage/managers/yuanrong_manager.py index f76b47b2..a409cb66 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.asyncio + 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) From 30b83abb50b8285e0423a5c60ccd2973dd0197b3 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 29 Jul 2026 22:31:51 +0800 Subject: [PATCH 05/10] Revert "feat: share fixed ZMQ context pool across client requests" This reverts commit a8bfbd81c68226f0c679ce3673c467866d470881. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 49 +----------------- transfer_queue/client.py | 22 ++------ transfer_queue/storage/managers/base.py | 51 ++++--------------- .../storage/managers/mooncake_manager.py | 11 +--- .../storage/managers/ray_storage_manager.py | 15 +----- .../managers/simple_storage_manager.py | 9 +--- .../storage/managers/yuanrong_manager.py | 11 +--- transfer_queue/utils/zmq_utils.py | 22 +------- 8 files changed, 27 insertions(+), 163 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 808df2a3..37700189 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -21,13 +21,12 @@ 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 fixed-size context -pool and that the context is never terminated between calls, only when the client is closed. +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 @@ -149,50 +148,6 @@ def _spy_create(ctx, *args, **kwargs): client.close() -def test_context_uses_configured_fixed_io_thread_pool(echo_controller): - """All sockets from a client share the configured native ZMQ I/O-thread pool.""" - client = AsyncTransferQueueClient( - client_id="client_fixed_ctx_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_context_rejects_invalid_io_thread_pool_size(echo_controller): - with pytest.raises(ValueError, match="at least 1"): - AsyncTransferQueueClient( - client_id="client_invalid_ctx_pool", - controller_info=echo_controller.zmq_server_info, - zmq_io_threads=0, - ) - - -def test_client_shares_context_pool_with_storage_manager(echo_controller): - client = AsyncTransferQueueClient( - client_id="client_storage_ctx_pool", - controller_info=echo_controller.zmq_server_info, - zmq_io_threads=4, - ) - config = {"client_name": "unused"} - - with patch("transfer_queue.client.StorageManagerFactory.create") as create_manager: - client.initialize_storage_manager("unused", config) - - create_manager.assert_called_once_with( - "unused", - controller_info=echo_controller.zmq_server_info, - config=config, - zmq_context=client.zmq_context, - ) - assert config == {"client_name": "unused"} - - 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).""" diff --git a/transfer_queue/client.py b/transfer_queue/client.py index fda39a0b..6b088785 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -31,7 +31,6 @@ ZMQMessage, ZMQRequestType, ZMQServerInfo, - create_zmq_context, with_zmq_socket, ) @@ -59,15 +58,12 @@ def __init__( self, client_id: str, controller_info: ZMQServerInfo, - 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 - zmq_io_threads: Size of the long-lived ZMQ context's native I/O-thread - pool. Defaults to ``TQ_ZMQ_IO_THREADS`` (8). """ if controller_info is None: raise ValueError("controller_info cannot be None") @@ -75,11 +71,10 @@ def __init__( raise TypeError(f"controller_info must be ZMQServerInfo, got {type(controller_info)}") self.client_id = client_id self._controller: ZMQServerInfo = controller_info - self._zmq_io_threads = zmq_io_threads - # One long-lived ZMQ context per client. Its fixed native I/O-thread pool is shared - # by all concurrent RPC sockets; sockets remain per-request because ZMQ sockets are - # not thread-safe. The context is terminated once in close(). - self.zmq_context = create_zmq_context(zmq_io_threads) + # Long-lived ZMQ context shared by all controller RPCs on this client. Contexts are + # thread-safe and event-loop-agnostic; each RPC creates and closes its own DEALER + # socket from this context (see with_controller_socket). Terminated once in close(). + self.zmq_context = zmq.asyncio.Context() logger.info(f"[{self.client_id}]: Registered Controller server {controller_info.id} at {controller_info.ip}") def initialize_storage_manager( @@ -98,10 +93,7 @@ def initialize_storage_manager( """ self.storage_manager = StorageManagerFactory.create( - manager_type, - controller_info=self._controller, - config=config, - zmq_context=self.zmq_context, + manager_type, controller_info=self._controller, config=config ) # ==================== Basic API ==================== @@ -1251,20 +1243,16 @@ def __init__( self, client_id: str, controller_info: ZMQServerInfo, - 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 - zmq_io_threads: Size of the long-lived ZMQ context's native I/O-thread - pool. Defaults to ``TQ_ZMQ_IO_THREADS`` (8). """ super().__init__( client_id, controller_info, - 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 a0effa6f..c6c959c4 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -14,7 +14,6 @@ # limitations under the License. import asyncio -import inspect import itertools import os import threading @@ -36,13 +35,7 @@ from transfer_queue.metadata import BatchMeta, extract_field_schema from transfer_queue.storage.clients.base import StorageClientFactory from transfer_queue.utils.logging_utils import get_logger -from transfer_queue.utils.zmq_utils import ( - ZMQMessage, - ZMQRequestType, - ZMQServerInfo, - create_zmq_context, - create_zmq_socket, -) +from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo, create_zmq_socket logger = get_logger(__name__) @@ -70,12 +63,7 @@ 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, - zmq_context: zmq.asyncio.Context | None = None, - ): + def __init__(self, controller_info: ZMQServerInfo, config: DictConfig): self.storage_manager_id = f"TQ_STORAGE_{uuid4().hex[:8]}" self.config = config self.controller_info = controller_info @@ -83,11 +71,7 @@ def __init__( # Handshake socket is sync (used only during initialization) self.controller_handshake_socket: zmq.Socket | None = None - # A manager created by TransferQueueClient borrows the client's context so - # controller and storage requests share one fixed native I/O-thread pool. - # Standalone managers create and own an equivalent long-lived context. - self._owns_zmq_context = zmq_context is None - self.zmq_context = zmq_context or create_zmq_context(config.get("zmq_io_threads", None)) + self.zmq_context = zmq.asyncio.Context() self._connect_to_controller() # Dedicated asyncio loop for ZMQ notify traffic, isolated from the caller's loop @@ -407,10 +391,9 @@ def close(self) -> None: else: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") - 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) + # 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.""" @@ -441,21 +424,12 @@ def decorator(manager_cls: type[StorageManager]): return decorator @classmethod - def create( - cls, - manager_type: str, - controller_info: ZMQServerInfo, - config: dict[str, Any], - zmq_context: zmq.asyncio.Context | None = None, - ) -> StorageManager: + def create(cls, manager_type: str, controller_info: ZMQServerInfo, config: dict[str, Any]) -> 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())}" ) - manager_cls = cls._registry[manager_type] - if zmq_context is not None and "zmq_context" in inspect.signature(manager_cls).parameters: - return manager_cls(controller_info, config, zmq_context=zmq_context) - return manager_cls(controller_info, config) + return cls._registry[manager_type](controller_info, config) class KVStorageManager(StorageManager): @@ -464,19 +438,14 @@ 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], - zmq_context: zmq.asyncio.Context | None = None, - ): + def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): """ Initialize the KVStorageManager with configuration. """ client_name = config.get("client_name", None) if client_name is None: raise ValueError("Missing client_name in config") - super().__init__(controller_info, config, zmq_context=zmq_context) + super().__init__(controller_info, config) self.storage_client = StorageClientFactory.create(client_name, config) self._multi_threads_executor: ThreadPoolExecutor | None = None self._executor_finalizer = weakref.finalize(self, self._shutdown_executor, self._multi_threads_executor) diff --git a/transfer_queue/storage/managers/mooncake_manager.py b/transfer_queue/storage/managers/mooncake_manager.py index a929d6b7..c3e8f5ce 100644 --- a/transfer_queue/storage/managers/mooncake_manager.py +++ b/transfer_queue/storage/managers/mooncake_manager.py @@ -15,8 +15,6 @@ from typing import Any -import zmq.asyncio - from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory from transfer_queue.utils.zmq_utils import ZMQServerInfo @@ -31,11 +29,6 @@ class MooncakeStorageManager(KVStorageManager): pybind bindings. """ - def __init__( - self, - controller_info: ZMQServerInfo, - config: dict[str, Any], - zmq_context: zmq.asyncio.Context | None = None, - ): + def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): config["client_name"] = "MooncakeStoreClient" - super().__init__(controller_info, config, zmq_context=zmq_context) + super().__init__(controller_info, config) diff --git a/transfer_queue/storage/managers/ray_storage_manager.py b/transfer_queue/storage/managers/ray_storage_manager.py index f91f008a..0cc2a09c 100644 --- a/transfer_queue/storage/managers/ray_storage_manager.py +++ b/transfer_queue/storage/managers/ray_storage_manager.py @@ -15,8 +15,6 @@ from typing import Any -import zmq.asyncio - from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory from transfer_queue.utils.zmq_utils import ZMQServerInfo @@ -25,17 +23,8 @@ class RayStorageManager(KVStorageManager): """Storage manager for Ray-RDT backend.""" - def __init__( - self, - controller_info: ZMQServerInfo, - config: dict[str, Any], - zmq_context: zmq.asyncio.Context | None = None, - ): + def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): 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"}, - zmq_context=zmq_context, - ) + super().__init__(controller_info, {**config, "client_name": "RayStorageClient"}) diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index aec260b6..c6b20eff 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -73,13 +73,8 @@ class AsyncSimpleStorageManager(StorageManager): instances using ZMQ communication and dynamic socket management. """ - def __init__( - self, - controller_info: ZMQServerInfo, - config: DictConfig, - zmq_context: zmq.asyncio.Context | None = None, - ): - super().__init__(controller_info, config, zmq_context=zmq_context) + def __init__(self, controller_info: ZMQServerInfo, config: DictConfig): + super().__init__(controller_info, config) 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 a409cb66..f76b47b2 100644 --- a/transfer_queue/storage/managers/yuanrong_manager.py +++ b/transfer_queue/storage/managers/yuanrong_manager.py @@ -15,8 +15,6 @@ from typing import Any -import zmq.asyncio - 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 @@ -28,12 +26,7 @@ class YuanrongStorageManager(KVStorageManager): """Storage manager for Yuanrong backend.""" - def __init__( - self, - controller_info: ZMQServerInfo, - config: dict[str, Any], - zmq_context: zmq.asyncio.Context | None = None, - ): + def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]): worker_port = config.get("worker_port", None) client_name = config.get("client_name", None) @@ -45,4 +38,4 @@ def __init__( 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, zmq_context=zmq_context) + super().__init__(controller_info, config) diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 713e6396..09fc5fb5 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os import socket import time from collections.abc import Sequence @@ -306,22 +305,6 @@ def get_free_port(ip: str) -> int: return sock.getsockname()[1] -TQ_ZMQ_IO_THREADS = int(os.environ.get("TQ_ZMQ_IO_THREADS", 8)) - - -def create_zmq_context(io_threads: int | None = None) -> "zmq.asyncio.Context": - """Create a long-lived async ZMQ context with a fixed I/O-thread pool. - - A ZMQ context owns the native I/O-thread pool used by all sockets created from - that context. Keeping one context per owner lets concurrent request sockets share - the whole pool without creating or terminating contexts per request. - """ - pool_size = TQ_ZMQ_IO_THREADS if io_threads is None else io_threads - if pool_size < 1: - raise ValueError(f"ZMQ I/O thread pool size must be at least 1, got {pool_size}") - return zmq.asyncio.Context(io_threads=pool_size) - - def create_zmq_socket( ctx: zmq.Context, socket_type: Any, @@ -389,9 +372,8 @@ def with_zmq_socket( 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. The context's fixed native - I/O-thread pool is shared by all request sockets; each socket is created and fully used - within one awaited call on one loop. + 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``. From 9b63baec39ce764b2e88f8a17735e8bfbbc621b2 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 29 Jul 2026 22:38:38 +0800 Subject: [PATCH 06/10] feat: share fixed ZMQ context pool with SimpleStorage Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 85 +++++++++++++++++++ transfer_queue/client.py | 36 ++++++-- transfer_queue/storage/managers/base.py | 33 +++++-- .../managers/simple_storage_manager.py | 9 +- 4 files changed, 149 insertions(+), 14 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 37700189..a4dde94f 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -27,6 +27,7 @@ import asyncio from threading import Thread +from unittest.mock import patch import pytest import zmq @@ -34,6 +35,7 @@ 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 @@ -148,6 +150,89 @@ def _spy_create(ctx, *args, **kwargs): 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).""" diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 6b088785..656331cc 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -37,6 +37,7 @@ 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( @@ -58,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") @@ -71,10 +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 - # Long-lived ZMQ context shared by all controller RPCs on this client. Contexts are - # thread-safe and event-loop-agnostic; each RPC creates and closes its own DEALER - # socket from this context (see with_controller_socket). Terminated once in close(). - self.zmq_context = zmq.asyncio.Context() + # 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( @@ -92,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 ==================== @@ -1243,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 c6c959c4..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,9 +400,10 @@ def close(self) -> None: else: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") - # 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) + 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.""" @@ -424,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 c6b20eff..aec260b6 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -73,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) From 22f48c8a12dfe3eddf8198848b1213d04ee60696 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 31 Jul 2026 22:29:29 +0800 Subject: [PATCH 07/10] style: apply ruff-format to client.py Co-Authored-By: Claude Opus 4.8 Signed-off-by: OutstanderWang --- transfer_queue/client.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 656331cc..b711c08d 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -80,15 +80,10 @@ def __init__( # 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 + 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}" - ) + raise ValueError(f"SimpleStorage ZMQ I/O thread pool size must be at least 1, 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}") From ae7ebf7c376b8f3d26b080086d71ad5bbdc791ab Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 3 Aug 2026 17:30:31 +0800 Subject: [PATCH 08/10] fix: address PR #145 review comments - Guard zmq_context.destroy() against live threads. destroy() calls Socket.close() internally and is not thread-safe, but both TransferQueueClient.close() and StorageManager.close() joined their threads with a warn-only timeout and then destroyed anyway. Add a _can_destroy_zmq_context() veto that leaks the context with a loud warning instead, and document the load-bearing close() ordering. - Rename the I/O-thread knob to reflect its real scope: the context serves every backend's controller RPCs, not just SimpleStorage. TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS -> TQ_CLIENT_ZMQ_IO_THREADS, kwarg simple_storage_zmq_io_threads -> zmq_io_threads. - Expose zmq_max_sockets / TQ_CLIENT_ZMQ_MAX_SOCKETS. Sharing one context per client means all in-flight sockets share libzmq's 1023-per-context budget, where previously each call had a private one. Opt-in, validated against the build's ZMQ_SOCKET_LIMIT. - Keep StorageManagerFactory.create backend-agnostic: forward **kwargs and let each registered manager decide, instead of hard-coding "SimpleStorage" in the factory and the client. KV managers accept zmq_context and deliberately keep their own. - Align the ownership check on an explicit `is None` test. Co-Authored-By: Claude Opus 4.8 Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 93 ++++++++++++++- transfer_queue/client.py | 107 ++++++++++++++---- transfer_queue/storage/managers/base.py | 59 +++++++--- .../storage/managers/mooncake_manager.py | 11 +- .../storage/managers/ray_storage_manager.py | 11 +- .../storage/managers/yuanrong_manager.py | 11 +- 6 files changed, 248 insertions(+), 44 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index a4dde94f..d52e1057 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -33,8 +33,9 @@ import zmq import transfer_queue.utils.zmq_utils as zmq_utils -from transfer_queue.client import AsyncTransferQueueClient +from transfer_queue.client import AsyncTransferQueueClient, TransferQueueClient from transfer_queue.metadata import BatchMeta +from transfer_queue.storage.managers.base import KVStorageManager 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 @@ -154,7 +155,7 @@ 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, + zmq_io_threads=4, ) assert client.zmq_context.get(zmq.IO_THREADS) == 4 @@ -167,7 +168,7 @@ def test_client_rejects_invalid_context_pool_size(echo_controller): AsyncTransferQueueClient( client_id="client_invalid_context_pool", controller_info=echo_controller.zmq_server_info, - simple_storage_zmq_io_threads=0, + zmq_io_threads=0, ) @@ -214,7 +215,8 @@ def test_simple_storage_does_not_destroy_borrowed_context(echo_controller): assert client.zmq_context.closed -def test_other_backends_do_not_borrow_client_context(echo_controller): +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, @@ -228,11 +230,42 @@ def test_other_backends_do_not_borrow_client_context(echo_controller): "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() + + @pytest.mark.asyncio async def test_close_destroys_context(echo_controller): """close() must terminate the shared context exactly once (no leak, no hang).""" @@ -251,3 +284,55 @@ async def test_close_destroys_context(echo_controller): 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_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) diff --git a/transfer_queue/client.py b/transfer_queue/client.py index b711c08d..d048a7ec 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -37,7 +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)) +# 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 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") # Pre-bound decorator for controller socket operations. with_controller_socket = with_zmq_socket( @@ -59,16 +66,22 @@ def __init__( self, client_id: str, controller_info: ZMQServerInfo, - simple_storage_zmq_io_threads: int | None = None, + 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 - 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). + 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") @@ -79,12 +92,24 @@ def __init__( # 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 - ) + io_threads = TQ_CLIENT_ZMQ_IO_THREADS if zmq_io_threads is None else zmq_io_threads if io_threads < 1: - raise ValueError(f"SimpleStorage ZMQ I/O thread pool size must be at least 1, got {io_threads}") + 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 = ( + int(TQ_CLIENT_ZMQ_MAX_SOCKETS) + if zmq_max_sockets is None and TQ_CLIENT_ZMQ_MAX_SOCKETS is not None + else zmq_max_sockets + ) + 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( @@ -94,6 +119,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. @@ -102,14 +131,11 @@ 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, - **create_kwargs, + zmq_context=self.zmq_context, ) # ==================== Basic API ==================== @@ -1108,7 +1134,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"): @@ -1118,12 +1152,27 @@ def close(self) -> None: # 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. + + Always true here: this class owns no background thread, so the caller's + quiescence contract (documented on :meth:`close`) is the only requirement. + Subclasses that run their own loop thread override this. + """ + return True + # ==================== Checkpoint API ==================== @with_controller_socket async def async_save_controller_checkpoint( @@ -1259,21 +1308,25 @@ def __init__( self, client_id: str, controller_info: ZMQServerInfo, - simple_storage_zmq_io_threads: int | None = None, + 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 - 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). + 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, - simple_storage_zmq_io_threads, + zmq_io_threads, + zmq_max_sockets, ) # create new event loop in a separate thread @@ -1807,7 +1860,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) @@ -1823,3 +1885,8 @@ 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.""" + thread = getattr(self, "_thread", None) + return thread is None or not thread.is_alive() diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index f435fd1e..51414543 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -76,11 +76,12 @@ def __init__( # Handshake socket is sync (used only during initialization) self.controller_handshake_socket: zmq.Socket | None = None - # 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. + # 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_context or zmq.asyncio.Context() + 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 @@ -393,17 +394,31 @@ 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.") 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) + # 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 __del__(self): """Destructor to ensure resources are cleaned up.""" @@ -439,16 +454,19 @@ def create( manager_type: str, controller_info: ZMQServerInfo, config: dict[str, Any], - zmq_context: zmq.asyncio.Context | None = None, + **kwargs: Any, ) -> StorageManager: - """Create and return a StorageManager instance.""" + """Create and return a StorageManager instance. + + Extra keyword arguments are forwarded verbatim 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). + """ assert manager_type in cls._registry, ( f"Unknown manager_type: {manager_type}. Supported managers include: {list(cls._registry.keys())}" ) - 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) + return cls._registry[manager_type](controller_info, config, **kwargs) class KVStorageManager(StorageManager): @@ -457,9 +475,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/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) From ffd0d282dd63d636fb043d3f4703d35702507fbb Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 3 Aug 2026 20:19:19 +0800 Subject: [PATCH 09/10] fix: harden TQ_CLIENT_ZMQ_MAX_SOCKETS env var and cover it with tests The socket-ceiling knob had tests only for the kwarg, never the env var -- which is the deployment-facing path, since the kwarg needs a code edit. Exercising it end to end surfaced two bugs: - An empty value (TQ_CLIENT_ZMQ_MAX_SOCKETS=, the usual shell idiom for clearing a variable) hit int('') and crashed the client instead of falling back to libzmq's default. Treat empty as unset. - A non-numeric value raised a bare int() error naming no variable, which is hard to trace in a worker log. Name the variable. Also add a regression test that drives the real StorageManagerFactory with an independently-registered third-party manager. The existing factory tests patch create() out, so nothing executed the **kwargs forwarding; verified the new test fails if the old `if manager_type == "SimpleStorage"` special-case is reintroduced, while the rest of the suite stays green. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 120 ++++++++++++++++++++++++++++++- transfer_queue/client.py | 22 +++--- 2 files changed, 132 insertions(+), 10 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index d52e1057..285bb516 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -35,7 +35,7 @@ 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 +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 @@ -266,6 +266,69 @@ def test_kv_backends_keep_own_context(echo_controller): 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).""" @@ -313,6 +376,61 @@ def test_client_rejects_max_sockets_above_build_limit(echo_controller): ) +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. diff --git a/transfer_queue/client.py b/transfer_queue/client.py index d048a7ec..9493f068 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -41,10 +41,10 @@ # 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 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") +# 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( @@ -97,11 +97,15 @@ def __init__( 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 = ( - int(TQ_CLIENT_ZMQ_MAX_SOCKETS) - if zmq_max_sockets is None and TQ_CLIENT_ZMQ_MAX_SOCKETS is not None - else zmq_max_sockets - ) + 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: From fded800da9c94560baa171fc5ad866c80fbf4474 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 3 Aug 2026 20:55:26 +0800 Subject: [PATCH 10/10] fix: close two shutdown/compatibility holes in the shared ZMQ context Both were reachable and both are now covered by tests verified to fail when the fix is reverted. 1. A borrowing manager's stuck notify thread could not veto destroy(). StorageManager.close() detected the failed join but only acted on it inside `if self._owns_zmq_context`, so a manager that borrowed the client's context warned and returned silently. The client's veto checked only its own loop thread, so it went on to destroy a context whose sockets that thread might still hold -- the documented non-thread-safe Socket.close() hazard. The manager now records the outcome unconditionally and exposes can_destroy_zmq_context(); the client consults it, but only for a manager that actually shares the context, and TransferQueueClient now combines that with its own loop-thread check instead of replacing it. 2. Externally registered managers on the old (controller_info, config) contract raised TypeError, because the client passes zmq_context unconditionally. Registration is an extension mechanism, so managers are not required to update in lockstep: the factory now drops keywords a constructor cannot accept, warning with the class and parameter name so the drop is never silent. Constructors taking **kwargs, and un-introspectable ones, pass through unchanged. Also fix __del__ reaching for storage_manager_id on a half-constructed object, which raised AttributeError and masked the real constructor error -- found while reproducing issue 2. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 131 ++++++++++++++++++++++++ transfer_queue/client.py | 26 ++++- transfer_queue/storage/managers/base.py | 67 ++++++++++-- 3 files changed, 213 insertions(+), 11 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 285bb516..f43dcd96 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -454,3 +454,134 @@ def test_close_skips_destroy_while_loop_thread_alive(echo_controller): # 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 9493f068..ba45b60d 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -1171,10 +1171,20 @@ def close(self) -> None: def _can_destroy_zmq_context(self) -> bool: """Whether it is safe to call ``destroy()`` on the shared context. - Always true here: this class owns no background thread, so the caller's - quiescence contract (documented on :meth:`close`) is the only requirement. - Subclasses that run their own loop thread override this. + 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 ==================== @@ -1891,6 +1901,12 @@ def close(self) -> None: super().close() def _can_destroy_zmq_context(self) -> bool: - """False while the loop thread that owns sockets on the context is still alive.""" + """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) - return thread is None or not thread.is_alive() + 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 51414543..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 @@ -403,6 +404,11 @@ def close(self) -> None: else: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") + # 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 @@ -420,12 +426,28 @@ def close(self) -> None: 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: @@ -458,15 +480,48 @@ def create( ) -> StorageManager: """Create and return a StorageManager instance. - Extra keyword arguments are forwarded verbatim 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). + 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, **kwargs) + 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):