From 2c34db3fc4094e0a477da07120028c79324d5ca5 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Thu, 23 Jul 2026 17:04:16 +0800 Subject: [PATCH 01/13] 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). 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 a47c3c526cb7baa1dd4a038a79a3df1b029e76a4 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 29 Jul 2026 22:10:36 +0800 Subject: [PATCH 02/13] 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 12ce1b0b5364c8850279cb52fb86d766df546102 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 29 Jul 2026 22:16:43 +0800 Subject: [PATCH 03/13] 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 b8907a5260fb6942f89cb8f9d4395e0ae56a230a Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 29 Jul 2026 22:28:45 +0800 Subject: [PATCH 04/13] 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 c82191c3fa21928bb3d2cf125d0b15ea8a3669be Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 29 Jul 2026 22:31:51 +0800 Subject: [PATCH 05/13] 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 2039e29157af659bdd13210eb2f9898744322ffe Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 29 Jul 2026 22:38:38 +0800 Subject: [PATCH 06/13] 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 589438395e39707b794932685052fd58e228c36d Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 31 Jul 2026 22:29:29 +0800 Subject: [PATCH 07/13] style: apply ruff-format to client.py 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 a00f670d06088aa2772f09b7bf232890b4d60d65 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 3 Aug 2026 17:30:31 +0800 Subject: [PATCH 08/13] 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. 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 02ce20eac78d66f851f5c4dfa83e4d4b890a984f Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 3 Aug 2026 20:19:19 +0800 Subject: [PATCH 09/13] 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 bae821941ffc320bfd817031df1578e0b04c1dcb Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 3 Aug 2026 20:55:26 +0800 Subject: [PATCH 10/13] 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): From aca4d2772270aebd9e716dd1e6795f8d7673e8de Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 4 Aug 2026 20:27:08 +0800 Subject: [PATCH 11/13] fix: address second round of PR #145 review comments Raise the client ZMQ socket ceiling by default, drop a dead flag, and add a finalizer backstop for clients that are never closed. - max sockets: default to DEFAULT_CLIENT_ZMQ_MAX_SOCKETS (8192) instead of inheriting libzmq's 1023, which large-scale training can exhaust since a borrowing storage manager draws on the same per-client budget. Validation now splits on provenance: an explicitly requested value (kwarg or env var) still raises when out of range, while the built-in default clamps to the build's ZMQ_SOCKET_LIMIT rather than making the client fail to construct on a build whose limit is lower. - StorageManager.close(): remove the write to self._notify_thread_stopped. No code read it -- can_destroy_zmq_context() checks the thread liveness directly so it stays correct when called before close(). - AsyncTransferQueueClient: register a weakref.finalize to destroy the context when close() is never called, so the context and its I/O threads do not leak for the process lifetime. close() detaches it, including on the path that deliberately leaks the context, so an unsafe destroy() cannot run later. Note TransferQueueClient's loop thread references the client, deferring collection to interpreter exit; close() remains the supported path. - tests: drop the factory/backend-agnosticism, max-sockets config, and legacy signature cases per review. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 239 +----------------------- transfer_queue/client.py | 75 ++++++-- transfer_queue/storage/managers/base.py | 5 - 3 files changed, 65 insertions(+), 254 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index f43dcd96..4c082703 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, StorageManager, StorageManagerFactory +from transfer_queue.storage.managers.base import StorageManager from transfer_queue.storage.managers.simple_storage_manager import AsyncSimpleStorageManager from transfer_queue.utils.enum_utils import Role from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo @@ -215,120 +215,6 @@ def test_simple_storage_does_not_destroy_borrowed_context(echo_controller): assert client.zmq_context.closed -def test_factory_is_backend_agnostic(echo_controller): - """The client offers its context to every backend uniformly, naming none of them.""" - client = AsyncTransferQueueClient( - client_id="client_other_storage_context", - controller_info=echo_controller.zmq_server_info, - ) - config = {"client_name": "unused"} - - with patch("transfer_queue.client.StorageManagerFactory.create") as create_manager: - client.initialize_storage_manager("OtherStorage", config) - - create_manager.assert_called_once_with( - "OtherStorage", - controller_info=echo_controller.zmq_server_info, - config=config, - zmq_context=client.zmq_context, - ) - - client.close() - - -def test_kv_backends_keep_own_context(echo_controller): - """KV managers accept the shared context but deliberately keep an independent one. - - They move bulk data through their own SDKs and use ZMQ only for the controller - notify/handshake path, so they must not draw on the client's socket budget. - """ - client = AsyncTransferQueueClient( - client_id="client_kv_own_context", - controller_info=echo_controller.zmq_server_info, - ) - - with ( - patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"), - patch("transfer_queue.storage.managers.base.StorageClientFactory.create"), - ): - manager = KVStorageManager( - echo_controller.zmq_server_info, - {"client_name": "unused"}, - zmq_context=client.zmq_context, - ) - - assert manager.zmq_context is not client.zmq_context - assert manager._owns_zmq_context - - manager.close() - assert not client.zmq_context.closed - - client.close() - - -def test_factory_forwards_kwargs_to_registered_manager(echo_controller): - """The real factory forwards **kwargs verbatim, knowing no backend by name. - - The other factory test patches ``create`` out, so this one exercises the real - dispatch: a third-party manager registered from outside this package must receive - ``zmq_context`` without the factory special-casing its name. A regression to the old - ``if manager_type == "SimpleStorage"`` branch would silently drop the kwarg here, and - a manager whose signature drifts would raise TypeError -- neither of which mypy can - catch through ``**kwargs: Any``. - """ - received = {} - - @StorageManagerFactory.register("THIRD_PARTY_PROBE") - class ThirdPartyManager(StorageManager): - def __init__(self, controller_info, config, zmq_context=None): - received["zmq_context"] = zmq_context - received["config"] = config - super().__init__(controller_info, config, zmq_context=zmq_context) - - def _connect_to_controller(self): - pass - - def _do_handshake_with_controller(self): - pass - - async def put_data(self, *args, **kwargs): - return None - - async def get_data(self, *args, **kwargs): - return None - - async def clear_data(self, *args, **kwargs): - return None - - async def notify_data_update(self, *args, **kwargs): - return None - - try: - client = AsyncTransferQueueClient( - client_id="client_third_party_factory", - controller_info=echo_controller.zmq_server_info, - ) - config = {"marker": "forwarded"} - - client.initialize_storage_manager("THIRD_PARTY_PROBE", config) - - # The kwarg survived dispatch through the unpatched factory... - assert received["zmq_context"] is client.zmq_context - assert received["config"] == config - # ...and a manager that opts in genuinely borrows rather than re-creating. - assert client.storage_manager.zmq_context is client.zmq_context - assert not client.storage_manager._owns_zmq_context - - # A borrower must not tear down a context it does not own. - client.storage_manager.close() - assert not client.zmq_context.closed - - client.close() - assert client.zmq_context.closed - finally: - StorageManagerFactory._registry.pop("THIRD_PARTY_PROBE", None) - - @pytest.mark.asyncio async def test_close_destroys_context(echo_controller): """close() must terminate the shared context exactly once (no leak, no hang).""" @@ -349,78 +235,6 @@ async def test_close_destroys_context(echo_controller): assert client.zmq_context.closed -def test_client_applies_max_sockets(echo_controller): - """The per-context socket ceiling is configurable, since it is now shared per client.""" - client = AsyncTransferQueueClient( - client_id="client_max_sockets", - controller_info=echo_controller.zmq_server_info, - zmq_max_sockets=2048, - ) - - assert client.zmq_context.get(zmq.MAX_SOCKETS) == 2048 - - client.close() - - -def test_client_rejects_max_sockets_above_build_limit(echo_controller): - """Values above this libzmq build's ZMQ_SOCKET_LIMIT are rejected up front.""" - probe = zmq.Context() - socket_limit = probe.get(zmq.SOCKET_LIMIT) - probe.term() - - with pytest.raises(ValueError, match="ZMQ_SOCKET_LIMIT"): - AsyncTransferQueueClient( - client_id="client_max_sockets_too_big", - controller_info=echo_controller.zmq_server_info, - zmq_max_sockets=socket_limit + 1, - ) - - -def test_max_sockets_from_env_var(echo_controller): - """TQ_CLIENT_ZMQ_MAX_SOCKETS configures the ceiling without touching call sites. - - The env var is the deployment-facing knob (the kwarg requires editing code), so it - needs its own coverage. Patched at the module constant because it is read at import. - """ - with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", "4096"): - client = AsyncTransferQueueClient( - client_id="client_max_sockets_env", - controller_info=echo_controller.zmq_server_info, - ) - - assert client.zmq_context.get(zmq.MAX_SOCKETS) == 4096 - client.close() - - -def test_explicit_max_sockets_overrides_env_var(echo_controller): - """An explicit kwarg wins over the env var, matching the io_threads precedence.""" - with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", "4096"): - client = AsyncTransferQueueClient( - client_id="client_max_sockets_precedence", - controller_info=echo_controller.zmq_server_info, - zmq_max_sockets=2048, - ) - - assert client.zmq_context.get(zmq.MAX_SOCKETS) == 2048 - client.close() - - -def test_unset_max_sockets_leaves_libzmq_default(echo_controller): - """Opt-in only: with nothing configured, libzmq's own default must be untouched.""" - probe = zmq.Context() - default = probe.get(zmq.MAX_SOCKETS) - probe.term() - - with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", None): - client = AsyncTransferQueueClient( - client_id="client_max_sockets_unset", - controller_info=echo_controller.zmq_server_info, - ) - - assert client.zmq_context.get(zmq.MAX_SOCKETS) == default - client.close() - - def test_non_numeric_max_sockets_env_var_names_the_variable(echo_controller): """A typo'd value must say which env var is wrong, not raise a bare int() error.""" with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", "not-a-number"): @@ -534,54 +348,3 @@ def test_manager_with_own_context_does_not_veto(echo_controller): 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 ba45b60d..28c717db 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -16,6 +16,7 @@ import asyncio import os import threading +import weakref from typing import Any, Callable import torch @@ -41,10 +42,14 @@ # controller RPC (all backends) and, for SimpleStorage, storage-unit requests too, # so this knob is client-scoped rather than backend-scoped. TQ_CLIENT_ZMQ_IO_THREADS = int(os.environ.get("TQ_CLIENT_ZMQ_IO_THREADS", 8)) -# Per-context socket ceiling. Unset (or empty) means libzmq's default (1023). Because the -# context is now shared per client instead of created per call, all in-flight sockets draw -# on a single budget; raise this if a large num_data_storage_units fan-out exhausts it. +# Per-context socket ceiling. Because the context is shared per client instead of created +# per call, all in-flight sockets draw on a single budget: a borrowing storage manager adds +# one socket per storage unit per in-flight put/get, so libzmq's own default (1023) is +# reachable in large-scale training. Default well above it and clamp to the build's +# ZMQ_SOCKET_LIMIT. Sockets are still created lazily -- this raises the ceiling, it does +# not preallocate -- but the process also needs file descriptors for them (``ulimit -n``). TQ_CLIENT_ZMQ_MAX_SOCKETS = os.environ.get("TQ_CLIENT_ZMQ_MAX_SOCKETS") or None +DEFAULT_CLIENT_ZMQ_MAX_SOCKETS = 8192 # Pre-bound decorator for controller socket operations. with_controller_socket = with_zmq_socket( @@ -77,11 +82,12 @@ def __init__( 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``). + at once. Defaults to ``TQ_CLIENT_ZMQ_MAX_SOCKETS``, and to + ``DEFAULT_CLIENT_ZMQ_MAX_SOCKETS`` (8192) when that is unset -- well above + libzmq's own 1023, which large-scale training can exhaust. This ceiling is + per-client: a single ``put``/``get`` fans out one socket per storage unit, + and a borrowing storage manager shares the same budget. Raising it also + requires enough file descriptors (``ulimit -n``). """ if controller_info is None: raise ValueError("controller_info cannot be None") @@ -98,6 +104,7 @@ def __init__( self.zmq_context = zmq.asyncio.Context(io_threads=io_threads) max_sockets = zmq_max_sockets + explicitly_requested = max_sockets is not None if max_sockets is None and TQ_CLIENT_ZMQ_MAX_SOCKETS is not None: try: max_sockets = int(TQ_CLIENT_ZMQ_MAX_SOCKETS) @@ -106,16 +113,50 @@ def __init__( 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) + explicitly_requested = True + if max_sockets is None: + max_sockets = DEFAULT_CLIENT_ZMQ_MAX_SOCKETS + + socket_limit = self.zmq_context.get(zmq.SOCKET_LIMIT) + if explicitly_requested: + # A value the caller asked for must not be silently reinterpreted. if not 1 <= max_sockets <= socket_limit: raise ValueError( f"Client ZMQ max sockets must be between 1 and this build's " f"ZMQ_SOCKET_LIMIT ({socket_limit}), got {max_sockets}" ) - self.zmq_context.set(zmq.MAX_SOCKETS, max_sockets) + elif max_sockets > socket_limit: + # Nobody asked for the default, so clamp instead of failing to construct on a + # build whose ZMQ_SOCKET_LIMIT is below it. + logger.debug( + f"[{client_id}]: Clamping default ZMQ max sockets {max_sockets} to this " + f"build's ZMQ_SOCKET_LIMIT ({socket_limit})." + ) + max_sockets = socket_limit + self.zmq_context.set(zmq.MAX_SOCKETS, max_sockets) + + # Users forget to call close(); without this the context and its I/O threads leak + # for the process lifetime. finalize() (not __del__) so it also runs at interpreter + # exit, and it is idempotent with an explicit close(). Note that for + # TransferQueueClient the running loop thread references the client, so collection + # is deferred until interpreter exit -- close() is still the supported path. + self._finalizer = weakref.finalize(self, self._release_zmq_context, self.zmq_context, client_id) logger.info(f"[{self.client_id}]: Registered Controller server {controller_info.id} at {controller_info.ip}") + @staticmethod + def _release_zmq_context(context: "zmq.asyncio.Context", client_id: str) -> None: + """Destroy *context* if it is still open. + + Must stay a staticmethod taking the context explicitly: a bound method would keep + the client alive and the finalizer would never fire. Only reached for a client that + was never closed, so the owning threads are gone (interpreter exit) or unreachable. + """ + try: + if context is not None and not context.closed: + context.destroy(linger=0) + except Exception as e: + logger.warning(f"[{client_id}]: Error destroying zmq_context in finalizer: {e}") + def initialize_storage_manager( self, manager_type: str, @@ -1157,6 +1198,9 @@ 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(): + # Detach before returning: the finalizer would otherwise fire later and run the + # very destroy() this branch judged unsafe, defeating the deliberate leak. + self._detach_finalizer() logger.warning( f"[{self.client_id}]: Skipping zmq_context.destroy() because a thread owning " f"sockets on it is still alive; destroy() is not thread-safe. Leaking the context." @@ -1167,6 +1211,15 @@ def close(self) -> None: self.zmq_context.destroy(linger=0) except Exception as e: logger.warning(f"[{self.client_id}]: Error terminating zmq_context: {e}") + finally: + # This close() superseded the finalizer; disarm it so it cannot run twice. + self._detach_finalizer() + + def _detach_finalizer(self) -> None: + """Disarm the context finalizer, if one was armed.""" + finalizer = getattr(self, "_finalizer", None) + if finalizer is not None: + finalizer.detach() def _can_destroy_zmq_context(self) -> bool: """Whether it is safe to call ``destroy()`` on the shared context. diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index aafa0691..6d0bc9c5 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -404,11 +404,6 @@ 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 From ff301c436bc5a522049e6bb3615dcdb2f81c1e77 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 4 Aug 2026 21:18:11 +0800 Subject: [PATCH 12/13] style: tighten comments added by the shared ZMQ context change Inline comment blocks in this PR ran 4-6 lines while the surrounding files use 1-2, and several docstrings had grown to full paragraphs. Condense both to the essential why, keeping every load-bearing caveat (destroy() thread-safety, per-call context churn, ownership rules). No logic changes. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 18 ++--- transfer_queue/client.py | 66 +++++++------------ transfer_queue/storage/managers/base.py | 40 ++++------- .../managers/simple_storage_manager.py | 5 +- transfer_queue/utils/zmq_utils.py | 23 +++---- 5 files changed, 54 insertions(+), 98 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 4c082703..b31151a1 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -15,14 +15,10 @@ """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. +with_zmq_socket used to create and term() a context per RPC call, which churned libzmq +signaler file descriptors and crashed the process under concurrency (Bad file descriptor +-> SIGABRT). It now reuses the owner's context and only creates the socket per call, so +these tests assert every concurrent call sees the SAME context, alive until close(). """ import asyncio @@ -298,10 +294,8 @@ async def notify_data_update(self, *args, **kwargs): 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. + The manager has no destroy() of its own to skip, so staying silent would let the client + destroy a context whose sockets that thread may still hold. It must be asked first. """ client = AsyncTransferQueueClient( client_id="client_notify_thread_veto", diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 28c717db..65559bea 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -38,16 +38,12 @@ logger = get_logger(__name__) TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8)) -# Size of the client context's native I/O-thread pool. The context serves every -# controller RPC (all backends) and, for SimpleStorage, storage-unit requests too, -# so this knob is client-scoped rather than backend-scoped. +# Size of the client context's native I/O-thread pool, shared by all controller RPCs and, +# for SimpleStorage, storage-unit requests -- so this knob is client- not backend-scoped. TQ_CLIENT_ZMQ_IO_THREADS = int(os.environ.get("TQ_CLIENT_ZMQ_IO_THREADS", 8)) -# Per-context socket ceiling. Because the context is shared per client instead of created -# per call, all in-flight sockets draw on a single budget: a borrowing storage manager adds -# one socket per storage unit per in-flight put/get, so libzmq's own default (1023) is -# reachable in large-scale training. Default well above it and clamp to the build's -# ZMQ_SOCKET_LIMIT. Sockets are still created lazily -- this raises the ceiling, it does -# not preallocate -- but the process also needs file descriptors for them (``ulimit -n``). +# Per-context socket ceiling, shared by every in-flight socket on the client's context +# (including a borrowing storage manager's), so libzmq's own 1023 is reachable at scale. +# Raising it also needs enough file descriptors (``ulimit -n``). TQ_CLIENT_ZMQ_MAX_SOCKETS = os.environ.get("TQ_CLIENT_ZMQ_MAX_SOCKETS") or None DEFAULT_CLIENT_ZMQ_MAX_SOCKETS = 8192 @@ -95,9 +91,8 @@ def __init__( raise TypeError(f"controller_info must be ZMQServerInfo, got {type(controller_info)}") self.client_id = client_id self._controller: ZMQServerInfo = controller_info - # One long-lived context per client. Its fixed native I/O-thread pool is shared - # by all controller RPCs and, for SimpleStorage only, storage-unit requests. - # Sockets remain per-request because ZMQ sockets are not thread-safe. + # One long-lived context per client; sockets stay per-request because ZMQ sockets + # are not thread-safe. io_threads = TQ_CLIENT_ZMQ_IO_THREADS if zmq_io_threads is None else zmq_io_threads if io_threads < 1: raise ValueError(f"Client ZMQ I/O thread pool size must be at least 1, got {io_threads}") @@ -135,11 +130,9 @@ def __init__( max_sockets = socket_limit self.zmq_context.set(zmq.MAX_SOCKETS, max_sockets) - # Users forget to call close(); without this the context and its I/O threads leak - # for the process lifetime. finalize() (not __del__) so it also runs at interpreter - # exit, and it is idempotent with an explicit close(). Note that for - # TransferQueueClient the running loop thread references the client, so collection - # is deferred until interpreter exit -- close() is still the supported path. + # Backstop for a client that is never closed, so the context and its I/O threads do + # not leak for the process lifetime. finalize() (not __del__) also runs at + # interpreter exit; close() detaches it. See _release_zmq_context(). self._finalizer = weakref.finalize(self, self._release_zmq_context, self.zmq_context, client_id) logger.info(f"[{self.client_id}]: Registered Controller server {controller_info.id} at {controller_info.ip}") @@ -147,9 +140,9 @@ def __init__( def _release_zmq_context(context: "zmq.asyncio.Context", client_id: str) -> None: """Destroy *context* if it is still open. - Must stay a staticmethod taking the context explicitly: a bound method would keep - the client alive and the finalizer would never fire. Only reached for a client that - was never closed, so the owning threads are gone (interpreter exit) or unreachable. + Must stay a staticmethod taking the context explicitly, or the bound method would + keep the client alive and never fire. TransferQueueClient's loop thread references + the client, deferring this to interpreter exit -- close() remains the supported path. """ try: if context is not None and not context.closed: @@ -1181,12 +1174,9 @@ async def async_kv_list( def close(self) -> None: """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`). + The caller must ensure no RPCs are in flight: this destroys the shared context, and + ``destroy()`` calls ``Socket.close()`` internally, which is **not** thread-safe. + Subclasses owning a loop thread must join it first (see TransferQueueClient.close). """ try: if hasattr(self, "storage_manager") and self.storage_manager: @@ -1195,11 +1185,9 @@ 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. + # Tear down the shared context last; linger=0 so shutdown cannot hang. if not self._can_destroy_zmq_context(): - # Detach before returning: the finalizer would otherwise fire later and run the - # very destroy() this branch judged unsafe, defeating the deliberate leak. + # Detach first, or the finalizer would later run the destroy() judged unsafe here. self._detach_finalizer() logger.warning( f"[{self.client_id}]: Skipping zmq_context.destroy() because a thread owning " @@ -1224,11 +1212,9 @@ def _detach_finalizer(self) -> None: def _can_destroy_zmq_context(self) -> bool: """Whether it is safe to call ``destroy()`` on the shared context. - This class owns no background thread of its own, so the caller's quiescence - contract (documented on :meth:`close`) covers the client side. But a storage - manager that *borrowed* this context runs a notify thread the client cannot see, - and that thread holds sockets on it -- so ask the manager whether it finished - shutting down. Subclasses that run their own loop thread extend this. + A storage manager that *borrowed* this context runs a notify thread the client + cannot see, and that thread holds sockets on it, so ask the manager first. + Subclasses running their own loop thread extend this. """ manager = getattr(self, "storage_manager", None) if manager is not None: @@ -1929,13 +1915,9 @@ def load_storage_checkpoint(self, checkpoint_dir: str) -> None: def close(self) -> None: """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. + Ordering is load-bearing: the loop thread is joined *before* the base close() + destroys the context, since ``destroy()`` is not thread-safe. If the join times + out, ``_can_destroy_zmq_context()`` returns False and the context is leaked instead. """ if hasattr(self, "_loop") and self._loop is not None: diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 6d0bc9c5..9c32e213 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -77,10 +77,8 @@ def __init__( # Handshake socket is sync (used only during initialization) self.controller_handshake_socket: zmq.Socket | None = None - # 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()). + # A manager may borrow a caller-owned context (SimpleStorage does) or own the one it + # creates when handed nothing. Only an owner tears its context down (see close()). self._owns_zmq_context = zmq_context is None self.zmq_context = zmq.asyncio.Context() if zmq_context is None else zmq_context self._connect_to_controller() @@ -405,14 +403,12 @@ def close(self) -> None: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") if self._owns_zmq_context: - # Ordering below is load-bearing: destroy() calls Socket.close() internally, - # which is NOT thread-safe, so it must run only after the notify thread that - # owns sockets on this context is gone. If that thread outlived its join - # timeout, leak the context rather than risk a crash during shutdown -- the - # process is terminating anyway, so a leaked context is the cheaper outcome. + # destroy() calls Socket.close(), which is not thread-safe, so it must run only + # after the notify thread holding sockets is gone. If that thread outlived its + # join, leak the context rather than risk a crash on a terminating process. if notify_thread_stopped: - # 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(). + # linger=0 force-closes sockets left by an interrupted request, so this + # cannot hang on term(). self.zmq_context.destroy(linger=0) else: logger.warning( @@ -424,12 +420,9 @@ def close(self) -> None: 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. + ``destroy()`` is not thread-safe, so the notify thread must be gone first, and only + this manager can see it. Checks the thread directly so it stays correct when called + before ``close()`` or if the thread exited late. """ thread = getattr(self, "_notify_thread", None) return thread is None or not thread.is_alive() @@ -475,15 +468,10 @@ def create( ) -> StorageManager: """Create and return a StorageManager instance. - Extra keyword arguments are forwarded to the registered class. The factory - deliberately knows nothing about any individual backend: each manager decides for - itself what to do with what it receives (e.g. whether to borrow a caller-supplied - ``zmq_context`` or keep its own). - - Registration is an extension mechanism, so managers written against the older - ``(controller_info, config)`` contract must keep working without being updated in - lockstep. Any keyword the constructor does not accept is therefore dropped, with a - warning, rather than raising ``TypeError``. + Extra keywords are forwarded verbatim; the factory knows no backend by name, so each + manager decides what to do with what it receives. Keywords a constructor cannot + accept are dropped with a warning, keeping older ``(controller_info, config)`` + managers working without a lockstep update. """ assert manager_type in cls._registry, ( f"Unknown manager_type: {manager_type}. Supported managers include: {list(cls._registry.keys())}" diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index aec260b6..acc73a28 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -49,9 +49,8 @@ "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. + # Long-lived context from the base StorageManager, shared with the notify path. Safe + # because the context is loop-agnostic and each socket stays per-call. get_context=lambda self: self.zmq_context, resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"), timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 09fc5fb5..2ca63d7f 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -362,18 +362,12 @@ def with_zmq_socket( ): """Create a reusable async decorator for request sockets. - This decorator encapsulates the common socket lifecycle used by both - client-side and storage-manager-side request paths: - 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. + Lifecycle: get owner's shared context -> create/connect socket -> inject -> close socket. + + The context comes from ``self`` via ``get_context`` and is long-lived; only the DEALER + socket is per-call. Do NOT create or terminate a context here -- per-call churn corrupts + libzmq's signaler file descriptors under concurrency (Bad file descriptor / SIGABRT) and + can hang on term(). Contexts are thread-safe and loop-agnostic, so sharing one is safe. Args: socket_name: Socket port key in ``ZMQServerInfo.ports``. @@ -427,9 +421,8 @@ async def wrapper(self, *args, **kwargs): kwargs["socket"] = sock return await func(self, *args, **kwargs) finally: - # 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. + # Close the per-call socket only; the context outlives the call. linger=0 + # drops unsent frames so close never blocks the event loop. if sock is not None and not sock.closed: sock.close(linger=0) From 309c0aaa7ac1892a6402a65d69fbfb6c5c47dbe5 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 4 Aug 2026 21:27:23 +0800 Subject: [PATCH 13/13] fix: drop dead code and a stale default from the shared ZMQ context change - TransferQueueClient.__init__ docstring still documented libzmq's 1023 as the zmq_max_sockets default after DEFAULT_CLIENT_ZMQ_MAX_SOCKETS (8192) landed. - _release_zmq_context: drop the `context is not None` check. The finalizer is armed only after the context is constructed, so it can never be None. - test Borrower stub: drop the _do_handshake_with_controller and notify_data_update overrides. Neither is abstract, and the handshake is reachable only from _connect_to_controller, which the stub already overrides to pass. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 6 ------ transfer_queue/client.py | 6 +++--- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index b31151a1..2bd9a6c7 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -273,9 +273,6 @@ 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 @@ -285,9 +282,6 @@ async def get_data(self, *args, **kwargs): 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) diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 65559bea..a0b82f92 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -145,7 +145,7 @@ def _release_zmq_context(context: "zmq.asyncio.Context", client_id: str) -> None the client, deferring this to interpreter exit -- close() remains the supported path. """ try: - if context is not None and not context.closed: + if not context.closed: context.destroy(linger=0) except Exception as e: logger.warning(f"[{client_id}]: Error destroying zmq_context in finalizer: {e}") @@ -1372,8 +1372,8 @@ def __init__( 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. + at once. Defaults to ``TQ_CLIENT_ZMQ_MAX_SOCKETS``, and to + ``DEFAULT_CLIENT_ZMQ_MAX_SOCKETS`` (8192) when that is unset. """ super().__init__( client_id,