From fe97e065a5d165c6fe284b6419f781376db8ac42 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 10 Aug 2026 16:40:47 +0800 Subject: [PATCH 01/28] [perf, fix] Pool and reuse ZMQ request sockets instead of one per request Every request path built a fresh DEALER socket -- create, connect, send, recv, close -- paying a TCP + ZMTP handshake per call, minting a new ROUTER peer identity each time, and calling psutil.virtual_memory() per socket. Sockets are now leased from a long-lived pool: 0.43ms -> 0.19ms per round trip (2.2x), and in-flight socket count tracks concurrency rather than request count, which relieves pressure on TQ_CLIENT_ZMQ_MAX_SOCKETS during large fan-outs. A lease is exclusive. Responses do not echo their request's request_id, so a reply is matched to its request only by arrival order; two concurrent users of one socket would read each other's replies. A socket therefore returns to the pool only after a clean send+recv, and is closed on timeout, error, or cancellation -- a request already on the wire may still get a reply, and the next lessee would read it as its own. asyncio.gather cancels siblings on the first failure, so cancellation is a live path, not a theoretical one. Buckets are keyed by event loop (or thread, outside one), since pyzmq rebinds a socket to whatever loop it next sees and one bound to a closed loop is the documented "Bad file descriptor / SIGABRT" hazard. Finished owners are evicted on the next lease: a pooled async socket references its own loop, so these cannot be reclaimed by garbage collection. The keying also lets the synchronous metrics collector share one pool implementation with the async callers. This fixes a live bug in metrics collection. _query_storage_unit caught zmq.error.Again before the branch that evicted the socket, and Again is an Exception subclass, so after any timeout the next cycle read the previous cycle's reply and attributed it to the wrong cycle. The pool's poison rule removes the bug along with the bookkeeping. Pool and context ownership are now enforced to arrive together: a pool built over a borrowed context would mint sockets the lender may destroy underneath it. Verified: 582 passed, 10 skipped. The 8 remaining errors are pre-existing and require the optional Yuanrong datasystem SDK. The three poison-rule tests fail against a pool without the rule, confirming they are real guards. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 51 ++- tests/test_zmq_socket_pool.py | 348 ++++++++++++++++++ transfer_queue/client.py | 25 +- transfer_queue/metrics.py | 59 ++- transfer_queue/storage/managers/base.py | 87 +++-- .../storage/managers/mooncake_manager.py | 5 +- .../storage/managers/ray_storage_manager.py | 10 +- .../managers/simple_storage_manager.py | 11 +- .../storage/managers/yuanrong_manager.py | 5 +- transfer_queue/utils/zmq_utils.py | 205 +++++++++-- 10 files changed, 659 insertions(+), 147 deletions(-) create mode 100644 tests/test_zmq_socket_pool.py diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 2bd9a6c7..34f317ad 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -17,8 +17,9 @@ 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 +-> SIGABRT). It now reuses the owner's context and leases sockets from a shared pool, so these tests assert every concurrent call sees the SAME context, alive until close(). +Pool-specific behavior is covered in test_zmq_socket_pool.py. """ import asyncio @@ -183,11 +184,41 @@ def test_simple_storage_borrows_client_context(echo_controller): controller_info=echo_controller.zmq_server_info, config=config, zmq_context=client.zmq_context, + zmq_socket_pool=client.zmq_socket_pool, ) client.close() +def test_simple_storage_borrows_client_socket_pool(echo_controller): + """A manager borrowing the client's context must borrow its pool too. + + A pool over a context it does not own would either outlive its sockets' context or be + closed while the lender is still leasing from it. + """ + client = AsyncTransferQueueClient( + client_id="client_borrowed_pool", + 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, + zmq_socket_pool=client.zmq_socket_pool, + ) + + assert manager.zmq_socket_pool is client.zmq_socket_pool + assert not manager._owns_zmq_context + + # Closing the borrower must leave the lender's pool usable. + manager.close() + assert not client.zmq_context.closed + + client.close() + + def test_simple_storage_does_not_destroy_borrowed_context(echo_controller): client = AsyncTransferQueueClient( client_id="client_borrowed_context_lifecycle", @@ -199,6 +230,7 @@ def test_simple_storage_does_not_destroy_borrowed_context(echo_controller): echo_controller.zmq_server_info, {"zmq_info": {"storage_0": echo_controller.zmq_server_info}}, zmq_context=client.zmq_context, + zmq_socket_pool=client.zmq_socket_pool, ) assert manager.zmq_context is client.zmq_context @@ -266,8 +298,11 @@ def test_close_skips_destroy_while_loop_thread_alive(echo_controller): context.destroy(linger=0) -def _make_borrowing_manager(zmq_context): - """A minimal manager that borrows a caller's context, like SimpleStorage does.""" +def _make_borrowing_manager(client=None): + """A minimal manager borrowing *client*'s context and pool, like SimpleStorage does. + + With no client it creates its own, since the two must always arrive together. + """ class Borrower(StorageManager): def _connect_to_controller(self): @@ -282,7 +317,9 @@ async def get_data(self, *args, **kwargs): async def clear_data(self, *args, **kwargs): return None - return Borrower(None, {}, zmq_context=zmq_context) + if client is None: + return Borrower(None, {}) + return Borrower(None, {}, zmq_context=client.zmq_context, zmq_socket_pool=client.zmq_socket_pool) def test_stuck_notify_thread_vetoes_destroy_of_borrowed_context(echo_controller): @@ -295,7 +332,7 @@ def test_stuck_notify_thread_vetoes_destroy_of_borrowed_context(echo_controller) client_id="client_notify_thread_veto", controller_info=echo_controller.zmq_server_info, ) - client.storage_manager = _make_borrowing_manager(client.zmq_context) + client.storage_manager = _make_borrowing_manager(client) context = client.zmq_context with patch.object(client.storage_manager._notify_thread, "is_alive", return_value=True): @@ -314,7 +351,7 @@ def test_healthy_notify_thread_does_not_block_destroy(echo_controller): client_id="client_notify_thread_clean", controller_info=echo_controller.zmq_server_info, ) - client.storage_manager = _make_borrowing_manager(client.zmq_context) + client.storage_manager = _make_borrowing_manager(client) client.close() assert client.zmq_context.closed @@ -326,7 +363,7 @@ def test_manager_with_own_context_does_not_veto(echo_controller): client_id="client_independent_manager", controller_info=echo_controller.zmq_server_info, ) - client.storage_manager = _make_borrowing_manager(None) # creates its own context + client.storage_manager = _make_borrowing_manager() # 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. diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py new file mode 100644 index 00000000..f3305ef9 --- /dev/null +++ b/tests/test_zmq_socket_pool.py @@ -0,0 +1,348 @@ +# 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. + +"""Tests for ZMQSocketPool. + +Sockets used to be created and closed per request. The pool reuses them, which is only +safe because a lease is exclusive and a socket that did not complete a clean send/recv is +discarded: replies carry no request id (see ZMQMessage.create), so a reply left in flight +by a timed-out or cancelled request would be read by the next user of that socket as its +own. test_timed_out_socket_is_not_reused pins exactly that. +""" + +import asyncio +import threading + +import pytest +import zmq +import zmq.asyncio + +import transfer_queue.utils.zmq_utils as zmq_utils +from transfer_queue.utils.enum_utils import Role +from transfer_queue.utils.zmq_utils import ZMQServerInfo, ZMQSocketPool + + +class _Peer: + """A ROUTER that echoes one reply per request, optionally after a delay.""" + + def __init__(self, delay_first_reply: float = 0.0): + self.context = zmq.Context() + self.socket = self.context.socket(zmq.ROUTER) + port = self.socket.bind_to_random_port("tcp://127.0.0.1") + self.info = ZMQServerInfo(role=Role.STORAGE, id="peer_0", ip="127.0.0.1", ports={"put_get_socket": port}) + self._delay_first_reply = delay_first_reply + self._replies = 0 + self.running = True + self.thread = threading.Thread(target=self._serve, daemon=True) + self.thread.start() + + def _serve(self): + poller = zmq.Poller() + poller.register(self.socket, zmq.POLLIN) + while self.running: + if not dict(poller.poll(50)): + continue + identity, request = self.socket.recv_multipart() + if self._replies == 0 and self._delay_first_reply: + # Reply late enough that the requester has already timed out, leaving this + # reply in flight -- the poisoned-socket scenario. + time_left = self._delay_first_reply + while time_left > 0 and self.running: + sleep = min(0.05, time_left) + threading.Event().wait(sleep) + time_left -= sleep + self._replies += 1 + self.socket.send_multipart([identity, b"reply-to-" + request]) + + def stop(self): + self.running = False + self.thread.join(timeout=2.0) + self.socket.close(linger=0) + self.context.term() + + +@pytest.fixture +def peer(): + p = _Peer() + yield p + p.stop() + + +def _idle_sockets(pool: ZMQSocketPool) -> list: + """Every socket currently parked in the pool, across all owners and buckets.""" + return [s for buckets in pool._idle.values() for bucket in buckets.values() for s in bucket] + + +async def _round_trip(pool, peer_info, payload=b"req", timeout=None): + with pool.lease(peer_info, "put_get_socket", timeout=timeout) as sock: + await sock.send_multipart([payload]) + return (await sock.recv_multipart())[0] + + +@pytest.mark.asyncio +async def test_socket_is_reused_across_requests(peer): + """Sequential requests to one peer must share a single socket.""" + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + + for i in range(10): + assert await _round_trip(pool, peer.info, f"req{i}".encode()) == f"reply-to-req{i}".encode() + + idle = _idle_sockets(pool) + assert len(idle) == 1, "each request opened its own socket instead of reusing one" + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_timed_out_socket_is_not_reused(): + """A timed-out request must not leave its socket -- or its late reply -- in the pool. + + Regression guard for the core hazard: with the socket pooled, the *next* request would + receive the previous request's reply, silently attributing one response to another. + """ + # Reply to the first request only after its 1s timeout has expired, so the reply is + # still in flight when the socket would otherwise be handed to the next caller. + peer = _Peer(delay_first_reply=2.0) + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + try: + with pytest.raises(zmq.error.Again): + await _round_trip(pool, peer.info, b"first", timeout=1) + + assert _idle_sockets(pool) == [], "a timed-out socket was returned to the pool" + + # The late reply to "first" must not surface as the answer to "second". + assert await _round_trip(pool, peer.info, b"second", timeout=10) == b"reply-to-second" + finally: + pool.close() + ctx.destroy(linger=0) + peer.stop() + + +@pytest.mark.asyncio +async def test_cancelled_lease_discards_socket(peer): + """Cancellation mid-recv poisons the socket: asyncio.gather cancels siblings routinely.""" + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + leased = [] + + async def never_answered(): + with pool.lease(peer.info, "put_get_socket") as sock: + leased.append(sock) + await asyncio.sleep(60) # cancelled here, after the lease was handed out + + task = asyncio.create_task(never_answered()) + await asyncio.sleep(0.1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert leased and leased[0].closed + assert _idle_sockets(pool) == [] + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_failed_lease_discards_socket(peer): + """Any exception in the body poisons the socket, not just timeouts.""" + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + + with pytest.raises(RuntimeError): + with pool.lease(peer.info, "put_get_socket") as sock: + await sock.send_multipart([b"req"]) + raise RuntimeError("handler blew up") + + assert _idle_sockets(pool) == [] + + pool.close() + ctx.destroy(linger=0) + + +def test_sockets_are_not_reused_across_event_loops(peer): + """A socket bound to a finished loop must never be handed to another one. + + pyzmq rebinds an async socket to whatever loop it next sees; one bound to a *closed* + loop is the "Bad file descriptor / SIGABRT" failure this keying exists to prevent. + """ + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + leased = [] + + async def lease_twice(tag): + # Twice per loop, so a socket IS reused within a loop -- which is what makes the + # cross-loop comparison below meaningful rather than trivially true. + for i in range(2): + with pool.lease(peer.info, "put_get_socket") as sock: + leased.append(sock) + await sock.send_multipart([f"{tag}{i}".encode()]) + await sock.recv_multipart() + + asyncio.run(lease_twice("a")) + first = [s for s in leased] + asyncio.run(lease_twice("b")) + second = [s for s in leased if s not in first] + + assert len({id(s) for s in first}) == 1, "a socket should be reused within one loop" + assert len({id(s) for s in second}) == 1 + assert not ({id(s) for s in first} & {id(s) for s in second}), "a socket crossed event loops" + + pool.close() + ctx.destroy(linger=0) + + +def test_finished_loop_releases_its_sockets(peer): + """A finished loop's sockets must be closed, not left parked bound to a dead loop. + + ``asyncio.run()`` per call -- the pattern the client docstrings show -- creates one loop + per call. A pooled async socket keeps its own loop referenced, so these cannot be reaped + by garbage collection; the pool evicts finished owners on the next lease instead. + """ + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + leased = [] + + async def once(): + with pool.lease(peer.info, "put_get_socket") as sock: + leased.append(sock) + await sock.send_multipart([b"q"]) + await sock.recv_multipart() + + for _ in range(5): + asyncio.run(once()) + + assert len(leased) == 5, "each fresh loop needs its own socket" + # The last loop's socket is still parked (nothing has leased since), but every earlier + # one must have been closed rather than accumulating. + assert all(sock.closed for sock in leased[:-1]), "a finished loop left an open socket behind" + assert len(_idle_sockets(pool)) == 1 + assert len(pool._idle) == 1, "finished owners must be evicted, not accumulated" + + pool.close() + ctx.destroy(linger=0) + + +def test_pooled_identities_are_unique_across_pools(peer): + """Two pools with the same owner_id must not collide on the wire. + + A ROUTER silently drops a second peer claiming an identity it already has, so colliding + identities would blackhole one process's traffic. Client ids are pid-derived and pids + repeat across nodes, so owner_id alone cannot carry uniqueness. + """ + ctx = zmq.asyncio.Context() + # Same owner_id, as two processes on different nodes with equal pids would produce. + a, b = ZMQSocketPool(ctx, "TransferQueueClient_1234"), ZMQSocketPool(ctx, "TransferQueueClient_1234") + + async def identity_of(pool): + with pool.lease(peer.info, "put_get_socket") as sock: + return sock.getsockopt(zmq.IDENTITY) + + first, second = asyncio.run(identity_of(a)), asyncio.run(identity_of(b)) + assert first != second, "two pools minted the same ZMQ identity" + + a.close() + b.close() + ctx.destroy(linger=0) + + +def test_connect_failure_does_not_leak_a_socket(peer): + """A socket is nobody's responsibility until it reaches a lease, so _connect closes it.""" + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + bad = ZMQServerInfo(role=Role.STORAGE, id="bad", ip="127.0.0.1", ports={"put_get_socket": -1}) + + created = [] + original = zmq_utils.create_zmq_socket + + def spy(*args, **kwargs): + sock = original(*args, **kwargs) + created.append(sock) + return sock + + zmq_utils.create_zmq_socket = spy + try: + with pytest.raises(zmq.ZMQError): + with pool.lease(bad, "put_get_socket"): + pass + finally: + zmq_utils.create_zmq_socket = original + + assert created and created[0].closed, "a socket that failed to connect was left open" + + pool.close() + ctx.destroy(linger=0) + ctx.destroy(linger=0) + + +def test_sync_caller_can_lease(peer): + """The metrics collector leases from a plain thread, with no event loop running.""" + ctx = zmq.Context() + pool = ZMQSocketPool(ctx, "metrics_collector") + + for i in range(3): + with pool.lease(peer.info, "put_get_socket", timeout=5) as sock: + sock.send_multipart([f"m{i}".encode()]) + assert sock.recv_multipart()[0] == f"reply-to-m{i}".encode() + + assert len(_idle_sockets(pool)) == 1 + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_pool_size_is_a_soft_cap(peer): + """Concurrency above maxsize still gets sockets; only the steady state is bounded.""" + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", maxsize=2) + + results = await asyncio.gather(*[_round_trip(pool, peer.info, f"c{i}".encode()) for i in range(8)]) + assert len(results) == 8, "a burst beyond maxsize must not be refused or blocked" + assert len(_idle_sockets(pool)) == 2, "excess sockets must be closed on return, not parked" + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_unknown_socket_name_is_reported(peer): + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + + with pytest.raises(RuntimeError, match="not configured"): + with pool.lease(peer.info, "no_such_socket"): + pass + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_close_is_idempotent_and_survives_dead_context(peer): + """Teardown ordering is not guaranteed, so close() must tolerate a destroyed context.""" + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + await _round_trip(pool, peer.info) + + pool.close() + assert _idle_sockets(pool) == [] + pool.close() # twice + + ctx.destroy(linger=0) + pool.close() # after the context is gone diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 01e50b85..a4dc33d4 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -32,6 +32,7 @@ ZMQMessage, ZMQRequestType, ZMQServerInfo, + ZMQSocketPool, with_zmq_socket, ) @@ -46,13 +47,15 @@ # 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 +# Idle sockets kept per (loop, peer, timeout) bucket. A soft cap: bursts beyond it still +# get sockets, so this bounds the steady state rather than the peak. +TQ_CLIENT_ZMQ_POOL_SIZE = int(os.environ.get("TQ_CLIENT_ZMQ_POOL_SIZE", 8)) # Pre-bound decorator for controller socket operations. with_controller_socket = with_zmq_socket( "request_handle_socket", - get_identity=lambda self: self.client_id, get_peer=lambda self, target: self._controller, - get_context=lambda self: self.zmq_context, + get_pool=lambda self: self.zmq_socket_pool, ) @@ -91,8 +94,9 @@ def __init__( raise TypeError(f"controller_info must be ZMQServerInfo, got {type(controller_info)}") self.client_id = client_id self._controller: ZMQServerInfo = controller_info - # One long-lived context per client; sockets stay per-request because ZMQ sockets - # are not thread-safe. + # One long-lived context per client, with sockets leased from a pool over it rather + # than built per request; a lease is exclusive because ZMQ sockets are not + # thread-safe and replies are matched to requests by arrival order. 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}") @@ -129,6 +133,10 @@ def __init__( ) max_sockets = socket_limit self.zmq_context.set(zmq.MAX_SOCKETS, max_sockets) + # Sockets are leased from this pool and reused across requests, so the context's + # socket budget above is consumed by the concurrency high-water mark, not by + # request count. Lent to a borrowing storage manager alongside the context. + self.zmq_socket_pool = ZMQSocketPool(self.zmq_context, client_id, maxsize=TQ_CLIENT_ZMQ_POOL_SIZE) # 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 @@ -157,9 +165,9 @@ 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. + The client's long-lived ZMQ context and socket pool are offered to every backend + uniformly; each registered manager decides whether to borrow them 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: @@ -174,6 +182,7 @@ def initialize_storage_manager( controller_info=self._controller, config=config, zmq_context=self.zmq_context, + zmq_socket_pool=self.zmq_socket_pool, ) async def _request_controller( @@ -1053,6 +1062,8 @@ def close(self) -> None: ) return try: + # Close pooled sockets before the context that owns them. + self.zmq_socket_pool.close() if hasattr(self, "zmq_context") and self.zmq_context is not None: self.zmq_context.destroy(linger=0) except Exception as e: diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index e90e9c59..03690270 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -18,7 +18,6 @@ from contextlib import contextmanager from threading import Thread from typing import Any -from uuid import uuid4 import psutil import zmq @@ -30,8 +29,7 @@ ZMQMessage, ZMQRequestType, ZMQServerInfo, - create_zmq_socket, - format_zmq_address, + ZMQSocketPool, ) logger = get_logger(__name__) @@ -74,7 +72,7 @@ def __init__(self, role: str = "controller"): self._role = role self._storage_unit_infos: dict[str, ZMQServerInfo] = {} self._zmq_ctx: zmq.Context | None = None - self._zmq_sockets: dict[str, zmq.Socket] = {} + self._zmq_socket_pool: ZMQSocketPool | None = None self._known_partition_ids: set[str] = set() self._known_production_labels: set[tuple[str, str]] = set() self._known_consumption_labels: set[tuple[str, str]] = set() @@ -372,49 +370,36 @@ def collect_storage_metrics(self) -> None: except Exception as e: logger.warning(f"Failed to collect metrics from storage unit {su_id}: {e}") - def _get_or_create_socket(self, su_id: str, su_info: ZMQServerInfo) -> zmq.Socket: - """Return a cached ZMQ DEALER socket for *su_id*, creating one if needed.""" - if self._zmq_ctx is None: + def _get_socket_pool(self) -> ZMQSocketPool: + """Return the lazily-created socket pool for storage-unit queries.""" + if self._zmq_socket_pool is None: self._zmq_ctx = zmq.Context() - - sock = self._zmq_sockets.get(su_id) - if sock is not None and not sock.closed: - return sock - - identity = f"{METRICS_COLLECTOR_IDENTITY_PREFIX}{uuid4().hex[:8]}".encode() - sock = create_zmq_socket(self._zmq_ctx, zmq.DEALER, su_info.ip, identity) - timeout_ms = TQ_METRICS_STORAGE_TIMEOUT * 1000 - address = format_zmq_address(su_info.ip, su_info.ports["put_get_socket"]) - sock.connect(address) - sock.setsockopt(zmq.RCVTIMEO, timeout_ms) - sock.setsockopt(zmq.SNDTIMEO, timeout_ms) - self._zmq_sockets[su_id] = sock - return sock + self._zmq_socket_pool = ZMQSocketPool(self._zmq_ctx, "metrics_collector") + return self._zmq_socket_pool def _query_storage_unit(self, su_info: ZMQServerInfo, su_id: str) -> dict[str, Any] | None: """Send a synchronous GET_METRICS request to a single storage unit.""" try: - sock = self._get_or_create_socket(su_id, su_info) - request_msg = ZMQMessage.create( - request_type=ZMQRequestType.GET_METRICS, - sender_id="metrics_collector", - body={}, - ) - sock.send_multipart(request_msg.serialize()) - response_frames = sock.recv_multipart(copy=False) - response_msg = ZMQMessage.deserialize(response_frames) - if response_msg.request_type == ZMQRequestType.METRICS_RESPONSE: - return response_msg.body - return None + pool = self._get_socket_pool() + with pool.lease(su_info, "put_get_socket", timeout=TQ_METRICS_STORAGE_TIMEOUT) as sock: + request_msg = ZMQMessage.create( + request_type=ZMQRequestType.GET_METRICS, + sender_id="metrics_collector", + body={}, + ) + sock.send_multipart(request_msg.serialize()) + response_frames = sock.recv_multipart(copy=False) + response_msg = ZMQMessage.deserialize(response_frames) + if response_msg.request_type == ZMQRequestType.METRICS_RESPONSE: + return response_msg.body + return None except zmq.error.Again: + # The pool discarded the socket, so the next cycle starts clean. Reusing it + # would read this reply, if it lands late, as the answer to that cycle's query. logger.debug(f"Timeout querying metrics from {su_id}") return None except Exception as e: logger.warning(f"Error querying metrics from {su_id}: {e}") - # Close broken socket so it gets recreated next cycle - sock = self._zmq_sockets.pop(su_id, None) - if sock and not sock.closed: - sock.close(linger=0) return None def start(self, node_ip: str = "0.0.0.0", port: int = 0) -> str: diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 14790523..5067b1de 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -37,10 +37,10 @@ from transfer_queue.storage.clients.base import StorageClientFactory from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( - STORAGE_MANAGER_IDENTITY_PREFIX, ZMQMessage, ZMQRequestType, ZMQServerInfo, + ZMQSocketPool, create_zmq_socket, ) @@ -75,6 +75,7 @@ def __init__( controller_info: ZMQServerInfo, config: DictConfig, zmq_context: zmq.asyncio.Context | None = None, + zmq_socket_pool: ZMQSocketPool | None = None, ): self.storage_manager_id = f"{STORAGE_MANAGER_IDENTITY_PREFIX}{uuid4().hex[:8]}" self.config = config @@ -85,8 +86,15 @@ def __init__( # 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()). + # The pool must follow the context: one built over a borrowed context would mint + # sockets the lender may destroy underneath it, so the two arrive together. + if (zmq_context is None) != (zmq_socket_pool is None): + raise ValueError("zmq_context and zmq_socket_pool must be passed together, or neither") self._owns_zmq_context = zmq_context is None self.zmq_context = zmq.asyncio.Context() if zmq_context is None else zmq_context + self.zmq_socket_pool = ( + ZMQSocketPool(self.zmq_context, self.storage_manager_id) if zmq_socket_pool is None else zmq_socket_pool + ) self._connect_to_controller() # Dedicated asyncio loop for ZMQ notify traffic, isolated from the caller's loop @@ -271,56 +279,41 @@ async def notify_data_update( async def _notify_and_wait(self, request_msg: list) -> None: """Send a data status notification to the controller and block until ACK is received.""" - identity = f"{self.storage_manager_id}-notify-{uuid4().hex[:8]}".encode() - sock = create_zmq_socket( - ctx=self.zmq_context, socket_type=zmq.DEALER, ip=self.controller_info.ip, identity=identity - ) - sock.setsockopt(zmq.LINGER, 0) - sock.connect(self.controller_info.to_addr("request_handle_socket")) - - try: - await sock.send_multipart(request_msg) - logger.debug( - f"[{self.storage_manager_id}]: Sent data status update request " - f"to controller id #{self.controller_info.id} successfully." - ) - - response_received = False - timeout = TQ_DATA_UPDATE_RESPONSE_TIMEOUT + # Acquiring the lease sits outside the handler below: a missing socket name or a dead + # context is a configuration/lifecycle fault the caller must see, not a slow ACK. + with self.zmq_socket_pool.lease(self.controller_info, "request_handle_socket") as sock: + try: + await sock.send_multipart(request_msg) + logger.debug( + f"[{self.storage_manager_id}]: Sent data status update request " + f"to controller id #{self.controller_info.id} successfully." + ) - while not response_received and timeout > 0: - try: - poll_interval = min(TQ_STORAGE_POLLER_TIMEOUT, timeout) - messages = await asyncio.wait_for( - sock.recv_multipart(copy=False), - timeout=poll_interval, - ) + # One deadline for the whole wait, so unrelated traffic on this socket cannot + # extend it and a quiet controller still gets the full budget. + deadline = time.monotonic() + TQ_DATA_UPDATE_RESPONSE_TIMEOUT + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + f"no ACK from controller {self.controller_info.id} after {TQ_DATA_UPDATE_RESPONSE_TIMEOUT}s" + ) + messages = await asyncio.wait_for(sock.recv_multipart(copy=False), timeout=remaining) response_msg = ZMQMessage.deserialize(messages) if response_msg.request_type == ZMQRequestType.NOTIFY_DATA_UPDATE_ACK: # type: ignore[arg-type] - response_received = True logger.debug( f"[{self.storage_manager_id}]: Get data status update ACK response " f"from controller id #{response_msg.sender_id} successfully." ) - break - except asyncio.TimeoutError: - timeout -= poll_interval - except Exception as e: - logger.warning(f"[{self.storage_manager_id}]: Error receiving response: {e}") - break - - if not response_received: - logger.error( - f"[{self.storage_manager_id}]: Timeout waiting for data status update ACK " - f"from controller after {TQ_DATA_UPDATE_RESPONSE_TIMEOUT}s." - ) - finally: - try: - if not sock.closed: - sock.close(linger=0) - except Exception: - pass + return + except Exception as e: + # Notification failure has always been logged rather than raised, so a slow + # controller does not fail the put that triggered it. Close the socket instead + # of reusing it: an ACK may still be in flight, and the next lessee would read + # it as its own reply. The pool discards a closed socket on its next lease. + logger.error(f"[{self.storage_manager_id}]: Data status update failed: {type(e).__name__}: {e}") + sock.close(linger=0) @abstractmethod async def put_data( @@ -413,8 +406,10 @@ def close(self) -> None: # after the notify thread holding sockets is gone. If that thread outlived its # join, leak the context rather than risk a crash on a terminating process. if notify_thread_stopped: - # linger=0 force-closes sockets left by an interrupted request, so this - # cannot hang on term(). + # Pooled sockets first: they live on the context destroyed next. linger=0 + # force-closes sockets left by an interrupted request, so this cannot hang + # on term(). + self.zmq_socket_pool.close() self.zmq_context.destroy(linger=0) else: logger.warning( @@ -524,6 +519,7 @@ def __init__( controller_info: ZMQServerInfo, config: dict[str, Any], zmq_context: zmq.asyncio.Context | None = None, + zmq_socket_pool: ZMQSocketPool | None = None, ): """ Initialize the KVStorageManager with configuration. @@ -535,6 +531,7 @@ def __init__( 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. + zmq_socket_pool: Ignored for the same reason; the pool must follow the context. """ 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 48fe6280..f9827393 100644 --- a/transfer_queue/storage/managers/mooncake_manager.py +++ b/transfer_queue/storage/managers/mooncake_manager.py @@ -18,7 +18,7 @@ import zmq from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory -from transfer_queue.utils.zmq_utils import ZMQServerInfo +from transfer_queue.utils.zmq_utils import ZMQServerInfo, ZMQSocketPool @StorageManagerFactory.register("MooncakeStore") @@ -36,6 +36,7 @@ def __init__( controller_info: ZMQServerInfo, config: dict[str, Any], zmq_context: zmq.asyncio.Context | None = None, + zmq_socket_pool: ZMQSocketPool | None = None, ): config["client_name"] = "MooncakeStoreClient" - super().__init__(controller_info, config, zmq_context=zmq_context) + super().__init__(controller_info, config, zmq_context=zmq_context, zmq_socket_pool=zmq_socket_pool) diff --git a/transfer_queue/storage/managers/ray_storage_manager.py b/transfer_queue/storage/managers/ray_storage_manager.py index a48176e4..09ecf884 100644 --- a/transfer_queue/storage/managers/ray_storage_manager.py +++ b/transfer_queue/storage/managers/ray_storage_manager.py @@ -18,7 +18,7 @@ import zmq from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory -from transfer_queue.utils.zmq_utils import ZMQServerInfo +from transfer_queue.utils.zmq_utils import ZMQServerInfo, ZMQSocketPool @StorageManagerFactory.register("RayStore") @@ -30,8 +30,14 @@ def __init__( controller_info: ZMQServerInfo, config: dict[str, Any], zmq_context: zmq.asyncio.Context | None = None, + zmq_socket_pool: ZMQSocketPool | 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"}, zmq_context=zmq_context) + super().__init__( + controller_info, + {**config, "client_name": "RayStorageClient"}, + zmq_context=zmq_context, + zmq_socket_pool=zmq_socket_pool, + ) diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 88340006..ff3965ed 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -35,6 +35,7 @@ ZMQMessage, ZMQRequestType, ZMQServerInfo, + ZMQSocketPool, with_zmq_socket, ) @@ -48,11 +49,10 @@ # Pre-bound decorator for storage-unit socket operations. with_storage_unit_socket = with_zmq_socket( "put_get_socket", - get_identity=lambda self: self.storage_manager_id, get_peer=lambda self, target: self.storage_unit_infos[target], - # Long-lived context from the base StorageManager, shared with the notify path. Safe - # because the context is loop-agnostic and each socket stays per-call. - get_context=lambda self: self.zmq_context, + # Long-lived pool from the base StorageManager, shared with the notify path. Safe + # because leases are keyed by event loop and are exclusive for their duration. + get_pool=lambda self: self.zmq_socket_pool, resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"), timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, ) @@ -78,8 +78,9 @@ def __init__( controller_info: ZMQServerInfo, config: DictConfig, zmq_context: zmq.asyncio.Context | None = None, + zmq_socket_pool: ZMQSocketPool | None = None, ): - super().__init__(controller_info, config, zmq_context=zmq_context) + super().__init__(controller_info, config, zmq_context=zmq_context, zmq_socket_pool=zmq_socket_pool) 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 26c30c4b..4930e727 100644 --- a/transfer_queue/storage/managers/yuanrong_manager.py +++ b/transfer_queue/storage/managers/yuanrong_manager.py @@ -19,7 +19,7 @@ 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 +from transfer_queue.utils.zmq_utils import ZMQServerInfo, ZMQSocketPool logger = get_logger(__name__) @@ -33,6 +33,7 @@ def __init__( controller_info: ZMQServerInfo, config: dict[str, Any], zmq_context: zmq.asyncio.Context | None = None, + zmq_socket_pool: ZMQSocketPool | None = None, ): worker_port = config.get("worker_port", None) client_name = config.get("client_name", None) @@ -45,4 +46,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, zmq_context=zmq_context, zmq_socket_pool=zmq_socket_pool) diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index c4b04ce7..792ff962 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -13,9 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio +import itertools import socket +import threading import time -from collections.abc import Sequence +from collections.abc import Iterator, Sequence +from contextlib import contextmanager from dataclasses import dataclass from functools import wraps from typing import Any, Callable, TypeAlias @@ -360,34 +364,177 @@ def create_zmq_socket( return socket +def _lease_owner() -> Any: + """The running event loop, or the current thread outside one. + + A ZMQ socket is safe on neither a second thread nor a second event loop, so leases + never cross either. pyzmq silently rebinds an async socket to whatever loop it next + sees, and one bound to a *closed* loop is the "Bad file descriptor / SIGABRT" failure + that per-call context churn used to cause. Returning the thread outside a loop lets + synchronous callers (the metrics collector) share this pool unchanged. + """ + try: + return asyncio.get_running_loop() + except RuntimeError: + return threading.current_thread() + + +def _owner_finished(owner: Any) -> bool: + """Whether a lease owner can no longer serve its sockets.""" + if isinstance(owner, threading.Thread): + return not owner.is_alive() + return owner.is_closed() + + +class ZMQSocketPool: + """Lends connected DEALER sockets, reusing them across requests. + + A lease is exclusive. Responses do not echo their request's ``request_id`` (see + ``ZMQMessage.create``), so a reply is matched to its request only by arrival order: + two concurrent users of one socket would read each other's replies. Each caller + therefore gets a socket to itself and returns it only after a clean send/recv. + + Reuse avoids a TCP + ZMTP handshake per request and stops the peer's ROUTER from + accreting a fresh identity every call. + """ + + def __init__(self, ctx: zmq.Context, owner_id: str, maxsize: int = 8): + """ + Args: + ctx: Long-lived context to open sockets on, sync or async. The pool borrows it + and never terminates it; the owner outlives the pool. + owner_id: Identity prefix for pooled sockets, for readable peer-side logs. + maxsize: Idle sockets kept per bucket. A soft cap: a burst beyond it still gets + sockets, and the excess is closed on return rather than made to wait. + """ + self._ctx = ctx + self._owner_id = owner_id + self._maxsize = maxsize + # Keyed by lease owner (see _lease_owner), then by (peer, socket name, timeout) -- + # the same peer is dialed with different timeouts. + self._idle: dict[Any, dict[tuple, list[zmq.Socket]]] = {} + self._lock = threading.Lock() + # A ROUTER silently drops a second peer claiming an identity it already has, so + # identities must not collide between processes: owner_id alone does not suffice + # (client ids are pid-derived, and pids repeat across nodes). + self._identity_prefix = f"{owner_id}_{uuid4().hex[:8]}" + self._counter = itertools.count() + + @contextmanager + def lease(self, peer: ZMQServerInfo, socket_name: str, timeout: int | None = None) -> Iterator[zmq.Socket]: + """Yield a socket connected to ``peer``, returning it to the pool only on success. + + A plain (non-async) contextmanager on purpose: ``with`` still sees exceptions and + ``CancelledError`` raised across ``await``s in its body, so this one definition + serves both async and synchronous callers. + """ + port = peer.ports.get(socket_name) + if port is None: + raise RuntimeError(f"Socket '{socket_name}' not configured for server '{peer.id}'") + + key = (peer.id, socket_name, timeout) + sock = self._take(key) or self._connect(peer, port, timeout) + try: + yield sock + except BaseException: + # Poisoned: the request may already be on the wire, so its reply could still + # arrive and the next lessee would read it as its own. Timeouts and + # cancellation count -- asyncio.gather cancels siblings on the first failure. + sock.close(linger=0) + raise + else: + self._release(key, sock) + + def _owner_buckets(self) -> dict[tuple, list[zmq.Socket]]: + """Buckets for the current lease owner, first evicting any owner that has finished. + + Callers must hold ``self._lock``. A closed loop's sockets must not linger: a pooled + async socket keeps its loop referenced, so they would never be collected, and reusing + one is the SIGABRT hazard in _lease_owner. Sweeping here needs no background thread, + and the loop count stays tiny (one per client, plus one per notify thread). + """ + for owner in [o for o in self._idle if _owner_finished(o)]: + self._close_all(self._idle.pop(owner)) + return self._idle.setdefault(_lease_owner(), {}) + + def _take(self, key: tuple) -> zmq.Socket | None: + """Pop a live idle socket for *key*, discarding any found closed.""" + with self._lock: + bucket = self._owner_buckets().get(key) + while bucket: + sock = bucket.pop() + if not sock.closed: + return sock + return None + + def _release(self, key: tuple, sock: zmq.Socket) -> None: + """Return a cleanly-used socket, closing it if its bucket is already full.""" + with self._lock: + bucket = self._owner_buckets().setdefault(key, []) + if len(bucket) < self._maxsize: + bucket.append(sock) + return + sock.close(linger=0) + + def _connect(self, peer: ZMQServerInfo, port: int, timeout: int | None) -> zmq.Socket: + """Open and connect a new DEALER socket to *peer*.""" + identity = f"{self._identity_prefix}_to_{peer.id}_{next(self._counter)}".encode() + sock = create_zmq_socket(self._ctx, zmq.DEALER, peer.ip, identity=identity) + try: + if timeout is not None: + sock.setsockopt(zmq.RCVTIMEO, timeout * 1000) + sock.setsockopt(zmq.SNDTIMEO, timeout * 1000) + sock.connect(format_zmq_address(peer.ip, port)) + except BaseException: + # Nothing owns the socket until it is handed to a lease, so close it here or it + # leaks. connect() raises on a malformed endpoint or a terminating context. + sock.close(linger=0) + raise + return sock + + def close(self) -> None: + """Close every idle socket. Safe to call twice, and after the context is gone. + + The pool stays usable afterwards: a lease still outstanding returns its socket to a + fresh bucket. Both callers destroy the context right after, so nothing is reused. + """ + with self._lock: + owned = list(self._idle.values()) + self._idle = {} + for buckets in owned: + self._close_all(buckets) + + def _close_all(self, buckets: dict[tuple, list[zmq.Socket]]) -> None: + """Close every socket in *buckets*, tolerating an already-destroyed context.""" + for sock in itertools.chain.from_iterable(buckets.values()): + try: + if not sock.closed: + sock.close(linger=0) + except Exception as e: + logger.debug(f"[{self._owner_id}]: Error closing pooled socket: {e}") + + def with_zmq_socket( socket_name: str, *, - get_identity: Callable[[Any], str], get_peer: Callable[[Any, str | None], ZMQServerInfo], - get_context: Callable[[Any], "zmq.asyncio.Context"], + get_pool: Callable[[Any], ZMQSocketPool], resolve_target: Callable[[tuple, dict], str | None] | None = None, timeout: int | None = None, ): - """Create a reusable async decorator for request sockets. - - Lifecycle: get owner's shared context -> create/connect socket -> inject -> close socket. + """Create a reusable async decorator that injects a pooled request 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. + Lifecycle: resolve peer -> lease a socket from ``self``'s pool -> inject as the + ``socket`` kwarg -> return it to the pool if the call succeeded, else discard it. Args: socket_name: Socket port key in ``ZMQServerInfo.ports``. - get_identity: Callable that extracts owner identity from ``self``. - Example: ``lambda self: self.client_id`` get_peer: Callable that returns ``ZMQServerInfo`` for the target. 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`` + get_pool: Callable that returns the owner's ``ZMQSocketPool``. + Example: ``lambda self: self.zmq_socket_pool`` 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")`` @@ -397,10 +544,6 @@ def with_zmq_socket( def decorator(func: Callable): @wraps(func) async def wrapper(self, *args, **kwargs): - owner_id = get_identity(self) - if owner_id is None: - raise RuntimeError("get_identity returned None") - target_name: str | None = None if resolve_target is not None: target_name = resolve_target(args, kwargs) @@ -409,31 +552,13 @@ async def wrapper(self, *args, **kwargs): if server_info is None: raise RuntimeError(f"get_peer returned None for target '{target_name}'") - port = server_info.ports.get(socket_name) - if port is None: - raise RuntimeError(f"Socket '{socket_name}' not configured for server '{server_info.id}'") + pool = get_pool(self) + if pool is None: + raise RuntimeError("get_pool returned None") - # 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) - identity = f"{owner_id}_to_{server_info.id}_{uuid4().hex[:8]}".encode() - sock = create_zmq_socket(context, zmq.DEALER, server_info.ip, identity=identity) - sock.connect(address) - if timeout is not None: - sock.setsockopt(zmq.RCVTIMEO, timeout * 1000) - sock.setsockopt(zmq.SNDTIMEO, timeout * 1000) + with pool.lease(server_info, socket_name, timeout) as sock: kwargs["socket"] = sock return await func(self, *args, **kwargs) - finally: - # Close the per-call socket only; the context outlives the call. linger=0 - # drops unsent frames so close never blocks the event loop. - if sock is not None and not sock.closed: - sock.close(linger=0) return wrapper From b3635fac827609316e53d2f10fab74fbad1f6276 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 10 Aug 2026 19:23:41 +0800 Subject: [PATCH 02/28] [fix] Keep supporting a standalone zmq_context in StorageManager Requiring zmq_context and zmq_socket_pool to arrive together broke a previously supported signature. Two reachable paths raised at construction: a caller or third-party manager passing only zmq_context, and any registered manager whose __init__ predates zmq_socket_pool -- StorageManagerFactory filters the keyword it cannot accept and forwards the context alone. A lone context now builds a pool over it instead. Pool and context ownership are tracked separately, so such a manager closes only its own pooled sockets and never the lender's context. Passing both still shares the lender's pool rather than duplicating it. Teardown moves out of the _owns_zmq_context branch accordingly, or a manager-owned pool over a borrowed context would never have been closed. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 80 ++++++++++++++++++++++++- transfer_queue/storage/managers/base.py | 25 +++++--- 2 files changed, 95 insertions(+), 10 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 34f317ad..8544178d 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -32,7 +32,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 StorageManager +from transfer_queue.storage.managers.base import 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 @@ -211,6 +211,7 @@ def test_simple_storage_borrows_client_socket_pool(echo_controller): assert manager.zmq_socket_pool is client.zmq_socket_pool assert not manager._owns_zmq_context + assert not manager._owns_zmq_socket_pool # Closing the borrower must leave the lender's pool usable. manager.close() @@ -219,6 +220,83 @@ def test_simple_storage_borrows_client_socket_pool(echo_controller): client.close() +def test_manager_pools_over_a_context_lent_without_a_pool(echo_controller): + """Passing ``zmq_context`` alone stays supported and must not raise. + + Older callers and third-party managers pass only a context, and + StorageManagerFactory drops ``zmq_socket_pool`` for constructors that do not accept + it -- so the manager pools over the borrowed context and owns only that pool. + """ + client = AsyncTransferQueueClient( + client_id="client_context_without_pool", + 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 + # Its own pool, built over the borrowed context rather than shared with the lender. + assert manager._owns_zmq_socket_pool + assert manager.zmq_socket_pool is not client.zmq_socket_pool + assert manager.zmq_socket_pool._ctx is client.zmq_context + + # Closing it must release its own pool without touching the lender's context. + manager.close() + assert not client.zmq_context.closed + + client.close() + + +def test_factory_construction_without_pool_support_still_works(echo_controller): + """A registered manager whose __init__ predates ``zmq_socket_pool`` must still build. + + The factory filters the unsupported keyword and forwards the context alone, which used + to hit a construction-time error. + """ + + @StorageManagerFactory.register("_LegacyContextOnly") + class LegacyContextOnly(StorageManager): + def __init__(self, controller_info, config, zmq_context=None): + super().__init__(controller_info, config, zmq_context=zmq_context) + + def _connect_to_controller(self): + pass + + async def put_data(self, *args, **kwargs): + return None + + async def get_data(self, *args, **kwargs): + return None + + async def clear_data(self, *args, **kwargs): + return None + + client = AsyncTransferQueueClient( + client_id="client_legacy_factory", + controller_info=echo_controller.zmq_server_info, + ) + try: + manager = StorageManagerFactory.create( + "_LegacyContextOnly", + controller_info=echo_controller.zmq_server_info, + config={}, + zmq_context=client.zmq_context, + zmq_socket_pool=client.zmq_socket_pool, + ) + assert manager.zmq_socket_pool._ctx is client.zmq_context + manager.close() + assert not client.zmq_context.closed + finally: + StorageManagerFactory._registry.pop("_LegacyContextOnly", None) + client.close() + + def test_simple_storage_does_not_destroy_borrowed_context(echo_controller): client = AsyncTransferQueueClient( client_id="client_borrowed_context_lifecycle", diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 5067b1de..f46d4542 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -86,12 +86,14 @@ def __init__( # 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()). - # The pool must follow the context: one built over a borrowed context would mint - # sockets the lender may destroy underneath it, so the two arrive together. - if (zmq_context is None) != (zmq_socket_pool is None): - raise ValueError("zmq_context and zmq_socket_pool must be passed together, or neither") self._owns_zmq_context = zmq_context is None self.zmq_context = zmq.asyncio.Context() if zmq_context is None else zmq_context + # A caller that lends a context may also lend its pool, so sockets are shared rather + # than duplicated. Passing a context alone stays supported -- older callers and + # third-party managers do, and the factory drops the pool for constructors that do + # not accept it -- in which case this manager pools over the borrowed context and + # closes only its own sockets, never the lender's context. + self._owns_zmq_socket_pool = zmq_socket_pool is None self.zmq_socket_pool = ( ZMQSocketPool(self.zmq_context, self.storage_manager_id) if zmq_socket_pool is None else zmq_socket_pool ) @@ -401,15 +403,19 @@ def close(self) -> None: else: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") + # Close pooled sockets whenever this manager owns them, even over a borrowed context: + # the lender closes its own pool, not this one. Only after the notify thread is gone, + # since Socket.close() is not thread-safe and that thread holds leases. + if self._owns_zmq_socket_pool and notify_thread_stopped: + # linger=0 force-closes sockets left by an interrupted request, so this cannot + # hang, and it runs before any destroy() of the context they live on. + self.zmq_socket_pool.close() + if self._owns_zmq_context: # destroy() calls Socket.close(), which is not thread-safe, so it must run only # after the notify thread holding sockets is gone. If that thread outlived its # join, leak the context rather than risk a crash on a terminating process. if notify_thread_stopped: - # Pooled sockets first: they live on the context destroyed next. linger=0 - # force-closes sockets left by an interrupted request, so this cannot hang - # on term(). - self.zmq_socket_pool.close() self.zmq_context.destroy(linger=0) else: logger.warning( @@ -531,7 +537,8 @@ def __init__( 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. - zmq_socket_pool: Ignored for the same reason; the pool must follow the context. + zmq_socket_pool: Ignored for the same reason, so this manager pools over its own + context rather than the caller's. """ client_name = config.get("client_name", None) if client_name is None: From e48611d0ec068192903f7d5ff18d4daef494814e Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 10 Aug 2026 21:17:37 +0800 Subject: [PATCH 03/28] [fix] Key pooled sockets by endpoint and reject pool sizes below 1 The pool keyed reuse on (peer id, socket name, timeout), which assumed an id permanently identifies one address. A storage unit or controller restarted or re-registered keeps its id but gets a fresh port, so the pool kept leasing a socket still connected to the address the peer no longer answers on -- reachable via register_storage_units(), which updates ZMQServerInfo entries in place. The key now carries the formatted endpoint. Sockets for an address a peer has moved away from are dropped on the next miss, so past addresses cannot accumulate buckets that each keep libzmq reconnecting to a dead endpoint. maxsize was also unvalidated: TQ_CLIENT_ZMQ_POOL_SIZE=-1 or 0 built a client successfully but made the bucket bound permanently false, silently disabling reuse while still looking pooled. Both the pool and the client now reject sizes below 1, and the client names the environment variable. The key becomes a NamedTuple so address-vs-id and the eviction scan read clearly at the point they matter. Regression tests cover a peer moving to a new address, moving back, and not accumulating stale buckets, plus both rejection paths. Each was confirmed to fail against the specific defect it guards. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 15 +++++++ tests/test_zmq_socket_pool.py | 74 +++++++++++++++++++++++++++++-- transfer_queue/client.py | 10 ++++- transfer_queue/utils/zmq_utils.py | 63 +++++++++++++++++++------- 4 files changed, 140 insertions(+), 22 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 8544178d..8e6e0f9e 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -169,6 +169,21 @@ def test_client_rejects_invalid_context_pool_size(echo_controller): ) +def test_client_rejects_invalid_socket_pool_size(echo_controller): + """A bad TQ_CLIENT_ZMQ_POOL_SIZE must name the variable, not silently disable reuse. + + Below 1 nothing is ever parked, so every request pays a fresh connect while the client + still looks pooled. + """ + for bad in (-1, 0): + with patch("transfer_queue.client.TQ_CLIENT_ZMQ_POOL_SIZE", bad): + with pytest.raises(ValueError, match="TQ_CLIENT_ZMQ_POOL_SIZE must be at least 1"): + AsyncTransferQueueClient( + client_id="client_invalid_socket_pool", + controller_info=echo_controller.zmq_server_info, + ) + + def test_simple_storage_borrows_client_context(echo_controller): client = AsyncTransferQueueClient( client_id="client_simple_storage_context", diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py index f3305ef9..95bc83b1 100644 --- a/tests/test_zmq_socket_pool.py +++ b/tests/test_zmq_socket_pool.py @@ -35,13 +35,17 @@ class _Peer: - """A ROUTER that echoes one reply per request, optionally after a delay.""" + """A ROUTER that echoes one reply per request, optionally after a delay. - def __init__(self, delay_first_reply: float = 0.0): + ``tag`` distinguishes which endpoint answered, for the re-registration tests. + """ + + def __init__(self, delay_first_reply: float = 0.0, peer_id: str = "peer_0", tag: bytes = b"reply-to-"): self.context = zmq.Context() self.socket = self.context.socket(zmq.ROUTER) port = self.socket.bind_to_random_port("tcp://127.0.0.1") - self.info = ZMQServerInfo(role=Role.STORAGE, id="peer_0", ip="127.0.0.1", ports={"put_get_socket": port}) + self.info = ZMQServerInfo(role=Role.STORAGE, id=peer_id, ip="127.0.0.1", ports={"put_get_socket": port}) + self._tag = tag self._delay_first_reply = delay_first_reply self._replies = 0 self.running = True @@ -64,7 +68,7 @@ def _serve(self): threading.Event().wait(sleep) time_left -= sleep self._replies += 1 - self.socket.send_multipart([identity, b"reply-to-" + request]) + self.socket.send_multipart([identity, self._tag + request]) def stop(self): self.running = False @@ -261,6 +265,68 @@ async def identity_of(pool): ctx.destroy(linger=0) +@pytest.mark.asyncio +async def test_reregistered_peer_is_not_served_a_stale_socket(): + """A peer that moves to a new address must not be answered by its old endpoint. + + A storage unit or controller restarted or re-registered keeps its id but gets a fresh + port, so keying reuse on the id alone would keep leasing a socket wired to the address + it no longer answers on -- silently talking to a dead or reassigned endpoint. All calls + share one event loop, as a real client or storage manager does; a loop per call would + discard the socket via owner eviction and hide the bug. + """ + old = _Peer(peer_id="su0", tag=b"OLD:") + new = _Peer(peer_id="su0", tag=b"NEW:") + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + try: + assert await _round_trip(pool, old.info) == b"OLD:req" + # Same id, different port -- exactly what re-registration produces. + assert await _round_trip(pool, new.info) == b"NEW:req" + # And a peer that moves back is still reachable. + assert await _round_trip(pool, old.info) == b"OLD:req" + finally: + pool.close() + ctx.destroy(linger=0) + old.stop() + new.stop() + + +@pytest.mark.asyncio +async def test_moved_peer_does_not_accumulate_stale_buckets(): + """Sockets for an address a peer no longer answers on must be dropped, not kept. + + Keying by endpoint alone would leave one bucket per past address, each holding an open + socket that libzmq keeps trying to reconnect. + """ + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + moved = [_Peer(peer_id="su_moves", tag=b"E%d:" % i) for i in range(4)] + try: + for m in moved: + await _round_trip(pool, m.info) + buckets = [b for owner in pool._idle.values() for b in owner] + assert len(buckets) == 1, f"stale endpoint buckets accumulated: {len(buckets)}" + assert len(_idle_sockets(pool)) == 1 + finally: + pool.close() + ctx.destroy(linger=0) + for m in moved: + m.stop() + + +def test_pool_size_below_one_is_rejected(): + """A size under 1 parks nothing, so reuse is silently off while still looking pooled.""" + ctx = zmq.asyncio.Context() + try: + for bad in (-1, 0): + with pytest.raises(ValueError, match="at least 1"): + ZMQSocketPool(ctx, "owner", maxsize=bad) + ZMQSocketPool(ctx, "owner", maxsize=1) # the boundary is valid + finally: + ctx.destroy(linger=0) + + def test_connect_failure_does_not_leak_a_socket(peer): """A socket is nobody's responsibility until it reaches a lease, so _connect closes it.""" ctx = zmq.asyncio.Context() diff --git a/transfer_queue/client.py b/transfer_queue/client.py index a4dc33d4..4c271994 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -47,8 +47,8 @@ # 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 -# Idle sockets kept per (loop, peer, timeout) bucket. A soft cap: bursts beyond it still -# get sockets, so this bounds the steady state rather than the peak. +# Idle sockets kept per (loop, endpoint, timeout) bucket, at least 1. A soft cap: bursts +# beyond it still get sockets, so this bounds the steady state rather than the peak. TQ_CLIENT_ZMQ_POOL_SIZE = int(os.environ.get("TQ_CLIENT_ZMQ_POOL_SIZE", 8)) # Pre-bound decorator for controller socket operations. @@ -136,6 +136,12 @@ def __init__( # Sockets are leased from this pool and reused across requests, so the context's # socket budget above is consumed by the concurrency high-water mark, not by # request count. Lent to a borrowing storage manager alongside the context. + if TQ_CLIENT_ZMQ_POOL_SIZE < 1: + # Name the variable: the pool's own error cannot say which knob supplied the value. + raise ValueError( + f"TQ_CLIENT_ZMQ_POOL_SIZE must be at least 1, got {TQ_CLIENT_ZMQ_POOL_SIZE}. " + f"The pool always reuses at least one socket per endpoint; it cannot be disabled." + ) self.zmq_socket_pool = ZMQSocketPool(self.zmq_context, client_id, maxsize=TQ_CLIENT_ZMQ_POOL_SIZE) # Backstop for a client that is never closed, so the context and its I/O threads do diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 792ff962..99ae1f7d 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -22,7 +22,7 @@ from contextlib import contextmanager from dataclasses import dataclass from functools import wraps -from typing import Any, Callable, TypeAlias +from typing import Any, Callable, NamedTuple, TypeAlias from uuid import uuid4 import psutil @@ -386,6 +386,22 @@ def _owner_finished(owner: Any) -> bool: return owner.is_closed() +class _PoolKey(NamedTuple): + """Identifies one interchangeable set of pooled sockets. + + ``address`` rather than ``peer_id`` alone decides reuse: a peer restarted or + re-registered under the same id at a new address must not be handed a socket still + connected to the old one. ``peer_id`` rides along so a moved peer's now-unreachable + sockets can be found and dropped. ``timeout`` is part of the identity because one peer + is dialed with different send/recv timeouts. + """ + + address: str + socket_name: str + timeout: int | None + peer_id: str + + class ZMQSocketPool: """Lends connected DEALER sockets, reusing them across requests. @@ -404,15 +420,22 @@ def __init__(self, ctx: zmq.Context, owner_id: str, maxsize: int = 8): ctx: Long-lived context to open sockets on, sync or async. The pool borrows it and never terminates it; the owner outlives the pool. owner_id: Identity prefix for pooled sockets, for readable peer-side logs. - maxsize: Idle sockets kept per bucket. A soft cap: a burst beyond it still gets - sockets, and the excess is closed on return rather than made to wait. + maxsize: Idle sockets kept per bucket, at least 1. A soft cap: a burst beyond it + still gets sockets, and the excess is closed on return rather than made to + wait. """ + if maxsize < 1: + # Below 1 nothing is ever parked, so every request pays a fresh connect while + # still looking pooled. Reject it rather than silently disable reuse. + raise ValueError(f"ZMQ socket pool size must be at least 1, got {maxsize}") self._ctx = ctx self._owner_id = owner_id self._maxsize = maxsize - # Keyed by lease owner (see _lease_owner), then by (peer, socket name, timeout) -- - # the same peer is dialed with different timeouts. - self._idle: dict[Any, dict[tuple, list[zmq.Socket]]] = {} + # Keyed by lease owner (see _lease_owner), then by (endpoint, socket name, timeout). + # The endpoint, not the peer id: a peer restarted or re-registered under the same id + # at a new address must not be handed a socket still connected to the old one. The + # timeout is in the key because one peer is dialed with different timeouts. + self._idle: dict[Any, dict[_PoolKey, list[zmq.Socket]]] = {} self._lock = threading.Lock() # A ROUTER silently drops a second peer claiming an identity it already has, so # identities must not collide between processes: owner_id alone does not suffice @@ -432,8 +455,9 @@ def lease(self, peer: ZMQServerInfo, socket_name: str, timeout: int | None = Non if port is None: raise RuntimeError(f"Socket '{socket_name}' not configured for server '{peer.id}'") - key = (peer.id, socket_name, timeout) - sock = self._take(key) or self._connect(peer, port, timeout) + address = format_zmq_address(peer.ip, port) + key = _PoolKey(address=address, socket_name=socket_name, timeout=timeout, peer_id=peer.id) + sock = self._take(key) or self._connect(peer, address, timeout) try: yield sock except BaseException: @@ -445,7 +469,7 @@ def lease(self, peer: ZMQServerInfo, socket_name: str, timeout: int | None = Non else: self._release(key, sock) - def _owner_buckets(self) -> dict[tuple, list[zmq.Socket]]: + def _owner_buckets(self) -> dict[_PoolKey, list[zmq.Socket]]: """Buckets for the current lease owner, first evicting any owner that has finished. Callers must hold ``self._lock``. A closed loop's sockets must not linger: a pooled @@ -457,17 +481,24 @@ def _owner_buckets(self) -> dict[tuple, list[zmq.Socket]]: self._close_all(self._idle.pop(owner)) return self._idle.setdefault(_lease_owner(), {}) - def _take(self, key: tuple) -> zmq.Socket | None: + def _take(self, key: "_PoolKey") -> zmq.Socket | None: """Pop a live idle socket for *key*, discarding any found closed.""" with self._lock: - bucket = self._owner_buckets().get(key) + buckets = self._owner_buckets() + bucket = buckets.get(key) while bucket: sock = bucket.pop() if not sock.closed: return sock + # Missed, so this peer may have just moved. Drop any socket still connected to an + # address it used to answer on: nothing will ask for those again, and each one + # keeps libzmq reconnecting to a dead endpoint in the background. + stale = [k for k in buckets if k.peer_id == key.peer_id and k.address != key.address] + for k in stale: + self._close_all({k: buckets.pop(k)}) return None - def _release(self, key: tuple, sock: zmq.Socket) -> None: + def _release(self, key: "_PoolKey", sock: zmq.Socket) -> None: """Return a cleanly-used socket, closing it if its bucket is already full.""" with self._lock: bucket = self._owner_buckets().setdefault(key, []) @@ -476,15 +507,15 @@ def _release(self, key: tuple, sock: zmq.Socket) -> None: return sock.close(linger=0) - def _connect(self, peer: ZMQServerInfo, port: int, timeout: int | None) -> zmq.Socket: - """Open and connect a new DEALER socket to *peer*.""" + def _connect(self, peer: ZMQServerInfo, address: str, timeout: int | None) -> zmq.Socket: + """Open and connect a new DEALER socket to *address*.""" identity = f"{self._identity_prefix}_to_{peer.id}_{next(self._counter)}".encode() sock = create_zmq_socket(self._ctx, zmq.DEALER, peer.ip, identity=identity) try: if timeout is not None: sock.setsockopt(zmq.RCVTIMEO, timeout * 1000) sock.setsockopt(zmq.SNDTIMEO, timeout * 1000) - sock.connect(format_zmq_address(peer.ip, port)) + sock.connect(address) except BaseException: # Nothing owns the socket until it is handed to a lease, so close it here or it # leaks. connect() raises on a malformed endpoint or a terminating context. @@ -504,7 +535,7 @@ def close(self) -> None: for buckets in owned: self._close_all(buckets) - def _close_all(self, buckets: dict[tuple, list[zmq.Socket]]) -> None: + def _close_all(self, buckets: dict[_PoolKey, list[zmq.Socket]]) -> None: """Close every socket in *buckets*, tolerating an already-destroyed context.""" for sock in itertools.chain.from_iterable(buckets.values()): try: From f1d3026444bdde9e8e2b7b3c943f37c6ea3c5bfc Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 10 Aug 2026 21:32:40 +0800 Subject: [PATCH 04/28] [fix] Discard late socket returns from a superseded endpoint Endpoint migration cleanup only scanned idle buckets, so a lease already in flight to the old address escaped it. When that request finished, _release() recreated the bucket and parked a socket wired to the obsolete endpoint indefinitely. The pool now records the address most recently leased per (peer, socket name). Both paths consult it: _take() drops superseded buckets as before, and _release() declines to park a socket whose address has since been retired. A single predicate covers the race in either direction, and the map is bounded by peer count and cleared on close(). An in-flight request still completes normally -- only its socket is discarded rather than reused. Also move the TQ_CLIENT_ZMQ_POOL_SIZE check ahead of context construction. It ran after the context was built and configured but before the finalizer was armed, so an invalid value raised while leaving the context and its native I/O threads open. All cheap validation now happens before anything is allocated. Both regression tests were confirmed to fail against the pre-fix code. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 29 +++++++++++++++++----- tests/test_zmq_socket_pool.py | 32 ++++++++++++++++++++++++ transfer_queue/client.py | 14 ++++++----- transfer_queue/utils/zmq_utils.py | 41 ++++++++++++++++++++++--------- 4 files changed, 93 insertions(+), 23 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 8e6e0f9e..a6587d0e 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -173,15 +173,32 @@ def test_client_rejects_invalid_socket_pool_size(echo_controller): """A bad TQ_CLIENT_ZMQ_POOL_SIZE must name the variable, not silently disable reuse. Below 1 nothing is ever parked, so every request pays a fresh connect while the client - still looks pooled. + still looks pooled. Validation must also run before the context is built: the finalizer + is not armed until __init__ finishes, so a raise afterwards would leak the context and + its native I/O threads. """ for bad in (-1, 0): + created = [] + real_context = zmq.asyncio.Context + + def _spy(*args, _real=real_context, _seen=created, **kwargs): + ctx = _real(*args, **kwargs) + _seen.append(ctx) + return ctx + with patch("transfer_queue.client.TQ_CLIENT_ZMQ_POOL_SIZE", bad): - with pytest.raises(ValueError, match="TQ_CLIENT_ZMQ_POOL_SIZE must be at least 1"): - AsyncTransferQueueClient( - client_id="client_invalid_socket_pool", - controller_info=echo_controller.zmq_server_info, - ) + with patch("zmq.asyncio.Context", side_effect=_spy): + with pytest.raises(ValueError, match="TQ_CLIENT_ZMQ_POOL_SIZE must be at least 1"): + AsyncTransferQueueClient( + client_id="client_invalid_socket_pool", + controller_info=echo_controller.zmq_server_info, + ) + try: + assert created == [], "the context must not be allocated before validation" + finally: + for ctx in created: + if not ctx.closed: + ctx.destroy(linger=0) def test_simple_storage_borrows_client_context(echo_controller): diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py index 95bc83b1..bc8799aa 100644 --- a/tests/test_zmq_socket_pool.py +++ b/tests/test_zmq_socket_pool.py @@ -315,6 +315,38 @@ async def test_moved_peer_does_not_accumulate_stale_buckets(): m.stop() +@pytest.mark.asyncio +async def test_lease_in_flight_when_peer_moves_is_not_parked(): + """A lease already out on loan when its peer moves must not re-enter the pool. + + _take's sweep only sees idle sockets, so a request in flight to the old address escapes + it and would recreate that bucket on return -- leaving a socket wired to an obsolete + endpoint parked indefinitely. + """ + # Answers slowly, so its lease is still out when the new endpoint is first used. + old = _Peer(delay_first_reply=1.0, peer_id="su0", tag=b"OLD:") + new = _Peer(peer_id="su0", tag=b"NEW:") + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + try: + in_flight = asyncio.create_task(_round_trip(pool, old.info)) + await asyncio.sleep(0.2) # let it reach recv before the peer "moves" + + assert await _round_trip(pool, new.info) == b"NEW:req" + assert await in_flight == b"OLD:req", "the in-flight request should still complete" + + keys = [k for owner in pool._idle.values() for k in owner] + assert len(keys) == 1, f"a superseded endpoint was parked: {[k.address for k in keys]}" + assert keys[0].address == new.info.to_addr("put_get_socket") + # And the next request still reaches the current endpoint. + assert await _round_trip(pool, new.info) == b"NEW:req" + finally: + pool.close() + ctx.destroy(linger=0) + old.stop() + new.stop() + + def test_pool_size_below_one_is_rejected(): """A size under 1 parks nothing, so reuse is silently off while still looking pooled.""" ctx = zmq.asyncio.Context() diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 4c271994..fe1e75a1 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -97,9 +97,17 @@ def __init__( # One long-lived context per client, with sockets leased from a pool over it rather # than built per request; a lease is exclusive because ZMQ sockets are not # thread-safe and replies are matched to requests by arrival order. + # Validate every knob before allocating: a raise after the context exists would leak + # it and its native I/O threads, since the finalizer is not armed until the end. 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}") + if TQ_CLIENT_ZMQ_POOL_SIZE < 1: + # Name the variable: the pool's own error cannot say which knob supplied the value. + raise ValueError( + f"TQ_CLIENT_ZMQ_POOL_SIZE must be at least 1, got {TQ_CLIENT_ZMQ_POOL_SIZE}. " + f"The pool always reuses at least one socket per endpoint; it cannot be disabled." + ) self.zmq_context = zmq.asyncio.Context(io_threads=io_threads) max_sockets = zmq_max_sockets @@ -136,12 +144,6 @@ def __init__( # Sockets are leased from this pool and reused across requests, so the context's # socket budget above is consumed by the concurrency high-water mark, not by # request count. Lent to a borrowing storage manager alongside the context. - if TQ_CLIENT_ZMQ_POOL_SIZE < 1: - # Name the variable: the pool's own error cannot say which knob supplied the value. - raise ValueError( - f"TQ_CLIENT_ZMQ_POOL_SIZE must be at least 1, got {TQ_CLIENT_ZMQ_POOL_SIZE}. " - f"The pool always reuses at least one socket per endpoint; it cannot be disabled." - ) self.zmq_socket_pool = ZMQSocketPool(self.zmq_context, client_id, maxsize=TQ_CLIENT_ZMQ_POOL_SIZE) # Backstop for a client that is never closed, so the context and its I/O threads do diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 99ae1f7d..028a3369 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -436,6 +436,10 @@ def __init__(self, ctx: zmq.Context, owner_id: str, maxsize: int = 8): # at a new address must not be handed a socket still connected to the old one. The # timeout is in the key because one peer is dialed with different timeouts. self._idle: dict[Any, dict[_PoolKey, list[zmq.Socket]]] = {} + # (peer id, socket name) -> the address most recently leased for it. Bounded by the + # peer count, and what lets a lease returning from a superseded address be told apart + # from a current one -- _idle alone cannot, since an in-flight socket is not in it. + self._endpoints: dict[tuple[str, str], str] = {} self._lock = threading.Lock() # A ROUTER silently drops a second peer claiming an identity it already has, so # identities must not collide between processes: owner_id alone does not suffice @@ -485,26 +489,40 @@ def _take(self, key: "_PoolKey") -> zmq.Socket | None: """Pop a live idle socket for *key*, discarding any found closed.""" with self._lock: buckets = self._owner_buckets() + # Leasing this address makes it the peer's current one, retiring any earlier + # address. Recorded on acquire so a lease already in flight to a superseded + # address is recognised as stale when it comes back (see _release). + self._endpoints[key.peer_id, key.socket_name] = key.address + for stale in [k for k in buckets if self._superseded(k)]: + # Nothing will ask for these again, and each keeps libzmq reconnecting to a + # dead endpoint in the background. + self._close_all({stale: buckets.pop(stale)}) bucket = buckets.get(key) while bucket: sock = bucket.pop() if not sock.closed: return sock - # Missed, so this peer may have just moved. Drop any socket still connected to an - # address it used to answer on: nothing will ask for those again, and each one - # keeps libzmq reconnecting to a dead endpoint in the background. - stale = [k for k in buckets if k.peer_id == key.peer_id and k.address != key.address] - for k in stale: - self._close_all({k: buckets.pop(k)}) return None + def _superseded(self, key: "_PoolKey") -> bool: + """Whether *key* names an address its peer has since moved away from. + + Callers must hold ``self._lock``. Unknown peers are not superseded, so a socket is + only ever retired because a newer address was actually seen. + """ + return self._endpoints.get((key.peer_id, key.socket_name), key.address) != key.address + def _release(self, key: "_PoolKey", sock: zmq.Socket) -> None: - """Return a cleanly-used socket, closing it if its bucket is already full.""" + """Return a cleanly-used socket, closing it if superseded or its bucket is full.""" with self._lock: - bucket = self._owner_buckets().setdefault(key, []) - if len(bucket) < self._maxsize: - bucket.append(sock) - return + # A lease that was already in flight when the peer moved must not be parked: it is + # wired to an address nobody will ask for again, and _take's sweep cannot see a + # socket that is out on loan. + if not self._superseded(key): + bucket = self._owner_buckets().setdefault(key, []) + if len(bucket) < self._maxsize: + bucket.append(sock) + return sock.close(linger=0) def _connect(self, peer: ZMQServerInfo, address: str, timeout: int | None) -> zmq.Socket: @@ -532,6 +550,7 @@ def close(self) -> None: with self._lock: owned = list(self._idle.values()) self._idle = {} + self._endpoints = {} for buckets in owned: self._close_all(buckets) From 7f7170f450685eff0fb807a55fc3f37a99f67e0c Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 10 Aug 2026 21:55:02 +0800 Subject: [PATCH 05/28] [fix] Retire superseded endpoints in every pool owner, and stop leaking the context _endpoints is pool-wide but the migration sweep only walked the current owner's buckets. The pool is deliberately shared between a client's RPC loop and a storage manager's notify loop, so several owners can hold idle sockets for one peer: when a migration was observed by one, the others kept sockets open against the retired address until they happened to lease again, or forever once they went quiet. The sweep now runs wherever the current endpoint is recorded and covers every owner's buckets, closing those sockets without ever handing one across owners. It is also skipped entirely when the address has not changed, which is the common path. Client construction still leaked the context on two max-sockets paths: TQ_CLIENT_ZMQ_MAX_SOCKETS parsing and the caller-supplied lower-bound check both ran after zmq.asyncio.Context() but before the finalizer was armed. Both are checkable without a context and now run before allocation. Only the upper bound needs ZMQ_SOCKET_LIMIT, so that check stays after construction and destroys the context if it fails. Regression tests cover three concurrently live owners across a migration and every invalid max-sockets input; both were confirmed to fail against the pre-fix code. The context-leak assertion is now a shared helper, since three tests need it. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 82 ++++++++++++++++++++++++------- tests/test_zmq_socket_pool.py | 49 ++++++++++++++++++ transfer_queue/client.py | 49 +++++++++++------- transfer_queue/utils/zmq_utils.py | 27 +++++++--- 4 files changed, 162 insertions(+), 45 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index a6587d0e..3ea33cb2 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -23,6 +23,7 @@ """ import asyncio +from contextlib import contextmanager from threading import Thread from unittest.mock import patch @@ -169,36 +170,81 @@ def test_client_rejects_invalid_context_pool_size(echo_controller): ) +@contextmanager +def _no_context_left_open(): + """Assert every ZMQ context built inside the block is closed by the time it exits. + + __init__ arms the finalizer only on success, so an invalid argument must either be + rejected before the context is built or destroy it on the way out -- otherwise the + context and its native I/O threads leak for the process lifetime. + """ + created = [] + real_context = zmq.asyncio.Context + + def _spy(*args, **kwargs): + ctx = real_context(*args, **kwargs) + created.append(ctx) + return ctx + + with patch("zmq.asyncio.Context", side_effect=_spy): + try: + yield created + finally: + leaked = [ctx for ctx in created if not ctx.closed] + for ctx in leaked: + ctx.destroy(linger=0) + assert not leaked, f"{len(leaked)} ZMQ context(s) left open on the failure path" + + def test_client_rejects_invalid_socket_pool_size(echo_controller): """A bad TQ_CLIENT_ZMQ_POOL_SIZE must name the variable, not silently disable reuse. Below 1 nothing is ever parked, so every request pays a fresh connect while the client - still looks pooled. Validation must also run before the context is built: the finalizer - is not armed until __init__ finishes, so a raise afterwards would leak the context and - its native I/O threads. + still looks pooled. """ for bad in (-1, 0): - created = [] - real_context = zmq.asyncio.Context - - def _spy(*args, _real=real_context, _seen=created, **kwargs): - ctx = _real(*args, **kwargs) - _seen.append(ctx) - return ctx - with patch("transfer_queue.client.TQ_CLIENT_ZMQ_POOL_SIZE", bad): - with patch("zmq.asyncio.Context", side_effect=_spy): + with _no_context_left_open() as created: with pytest.raises(ValueError, match="TQ_CLIENT_ZMQ_POOL_SIZE must be at least 1"): AsyncTransferQueueClient( client_id="client_invalid_socket_pool", controller_info=echo_controller.zmq_server_info, ) - try: - assert created == [], "the context must not be allocated before validation" - finally: - for ctx in created: - if not ctx.closed: - ctx.destroy(linger=0) + assert created == [], "checkable without a context, so none should be built" + + +def test_invalid_max_sockets_does_not_leak_a_context(echo_controller): + """No invalid max-sockets input may leave the context or its I/O threads behind.""" + # Rejectable without a live context. + with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", "not-a-number"): + with _no_context_left_open() as created: + with pytest.raises(ValueError, match="TQ_CLIENT_ZMQ_MAX_SOCKETS must be an integer"): + AsyncTransferQueueClient( + client_id="client_garbage_max_sockets", + controller_info=echo_controller.zmq_server_info, + ) + assert created == [], "parsing needs no context, so none should be built" + + for bad in (0, -5): + with _no_context_left_open() as created: + with pytest.raises(ValueError, match="at least 1"): + AsyncTransferQueueClient( + client_id="client_low_max_sockets", + controller_info=echo_controller.zmq_server_info, + zmq_max_sockets=bad, + ) + assert created == [], "the lower bound needs no context, so none should be built" + + # Above ZMQ_SOCKET_LIMIT: this one genuinely needs a live context, so it must be + # destroyed rather than hoisted. + with _no_context_left_open() as created: + with pytest.raises(ValueError, match="ZMQ_SOCKET_LIMIT"): + AsyncTransferQueueClient( + client_id="client_huge_max_sockets", + controller_info=echo_controller.zmq_server_info, + zmq_max_sockets=10**9, + ) + assert len(created) == 1, "the limit check requires a built context" def test_simple_storage_borrows_client_context(echo_controller): diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py index bc8799aa..5f1e5a39 100644 --- a/tests/test_zmq_socket_pool.py +++ b/tests/test_zmq_socket_pool.py @@ -315,6 +315,55 @@ async def test_moved_peer_does_not_accumulate_stale_buckets(): m.stop() +@pytest.mark.asyncio +async def test_superseded_sockets_are_evicted_from_every_owner(): + """A migration seen by one owner must retire the old address in all of them. + + The pool is deliberately shared between a client's RPC loop and a storage manager's + notify loop, so several owners can hold idle sockets for the same peer. Sweeping only + the current owner leaves the others reconnecting to the retired address until they + happen to lease again -- or forever, if they go quiet. + """ + old = _Peer(peer_id="su0", tag=b"OLD:") + new = _Peer(peer_id="su0", tag=b"NEW:") + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner") + loops = [] + + def spawn_owner(): + loop = asyncio.new_event_loop() + threading.Thread(target=loop.run_forever, daemon=True).start() + loops.append(loop) + return loop + + def run_on(loop, info): + return asyncio.run_coroutine_threadsafe(_round_trip(pool, info), loop).result(timeout=10) + + try: + # Two owners park a socket to the old address and stay alive. + parked = [spawn_owner(), spawn_owner()] + for loop in parked: + assert run_on(loop, old.info) == b"OLD:req" + assert len(pool._idle) == 2, "each live loop should hold its own bucket" + + # A third owner observes the migration. + assert run_on(spawn_owner(), new.info) == b"NEW:req" + + addresses = {k.address for buckets in pool._idle.values() for k in buckets} + assert addresses == {new.info.to_addr("put_get_socket")}, f"stale addresses remain: {addresses}" + + # Every owner still reaches the current endpoint afterwards. + for loop in parked: + assert run_on(loop, new.info) == b"NEW:req" + finally: + for loop in loops: + loop.call_soon_threadsafe(loop.stop) + pool.close() + ctx.destroy(linger=0) + old.stop() + new.stop() + + @pytest.mark.asyncio async def test_lease_in_flight_when_peer_moves_is_not_parked(): """A lease already out on loan when its peer moves must not re-enter the pool. diff --git a/transfer_queue/client.py b/transfer_queue/client.py index fe1e75a1..6e112046 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -97,8 +97,9 @@ def __init__( # One long-lived context per client, with sockets leased from a pool over it rather # than built per request; a lease is exclusive because ZMQ sockets are not # thread-safe and replies are matched to requests by arrival order. - # Validate every knob before allocating: a raise after the context exists would leak - # it and its native I/O threads, since the finalizer is not armed until the end. + # Everything checkable without the context is checked first: the finalizer is not + # armed until __init__ returns, so a raise after allocation leaks the context and its + # native I/O 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"Client ZMQ I/O thread pool size must be at least 1, got {io_threads}") @@ -108,7 +109,6 @@ def __init__( f"TQ_CLIENT_ZMQ_POOL_SIZE must be at least 1, got {TQ_CLIENT_ZMQ_POOL_SIZE}. " f"The pool always reuses at least one socket per endpoint; it cannot be disabled." ) - self.zmq_context = zmq.asyncio.Context(io_threads=io_threads) max_sockets = zmq_max_sockets explicitly_requested = max_sockets is not None @@ -121,26 +121,37 @@ def __init__( f"TQ_CLIENT_ZMQ_MAX_SOCKETS must be an integer, got {TQ_CLIENT_ZMQ_MAX_SOCKETS!r}" ) from e explicitly_requested = True + if explicitly_requested and max_sockets < 1: + # The upper bound needs ZMQ_SOCKET_LIMIT, hence a live context, but the lower one + # does not -- so reject it before allocating anything. + raise ValueError(f"Client ZMQ max sockets must be at least 1, got {max_sockets}") 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 = zmq.asyncio.Context(io_threads=io_threads) + try: + # ZMQ_SOCKET_LIMIT is a property of the built context, so this last check cannot + # be hoisted above it; destroy the context rather than leak it on failure. + socket_limit = self.zmq_context.get(zmq.SOCKET_LIMIT) + if explicitly_requested: + # A value the caller asked for must not be silently reinterpreted. + if max_sockets > socket_limit: + raise ValueError( + f"Client ZMQ max sockets must be between 1 and this build's " + f"ZMQ_SOCKET_LIMIT ({socket_limit}), got {max_sockets}" + ) + elif max_sockets > socket_limit: + # Nobody asked for the default, so clamp instead of failing to construct on a + # build whose ZMQ_SOCKET_LIMIT is below it. + logger.debug( + f"[{client_id}]: Clamping default ZMQ max sockets {max_sockets} to this " + f"build's ZMQ_SOCKET_LIMIT ({socket_limit})." ) - 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) + max_sockets = socket_limit + self.zmq_context.set(zmq.MAX_SOCKETS, max_sockets) + except BaseException: + self.zmq_context.destroy(linger=0) + raise # Sockets are leased from this pool and reused across requests, so the context's # socket budget above is consumed by the concurrency high-water mark, not by # request count. Lent to a borrowing storage manager alongside the context. diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 028a3369..ef282e53 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -488,15 +488,9 @@ def _owner_buckets(self) -> dict[_PoolKey, list[zmq.Socket]]: def _take(self, key: "_PoolKey") -> zmq.Socket | None: """Pop a live idle socket for *key*, discarding any found closed.""" with self._lock: + # Drop finished owners first, so the endpoint sweep never walks their buckets. buckets = self._owner_buckets() - # Leasing this address makes it the peer's current one, retiring any earlier - # address. Recorded on acquire so a lease already in flight to a superseded - # address is recognised as stale when it comes back (see _release). - self._endpoints[key.peer_id, key.socket_name] = key.address - for stale in [k for k in buckets if self._superseded(k)]: - # Nothing will ask for these again, and each keeps libzmq reconnecting to a - # dead endpoint in the background. - self._close_all({stale: buckets.pop(stale)}) + self._mark_current(key) bucket = buckets.get(key) while bucket: sock = bucket.pop() @@ -504,6 +498,23 @@ def _take(self, key: "_PoolKey") -> zmq.Socket | None: return sock return None + def _mark_current(self, key: "_PoolKey") -> None: + """Record *key*'s address as its peer's current one and retire every older address. + + Callers must hold ``self._lock``. Recorded on acquire so a lease already in flight to + a superseded address is recognised when it returns (see _release). The sweep covers + every owner, not just this one: the pool is deliberately shared between a client's RPC + loop and a storage manager's notify loop, and an idle socket parked by an owner that + goes quiet would otherwise keep reconnecting to the retired address forever. Sockets + are only ever closed here, never handed across owners. + """ + if self._endpoints.get((key.peer_id, key.socket_name)) == key.address: + return # unchanged, so nothing to retire + self._endpoints[key.peer_id, key.socket_name] = key.address + for buckets in self._idle.values(): + for stale in [k for k in buckets if self._superseded(k)]: + self._close_all({stale: buckets.pop(stale)}) + def _superseded(self, key: "_PoolKey") -> bool: """Whether *key* names an address its peer has since moved away from. From 5b7b644e992896f0c6caf165a19ec50f928cac62 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 10 Aug 2026 22:44:09 +0800 Subject: [PATCH 06/28] [refactor] Give each request scenario its own socket pool One pool served controller RPC, storage RPC and notify, but it never actually shared a socket between them: buckets are keyed by owning loop, and the three differ by loop, socket name, or timeout, so every scenario already had its own sockets. The sharing was cost without benefit -- it required lending a pool alongside a context, ownership rules for the borrowed case, and a key wide enough to separate scenarios inside one pool. Each scenario now builds its own pool: controller_rpc_pool on the client, storage_rpc_pool on the SimpleStorage manager, notify_pool on the base manager, and one in the metrics exporter. They still share the client's context, so the socket budget stays client-wide. Because a pool now serves exactly one scenario, socket_name and timeout move to its constructor, and the key collapses from a four-field _PoolKey to the address. The zmq_socket_pool= parameter disappears from five constructors -- including the three KV managers, which only ever accepted it to keep the factory from warning -- along with the pool-vs-context ownership rules and their divergence guard. Endpoint-migration tracking becomes opt-in via follow_endpoint_changes, enabled only for metrics. Verified this is the sole endpoint mapping mutated at runtime (register_storage_units uses dict.update); client._controller, manager.controller_info and manager.storage_unit_infos are each assigned once at construction, so the other three pools no longer pay for a case they cannot hit. Net 68 lines lighter, with the same 2x speedup on the RPC round trip. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 116 +++----------- tests/test_zmq_socket_pool.py | 78 +++++----- transfer_queue/client.py | 27 ++-- transfer_queue/metrics.py | 12 +- transfer_queue/storage/managers/base.py | 31 ++-- .../storage/managers/mooncake_manager.py | 5 +- .../storage/managers/ray_storage_manager.py | 10 +- .../managers/simple_storage_manager.py | 20 ++- .../storage/managers/yuanrong_manager.py | 5 +- transfer_queue/utils/zmq_utils.py | 142 +++++++++--------- 10 files changed, 189 insertions(+), 257 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 3ea33cb2..7e6b7143 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -33,7 +33,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 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 @@ -262,20 +262,20 @@ def test_simple_storage_borrows_client_context(echo_controller): controller_info=echo_controller.zmq_server_info, config=config, zmq_context=client.zmq_context, - zmq_socket_pool=client.zmq_socket_pool, ) client.close() -def test_simple_storage_borrows_client_socket_pool(echo_controller): - """A manager borrowing the client's context must borrow its pool too. +def test_each_scenario_gets_its_own_pool(echo_controller): + """Controller RPC, storage RPC and notify must each hold a separate pool. - A pool over a context it does not own would either outlive its sockets' context or be - closed while the lender is still leasing from it. + They share one context but never one pool: each dials a different peer or socket name, + so a shared pool reused nothing while letting one scenario's sockets be swept by + another's. Separate pools keep each scenario's failures to itself. """ client = AsyncTransferQueueClient( - client_id="client_borrowed_pool", + client_id="client_scenario_pools", controller_info=echo_controller.zmq_server_info, ) @@ -284,97 +284,22 @@ def test_simple_storage_borrows_client_socket_pool(echo_controller): echo_controller.zmq_server_info, {"zmq_info": {"storage_0": echo_controller.zmq_server_info}}, zmq_context=client.zmq_context, - zmq_socket_pool=client.zmq_socket_pool, ) - assert manager.zmq_socket_pool is client.zmq_socket_pool - assert not manager._owns_zmq_context - assert not manager._owns_zmq_socket_pool + pools = [client.controller_rpc_pool, manager.storage_rpc_pool, manager.notify_pool] + assert len({id(pool) for pool in pools}) == 3, "scenarios must not share a pool" + # All three live on the one shared context, so the socket budget stays client-wide. + assert all(pool._ctx is client.zmq_context for pool in pools) + # Each dials the socket its scenario needs. + assert client.controller_rpc_pool._socket_name == "request_handle_socket" + assert manager.notify_pool._socket_name == "request_handle_socket" + assert manager.storage_rpc_pool._socket_name == "put_get_socket" + assert manager.storage_rpc_pool._timeout is not None, "storage RPC keeps its send/recv timeout" - # Closing the borrower must leave the lender's pool usable. manager.close() - assert not client.zmq_context.closed - - client.close() - - -def test_manager_pools_over_a_context_lent_without_a_pool(echo_controller): - """Passing ``zmq_context`` alone stays supported and must not raise. - - Older callers and third-party managers pass only a context, and - StorageManagerFactory drops ``zmq_socket_pool`` for constructors that do not accept - it -- so the manager pools over the borrowed context and owns only that pool. - """ - client = AsyncTransferQueueClient( - client_id="client_context_without_pool", - 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 - # Its own pool, built over the borrowed context rather than shared with the lender. - assert manager._owns_zmq_socket_pool - assert manager.zmq_socket_pool is not client.zmq_socket_pool - assert manager.zmq_socket_pool._ctx is client.zmq_context - - # Closing it must release its own pool without touching the lender's context. - manager.close() - assert not client.zmq_context.closed - client.close() -def test_factory_construction_without_pool_support_still_works(echo_controller): - """A registered manager whose __init__ predates ``zmq_socket_pool`` must still build. - - The factory filters the unsupported keyword and forwards the context alone, which used - to hit a construction-time error. - """ - - @StorageManagerFactory.register("_LegacyContextOnly") - class LegacyContextOnly(StorageManager): - def __init__(self, controller_info, config, zmq_context=None): - super().__init__(controller_info, config, zmq_context=zmq_context) - - def _connect_to_controller(self): - pass - - async def put_data(self, *args, **kwargs): - return None - - async def get_data(self, *args, **kwargs): - return None - - async def clear_data(self, *args, **kwargs): - return None - - client = AsyncTransferQueueClient( - client_id="client_legacy_factory", - controller_info=echo_controller.zmq_server_info, - ) - try: - manager = StorageManagerFactory.create( - "_LegacyContextOnly", - controller_info=echo_controller.zmq_server_info, - config={}, - zmq_context=client.zmq_context, - zmq_socket_pool=client.zmq_socket_pool, - ) - assert manager.zmq_socket_pool._ctx is client.zmq_context - manager.close() - assert not client.zmq_context.closed - finally: - StorageManagerFactory._registry.pop("_LegacyContextOnly", None) - client.close() - - def test_simple_storage_does_not_destroy_borrowed_context(echo_controller): client = AsyncTransferQueueClient( client_id="client_borrowed_context_lifecycle", @@ -386,7 +311,6 @@ def test_simple_storage_does_not_destroy_borrowed_context(echo_controller): echo_controller.zmq_server_info, {"zmq_info": {"storage_0": echo_controller.zmq_server_info}}, zmq_context=client.zmq_context, - zmq_socket_pool=client.zmq_socket_pool, ) assert manager.zmq_context is client.zmq_context @@ -455,9 +379,9 @@ def test_close_skips_destroy_while_loop_thread_alive(echo_controller): def _make_borrowing_manager(client=None): - """A minimal manager borrowing *client*'s context and pool, like SimpleStorage does. + """A minimal manager borrowing *client*'s context, like SimpleStorage does. - With no client it creates its own, since the two must always arrive together. + With no client it creates its own context. Either way it builds its own notify pool. """ class Borrower(StorageManager): @@ -473,9 +397,7 @@ async def get_data(self, *args, **kwargs): async def clear_data(self, *args, **kwargs): return None - if client is None: - return Borrower(None, {}) - return Borrower(None, {}, zmq_context=client.zmq_context, zmq_socket_pool=client.zmq_socket_pool) + return Borrower(None, {}) if client is None else Borrower(None, {}, zmq_context=client.zmq_context) def test_stuck_notify_thread_vetoes_destroy_of_borrowed_context(echo_controller): diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py index 5f1e5a39..4bf81f8e 100644 --- a/tests/test_zmq_socket_pool.py +++ b/tests/test_zmq_socket_pool.py @@ -89,8 +89,8 @@ def _idle_sockets(pool: ZMQSocketPool) -> list: return [s for buckets in pool._idle.values() for bucket in buckets.values() for s in bucket] -async def _round_trip(pool, peer_info, payload=b"req", timeout=None): - with pool.lease(peer_info, "put_get_socket", timeout=timeout) as sock: +async def _round_trip(pool, peer_info, payload=b"req"): + with pool.lease(peer_info) as sock: await sock.send_multipart([payload]) return (await sock.recv_multipart())[0] @@ -99,7 +99,7 @@ async def _round_trip(pool, peer_info, payload=b"req", timeout=None): async def test_socket_is_reused_across_requests(peer): """Sequential requests to one peer must share a single socket.""" ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") for i in range(10): assert await _round_trip(pool, peer.info, f"req{i}".encode()) == f"reply-to-req{i}".encode() @@ -118,19 +118,20 @@ async def test_timed_out_socket_is_not_reused(): Regression guard for the core hazard: with the socket pooled, the *next* request would receive the previous request's reply, silently attributing one response to another. """ - # Reply to the first request only after its 1s timeout has expired, so the reply is - # still in flight when the socket would otherwise be handed to the next caller. - peer = _Peer(delay_first_reply=2.0) + # Replies to the first request only after the pool's 1s timeout has expired, so that + # reply is still in flight when the socket would otherwise go to the next caller. The + # peer serves serially, so the delay stays well inside the second request's own timeout. + peer = _Peer(delay_first_reply=1.4) ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", timeout=1) try: with pytest.raises(zmq.error.Again): - await _round_trip(pool, peer.info, b"first", timeout=1) + await _round_trip(pool, peer.info, b"first") assert _idle_sockets(pool) == [], "a timed-out socket was returned to the pool" # The late reply to "first" must not surface as the answer to "second". - assert await _round_trip(pool, peer.info, b"second", timeout=10) == b"reply-to-second" + assert await _round_trip(pool, peer.info, b"second") == b"reply-to-second" finally: pool.close() ctx.destroy(linger=0) @@ -141,11 +142,11 @@ async def test_timed_out_socket_is_not_reused(): async def test_cancelled_lease_discards_socket(peer): """Cancellation mid-recv poisons the socket: asyncio.gather cancels siblings routinely.""" ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") leased = [] async def never_answered(): - with pool.lease(peer.info, "put_get_socket") as sock: + with pool.lease(peer.info) as sock: leased.append(sock) await asyncio.sleep(60) # cancelled here, after the lease was handed out @@ -166,10 +167,10 @@ async def never_answered(): async def test_failed_lease_discards_socket(peer): """Any exception in the body poisons the socket, not just timeouts.""" ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") with pytest.raises(RuntimeError): - with pool.lease(peer.info, "put_get_socket") as sock: + with pool.lease(peer.info) as sock: await sock.send_multipart([b"req"]) raise RuntimeError("handler blew up") @@ -186,14 +187,14 @@ def test_sockets_are_not_reused_across_event_loops(peer): loop is the "Bad file descriptor / SIGABRT" failure this keying exists to prevent. """ ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") leased = [] async def lease_twice(tag): # Twice per loop, so a socket IS reused within a loop -- which is what makes the # cross-loop comparison below meaningful rather than trivially true. for i in range(2): - with pool.lease(peer.info, "put_get_socket") as sock: + with pool.lease(peer.info) as sock: leased.append(sock) await sock.send_multipart([f"{tag}{i}".encode()]) await sock.recv_multipart() @@ -219,11 +220,11 @@ def test_finished_loop_releases_its_sockets(peer): by garbage collection; the pool evicts finished owners on the next lease instead. """ ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") leased = [] async def once(): - with pool.lease(peer.info, "put_get_socket") as sock: + with pool.lease(peer.info) as sock: leased.append(sock) await sock.send_multipart([b"q"]) await sock.recv_multipart() @@ -251,10 +252,13 @@ def test_pooled_identities_are_unique_across_pools(peer): """ ctx = zmq.asyncio.Context() # Same owner_id, as two processes on different nodes with equal pids would produce. - a, b = ZMQSocketPool(ctx, "TransferQueueClient_1234"), ZMQSocketPool(ctx, "TransferQueueClient_1234") + a, b = ( + ZMQSocketPool(ctx, "TransferQueueClient_1234", "put_get_socket"), + ZMQSocketPool(ctx, "TransferQueueClient_1234", "put_get_socket"), + ) async def identity_of(pool): - with pool.lease(peer.info, "put_get_socket") as sock: + with pool.lease(peer.info) as sock: return sock.getsockopt(zmq.IDENTITY) first, second = asyncio.run(identity_of(a)), asyncio.run(identity_of(b)) @@ -278,7 +282,7 @@ async def test_reregistered_peer_is_not_served_a_stale_socket(): old = _Peer(peer_id="su0", tag=b"OLD:") new = _Peer(peer_id="su0", tag=b"NEW:") ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", follow_endpoint_changes=True) try: assert await _round_trip(pool, old.info) == b"OLD:req" # Same id, different port -- exactly what re-registration produces. @@ -300,7 +304,7 @@ async def test_moved_peer_does_not_accumulate_stale_buckets(): socket that libzmq keeps trying to reconnect. """ ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", follow_endpoint_changes=True) moved = [_Peer(peer_id="su_moves", tag=b"E%d:" % i) for i in range(4)] try: for m in moved: @@ -327,7 +331,7 @@ async def test_superseded_sockets_are_evicted_from_every_owner(): old = _Peer(peer_id="su0", tag=b"OLD:") new = _Peer(peer_id="su0", tag=b"NEW:") ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", follow_endpoint_changes=True) loops = [] def spawn_owner(): @@ -349,7 +353,7 @@ def run_on(loop, info): # A third owner observes the migration. assert run_on(spawn_owner(), new.info) == b"NEW:req" - addresses = {k.address for buckets in pool._idle.values() for k in buckets} + addresses = {addr for buckets in pool._idle.values() for addr in buckets} assert addresses == {new.info.to_addr("put_get_socket")}, f"stale addresses remain: {addresses}" # Every owner still reaches the current endpoint afterwards. @@ -376,7 +380,7 @@ async def test_lease_in_flight_when_peer_moves_is_not_parked(): old = _Peer(delay_first_reply=1.0, peer_id="su0", tag=b"OLD:") new = _Peer(peer_id="su0", tag=b"NEW:") ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", follow_endpoint_changes=True) try: in_flight = asyncio.create_task(_round_trip(pool, old.info)) await asyncio.sleep(0.2) # let it reach recv before the peer "moves" @@ -384,9 +388,9 @@ async def test_lease_in_flight_when_peer_moves_is_not_parked(): assert await _round_trip(pool, new.info) == b"NEW:req" assert await in_flight == b"OLD:req", "the in-flight request should still complete" - keys = [k for owner in pool._idle.values() for k in owner] - assert len(keys) == 1, f"a superseded endpoint was parked: {[k.address for k in keys]}" - assert keys[0].address == new.info.to_addr("put_get_socket") + parked = [addr for owner in pool._idle.values() for addr in owner] + assert len(parked) == 1, f"a superseded endpoint was parked: {parked}" + assert parked[0] == new.info.to_addr("put_get_socket") # And the next request still reaches the current endpoint. assert await _round_trip(pool, new.info) == b"NEW:req" finally: @@ -402,8 +406,8 @@ def test_pool_size_below_one_is_rejected(): try: for bad in (-1, 0): with pytest.raises(ValueError, match="at least 1"): - ZMQSocketPool(ctx, "owner", maxsize=bad) - ZMQSocketPool(ctx, "owner", maxsize=1) # the boundary is valid + ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=bad) + ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=1) # the boundary is valid finally: ctx.destroy(linger=0) @@ -411,7 +415,7 @@ def test_pool_size_below_one_is_rejected(): def test_connect_failure_does_not_leak_a_socket(peer): """A socket is nobody's responsibility until it reaches a lease, so _connect closes it.""" ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") bad = ZMQServerInfo(role=Role.STORAGE, id="bad", ip="127.0.0.1", ports={"put_get_socket": -1}) created = [] @@ -425,7 +429,7 @@ def spy(*args, **kwargs): zmq_utils.create_zmq_socket = spy try: with pytest.raises(zmq.ZMQError): - with pool.lease(bad, "put_get_socket"): + with pool.lease(bad): pass finally: zmq_utils.create_zmq_socket = original @@ -440,10 +444,10 @@ def spy(*args, **kwargs): def test_sync_caller_can_lease(peer): """The metrics collector leases from a plain thread, with no event loop running.""" ctx = zmq.Context() - pool = ZMQSocketPool(ctx, "metrics_collector") + pool = ZMQSocketPool(ctx, "metrics_collector", "put_get_socket", timeout=5) for i in range(3): - with pool.lease(peer.info, "put_get_socket", timeout=5) as sock: + with pool.lease(peer.info) as sock: sock.send_multipart([f"m{i}".encode()]) assert sock.recv_multipart()[0] == f"reply-to-m{i}".encode() @@ -457,7 +461,7 @@ def test_sync_caller_can_lease(peer): async def test_pool_size_is_a_soft_cap(peer): """Concurrency above maxsize still gets sockets; only the steady state is bounded.""" ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner", maxsize=2) + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=2) results = await asyncio.gather(*[_round_trip(pool, peer.info, f"c{i}".encode()) for i in range(8)]) assert len(results) == 8, "a burst beyond maxsize must not be refused or blocked" @@ -470,10 +474,10 @@ async def test_pool_size_is_a_soft_cap(peer): @pytest.mark.asyncio async def test_unknown_socket_name_is_reported(peer): ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "no_such_socket") with pytest.raises(RuntimeError, match="not configured"): - with pool.lease(peer.info, "no_such_socket"): + with pool.lease(peer.info): pass pool.close() @@ -484,7 +488,7 @@ async def test_unknown_socket_name_is_reported(peer): async def test_close_is_idempotent_and_survives_dead_context(peer): """Teardown ordering is not guaranteed, so close() must tolerate a destroyed context.""" ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner") + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") await _round_trip(pool, peer.info) pool.close() diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 6e112046..866dc707 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -47,15 +47,14 @@ # 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 -# Idle sockets kept per (loop, endpoint, timeout) bucket, at least 1. A soft cap: bursts -# beyond it still get sockets, so this bounds the steady state rather than the peak. +# Idle sockets kept per (loop, endpoint) bucket, at least 1. A soft cap: bursts beyond it +# still get sockets, so this bounds the steady state rather than the peak. TQ_CLIENT_ZMQ_POOL_SIZE = int(os.environ.get("TQ_CLIENT_ZMQ_POOL_SIZE", 8)) # Pre-bound decorator for controller socket operations. with_controller_socket = with_zmq_socket( - "request_handle_socket", get_peer=lambda self, target: self._controller, - get_pool=lambda self: self.zmq_socket_pool, + get_pool=lambda self: self.controller_rpc_pool, ) @@ -154,8 +153,14 @@ def __init__( raise # Sockets are leased from this pool and reused across requests, so the context's # socket budget above is consumed by the concurrency high-water mark, not by - # request count. Lent to a borrowing storage manager alongside the context. - self.zmq_socket_pool = ZMQSocketPool(self.zmq_context, client_id, maxsize=TQ_CLIENT_ZMQ_POOL_SIZE) + # request count. Controller RPC only -- the storage backend keeps its own pools, so + # neither scenario can disturb the other's sockets. + self.controller_rpc_pool = ZMQSocketPool( + self.zmq_context, + client_id, + "request_handle_socket", + maxsize=TQ_CLIENT_ZMQ_POOL_SIZE, + ) # 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 @@ -184,9 +189,10 @@ def initialize_storage_manager( ): """Initialize the storage manager. - The client's long-lived ZMQ context and socket pool are offered to every backend - uniformly; each registered manager decides whether to borrow them or keep its own, - so the client needs no knowledge of specific backend names. + 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. Managers build their own pools over + whichever context they end up with, one per request scenario. Args: manager_type: Type of storage manager to create. Supported types include: @@ -201,7 +207,6 @@ def initialize_storage_manager( controller_info=self._controller, config=config, zmq_context=self.zmq_context, - zmq_socket_pool=self.zmq_socket_pool, ) async def _request_controller( @@ -1082,7 +1087,7 @@ def close(self) -> None: return try: # Close pooled sockets before the context that owns them. - self.zmq_socket_pool.close() + self.controller_rpc_pool.close() if hasattr(self, "zmq_context") and self.zmq_context is not None: self.zmq_context.destroy(linger=0) except Exception as e: diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index 03690270..10fe5de7 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -374,14 +374,22 @@ def _get_socket_pool(self) -> ZMQSocketPool: """Return the lazily-created socket pool for storage-unit queries.""" if self._zmq_socket_pool is None: self._zmq_ctx = zmq.Context() - self._zmq_socket_pool = ZMQSocketPool(self._zmq_ctx, "metrics_collector") + self._zmq_socket_pool = ZMQSocketPool( + self._zmq_ctx, + "metrics_collector", + "put_get_socket", + timeout=TQ_METRICS_STORAGE_TIMEOUT, + # register_storage_units() can remap a storage unit id onto a new address, so + # sockets left at the address it moved off must be closed rather than linger. + follow_endpoint_changes=True, + ) return self._zmq_socket_pool def _query_storage_unit(self, su_info: ZMQServerInfo, su_id: str) -> dict[str, Any] | None: """Send a synchronous GET_METRICS request to a single storage unit.""" try: pool = self._get_socket_pool() - with pool.lease(su_info, "put_get_socket", timeout=TQ_METRICS_STORAGE_TIMEOUT) as sock: + with pool.lease(su_info) as sock: request_msg = ZMQMessage.create( request_type=ZMQRequestType.GET_METRICS, sender_id="metrics_collector", diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index f46d4542..8faccc02 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -75,7 +75,6 @@ def __init__( controller_info: ZMQServerInfo, config: DictConfig, zmq_context: zmq.asyncio.Context | None = None, - zmq_socket_pool: ZMQSocketPool | None = None, ): self.storage_manager_id = f"{STORAGE_MANAGER_IDENTITY_PREFIX}{uuid4().hex[:8]}" self.config = config @@ -88,15 +87,10 @@ def __init__( # 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 - # A caller that lends a context may also lend its pool, so sockets are shared rather - # than duplicated. Passing a context alone stays supported -- older callers and - # third-party managers do, and the factory drops the pool for constructors that do - # not accept it -- in which case this manager pools over the borrowed context and - # closes only its own sockets, never the lender's context. - self._owns_zmq_socket_pool = zmq_socket_pool is None - self.zmq_socket_pool = ( - ZMQSocketPool(self.zmq_context, self.storage_manager_id) if zmq_socket_pool is None else zmq_socket_pool - ) + # Notify traffic gets its own pool. It runs on a dedicated loop, so it could not share + # sockets with another scenario in any case, and keeping it separate means no other + # scenario's sockets are reachable from here. + self.notify_pool = ZMQSocketPool(self.zmq_context, self.storage_manager_id, "request_handle_socket") self._connect_to_controller() # Dedicated asyncio loop for ZMQ notify traffic, isolated from the caller's loop @@ -283,7 +277,7 @@ async def _notify_and_wait(self, request_msg: list) -> None: """Send a data status notification to the controller and block until ACK is received.""" # Acquiring the lease sits outside the handler below: a missing socket name or a dead # context is a configuration/lifecycle fault the caller must see, not a slow ACK. - with self.zmq_socket_pool.lease(self.controller_info, "request_handle_socket") as sock: + with self.notify_pool.lease(self.controller_info) as sock: try: await sock.send_multipart(request_msg) logger.debug( @@ -403,13 +397,11 @@ def close(self) -> None: else: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") - # Close pooled sockets whenever this manager owns them, even over a borrowed context: - # the lender closes its own pool, not this one. Only after the notify thread is gone, - # since Socket.close() is not thread-safe and that thread holds leases. - if self._owns_zmq_socket_pool and notify_thread_stopped: - # linger=0 force-closes sockets left by an interrupted request, so this cannot - # hang, and it runs before any destroy() of the context they live on. - self.zmq_socket_pool.close() + # This manager always owns its notify pool, even over a borrowed context. Only after + # the notify thread is gone, since Socket.close() is not thread-safe and that thread + # holds the leases; linger=0 means this cannot hang. + if notify_thread_stopped: + self.notify_pool.close() if self._owns_zmq_context: # destroy() calls Socket.close(), which is not thread-safe, so it must run only @@ -525,7 +517,6 @@ def __init__( controller_info: ZMQServerInfo, config: dict[str, Any], zmq_context: zmq.asyncio.Context | None = None, - zmq_socket_pool: ZMQSocketPool | None = None, ): """ Initialize the KVStorageManager with configuration. @@ -537,8 +528,6 @@ def __init__( 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. - zmq_socket_pool: Ignored for the same reason, so this manager pools over its own - context rather than the caller's. """ 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 f9827393..48fe6280 100644 --- a/transfer_queue/storage/managers/mooncake_manager.py +++ b/transfer_queue/storage/managers/mooncake_manager.py @@ -18,7 +18,7 @@ import zmq from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory -from transfer_queue.utils.zmq_utils import ZMQServerInfo, ZMQSocketPool +from transfer_queue.utils.zmq_utils import ZMQServerInfo @StorageManagerFactory.register("MooncakeStore") @@ -36,7 +36,6 @@ def __init__( controller_info: ZMQServerInfo, config: dict[str, Any], zmq_context: zmq.asyncio.Context | None = None, - zmq_socket_pool: ZMQSocketPool | None = None, ): config["client_name"] = "MooncakeStoreClient" - super().__init__(controller_info, config, zmq_context=zmq_context, zmq_socket_pool=zmq_socket_pool) + 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 09ecf884..a48176e4 100644 --- a/transfer_queue/storage/managers/ray_storage_manager.py +++ b/transfer_queue/storage/managers/ray_storage_manager.py @@ -18,7 +18,7 @@ import zmq from transfer_queue.storage.managers.base import KVStorageManager, StorageManagerFactory -from transfer_queue.utils.zmq_utils import ZMQServerInfo, ZMQSocketPool +from transfer_queue.utils.zmq_utils import ZMQServerInfo @StorageManagerFactory.register("RayStore") @@ -30,14 +30,8 @@ def __init__( controller_info: ZMQServerInfo, config: dict[str, Any], zmq_context: zmq.asyncio.Context | None = None, - zmq_socket_pool: ZMQSocketPool | 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"}, - zmq_context=zmq_context, - zmq_socket_pool=zmq_socket_pool, - ) + 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 ff3965ed..bdd0d8d7 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -48,13 +48,11 @@ # Pre-bound decorator for storage-unit socket operations. with_storage_unit_socket = with_zmq_socket( - "put_get_socket", get_peer=lambda self, target: self.storage_unit_infos[target], - # Long-lived pool from the base StorageManager, shared with the notify path. Safe - # because leases are keyed by event loop and are exclusive for their duration. - get_pool=lambda self: self.zmq_socket_pool, + # Storage RPC has its own pool, separate from the notify pool on the same context: the + # two dial different peers with different timeouts and could never share a socket. + get_pool=lambda self: self.storage_rpc_pool, resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"), - timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, ) @@ -78,9 +76,15 @@ def __init__( controller_info: ZMQServerInfo, config: DictConfig, zmq_context: zmq.asyncio.Context | None = None, - zmq_socket_pool: ZMQSocketPool | None = None, ): - super().__init__(controller_info, config, zmq_context=zmq_context, zmq_socket_pool=zmq_socket_pool) + super().__init__(controller_info, config, zmq_context=zmq_context) + # Storage-unit RPC, on whichever context the base class settled on. + self.storage_rpc_pool = ZMQSocketPool( + self.zmq_context, + self.storage_manager_id, + "put_get_socket", + timeout=TQ_SIMPLE_STORAGE_SEND_RECV_TIMEOUT, + ) self.config = config server_infos: ZMQServerInfo | dict[str, ZMQServerInfo] | None = config.get("zmq_info", None) @@ -660,4 +664,6 @@ async def load_checkpoint(self, checkpoint_dir: str) -> None: def close(self) -> None: """Close all ZMQ sockets and context to prevent resource leaks.""" + # Before super(), which may destroy the context these sockets live on. + self.storage_rpc_pool.close() super().close() diff --git a/transfer_queue/storage/managers/yuanrong_manager.py b/transfer_queue/storage/managers/yuanrong_manager.py index 4930e727..26c30c4b 100644 --- a/transfer_queue/storage/managers/yuanrong_manager.py +++ b/transfer_queue/storage/managers/yuanrong_manager.py @@ -19,7 +19,7 @@ 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, ZMQSocketPool +from transfer_queue.utils.zmq_utils import ZMQServerInfo logger = get_logger(__name__) @@ -33,7 +33,6 @@ def __init__( controller_info: ZMQServerInfo, config: dict[str, Any], zmq_context: zmq.asyncio.Context | None = None, - zmq_socket_pool: ZMQSocketPool | None = None, ): worker_port = config.get("worker_port", None) client_name = config.get("client_name", None) @@ -46,4 +45,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, zmq_socket_pool=zmq_socket_pool) + 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 ef282e53..94731b81 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -22,7 +22,7 @@ from contextlib import contextmanager from dataclasses import dataclass from functools import wraps -from typing import Any, Callable, NamedTuple, TypeAlias +from typing import Any, Callable, TypeAlias from uuid import uuid4 import psutil @@ -386,24 +386,14 @@ def _owner_finished(owner: Any) -> bool: return owner.is_closed() -class _PoolKey(NamedTuple): - """Identifies one interchangeable set of pooled sockets. - - ``address`` rather than ``peer_id`` alone decides reuse: a peer restarted or - re-registered under the same id at a new address must not be handed a socket still - connected to the old one. ``peer_id`` rides along so a moved peer's now-unreachable - sockets can be found and dropped. ``timeout`` is part of the identity because one peer - is dialed with different send/recv timeouts. - """ - - address: str - socket_name: str - timeout: int | None - peer_id: str - - class ZMQSocketPool: - """Lends connected DEALER sockets, reusing them across requests. + """Lends connected DEALER sockets for one request scenario, reusing them across calls. + + One pool serves one kind of request -- controller RPC, storage RPC, notify, metrics -- + so ``socket_name`` and ``timeout`` are fixed here and every socket it holds is + interchangeable but for the address it is connected to. Separate pools per scenario also + keep them from interfering: a pool shared across scenarios reused nothing anyway, since + scenarios differ by owning loop, socket name, or timeout. A lease is exclusive. Responses do not echo their request's ``request_id`` (see ``ZMQMessage.create``), so a reply is matched to its request only by arrival order: @@ -414,15 +404,31 @@ class ZMQSocketPool: accreting a fresh identity every call. """ - def __init__(self, ctx: zmq.Context, owner_id: str, maxsize: int = 8): + def __init__( + self, + ctx: zmq.Context, + owner_id: str, + socket_name: str, + *, + timeout: int | None = None, + maxsize: int = 8, + follow_endpoint_changes: bool = False, + ): """ Args: ctx: Long-lived context to open sockets on, sync or async. The pool borrows it and never terminates it; the owner outlives the pool. owner_id: Identity prefix for pooled sockets, for readable peer-side logs. + socket_name: Port key in ``ZMQServerInfo.ports`` that every lease dials. + timeout: Send/recv timeout in seconds applied to every socket, or None for none. maxsize: Idle sockets kept per bucket, at least 1. A soft cap: a burst beyond it still gets sockets, and the excess is closed on return rather than made to wait. + follow_endpoint_changes: Set when a peer's address can change under a stable id, + as re-registration does. Sockets are keyed by address either way, so a moved + peer is always dialed correctly; this additionally closes the sockets left + behind at the address it abandoned instead of letting them linger. Only the + metrics collector needs it -- elsewhere endpoints are fixed at construction. """ if maxsize < 1: # Below 1 nothing is ever parked, so every request pays a fresh connect while @@ -430,16 +436,18 @@ def __init__(self, ctx: zmq.Context, owner_id: str, maxsize: int = 8): raise ValueError(f"ZMQ socket pool size must be at least 1, got {maxsize}") self._ctx = ctx self._owner_id = owner_id + self._socket_name = socket_name + self._timeout = timeout self._maxsize = maxsize - # Keyed by lease owner (see _lease_owner), then by (endpoint, socket name, timeout). - # The endpoint, not the peer id: a peer restarted or re-registered under the same id - # at a new address must not be handed a socket still connected to the old one. The - # timeout is in the key because one peer is dialed with different timeouts. - self._idle: dict[Any, dict[_PoolKey, list[zmq.Socket]]] = {} - # (peer id, socket name) -> the address most recently leased for it. Bounded by the - # peer count, and what lets a lease returning from a superseded address be told apart - # from a current one -- _idle alone cannot, since an in-flight socket is not in it. - self._endpoints: dict[tuple[str, str], str] = {} + self._follow_endpoint_changes = follow_endpoint_changes + # Keyed by lease owner (see _lease_owner), then by address. The address, not the peer + # id: a peer restarted under the same id at a new address must not be handed a socket + # still connected to the old one. + self._idle: dict[Any, dict[str, list[zmq.Socket]]] = {} + # peer id -> the address most recently leased for it, tracked only when endpoints can + # move. What lets a lease returning from an abandoned address be told apart from a + # current one; _idle alone cannot, since an in-flight socket is not in it. + self._endpoints: dict[str, str] = {} self._lock = threading.Lock() # A ROUTER silently drops a second peer claiming an identity it already has, so # identities must not collide between processes: owner_id alone does not suffice @@ -448,20 +456,19 @@ def __init__(self, ctx: zmq.Context, owner_id: str, maxsize: int = 8): self._counter = itertools.count() @contextmanager - def lease(self, peer: ZMQServerInfo, socket_name: str, timeout: int | None = None) -> Iterator[zmq.Socket]: + def lease(self, peer: ZMQServerInfo) -> Iterator[zmq.Socket]: """Yield a socket connected to ``peer``, returning it to the pool only on success. A plain (non-async) contextmanager on purpose: ``with`` still sees exceptions and ``CancelledError`` raised across ``await``s in its body, so this one definition serves both async and synchronous callers. """ - port = peer.ports.get(socket_name) + port = peer.ports.get(self._socket_name) if port is None: - raise RuntimeError(f"Socket '{socket_name}' not configured for server '{peer.id}'") + raise RuntimeError(f"Socket '{self._socket_name}' not configured for server '{peer.id}'") address = format_zmq_address(peer.ip, port) - key = _PoolKey(address=address, socket_name=socket_name, timeout=timeout, peer_id=peer.id) - sock = self._take(key) or self._connect(peer, address, timeout) + sock = self._take(peer.id, address) or self._connect(peer, address) try: yield sock except BaseException: @@ -471,9 +478,9 @@ def lease(self, peer: ZMQServerInfo, socket_name: str, timeout: int | None = Non sock.close(linger=0) raise else: - self._release(key, sock) + self._release(peer.id, address, sock) - def _owner_buckets(self) -> dict[_PoolKey, list[zmq.Socket]]: + def _owner_buckets(self) -> dict[str, list[zmq.Socket]]: """Buckets for the current lease owner, first evicting any owner that has finished. Callers must hold ``self._lock``. A closed loop's sockets must not linger: a pooled @@ -485,65 +492,66 @@ def _owner_buckets(self) -> dict[_PoolKey, list[zmq.Socket]]: self._close_all(self._idle.pop(owner)) return self._idle.setdefault(_lease_owner(), {}) - def _take(self, key: "_PoolKey") -> zmq.Socket | None: - """Pop a live idle socket for *key*, discarding any found closed.""" + def _take(self, peer_id: str, address: str) -> zmq.Socket | None: + """Pop a live idle socket for *address*, discarding any found closed.""" with self._lock: # Drop finished owners first, so the endpoint sweep never walks their buckets. buckets = self._owner_buckets() - self._mark_current(key) - bucket = buckets.get(key) + self._mark_current(peer_id, address) + bucket = buckets.get(address) while bucket: sock = bucket.pop() if not sock.closed: return sock return None - def _mark_current(self, key: "_PoolKey") -> None: - """Record *key*'s address as its peer's current one and retire every older address. + def _mark_current(self, peer_id: str, address: str) -> None: + """Record *address* as *peer_id*'s current one and close sockets to the old one. Callers must hold ``self._lock``. Recorded on acquire so a lease already in flight to - a superseded address is recognised when it returns (see _release). The sweep covers - every owner, not just this one: the pool is deliberately shared between a client's RPC - loop and a storage manager's notify loop, and an idle socket parked by an owner that - goes quiet would otherwise keep reconnecting to the retired address forever. Sockets - are only ever closed here, never handed across owners. + an abandoned address is recognised when it returns (see _release). The sweep spans + every owner: an idle socket parked by an owner that then goes quiet would otherwise + keep reconnecting to the abandoned address forever. Sockets are only ever closed + here, never handed across owners. """ - if self._endpoints.get((key.peer_id, key.socket_name)) == key.address: + if not self._follow_endpoint_changes: + return + if self._endpoints.get(peer_id) == address: return # unchanged, so nothing to retire - self._endpoints[key.peer_id, key.socket_name] = key.address + self._endpoints[peer_id] = address for buckets in self._idle.values(): - for stale in [k for k in buckets if self._superseded(k)]: + for stale in [addr for addr in buckets if self._superseded(peer_id, addr)]: self._close_all({stale: buckets.pop(stale)}) - def _superseded(self, key: "_PoolKey") -> bool: - """Whether *key* names an address its peer has since moved away from. + def _superseded(self, peer_id: str, address: str) -> bool: + """Whether *address* is one its peer has since moved away from. Callers must hold ``self._lock``. Unknown peers are not superseded, so a socket is only ever retired because a newer address was actually seen. """ - return self._endpoints.get((key.peer_id, key.socket_name), key.address) != key.address + return self._endpoints.get(peer_id, address) != address - def _release(self, key: "_PoolKey", sock: zmq.Socket) -> None: + def _release(self, peer_id: str, address: str, sock: zmq.Socket) -> None: """Return a cleanly-used socket, closing it if superseded or its bucket is full.""" with self._lock: # A lease that was already in flight when the peer moved must not be parked: it is # wired to an address nobody will ask for again, and _take's sweep cannot see a # socket that is out on loan. - if not self._superseded(key): - bucket = self._owner_buckets().setdefault(key, []) + if not self._superseded(peer_id, address): + bucket = self._owner_buckets().setdefault(address, []) if len(bucket) < self._maxsize: bucket.append(sock) return sock.close(linger=0) - def _connect(self, peer: ZMQServerInfo, address: str, timeout: int | None) -> zmq.Socket: + def _connect(self, peer: ZMQServerInfo, address: str) -> zmq.Socket: """Open and connect a new DEALER socket to *address*.""" identity = f"{self._identity_prefix}_to_{peer.id}_{next(self._counter)}".encode() sock = create_zmq_socket(self._ctx, zmq.DEALER, peer.ip, identity=identity) try: - if timeout is not None: - sock.setsockopt(zmq.RCVTIMEO, timeout * 1000) - sock.setsockopt(zmq.SNDTIMEO, timeout * 1000) + if self._timeout is not None: + sock.setsockopt(zmq.RCVTIMEO, self._timeout * 1000) + sock.setsockopt(zmq.SNDTIMEO, self._timeout * 1000) sock.connect(address) except BaseException: # Nothing owns the socket until it is handed to a lease, so close it here or it @@ -556,7 +564,7 @@ def close(self) -> None: """Close every idle socket. Safe to call twice, and after the context is gone. The pool stays usable afterwards: a lease still outstanding returns its socket to a - fresh bucket. Both callers destroy the context right after, so nothing is reused. + fresh bucket. Callers destroy the context right after, so nothing is reused. """ with self._lock: owned = list(self._idle.values()) @@ -565,7 +573,7 @@ def close(self) -> None: for buckets in owned: self._close_all(buckets) - def _close_all(self, buckets: dict[_PoolKey, list[zmq.Socket]]) -> None: + def _close_all(self, buckets: dict[str, list[zmq.Socket]]) -> None: """Close every socket in *buckets*, tolerating an already-destroyed context.""" for sock in itertools.chain.from_iterable(buckets.values()): try: @@ -576,30 +584,28 @@ def _close_all(self, buckets: dict[_PoolKey, list[zmq.Socket]]) -> None: def with_zmq_socket( - socket_name: str, *, get_peer: Callable[[Any, str | None], ZMQServerInfo], get_pool: Callable[[Any], ZMQSocketPool], resolve_target: Callable[[tuple, dict], str | None] | None = None, - timeout: int | None = None, ): """Create a reusable async decorator that injects a pooled request socket. Lifecycle: resolve peer -> lease a socket from ``self``'s pool -> inject as the ``socket`` kwarg -> return it to the pool if the call succeeded, else discard it. + The socket name and timeout come from the pool, which serves one request scenario. + Args: - socket_name: Socket port key in ``ZMQServerInfo.ports``. get_peer: Callable that returns ``ZMQServerInfo`` for the target. 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_pool: Callable that returns the owner's ``ZMQSocketPool``. - Example: ``lambda self: self.zmq_socket_pool`` + get_pool: Callable that returns the pool for this scenario. + Example: ``lambda self: self.controller_rpc_pool`` 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")`` - timeout: Optional timeout (seconds) for both send/recv operations. """ def decorator(func: Callable): @@ -617,7 +623,7 @@ async def wrapper(self, *args, **kwargs): if pool is None: raise RuntimeError("get_pool returned None") - with pool.lease(server_info, socket_name, timeout) as sock: + with pool.lease(server_info) as sock: kwargs["socket"] = sock return await func(self, *args, **kwargs) From 09a57e10d0583956789cf5cc289eb6c78a48b43b Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 7 Sep 2026 11:39:00 +0800 Subject: [PATCH 07/28] [fix] Retire only the moved peer's own former address The endpoint-migration sweep filtered every idle bucket through _superseded(peer_id, addr), which reduces to "addr is not this peer's current address". Applied to the buckets of *other* peers that predicate is always true, so one peer moving closed the pooled sockets of every other peer -- and a peer seen for the first time did the same, since an unrecorded id also compares unequal. The pool that opts into this is the metrics collector's, which queries the whole storage fleet through a single pool, so a single unit re-registering dropped every unit's socket. It self-heals on the next collection cycle, so the cost is a reconnect storm rather than lost traffic, but the sweep was never meant to reach beyond the peer that moved. The sweep now pops the peer's recorded previous address and nothing else, and returns early when there is no previous address to retire. _superseded is left alone: _release calls it with the peer id of the lease being returned, where the predicate is the intended one. The regression test uses three peers, the shape metrics actually runs in, and was confirmed to fail against the pre-fix code. Signed-off-by: OutstanderWang --- tests/test_zmq_socket_pool.py | 34 +++++++++++++++++++++++++++++++ transfer_queue/utils/zmq_utils.py | 19 ++++++++++------- 2 files changed, 46 insertions(+), 7 deletions(-) diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py index 4bf81f8e..dea6e9f6 100644 --- a/tests/test_zmq_socket_pool.py +++ b/tests/test_zmq_socket_pool.py @@ -319,6 +319,40 @@ async def test_moved_peer_does_not_accumulate_stale_buckets(): m.stop() +@pytest.mark.asyncio +async def test_migration_of_one_peer_leaves_other_peers_alone(): + """Retiring a moved peer's old address must not close sockets belonging to other peers. + + The pool that tracks migrations is the metrics collector's, which queries every storage + unit from one pool. Sweeping every address that is not the mover's new one would drop + the whole fleet's sockets each time a single unit re-registers. + """ + old = _Peer(peer_id="su0", tag=b"OLD:") + new = _Peer(peer_id="su0", tag=b"NEW:") + others = [_Peer(peer_id=f"su{i}", tag=b"P%d:" % i) for i in (1, 2)] + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "metrics_collector", "put_get_socket", follow_endpoint_changes=True) + try: + # Two cycles, so every endpoint is recorded and each peer has a socket parked. + for _ in range(2): + for p in [old, *others]: + await _round_trip(pool, p.info) + parked_before = {id(sock) for sock in _idle_sockets(pool)} + assert len(parked_before) == 3, "each peer should hold one idle socket" + + assert await _round_trip(pool, new.info) == b"NEW:req" + + survivors = parked_before & {id(sock) for sock in _idle_sockets(pool)} + assert len(survivors) == 2, "a peer that did not move lost its pooled socket" + addresses = {addr for buckets in pool._idle.values() for addr in buckets} + assert old.info.to_addr("put_get_socket") not in addresses, "the mover's old address stayed" + finally: + pool.close() + ctx.destroy(linger=0) + for p in [old, new, *others]: + p.stop() + + @pytest.mark.asyncio async def test_superseded_sockets_are_evicted_from_every_owner(): """A migration seen by one owner must retire the old address in all of them. diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 94731b81..1b60f061 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -509,19 +509,24 @@ def _mark_current(self, peer_id: str, address: str) -> None: """Record *address* as *peer_id*'s current one and close sockets to the old one. Callers must hold ``self._lock``. Recorded on acquire so a lease already in flight to - an abandoned address is recognised when it returns (see _release). The sweep spans - every owner: an idle socket parked by an owner that then goes quiet would otherwise - keep reconnecting to the abandoned address forever. Sockets are only ever closed - here, never handed across owners. + an abandoned address is recognised when it returns (see _release). Only this peer's + own former address is retired; peers that did not move keep their sockets. The sweep + spans every owner: an idle socket parked by an owner that then goes quiet would + otherwise keep reconnecting to the abandoned address forever. Sockets are only ever + closed here, never handed across owners. """ if not self._follow_endpoint_changes: return - if self._endpoints.get(peer_id) == address: + previous = self._endpoints.get(peer_id) + if previous == address: return # unchanged, so nothing to retire self._endpoints[peer_id] = address + if previous is None: + return # first sighting, so this peer has left nothing behind for buckets in self._idle.values(): - for stale in [addr for addr in buckets if self._superseded(peer_id, addr)]: - self._close_all({stale: buckets.pop(stale)}) + retired = buckets.pop(previous, None) + if retired: + self._close_all({previous: retired}) def _superseded(self, peer_id: str, address: str) -> bool: """Whether *address* is one its peer has since moved away from. From 3504edb1ba5acfc2d1412c2fbbe9fe168dce10c0 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 7 Sep 2026 11:39:26 +0800 Subject: [PATCH 08/28] [fix] Do not park a socket the lease body already closed lease() treats a clean exit as proof the socket completed a send/recv and hands it to _release(), which parked it without checking whether it was still open. StorageManager._notify_and_wait exits exactly that way: on a missing ACK it closes the socket, so a reply still in flight cannot be read by the next lessee as its own, then swallows the error because a slow controller must not fail the put that triggered the notification. The closed socket went back into the pool and held a slot until the next lease popped and discarded it, costing a reconnect after every failed notification. _release now refuses a closed socket, which is the invariant the pool wanted all along: nothing closed is ever parked. Fixing it here rather than at the call site keeps the notify path able to distinguish a slow ACK from the configuration and lifecycle faults lease() itself raises, which it deliberately does not swallow. The regression test was confirmed to fail against the pre-fix code. Signed-off-by: OutstanderWang --- tests/test_zmq_socket_pool.py | 24 ++++++++++++++++++++++++ transfer_queue/utils/zmq_utils.py | 11 ++++++----- 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py index dea6e9f6..35fa3a95 100644 --- a/tests/test_zmq_socket_pool.py +++ b/tests/test_zmq_socket_pool.py @@ -180,6 +180,30 @@ async def test_failed_lease_discards_socket(peer): ctx.destroy(linger=0) +@pytest.mark.asyncio +async def test_socket_closed_by_the_caller_is_not_parked(peer): + """A socket the body closed without raising must not occupy a pool slot. + + StorageManager._notify_and_wait does exactly this: it closes the socket so a late ACK + cannot be read as the next request's reply, but swallows the error so a slow controller + does not fail the put that triggered it. The lease therefore exits cleanly. + """ + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") + + with pool.lease(peer.info) as sock: + await sock.send_multipart([b"req"]) + sock.close(linger=0) # closed, but the body returns normally + + assert _idle_sockets(pool) == [], "a closed socket was returned to the pool" + + # The pool still works, and the next request gets a live socket. + assert await _round_trip(pool, peer.info) == b"reply-to-req" + + pool.close() + ctx.destroy(linger=0) + + def test_sockets_are_not_reused_across_event_loops(peer): """A socket bound to a finished loop must never be handed to another one. diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 1b60f061..2c3b5d2b 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -537,12 +537,13 @@ def _superseded(self, peer_id: str, address: str) -> bool: return self._endpoints.get(peer_id, address) != address def _release(self, peer_id: str, address: str, sock: zmq.Socket) -> None: - """Return a cleanly-used socket, closing it if superseded or its bucket is full.""" + """Return a socket, closing it if already closed, superseded, or its bucket is full.""" with self._lock: - # A lease that was already in flight when the peer moved must not be parked: it is - # wired to an address nobody will ask for again, and _take's sweep cannot see a - # socket that is out on loan. - if not self._superseded(peer_id, address): + # A caller may close the socket itself without raising, as the notify path does to + # discard a possibly-late ACK. A lease already in flight when the peer moved must + # not be parked either: it is wired to an address nobody will ask for again, and + # _take's sweep cannot see a socket that is out on loan. + if not sock.closed and not self._superseded(peer_id, address): bucket = self._owner_buckets().setdefault(address, []) if len(bucket) < self._maxsize: bucket.append(sock) From 0801c74c3c3fc95aba872aaa0d962ce8373c3309 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 7 Sep 2026 11:49:31 +0800 Subject: [PATCH 09/28] [fix] Keep close() working when the base constructor raised storage_rpc_pool is assigned after super().__init__(), which runs the controller handshake and raises TimeoutError once TQ_STORAGE_HANDSHAKE_MAX_RETRIES is exhausted -- an ordinary outcome when the controller is slow or unreachable at startup. close() dereferenced the pool unconditionally, so on that path it raised AttributeError before reaching super().close(), leaving the notify pool and the context with its native I/O threads behind. __del__ swallows the error into a log line, so the leak surfaced only as a confusing message. Guarding the attribute matches the base class, whose close() already reaches for every teardown attribute through hasattr/getattr precisely because __del__ must cope with a partially built object. This override was the only one that did not. The regression test drives a handshake failure and asserts the context is destroyed; it was confirmed to fail against the pre-fix code. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 28 +++++++++++++++++++ .../managers/simple_storage_manager.py | 7 +++-- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 7e6b7143..c2b2e69d 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -300,6 +300,34 @@ def test_each_scenario_gets_its_own_pool(echo_controller): client.close() +def test_close_after_failed_handshake_still_releases_the_context(echo_controller): + """A manager whose base constructor raised must still tear down on close(). + + storage_rpc_pool is assigned after super().__init__(), so a handshake timeout leaves it + unset while the context and its native I/O threads are already allocated. __del__ calls + close() regardless, so dereferencing the pool there aborts the base teardown and leaks + the context. + """ + built = [] + + def _fail(self): + built.append(self) + raise TimeoutError("handshake failed") + + with patch.object(StorageManager, "_connect_to_controller", _fail): + with pytest.raises(TimeoutError): + AsyncSimpleStorageManager( + echo_controller.zmq_server_info, + {"zmq_info": {"storage_0": echo_controller.zmq_server_info}}, + ) + + manager = built[0] + assert not hasattr(manager, "storage_rpc_pool"), "the test no longer exercises the partial-build path" + + manager.close() + assert manager.zmq_context.closed, "the context outlived a failed construction" + + def test_simple_storage_does_not_destroy_borrowed_context(echo_controller): client = AsyncTransferQueueClient( client_id="client_borrowed_context_lifecycle", diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index bdd0d8d7..6fb5ef1b 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -664,6 +664,9 @@ async def load_checkpoint(self, checkpoint_dir: str) -> None: def close(self) -> None: """Close all ZMQ sockets and context to prevent resource leaks.""" - # Before super(), which may destroy the context these sockets live on. - self.storage_rpc_pool.close() + # Before super(), which may destroy the context these sockets live on. Absent when + # the base constructor raised (a failed handshake does), and __del__ still calls this. + pool = getattr(self, "storage_rpc_pool", None) + if pool is not None: + pool.close() super().close() From 947beac8ba65ba3b170969482df850d4823bc4a4 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 7 Sep 2026 11:51:05 +0800 Subject: [PATCH 10/28] [fix] Let close() tolerate a subclass that raised before super().__init__() KVStorageManager, RayStorageManager and YuanrongStorageManager all validate their config and raise before delegating to super().__init__(), so none of the attributes close() reaches for exist on the resulting object. __del__ calls close() anyway, and the very first line dereferenced controller_handshake_socket unguarded. The AttributeError was caught and logged, which buried the real ValueError under "Exception during __del__: object has no attribute controller_handshake_socket". Nothing leaks on this path, since super().__init__() is what allocates the context, the notify pool and the notify thread; the cost is a misleading error that sends readers after the wrong fault. The existing hasattr guards further down show close() was already meant to cope with a partially built object, just not with one where the base constructor never ran at all. Guarding on notify_pool covers the whole method in one check: it is the last attribute __init__ assigns before its first fallible step, so if it is present every attribute close() touches is too, and the existing guards already handle a handshake that failed after that point. Predates the socket pool work. The regression test was confirmed to fail against the pre-fix code. Signed-off-by: OutstanderWang --- tests/test_kv_storage_manager.py | 22 ++++++++++++++++++++++ transfer_queue/storage/managers/base.py | 7 +++++++ 2 files changed, 29 insertions(+) diff --git a/tests/test_kv_storage_manager.py b/tests/test_kv_storage_manager.py index 7ac3744d..14980fab 100644 --- a/tests/test_kv_storage_manager.py +++ b/tests/test_kv_storage_manager.py @@ -23,6 +23,28 @@ from transfer_queue.storage.managers.base import KVStorageManager +def test_close_is_quiet_when_config_is_rejected_before_super_init(): + """A config rejected before super().__init__() leaves nothing for close() to release. + + Every KV manager validates its config first, so none of the base attributes close() + reaches for exist yet. __del__ calls close() regardless, and an AttributeError there + buries the real ValueError under a spurious teardown error. + """ + built = [] + + class Probe(KVStorageManager): + def __init__(self, *args, **kwargs): + built.append(self) + super().__init__(*args, **kwargs) + + with pytest.raises(ValueError, match="Missing client_name"): + Probe(None, {}) + + manager = built[0] + assert not hasattr(manager, "notify_pool"), "the test no longer exercises the pre-super path" + manager.close() + + def get_meta(data, global_indexes=None): if not global_indexes: global_indexes = list(range(data.batch_size[0])) diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 8faccc02..537da503 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -378,6 +378,13 @@ async def load_checkpoint(self, checkpoint_dir: str) -> None: def close(self) -> None: """Close all ZMQ sockets/contexts and stop the notify loop.""" + # A subclass may reject its config before calling super().__init__(), as the KV + # managers do, leaving nothing here allocated. __del__ calls close() anyway, so + # return rather than burying the constructor's error under an AttributeError. + # notify_pool is the last thing __init__ sets before its first fallible step. + if not hasattr(self, "notify_pool"): + return + if self.controller_handshake_socket: try: if not self.controller_handshake_socket.closed: From b59bf68741cc5de0a03dec5be75bbfce5dbc78dc Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 7 Sep 2026 12:57:52 +0800 Subject: [PATCH 11/28] [refactor] Drop endpoint-migration tracking, which nothing can reach The pool carried a per-peer map of current addresses so that a peer moving to a new port under a stable id had the sockets at its old address retired. Only the metrics collector opted in, via register_storage_units(), and that runs exactly once: interface.init() calls it during first-time initialization, returns early through _init_from_existing() on every later call, and tq.close() ray.kill()s the controller actor rather than reconfiguring it. Nothing else remaps an endpoint at runtime -- client._controller, manager.controller_info and manager.storage_unit_infos are each assigned once at construction. Reuse was always keyed by address, never by peer id, so a moved peer is still dialed correctly without any of this; test_reregistered_peer_is_not_served_a_ stale_socket keeps that guarantee and passes unchanged. What the tracking added on top was closing the sockets left behind at the abandoned address, and its absence is now recorded where _idle is declared so that anyone adding runtime remapping knows to retire those buckets. Removing it takes with it the follow_endpoint_changes flag, the _endpoints map, _mark_current, _superseded, and the peer_id argument that only those two needed from _take and _release. Four tests covering the deleted state go too; the remaining eleven cover behaviour that is still reachable. Net 199 lines lighter, 60 of them in the pool itself. Signed-off-by: OutstanderWang --- tests/test_zmq_socket_pool.py | 140 +----------------------------- transfer_queue/metrics.py | 3 - transfer_queue/utils/zmq_utils.py | 70 +++------------ 3 files changed, 13 insertions(+), 200 deletions(-) diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py index 35fa3a95..38dffcc7 100644 --- a/tests/test_zmq_socket_pool.py +++ b/tests/test_zmq_socket_pool.py @@ -306,7 +306,7 @@ async def test_reregistered_peer_is_not_served_a_stale_socket(): old = _Peer(peer_id="su0", tag=b"OLD:") new = _Peer(peer_id="su0", tag=b"NEW:") ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner", "put_get_socket", follow_endpoint_changes=True) + pool = ZMQSocketPool(ctx, "owner", "put_get_socket") try: assert await _round_trip(pool, old.info) == b"OLD:req" # Same id, different port -- exactly what re-registration produces. @@ -320,144 +320,6 @@ async def test_reregistered_peer_is_not_served_a_stale_socket(): new.stop() -@pytest.mark.asyncio -async def test_moved_peer_does_not_accumulate_stale_buckets(): - """Sockets for an address a peer no longer answers on must be dropped, not kept. - - Keying by endpoint alone would leave one bucket per past address, each holding an open - socket that libzmq keeps trying to reconnect. - """ - ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner", "put_get_socket", follow_endpoint_changes=True) - moved = [_Peer(peer_id="su_moves", tag=b"E%d:" % i) for i in range(4)] - try: - for m in moved: - await _round_trip(pool, m.info) - buckets = [b for owner in pool._idle.values() for b in owner] - assert len(buckets) == 1, f"stale endpoint buckets accumulated: {len(buckets)}" - assert len(_idle_sockets(pool)) == 1 - finally: - pool.close() - ctx.destroy(linger=0) - for m in moved: - m.stop() - - -@pytest.mark.asyncio -async def test_migration_of_one_peer_leaves_other_peers_alone(): - """Retiring a moved peer's old address must not close sockets belonging to other peers. - - The pool that tracks migrations is the metrics collector's, which queries every storage - unit from one pool. Sweeping every address that is not the mover's new one would drop - the whole fleet's sockets each time a single unit re-registers. - """ - old = _Peer(peer_id="su0", tag=b"OLD:") - new = _Peer(peer_id="su0", tag=b"NEW:") - others = [_Peer(peer_id=f"su{i}", tag=b"P%d:" % i) for i in (1, 2)] - ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "metrics_collector", "put_get_socket", follow_endpoint_changes=True) - try: - # Two cycles, so every endpoint is recorded and each peer has a socket parked. - for _ in range(2): - for p in [old, *others]: - await _round_trip(pool, p.info) - parked_before = {id(sock) for sock in _idle_sockets(pool)} - assert len(parked_before) == 3, "each peer should hold one idle socket" - - assert await _round_trip(pool, new.info) == b"NEW:req" - - survivors = parked_before & {id(sock) for sock in _idle_sockets(pool)} - assert len(survivors) == 2, "a peer that did not move lost its pooled socket" - addresses = {addr for buckets in pool._idle.values() for addr in buckets} - assert old.info.to_addr("put_get_socket") not in addresses, "the mover's old address stayed" - finally: - pool.close() - ctx.destroy(linger=0) - for p in [old, new, *others]: - p.stop() - - -@pytest.mark.asyncio -async def test_superseded_sockets_are_evicted_from_every_owner(): - """A migration seen by one owner must retire the old address in all of them. - - The pool is deliberately shared between a client's RPC loop and a storage manager's - notify loop, so several owners can hold idle sockets for the same peer. Sweeping only - the current owner leaves the others reconnecting to the retired address until they - happen to lease again -- or forever, if they go quiet. - """ - old = _Peer(peer_id="su0", tag=b"OLD:") - new = _Peer(peer_id="su0", tag=b"NEW:") - ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner", "put_get_socket", follow_endpoint_changes=True) - loops = [] - - def spawn_owner(): - loop = asyncio.new_event_loop() - threading.Thread(target=loop.run_forever, daemon=True).start() - loops.append(loop) - return loop - - def run_on(loop, info): - return asyncio.run_coroutine_threadsafe(_round_trip(pool, info), loop).result(timeout=10) - - try: - # Two owners park a socket to the old address and stay alive. - parked = [spawn_owner(), spawn_owner()] - for loop in parked: - assert run_on(loop, old.info) == b"OLD:req" - assert len(pool._idle) == 2, "each live loop should hold its own bucket" - - # A third owner observes the migration. - assert run_on(spawn_owner(), new.info) == b"NEW:req" - - addresses = {addr for buckets in pool._idle.values() for addr in buckets} - assert addresses == {new.info.to_addr("put_get_socket")}, f"stale addresses remain: {addresses}" - - # Every owner still reaches the current endpoint afterwards. - for loop in parked: - assert run_on(loop, new.info) == b"NEW:req" - finally: - for loop in loops: - loop.call_soon_threadsafe(loop.stop) - pool.close() - ctx.destroy(linger=0) - old.stop() - new.stop() - - -@pytest.mark.asyncio -async def test_lease_in_flight_when_peer_moves_is_not_parked(): - """A lease already out on loan when its peer moves must not re-enter the pool. - - _take's sweep only sees idle sockets, so a request in flight to the old address escapes - it and would recreate that bucket on return -- leaving a socket wired to an obsolete - endpoint parked indefinitely. - """ - # Answers slowly, so its lease is still out when the new endpoint is first used. - old = _Peer(delay_first_reply=1.0, peer_id="su0", tag=b"OLD:") - new = _Peer(peer_id="su0", tag=b"NEW:") - ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner", "put_get_socket", follow_endpoint_changes=True) - try: - in_flight = asyncio.create_task(_round_trip(pool, old.info)) - await asyncio.sleep(0.2) # let it reach recv before the peer "moves" - - assert await _round_trip(pool, new.info) == b"NEW:req" - assert await in_flight == b"OLD:req", "the in-flight request should still complete" - - parked = [addr for owner in pool._idle.values() for addr in owner] - assert len(parked) == 1, f"a superseded endpoint was parked: {parked}" - assert parked[0] == new.info.to_addr("put_get_socket") - # And the next request still reaches the current endpoint. - assert await _round_trip(pool, new.info) == b"NEW:req" - finally: - pool.close() - ctx.destroy(linger=0) - old.stop() - new.stop() - - def test_pool_size_below_one_is_rejected(): """A size under 1 parks nothing, so reuse is silently off while still looking pooled.""" ctx = zmq.asyncio.Context() diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index 10fe5de7..49842259 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -379,9 +379,6 @@ def _get_socket_pool(self) -> ZMQSocketPool: "metrics_collector", "put_get_socket", timeout=TQ_METRICS_STORAGE_TIMEOUT, - # register_storage_units() can remap a storage unit id onto a new address, so - # sockets left at the address it moved off must be closed rather than linger. - follow_endpoint_changes=True, ) return self._zmq_socket_pool diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 2c3b5d2b..29df2f55 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -412,7 +412,6 @@ def __init__( *, timeout: int | None = None, maxsize: int = 8, - follow_endpoint_changes: bool = False, ): """ Args: @@ -424,11 +423,6 @@ def __init__( maxsize: Idle sockets kept per bucket, at least 1. A soft cap: a burst beyond it still gets sockets, and the excess is closed on return rather than made to wait. - follow_endpoint_changes: Set when a peer's address can change under a stable id, - as re-registration does. Sockets are keyed by address either way, so a moved - peer is always dialed correctly; this additionally closes the sockets left - behind at the address it abandoned instead of letting them linger. Only the - metrics collector needs it -- elsewhere endpoints are fixed at construction. """ if maxsize < 1: # Below 1 nothing is ever parked, so every request pays a fresh connect while @@ -439,15 +433,12 @@ def __init__( self._socket_name = socket_name self._timeout = timeout self._maxsize = maxsize - self._follow_endpoint_changes = follow_endpoint_changes # Keyed by lease owner (see _lease_owner), then by address. The address, not the peer # id: a peer restarted under the same id at a new address must not be handed a socket - # still connected to the old one. + # still connected to the old one. The bucket it moved off is then never asked for + # again and lingers until close(), which no supported path can reach today -- adding + # runtime endpoint remapping means retiring those buckets too. self._idle: dict[Any, dict[str, list[zmq.Socket]]] = {} - # peer id -> the address most recently leased for it, tracked only when endpoints can - # move. What lets a lease returning from an abandoned address be told apart from a - # current one; _idle alone cannot, since an in-flight socket is not in it. - self._endpoints: dict[str, str] = {} self._lock = threading.Lock() # A ROUTER silently drops a second peer claiming an identity it already has, so # identities must not collide between processes: owner_id alone does not suffice @@ -468,7 +459,7 @@ def lease(self, peer: ZMQServerInfo) -> Iterator[zmq.Socket]: raise RuntimeError(f"Socket '{self._socket_name}' not configured for server '{peer.id}'") address = format_zmq_address(peer.ip, port) - sock = self._take(peer.id, address) or self._connect(peer, address) + sock = self._take(address) or self._connect(peer, address) try: yield sock except BaseException: @@ -478,7 +469,7 @@ def lease(self, peer: ZMQServerInfo) -> Iterator[zmq.Socket]: sock.close(linger=0) raise else: - self._release(peer.id, address, sock) + self._release(address, sock) def _owner_buckets(self) -> dict[str, list[zmq.Socket]]: """Buckets for the current lease owner, first evicting any owner that has finished. @@ -492,58 +483,22 @@ def _owner_buckets(self) -> dict[str, list[zmq.Socket]]: self._close_all(self._idle.pop(owner)) return self._idle.setdefault(_lease_owner(), {}) - def _take(self, peer_id: str, address: str) -> zmq.Socket | None: + def _take(self, address: str) -> zmq.Socket | None: """Pop a live idle socket for *address*, discarding any found closed.""" with self._lock: - # Drop finished owners first, so the endpoint sweep never walks their buckets. - buckets = self._owner_buckets() - self._mark_current(peer_id, address) - bucket = buckets.get(address) + bucket = self._owner_buckets().get(address) while bucket: sock = bucket.pop() if not sock.closed: return sock return None - def _mark_current(self, peer_id: str, address: str) -> None: - """Record *address* as *peer_id*'s current one and close sockets to the old one. - - Callers must hold ``self._lock``. Recorded on acquire so a lease already in flight to - an abandoned address is recognised when it returns (see _release). Only this peer's - own former address is retired; peers that did not move keep their sockets. The sweep - spans every owner: an idle socket parked by an owner that then goes quiet would - otherwise keep reconnecting to the abandoned address forever. Sockets are only ever - closed here, never handed across owners. - """ - if not self._follow_endpoint_changes: - return - previous = self._endpoints.get(peer_id) - if previous == address: - return # unchanged, so nothing to retire - self._endpoints[peer_id] = address - if previous is None: - return # first sighting, so this peer has left nothing behind - for buckets in self._idle.values(): - retired = buckets.pop(previous, None) - if retired: - self._close_all({previous: retired}) - - def _superseded(self, peer_id: str, address: str) -> bool: - """Whether *address* is one its peer has since moved away from. - - Callers must hold ``self._lock``. Unknown peers are not superseded, so a socket is - only ever retired because a newer address was actually seen. - """ - return self._endpoints.get(peer_id, address) != address - - def _release(self, peer_id: str, address: str, sock: zmq.Socket) -> None: - """Return a socket, closing it if already closed, superseded, or its bucket is full.""" + def _release(self, address: str, sock: zmq.Socket) -> None: + """Return a socket, closing it if already closed or its bucket is full.""" with self._lock: - # A caller may close the socket itself without raising, as the notify path does to - # discard a possibly-late ACK. A lease already in flight when the peer moved must - # not be parked either: it is wired to an address nobody will ask for again, and - # _take's sweep cannot see a socket that is out on loan. - if not sock.closed and not self._superseded(peer_id, address): + # A caller may close the socket itself without raising, as the notify path does + # to discard a possibly-late ACK. + if not sock.closed: bucket = self._owner_buckets().setdefault(address, []) if len(bucket) < self._maxsize: bucket.append(sock) @@ -575,7 +530,6 @@ def close(self) -> None: with self._lock: owned = list(self._idle.values()) self._idle = {} - self._endpoints = {} for buckets in owned: self._close_all(buckets) From 341b8ec80030240c5650ebee972c72e321c14217 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 7 Sep 2026 13:22:47 +0800 Subject: [PATCH 12/28] [refactor] Borrow the controller's ZMQ context in the metrics exporter The exporter minted its own zmq.Context() the first time it queried a storage unit, and nothing ever terminated it: the exporter is created once per Ray actor by start_metrics(), which is idempotent, and the actor is torn down with ray.kill(), so there is no shutdown path to close it from. Adding one would have meant a graceful actor-exit protocol -- a stop signal for the collect loop, a shutdown for the Prometheus HTTP server, and an RPC in interface.close() ahead of ray.kill() -- to release a context in a process that is about to exit anyway. The controller already holds a long-lived synchronous context, which is exactly what these queries need, so the exporter now borrows it and the second context disappears along with the question of who closes it. This is the arrangement ZMQSocketPool already documents and the one the client uses when it lends its context to the SimpleStorage manager. The storage-role exporter passes nothing and needs nothing: only the controller role starts the collection loop, so only it ever builds a pool. Asking for a pool without a context now raises instead of quietly creating one, which keeps the lending explicit. Metrics sockets, one per storage unit, now count against the controller context's socket budget. That is libzmq's default 1023 and the context holds two ROUTERs today, so the headroom covers roughly a thousand storage units; the socket count itself is unchanged, only which context they belong to. Signed-off-by: OutstanderWang --- tests/test_metrics.py | 24 ++++++++++++++++++++++++ transfer_queue/controller.py | 3 ++- transfer_queue/metrics.py | 19 ++++++++++++++++--- 3 files changed, 42 insertions(+), 4 deletions(-) diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 15c43882..785d83dc 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -21,6 +21,8 @@ import pytest try: + import zmq + from transfer_queue.metrics import TQMetricsExporter _HAS_DEPS = True @@ -207,6 +209,28 @@ def test_multiple_ops_tracked_independently(self): # --------------------------------------------------------------------------- +class TestStorageQuerySocketPool: + def test_pool_borrows_the_owner_context(self): + """The exporter must query storage units over the context it was handed. + + The controller already holds a long-lived synchronous context. A second one would + add another native I/O thread and leave a context nobody closes, since the exporter + lives for the whole life of its Ray actor. + """ + ctx = zmq.Context() + try: + exporter = TQMetricsExporter(zmq_context=ctx) + assert exporter._get_socket_pool()._ctx is ctx + finally: + ctx.destroy(linger=0) + + def test_missing_context_is_reported(self): + """Without a context there is nothing to query over, so say so rather than mint one.""" + exporter = TQMetricsExporter() + with pytest.raises(RuntimeError, match="without a ZMQ context"): + exporter._get_socket_pool() + + class TestStorageMetricsCollection: def test_collect_with_no_storage_units(self): """No storage units registered — collect should be a no-op.""" diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index ef71bb4c..188c011c 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -2337,7 +2337,8 @@ def start_metrics(self, port: int = 0) -> str: return self._metrics_endpoint from transfer_queue.metrics import TQMetricsExporter - self._metrics = TQMetricsExporter() + # Lend the controller's context rather than let the exporter build a second one. + self._metrics = TQMetricsExporter(zmq_context=self.zmq_context) self._metrics_endpoint = self._metrics.start(node_ip=self._node_ip, port=port) # Launch a daemon thread that periodically pushes controller state # snapshots to the exporter, keeping them process-isolated. diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index 49842259..bb475561 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -66,12 +66,21 @@ class TQMetricsExporter: TQ_METRICS_STORAGE_TIMEOUT ZMQ timeout for storage queries (default 5s) """ - def __init__(self, role: str = "controller"): + def __init__(self, role: str = "controller", zmq_context: zmq.Context | None = None): + """ + Args: + role: Which process this exporter runs in; only "controller" collects from + storage units, so only that role needs a context. + zmq_context: The owner's long-lived synchronous context, borrowed for + storage-unit queries and never terminated here. Minting one instead would + add a second context and its native I/O thread with nobody to close them, + since the exporter lives as long as its Ray actor. + """ self._start_time = time.time() self._process = psutil.Process() self._role = role self._storage_unit_infos: dict[str, ZMQServerInfo] = {} - self._zmq_ctx: zmq.Context | None = None + self._zmq_ctx = zmq_context self._zmq_socket_pool: ZMQSocketPool | None = None self._known_partition_ids: set[str] = set() self._known_production_labels: set[tuple[str, str]] = set() @@ -373,7 +382,11 @@ def collect_storage_metrics(self) -> None: def _get_socket_pool(self) -> ZMQSocketPool: """Return the lazily-created socket pool for storage-unit queries.""" if self._zmq_socket_pool is None: - self._zmq_ctx = zmq.Context() + if self._zmq_ctx is None: + raise RuntimeError( + "TQMetricsExporter was built without a ZMQ context, so it cannot query " + "storage units; pass zmq_context= from the owning process." + ) self._zmq_socket_pool = ZMQSocketPool( self._zmq_ctx, "metrics_collector", From e6c48e1a2d4d2f7ca2d7106a1990319ec037c1fc Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 7 Sep 2026 15:21:54 +0800 Subject: [PATCH 13/28] [perf] Raise the default pooled-socket cap from 8 to 64 The cap bounds only how many idle sockets are kept per (owner, endpoint) bucket, not concurrency: a burst beyond it is still served, but the excess is closed on return and pays a fresh handshake next time. At 8 that made reuse fall off well below the concurrency these paths actually reach. Both defaults move together so all four pools follow: the client's controller RPC pool reads TQ_CLIENT_ZMQ_POOL_SIZE, while the storage RPC, notify, and metrics pools take the constructor default. Signed-off-by: OutstanderWang --- transfer_queue/client.py | 2 +- transfer_queue/utils/zmq_utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 866dc707..8d82b17b 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -49,7 +49,7 @@ DEFAULT_CLIENT_ZMQ_MAX_SOCKETS = 8192 # Idle sockets kept per (loop, endpoint) bucket, at least 1. A soft cap: bursts beyond it # still get sockets, so this bounds the steady state rather than the peak. -TQ_CLIENT_ZMQ_POOL_SIZE = int(os.environ.get("TQ_CLIENT_ZMQ_POOL_SIZE", 8)) +TQ_CLIENT_ZMQ_POOL_SIZE = int(os.environ.get("TQ_CLIENT_ZMQ_POOL_SIZE", 64)) # Pre-bound decorator for controller socket operations. with_controller_socket = with_zmq_socket( diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 29df2f55..b14cdc09 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -411,7 +411,7 @@ def __init__( socket_name: str, *, timeout: int | None = None, - maxsize: int = 8, + maxsize: int = 64, ): """ Args: From 2f2469c5b2d10cd6692a2982eeb50db8516efafc Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 7 Sep 2026 15:58:29 +0800 Subject: [PATCH 14/28] [fix] Narrow max_sockets by identity rather than by a parallel flag At this point explicitly_requested is true exactly when max_sockets is not None, so the two guards accept the same inputs. Testing the value itself lets a type checker see that the comparison below operates on an int. Signed-off-by: OutstanderWang --- transfer_queue/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 8d82b17b..d66911eb 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -120,7 +120,7 @@ def __init__( f"TQ_CLIENT_ZMQ_MAX_SOCKETS must be an integer, got {TQ_CLIENT_ZMQ_MAX_SOCKETS!r}" ) from e explicitly_requested = True - if explicitly_requested and max_sockets < 1: + if max_sockets is not None and max_sockets < 1: # The upper bound needs ZMQ_SOCKET_LIMIT, hence a live context, but the lower one # does not -- so reject it before allocating anything. raise ValueError(f"Client ZMQ max sockets must be at least 1, got {max_sockets}") From 71736a6f2e7228f08141984f5092f30d0bacaa63 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 7 Sep 2026 17:32:37 +0800 Subject: [PATCH 15/28] [fix] Scope the pooled-socket knob to the pool it configures TQ_CLIENT_ZMQ_POOL_SIZE reads as a client-wide budget but is applied at a single call site, the controller RPC pool. The storage RPC, notify, and metrics pools take ZMQSocketPool's own default, so the CLIENT prefix promises a reach the knob does not have. Rename it to TQ_CONTROLLER_RPC_POOL_SIZE, which names exactly what it sets. The knob is new in this branch and unreferenced outside it, so no configuration in the wild has to change. Document the two properties the name cannot carry: the cap counts per (owner, address) rather than per pool, so a pool dialling N peers may park N*maxsize idle sockets, and the pools that take the default are the ones whose concurrency does not warrant a knob -- notify serializes onto its own loop and metrics collects sequentially, each holding one socket at a time. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 6 +++--- transfer_queue/client.py | 14 ++++++++------ transfer_queue/utils/zmq_utils.py | 5 ++++- 3 files changed, 15 insertions(+), 10 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index c2b2e69d..aea8d8a1 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -197,15 +197,15 @@ def _spy(*args, **kwargs): def test_client_rejects_invalid_socket_pool_size(echo_controller): - """A bad TQ_CLIENT_ZMQ_POOL_SIZE must name the variable, not silently disable reuse. + """A bad TQ_CONTROLLER_RPC_POOL_SIZE must name the variable, not silently disable reuse. Below 1 nothing is ever parked, so every request pays a fresh connect while the client still looks pooled. """ for bad in (-1, 0): - with patch("transfer_queue.client.TQ_CLIENT_ZMQ_POOL_SIZE", bad): + with patch("transfer_queue.client.TQ_CONTROLLER_RPC_POOL_SIZE", bad): with _no_context_left_open() as created: - with pytest.raises(ValueError, match="TQ_CLIENT_ZMQ_POOL_SIZE must be at least 1"): + with pytest.raises(ValueError, match="TQ_CONTROLLER_RPC_POOL_SIZE must be at least 1"): AsyncTransferQueueClient( client_id="client_invalid_socket_pool", controller_info=echo_controller.zmq_server_info, diff --git a/transfer_queue/client.py b/transfer_queue/client.py index d66911eb..d634b220 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -47,9 +47,11 @@ # 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 -# Idle sockets kept per (loop, endpoint) bucket, at least 1. A soft cap: bursts beyond it -# still get sockets, so this bounds the steady state rather than the peak. -TQ_CLIENT_ZMQ_POOL_SIZE = int(os.environ.get("TQ_CLIENT_ZMQ_POOL_SIZE", 64)) +# Idle sockets kept per (loop, endpoint) bucket for controller RPC, at least 1. A soft cap: +# bursts beyond it still get sockets, so this bounds the steady state rather than the peak. +# Scoped to this one pool; the storage RPC, notify, and metrics pools take ZMQSocketPool's +# own default, as none of them reaches a concurrency worth tuning separately. +TQ_CONTROLLER_RPC_POOL_SIZE = int(os.environ.get("TQ_CONTROLLER_RPC_POOL_SIZE", 64)) # Pre-bound decorator for controller socket operations. with_controller_socket = with_zmq_socket( @@ -102,10 +104,10 @@ def __init__( 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}") - if TQ_CLIENT_ZMQ_POOL_SIZE < 1: + if TQ_CONTROLLER_RPC_POOL_SIZE < 1: # Name the variable: the pool's own error cannot say which knob supplied the value. raise ValueError( - f"TQ_CLIENT_ZMQ_POOL_SIZE must be at least 1, got {TQ_CLIENT_ZMQ_POOL_SIZE}. " + f"TQ_CONTROLLER_RPC_POOL_SIZE must be at least 1, got {TQ_CONTROLLER_RPC_POOL_SIZE}. " f"The pool always reuses at least one socket per endpoint; it cannot be disabled." ) @@ -159,7 +161,7 @@ def __init__( self.zmq_context, client_id, "request_handle_socket", - maxsize=TQ_CLIENT_ZMQ_POOL_SIZE, + maxsize=TQ_CONTROLLER_RPC_POOL_SIZE, ) # Backstop for a client that is never closed, so the context and its I/O threads do diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index b14cdc09..d6f99337 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -422,7 +422,10 @@ def __init__( timeout: Send/recv timeout in seconds applied to every socket, or None for none. maxsize: Idle sockets kept per bucket, at least 1. A soft cap: a burst beyond it still gets sockets, and the excess is closed on return rather than made to - wait. + wait. Counted per (owner, address), not per pool, so a pool dialling N peers + may hold N*maxsize idle sockets. Only the client's controller RPC pool tunes + this (``TQ_CONTROLLER_RPC_POOL_SIZE``); the storage RPC, notify, and metrics + pools take this default, the latter two holding one socket at a time anyway. """ if maxsize < 1: # Below 1 nothing is ever parked, so every request pays a fresh connect while From 2c9e45f5e4c777a8143a00d53fd1749c7a85b8f0 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 7 Sep 2026 17:33:05 +0800 Subject: [PATCH 16/28] [docs] Show the async client awaiting on one loop, not asyncio.run per call Every example on AsyncTransferQueueClient wrapped a single call in asyncio.run(), and the put/get walkthrough did so three times in a row. Pooled sockets are keyed by their owning event loop, so that pattern retires its socket at the end of each call and pays a fresh connect handshake on the next one -- the examples demonstrated precisely the shape that cannot reuse a connection. Await the calls inside one async def instead, which is both how the supported entry points drive the client and the shape the pool is built for. TransferQueueClient keeps a loop of its own for callers that have none, so state the reuse condition on the class itself rather than leaving it to be inferred from the pool internals. No behavior change; docstrings only. Signed-off-by: OutstanderWang --- transfer_queue/client.py | 214 +++++++++++++++++++++------------------ 1 file changed, 117 insertions(+), 97 deletions(-) diff --git a/transfer_queue/client.py b/transfer_queue/client.py index d634b220..eafad617 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -65,6 +65,11 @@ class AsyncTransferQueueClient: This client provides async methods for data transfer operations including getting metadata, reading data from storage, writing data to storage, and clearing data. + + Await these methods from one long-lived event loop, as the examples below do. Pooled + sockets are keyed by their owning loop, so a fresh ``asyncio.run()`` per call gets a + fresh socket every time and pays a connect handshake it cannot amortize. Callers with + no loop of their own should use ``TransferQueueClient``, which keeps one internally. """ def __init__( @@ -271,33 +276,35 @@ async def async_get_meta( RuntimeError: If communication fails or controller returns error response Example: - >>> # Example 1: Basic fetch metadata - >>> batch_meta = asyncio.run(client.async_get_meta( - ... data_fields=["input_ids", "attention_mask"], - ... batch_size=4, - ... partition_id="train_0", - ... mode="fetch", - ... task_name="generate_sequences" - ... )) - >>> print(batch_meta.is_ready) # True if all samples ready - >>> - >>> # Example 2: Fetch with self-defined samplers (using GRPOGroupNSampler as an example) - >>> batch_meta = asyncio.run(client.async_get_meta( - ... data_fields=["input_ids", "attention_mask"], - ... batch_size=8, - ... partition_id="train_0", - ... mode="fetch", - ... task_name="generate_sequences", - ... )) - >>> print(batch_meta.is_ready) # True if all samples ready - >>> - >>> # Example 3: Force fetch metadata (bypass production status check and Sampler, - >>> # so may include unready and already-consumed samples. No filtering by consumption status is applied.) - >>> batch_meta = asyncio.run(client.async_get_meta( - ... partition_id="train_0", # optional - ... mode="force_fetch", - ... )) - >>> print(batch_meta.is_ready) # May be False if some samples not ready + >>> async def main(): + ... # Example 1: Basic fetch metadata + ... batch_meta = await client.async_get_meta( + ... data_fields=["input_ids", "attention_mask"], + ... batch_size=4, + ... partition_id="train_0", + ... mode="fetch", + ... task_name="generate_sequences" + ... ) + ... print(batch_meta.is_ready) # True if all samples ready + ... + ... # Example 2: Fetch with self-defined samplers (using GRPOGroupNSampler as an example) + ... batch_meta = await client.async_get_meta( + ... data_fields=["input_ids", "attention_mask"], + ... batch_size=8, + ... partition_id="train_0", + ... mode="fetch", + ... task_name="generate_sequences", + ... ) + ... print(batch_meta.is_ready) # True if all samples ready + ... + ... # Example 3: Force fetch metadata (bypass production status check and Sampler, + ... # so may include unready and already-consumed samples. No filtering by + ... # consumption status is applied.) + ... batch_meta = await client.async_get_meta( + ... partition_id="train_0", # optional + ... mode="force_fetch", + ... ) + ... print(batch_meta.is_ready) # May be False if some samples not ready """ response_msg = await self._request_controller( socket=socket, @@ -337,10 +344,11 @@ async def async_set_custom_meta( RuntimeError: If communication fails or controller returns error response Example: - >>> # Create batch with custom metadata - >>> batch_meta = client.get_meta(data_fields=["input_ids"], batch_size=4, ...) - >>> batch_meta.update_custom_meta([{"score": 0.9}, {"score": 0.8}]) - >>> asyncio.run(client.async_set_custom_meta(batch_meta)) + >>> async def main(): + ... # Create batch with custom metadata + ... batch_meta = await client.async_get_meta(data_fields=["input_ids"], batch_size=4, ...) + ... batch_meta.update_custom_meta([{"score": 0.9}, {"score": 0.8}]) + ... await client.async_set_custom_meta(batch_meta) """ assert socket is not None @@ -422,30 +430,34 @@ async def async_put( >>> batch_size = 4 >>> seq_len = 16 >>> current_partition_id = "train_0" - >>> # Example 1: Normal usage with existing metadata - >>> batch_meta = asyncio.run(client.async_get_meta( - ... data_fields=["prompts", "attention_mask"], - ... batch_size=batch_size, - ... partition_id=current_partition_id, - ... mode="fetch", - ... task_name="generate_sequences", - ... )) - >>> batch = asyncio.run(client.async_get_data(batch_meta)) - >>> output = TensorDict({"response": torch.randn(batch_size, seq_len)}) - >>> asyncio.run(client.async_put(data=output, metadata=batch_meta)) - >>> - >>> # Example 2: Initial data insertion without pre-existing metadata - >>> # BE CAREFUL: this usage may overwrite any unconsumed data in the given partition_id! - >>> # Please make sure the corresponding partition_id is empty before calling the async_put() - >>> # without metadata. - >>> # Now we only support put all the data of the corresponding partition id in once. You should repeat with - >>> # interleave the initial data if n_sample > 1 before calling the async_put(). - >>> original_prompts = torch.randn(batch_size, seq_len) - >>> n_samples = 4 - >>> prompts_repeated = torch.repeat_interleave(original_prompts, n_samples, dim=0) - >>> prompts_repeated_batch = TensorDict({"prompts": prompts_repeated}) - >>> # This will create metadata in "insert" mode internally. - >>> metadata = asyncio.run(client.async_put(data=prompts_repeated_batch, partition_id=current_partition_id)) + >>> async def main(): + ... # Example 1: Normal usage with existing metadata + ... batch_meta = await client.async_get_meta( + ... data_fields=["prompts", "attention_mask"], + ... batch_size=batch_size, + ... partition_id=current_partition_id, + ... mode="fetch", + ... task_name="generate_sequences", + ... ) + ... batch = await client.async_get_data(batch_meta) + ... output = TensorDict({"response": torch.randn(batch_size, seq_len)}) + ... await client.async_put(data=output, metadata=batch_meta) + ... + ... # Example 2: Initial data insertion without pre-existing metadata + ... # BE CAREFUL: this usage may overwrite any unconsumed data in the given + ... # partition_id! Please make sure the corresponding partition_id is empty + ... # before calling the async_put() without metadata. + ... # Now we only support put all the data of the corresponding partition id in + ... # once. You should repeat with interleave the initial data if n_sample > 1 + ... # before calling the async_put(). + ... original_prompts = torch.randn(batch_size, seq_len) + ... n_samples = 4 + ... prompts_repeated = torch.repeat_interleave(original_prompts, n_samples, dim=0) + ... prompts_repeated_batch = TensorDict({"prompts": prompts_repeated}) + ... # This will create metadata in "insert" mode internally. + ... metadata = await client.async_put( + ... data=prompts_repeated_batch, partition_id=current_partition_id + ... ) """ if not hasattr(self, "storage_manager") or self.storage_manager is None: @@ -502,16 +514,18 @@ async def async_get_data(self, metadata: BatchMeta) -> TensorDict: - Requested data fields (e.g., "prompts", "attention_mask") Example: - >>> batch_meta = asyncio.run(client.async_get_meta( - ... data_fields=["prompts", "attention_mask"], - ... batch_size=4, - ... partition_id="train_0", - ... mode="fetch", - ... task_name="generate_sequences", - ... )) - >>> batch = asyncio.run(client.async_get_data(batch_meta)) - >>> print(batch) - >>> # TensorDict with fields "prompts", "attention_mask", and sample order matching metadata global_indexes + >>> async def main(): + ... batch_meta = await client.async_get_meta( + ... data_fields=["prompts", "attention_mask"], + ... batch_size=4, + ... partition_id="train_0", + ... mode="fetch", + ... task_name="generate_sequences", + ... ) + ... batch = await client.async_get_data(batch_meta) + ... print(batch) + ... # TensorDict with fields "prompts", "attention_mask", and sample order + ... # matching metadata global_indexes """ if not hasattr(self, "storage_manager") or self.storage_manager is None: @@ -712,12 +726,13 @@ async def async_get_consumption_status( RuntimeError: If communication fails or controller returns error response Example: - >>> # Get consumption status - >>> global_index, consumption_status = asyncio.run(client.async_get_consumption_status( - ... task_name="generate_sequences", - ... partition_id="train_0" - ... )) - >>> print(f"Global index: {global_index}, Consumption status: {consumption_status}") + >>> async def main(): + ... # Get consumption status + ... global_index, consumption_status = await client.async_get_consumption_status( + ... task_name="generate_sequences", + ... partition_id="train_0" + ... ) + ... print(f"Global index: {global_index}, Consumption status: {consumption_status}") """ try: @@ -759,12 +774,13 @@ async def async_get_production_status( RuntimeError: If communication fails or controller returns error response Example: - >>> # Get production status - >>> global_index, production_status = asyncio.run(client.async_get_production_status( - ... data_fields=["input_ids", "attention_mask"], - ... partition_id="train_0" - ... )) - >>> print(f"Global index: {global_index}, Production status: {production_status}") + >>> async def main(): + ... # Get production status + ... global_index, production_status = await client.async_get_production_status( + ... data_fields=["input_ids", "attention_mask"], + ... partition_id="train_0" + ... ) + ... print(f"Global index: {global_index}, Production status: {production_status}") """ try: response_msg = await self._request_controller( @@ -800,12 +816,13 @@ async def async_check_consumption_status( RuntimeError: If communication fails or controller returns error response Example: - >>> # Check if all samples have been consumed - >>> is_consumed = asyncio.run(client.async_check_consumption_status( - ... task_name="generate_sequences", - ... partition_id="train_0" - ... )) - >>> print(f"All samples consumed: {is_consumed}") + >>> async def main(): + ... # Check if all samples have been consumed + ... is_consumed = await client.async_check_consumption_status( + ... task_name="generate_sequences", + ... partition_id="train_0" + ... ) + ... print(f"All samples consumed: {is_consumed}") """ _, consumption_status = await self.async_get_consumption_status( @@ -836,12 +853,13 @@ async def async_check_production_status( RuntimeError: If communication fails or controller returns error response Example: - >>> # Check if all samples are ready for consumption - >>> is_ready = asyncio.run(client.async_check_production_status( - ... data_fields=["input_ids", "attention_mask"], - ... partition_id="train_0" - ... )) - >>> print(f"All samples ready: {is_ready}") + >>> async def main(): + ... # Check if all samples are ready for consumption + ... is_ready = await client.async_check_production_status( + ... data_fields=["input_ids", "attention_mask"], + ... partition_id="train_0" + ... ) + ... print(f"All samples ready: {is_ready}") """ _, production_status = await self.async_get_production_status( data_fields=data_fields, @@ -876,12 +894,13 @@ async def async_reset_consumption( RuntimeError: If communication fails or controller returns error response Example: - >>> # Reset consumption for train task to re-train on same data - >>> success = asyncio.run(client.async_reset_consumption( - ... partition_id="train_0", - ... task_name="train" - ... )) - >>> print(f"Reset successful: {success}") + >>> async def main(): + ... # Reset consumption for train task to re-train on same data + ... success = await client.async_reset_consumption( + ... partition_id="train_0", + ... task_name="train" + ... ) + ... print(f"Reset successful: {success}") """ body = {"partition_id": partition_id} if task_name is not None: @@ -914,8 +933,9 @@ async def async_get_partition_list( list[str]: List of partition ids managed by the controller Example: - >>> partition_ids = asyncio.run(client.get_partition_list()) - >>> print(f"Available partitions: {partition_ids}") + >>> async def main(): + ... partition_ids = await client.get_partition_list() + ... print(f"Available partitions: {partition_ids}") """ try: response_msg = await self._request_controller( From 505466a728f6cf51a8436b1f94ced366eb1d096b Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 7 Sep 2026 19:37:26 +0800 Subject: [PATCH 17/28] [test] Assert socket reuse from the peer, not from the pool's internals The pool tests read pool._idle through a helper and asserted on _ctx, _socket_name and _timeout, so they described the implementation rather than the behavior and broke on any refactor that kept the contract intact. Every pooled socket dials with its own ZMQ identity, so the peer can count distinct callers: one identity across many requests means the connection was reused, a new one means the old socket was discarded. That is the same property observed from the outside, and it needs no access to pool state. Drop the cases reachable only by construction -- identity collisions between same-pid pools, connect failure caught by patching create_zmq_socket, an unknown socket name, close() idempotency, and a size check already covered on the client -- and fold the cancelled and raising leases into one parametrized test. Eight tests remain, covering reuse, the timed-out socket that must not answer the next request, poisoned leases, loop keying, peer re-registration, synchronous callers, and bursts above maxsize. Signed-off-by: OutstanderWang --- tests/test_metrics.py | 6 +- tests/test_zmq_shared_context.py | 7 - tests/test_zmq_socket_pool.py | 254 ++++++------------------------- 3 files changed, 50 insertions(+), 217 deletions(-) diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 785d83dc..88401988 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -16,7 +16,7 @@ """Unit tests for the Prometheus metrics exporter (transfer_queue.metrics).""" import time -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest @@ -220,7 +220,9 @@ def test_pool_borrows_the_owner_context(self): ctx = zmq.Context() try: exporter = TQMetricsExporter(zmq_context=ctx) - assert exporter._get_socket_pool()._ctx is ctx + with patch("zmq.Context") as minted: + exporter._get_socket_pool() + minted.assert_not_called() finally: ctx.destroy(linger=0) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index aea8d8a1..c4d1ca0f 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -288,13 +288,6 @@ def test_each_scenario_gets_its_own_pool(echo_controller): pools = [client.controller_rpc_pool, manager.storage_rpc_pool, manager.notify_pool] assert len({id(pool) for pool in pools}) == 3, "scenarios must not share a pool" - # All three live on the one shared context, so the socket budget stays client-wide. - assert all(pool._ctx is client.zmq_context for pool in pools) - # Each dials the socket its scenario needs. - assert client.controller_rpc_pool._socket_name == "request_handle_socket" - assert manager.notify_pool._socket_name == "request_handle_socket" - assert manager.storage_rpc_pool._socket_name == "put_get_socket" - assert manager.storage_rpc_pool._timeout is not None, "storage RPC keeps its send/recv timeout" manager.close() client.close() diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py index 38dffcc7..fdb35e91 100644 --- a/tests/test_zmq_socket_pool.py +++ b/tests/test_zmq_socket_pool.py @@ -20,6 +20,10 @@ discarded: replies carry no request id (see ZMQMessage.create), so a reply left in flight by a timed-out or cancelled request would be read by the next user of that socket as its own. test_timed_out_socket_is_not_reused pins exactly that. + +Reuse is asserted from the peer's side rather than from the pool's internals: every socket +dials with its own ZMQ identity, so one identity across many requests means the connection +was reused, and a fresh identity means the old socket was discarded. """ import asyncio @@ -29,15 +33,14 @@ import zmq import zmq.asyncio -import transfer_queue.utils.zmq_utils as zmq_utils from transfer_queue.utils.enum_utils import Role from transfer_queue.utils.zmq_utils import ZMQServerInfo, ZMQSocketPool class _Peer: - """A ROUTER that echoes one reply per request, optionally after a delay. + """A ROUTER that echoes one reply per request, recording who dialled it. - ``tag`` distinguishes which endpoint answered, for the re-registration tests. + ``tag`` distinguishes which endpoint answered, for the re-registration test. """ def __init__(self, delay_first_reply: float = 0.0, peer_id: str = "peer_0", tag: bytes = b"reply-to-"): @@ -45,6 +48,7 @@ def __init__(self, delay_first_reply: float = 0.0, peer_id: str = "peer_0", tag: self.socket = self.context.socket(zmq.ROUTER) port = self.socket.bind_to_random_port("tcp://127.0.0.1") self.info = ZMQServerInfo(role=Role.STORAGE, id=peer_id, ip="127.0.0.1", ports={"put_get_socket": port}) + self.identities: list[bytes] = [] self._tag = tag self._delay_first_reply = delay_first_reply self._replies = 0 @@ -52,6 +56,11 @@ def __init__(self, delay_first_reply: float = 0.0, peer_id: str = "peer_0", tag: self.thread = threading.Thread(target=self._serve, daemon=True) self.thread.start() + @property + def callers(self) -> int: + """How many distinct sockets have dialled this peer.""" + return len(set(self.identities)) + def _serve(self): poller = zmq.Poller() poller.register(self.socket, zmq.POLLIN) @@ -59,6 +68,7 @@ def _serve(self): if not dict(poller.poll(50)): continue identity, request = self.socket.recv_multipart() + self.identities.append(identity) if self._replies == 0 and self._delay_first_reply: # Reply late enough that the requester has already timed out, leaving this # reply in flight -- the poisoned-socket scenario. @@ -84,11 +94,6 @@ def peer(): p.stop() -def _idle_sockets(pool: ZMQSocketPool) -> list: - """Every socket currently parked in the pool, across all owners and buckets.""" - return [s for buckets in pool._idle.values() for bucket in buckets.values() for s in bucket] - - async def _round_trip(pool, peer_info, payload=b"req"): with pool.lease(peer_info) as sock: await sock.send_multipart([payload]) @@ -104,8 +109,7 @@ async def test_socket_is_reused_across_requests(peer): for i in range(10): assert await _round_trip(pool, peer.info, f"req{i}".encode()) == f"reply-to-req{i}".encode() - idle = _idle_sockets(pool) - assert len(idle) == 1, "each request opened its own socket instead of reusing one" + assert peer.callers == 1, "each request opened its own socket instead of reusing one" pool.close() ctx.destroy(linger=0) @@ -128,10 +132,9 @@ async def test_timed_out_socket_is_not_reused(): with pytest.raises(zmq.error.Again): await _round_trip(pool, peer.info, b"first") - assert _idle_sockets(pool) == [], "a timed-out socket was returned to the pool" - # The late reply to "first" must not surface as the answer to "second". assert await _round_trip(pool, peer.info, b"second") == b"reply-to-second" + assert peer.callers == 2, "the timed-out socket was handed to the next request" finally: pool.close() ctx.destroy(linger=0) @@ -139,66 +142,36 @@ async def test_timed_out_socket_is_not_reused(): @pytest.mark.asyncio -async def test_cancelled_lease_discards_socket(peer): - """Cancellation mid-recv poisons the socket: asyncio.gather cancels siblings routinely.""" - ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner", "put_get_socket") - leased = [] - - async def never_answered(): - with pool.lease(peer.info) as sock: - leased.append(sock) - await asyncio.sleep(60) # cancelled here, after the lease was handed out - - task = asyncio.create_task(never_answered()) - await asyncio.sleep(0.1) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - assert leased and leased[0].closed - assert _idle_sockets(pool) == [] - - pool.close() - ctx.destroy(linger=0) - - -@pytest.mark.asyncio -async def test_failed_lease_discards_socket(peer): - """Any exception in the body poisons the socket, not just timeouts.""" - ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner", "put_get_socket") - - with pytest.raises(RuntimeError): - with pool.lease(peer.info) as sock: - await sock.send_multipart([b"req"]) - raise RuntimeError("handler blew up") - - assert _idle_sockets(pool) == [] - - pool.close() - ctx.destroy(linger=0) - - -@pytest.mark.asyncio -async def test_socket_closed_by_the_caller_is_not_parked(peer): - """A socket the body closed without raising must not occupy a pool slot. +@pytest.mark.parametrize("failure", ["exception", "cancellation"]) +async def test_poisoned_lease_is_discarded(peer, failure): + """A lease that did not complete cleanly must not return its socket to the pool. - StorageManager._notify_and_wait does exactly this: it closes the socket so a late ACK - cannot be read as the next request's reply, but swallows the error so a slow controller - does not fail the put that triggered it. The lease therefore exits cleanly. + Cancellation counts as well as a raise: asyncio.gather cancels its siblings on the + first failure, so this is the routine case rather than an exotic one. """ ctx = zmq.asyncio.Context() pool = ZMQSocketPool(ctx, "owner", "put_get_socket") - with pool.lease(peer.info) as sock: - await sock.send_multipart([b"req"]) - sock.close(linger=0) # closed, but the body returns normally + async def poisoned(): + with pool.lease(peer.info): + # Nothing is sent, so no reply is left in flight to confuse the next request; + # the socket is poisoned purely by the lease not completing. + if failure == "exception": + raise RuntimeError("handler blew up") + await asyncio.sleep(60) # cancelled here, after the lease was handed out - assert _idle_sockets(pool) == [], "a closed socket was returned to the pool" + if failure == "exception": + with pytest.raises(RuntimeError): + await poisoned() + else: + task = asyncio.create_task(poisoned()) + await asyncio.sleep(0.1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task - # The pool still works, and the next request gets a live socket. assert await _round_trip(pool, peer.info) == b"reply-to-req" + assert peer.callers == 1, "the discarded socket never reached the peer" pool.close() ctx.destroy(linger=0) @@ -212,87 +185,22 @@ def test_sockets_are_not_reused_across_event_loops(peer): """ ctx = zmq.asyncio.Context() pool = ZMQSocketPool(ctx, "owner", "put_get_socket") - leased = [] async def lease_twice(tag): # Twice per loop, so a socket IS reused within a loop -- which is what makes the - # cross-loop comparison below meaningful rather than trivially true. + # cross-loop comparison meaningful rather than trivially true. for i in range(2): - with pool.lease(peer.info) as sock: - leased.append(sock) - await sock.send_multipart([f"{tag}{i}".encode()]) - await sock.recv_multipart() + assert await _round_trip(pool, peer.info, f"{tag}{i}".encode()) == f"reply-to-{tag}{i}".encode() - asyncio.run(lease_twice("a")) - first = [s for s in leased] - asyncio.run(lease_twice("b")) - second = [s for s in leased if s not in first] + for tag in ("a", "b", "c"): + asyncio.run(lease_twice(tag)) - assert len({id(s) for s in first}) == 1, "a socket should be reused within one loop" - assert len({id(s) for s in second}) == 1 - assert not ({id(s) for s in first} & {id(s) for s in second}), "a socket crossed event loops" + assert peer.callers == 3, "one socket per loop, reused within it" pool.close() ctx.destroy(linger=0) -def test_finished_loop_releases_its_sockets(peer): - """A finished loop's sockets must be closed, not left parked bound to a dead loop. - - ``asyncio.run()`` per call -- the pattern the client docstrings show -- creates one loop - per call. A pooled async socket keeps its own loop referenced, so these cannot be reaped - by garbage collection; the pool evicts finished owners on the next lease instead. - """ - ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner", "put_get_socket") - leased = [] - - async def once(): - with pool.lease(peer.info) as sock: - leased.append(sock) - await sock.send_multipart([b"q"]) - await sock.recv_multipart() - - for _ in range(5): - asyncio.run(once()) - - assert len(leased) == 5, "each fresh loop needs its own socket" - # The last loop's socket is still parked (nothing has leased since), but every earlier - # one must have been closed rather than accumulating. - assert all(sock.closed for sock in leased[:-1]), "a finished loop left an open socket behind" - assert len(_idle_sockets(pool)) == 1 - assert len(pool._idle) == 1, "finished owners must be evicted, not accumulated" - - pool.close() - ctx.destroy(linger=0) - - -def test_pooled_identities_are_unique_across_pools(peer): - """Two pools with the same owner_id must not collide on the wire. - - A ROUTER silently drops a second peer claiming an identity it already has, so colliding - identities would blackhole one process's traffic. Client ids are pid-derived and pids - repeat across nodes, so owner_id alone cannot carry uniqueness. - """ - ctx = zmq.asyncio.Context() - # Same owner_id, as two processes on different nodes with equal pids would produce. - a, b = ( - ZMQSocketPool(ctx, "TransferQueueClient_1234", "put_get_socket"), - ZMQSocketPool(ctx, "TransferQueueClient_1234", "put_get_socket"), - ) - - async def identity_of(pool): - with pool.lease(peer.info) as sock: - return sock.getsockopt(zmq.IDENTITY) - - first, second = asyncio.run(identity_of(a)), asyncio.run(identity_of(b)) - assert first != second, "two pools minted the same ZMQ identity" - - a.close() - b.close() - ctx.destroy(linger=0) - - @pytest.mark.asyncio async def test_reregistered_peer_is_not_served_a_stale_socket(): """A peer that moves to a new address must not be answered by its old endpoint. @@ -320,47 +228,6 @@ async def test_reregistered_peer_is_not_served_a_stale_socket(): new.stop() -def test_pool_size_below_one_is_rejected(): - """A size under 1 parks nothing, so reuse is silently off while still looking pooled.""" - ctx = zmq.asyncio.Context() - try: - for bad in (-1, 0): - with pytest.raises(ValueError, match="at least 1"): - ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=bad) - ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=1) # the boundary is valid - finally: - ctx.destroy(linger=0) - - -def test_connect_failure_does_not_leak_a_socket(peer): - """A socket is nobody's responsibility until it reaches a lease, so _connect closes it.""" - ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner", "put_get_socket") - bad = ZMQServerInfo(role=Role.STORAGE, id="bad", ip="127.0.0.1", ports={"put_get_socket": -1}) - - created = [] - original = zmq_utils.create_zmq_socket - - def spy(*args, **kwargs): - sock = original(*args, **kwargs) - created.append(sock) - return sock - - zmq_utils.create_zmq_socket = spy - try: - with pytest.raises(zmq.ZMQError): - with pool.lease(bad): - pass - finally: - zmq_utils.create_zmq_socket = original - - assert created and created[0].closed, "a socket that failed to connect was left open" - - pool.close() - ctx.destroy(linger=0) - ctx.destroy(linger=0) - - def test_sync_caller_can_lease(peer): """The metrics collector leases from a plain thread, with no event loop running.""" ctx = zmq.Context() @@ -371,49 +238,20 @@ def test_sync_caller_can_lease(peer): sock.send_multipart([f"m{i}".encode()]) assert sock.recv_multipart()[0] == f"reply-to-m{i}".encode() - assert len(_idle_sockets(pool)) == 1 + assert peer.callers == 1, "a synchronous caller should reuse its socket too" pool.close() ctx.destroy(linger=0) @pytest.mark.asyncio -async def test_pool_size_is_a_soft_cap(peer): - """Concurrency above maxsize still gets sockets; only the steady state is bounded.""" +async def test_burst_beyond_pool_size_is_served(peer): + """Concurrency above maxsize must not be refused or blocked, only left unparked.""" ctx = zmq.asyncio.Context() pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=2) results = await asyncio.gather(*[_round_trip(pool, peer.info, f"c{i}".encode()) for i in range(8)]) - assert len(results) == 8, "a burst beyond maxsize must not be refused or blocked" - assert len(_idle_sockets(pool)) == 2, "excess sockets must be closed on return, not parked" + assert sorted(results) == sorted(f"reply-to-c{i}".encode() for i in range(8)) pool.close() ctx.destroy(linger=0) - - -@pytest.mark.asyncio -async def test_unknown_socket_name_is_reported(peer): - ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner", "no_such_socket") - - with pytest.raises(RuntimeError, match="not configured"): - with pool.lease(peer.info): - pass - - pool.close() - ctx.destroy(linger=0) - - -@pytest.mark.asyncio -async def test_close_is_idempotent_and_survives_dead_context(peer): - """Teardown ordering is not guaranteed, so close() must tolerate a destroyed context.""" - ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner", "put_get_socket") - await _round_trip(pool, peer.info) - - pool.close() - assert _idle_sockets(pool) == [] - pool.close() # twice - - ctx.destroy(linger=0) - pool.close() # after the context is gone From 6587f651e163f8f108d3a29897acead50ad37765 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Tue, 8 Sep 2026 13:06:02 +0800 Subject: [PATCH 18/28] [test] Assert each role's socket reuse where that role is tested Reuse was only covered against ZMQSocketPool directly, which is internal and not something a user reaches. Each pool belongs to a role, so the reuse each role depends on now lives with that role. test_client.py drives it through the public client: repeated get_partition_list() calls must reach the controller over one connection, which is the user-visible payoff of pooling and something the file asserted nowhere before. MockController records caller identities to make that observable; it answers requests exactly as before. test_metrics.py covers the collector, whose pool keys by thread because it runs with no event loop at all, against a ROUTER that counts callers. Both were checked by breaking _take() so no socket is ever handed back: each then fails with one identity per request, so neither passes vacuously. What stays in test_zmq_socket_pool.py needs control over reply timing and loop lifetime that a caller-level test cannot reach -- the late reply that must not answer the next request, poisoned leases, loop keying, peer re-registration, and bursts past maxsize. Its docstring now says where each role's own reuse is covered, so the split is not mistaken for an omission. Signed-off-by: OutstanderWang --- tests/test_client.py | 24 ++++++++++++++++++ tests/test_metrics.py | 48 +++++++++++++++++++++++++++++++++++ tests/test_zmq_socket_pool.py | 38 +++++---------------------- 3 files changed, 78 insertions(+), 32 deletions(-) diff --git a/tests/test_client.py b/tests/test_client.py index d287108e..aa887656 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -68,6 +68,7 @@ def __init__(self, controller_id="controller_0"): ) self.running = True + self.caller_identities: set[bytes] = set() self.request_thread = Thread(target=self._handle_requests, daemon=True) self.request_thread.start() @@ -85,6 +86,9 @@ def _handle_requests(self): if self.request_socket in socks: messages = self.request_socket.recv_multipart(copy=False) identity = messages.pop(0) + # Each pooled socket dials with its own identity, so this records how + # many distinct connections the client opened. + self.caller_identities.add(bytes(identity)) serialized_msg = messages request_msg = ZMQMessage.deserialize(serialized_msg) @@ -1330,6 +1334,26 @@ def test_kv_retrieve_keys_type_validation(self, client_setup): ) +# ===================================================== +# Controller RPC Socket Pool Tests +# ===================================================== + + +def test_controller_rpc_reuses_one_socket(client_setup): + """Repeated calls must share a single connection to the controller. + + This is the client-side payoff of the socket pool: TransferQueueClient runs every call + on the one event loop it builds in __init__, so all of them land in the same pool + bucket. Counted from the controller, which sees one ZMQ identity per socket dialled. + """ + client, mock_controller, _ = client_setup + + for _ in range(5): + assert client.get_partition_list() == ["partition_0", "partition_1", "test_partition"] + + assert len(mock_controller.caller_identities) == 1, "each request opened its own socket" + + # ===================================================== # Checkpoint Interface Tests # ===================================================== diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 88401988..62e9c25b 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -16,6 +16,7 @@ """Unit tests for the Prometheus metrics exporter (transfer_queue.metrics).""" import time +from threading import Thread from unittest.mock import MagicMock, patch import pytest @@ -24,6 +25,8 @@ import zmq from transfer_queue.metrics import TQMetricsExporter + from transfer_queue.utils.enum_utils import Role + from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo _HAS_DEPS = True except (ImportError, OSError): @@ -232,6 +235,51 @@ def test_missing_context_is_reported(self): with pytest.raises(RuntimeError, match="without a ZMQ context"): exporter._get_socket_pool() + def test_collector_reuses_one_socket_per_storage_unit(self): + """The collector runs on a plain thread with no event loop, and must still reuse. + + Its pool keys by thread rather than by loop, and the collect loop is one long-lived + daemon thread, so every cycle should reach a storage unit over the same connection. + Counted from the storage unit, which sees one ZMQ identity per socket dialled. + """ + identities: set[bytes] = set() + ctx_peer = zmq.Context() + router = ctx_peer.socket(zmq.ROUTER) + port = router.bind_to_random_port("tcp://127.0.0.1") + running = True + + def serve(): + poller = zmq.Poller() + poller.register(router, zmq.POLLIN) + while running: + if not dict(poller.poll(50)): + continue + identity, _ = router.recv_multipart() + identities.add(bytes(identity)) + response = ZMQMessage.create( + request_type=ZMQRequestType.METRICS_RESPONSE, + sender_id="storage_0", + body={}, + ) + router.send_multipart([identity, *response.serialize()]) + + server = Thread(target=serve, daemon=True) + server.start() + + su_info = ZMQServerInfo(role=Role.STORAGE, id="storage_0", ip="127.0.0.1", ports={"put_get_socket": port}) + ctx = zmq.Context() + try: + exporter = TQMetricsExporter(zmq_context=ctx) + for _ in range(3): + assert exporter._query_storage_unit(su_info, "storage_0") == {} + assert len(identities) == 1, "each collection cycle opened its own socket" + finally: + running = False + server.join(timeout=2.0) + ctx.destroy(linger=0) + router.close(linger=0) + ctx_peer.term() + class TestStorageMetricsCollection: def test_collect_with_no_storage_units(self): diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py index fdb35e91..ff74130f 100644 --- a/tests/test_zmq_socket_pool.py +++ b/tests/test_zmq_socket_pool.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for ZMQSocketPool. +"""Tests for ZMQSocketPool's concurrency invariants. Sockets used to be created and closed per request. The pool reuses them, which is only safe because a lease is exclusive and a socket that did not complete a clean send/recv is @@ -21,6 +21,11 @@ by a timed-out or cancelled request would be read by the next user of that socket as its own. test_timed_out_socket_is_not_reused pins exactly that. +Each role's own reuse is asserted where that role is tested -- the client's controller RPC +in test_client.py, the metrics collector in test_metrics.py, pool wiring and context +lifecycle in test_zmq_shared_context.py. What is left here needs control over reply timing +and loop lifetime that a caller-level test cannot reach. + Reuse is asserted from the peer's side rather than from the pool's internals: every socket dials with its own ZMQ identity, so one identity across many requests means the connection was reused, and a fresh identity means the old socket was discarded. @@ -100,21 +105,6 @@ async def _round_trip(pool, peer_info, payload=b"req"): return (await sock.recv_multipart())[0] -@pytest.mark.asyncio -async def test_socket_is_reused_across_requests(peer): - """Sequential requests to one peer must share a single socket.""" - ctx = zmq.asyncio.Context() - pool = ZMQSocketPool(ctx, "owner", "put_get_socket") - - for i in range(10): - assert await _round_trip(pool, peer.info, f"req{i}".encode()) == f"reply-to-req{i}".encode() - - assert peer.callers == 1, "each request opened its own socket instead of reusing one" - - pool.close() - ctx.destroy(linger=0) - - @pytest.mark.asyncio async def test_timed_out_socket_is_not_reused(): """A timed-out request must not leave its socket -- or its late reply -- in the pool. @@ -228,22 +218,6 @@ async def test_reregistered_peer_is_not_served_a_stale_socket(): new.stop() -def test_sync_caller_can_lease(peer): - """The metrics collector leases from a plain thread, with no event loop running.""" - ctx = zmq.Context() - pool = ZMQSocketPool(ctx, "metrics_collector", "put_get_socket", timeout=5) - - for i in range(3): - with pool.lease(peer.info) as sock: - sock.send_multipart([f"m{i}".encode()]) - assert sock.recv_multipart()[0] == f"reply-to-m{i}".encode() - - assert peer.callers == 1, "a synchronous caller should reuse its socket too" - - pool.close() - ctx.destroy(linger=0) - - @pytest.mark.asyncio async def test_burst_beyond_pool_size_is_served(peer): """Concurrency above maxsize must not be refused or blocked, only left unparked.""" From 6db3a70a5a7b33670abe692fae2b3ab84890916f Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 9 Sep 2026 10:56:33 +0800 Subject: [PATCH 19/28] [test, docs] Make the poisoned-lease test real, and trim example noise Review found test_poisoned_lease_is_discarded passing vacuously. The lease sent nothing before failing, so the peer never saw that socket and the caller count held whether the socket was discarded or wrongly parked. Confirmed by releasing the socket on the exception path instead of closing it: both parametrized cases still passed. Send the request and wait until the peer has it before failing, with the reply delayed so it is genuinely outstanding, then assert the next request both uses a different connection and receives its own reply. Sabotaging the discard now fails both cases on the reply belonging to the abandoned request, which is the misattribution the discard exists to prevent. Drop the async def main() wrapper from examples that make a single call; it only earns its keep where consecutive calls show reuse under one loop. Bare await matches the async examples already in interface.py. Shorten the close() guard comment to the reason it exists. Signed-off-by: OutstanderWang --- tests/test_zmq_socket_pool.py | 50 +++++++++++------- transfer_queue/client.py | 70 +++++++++++-------------- transfer_queue/storage/managers/base.py | 6 +-- 3 files changed, 65 insertions(+), 61 deletions(-) diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py index ff74130f..a23ef1e9 100644 --- a/tests/test_zmq_socket_pool.py +++ b/tests/test_zmq_socket_pool.py @@ -133,38 +133,50 @@ async def test_timed_out_socket_is_not_reused(): @pytest.mark.asyncio @pytest.mark.parametrize("failure", ["exception", "cancellation"]) -async def test_poisoned_lease_is_discarded(peer, failure): +async def test_poisoned_lease_is_discarded(failure): """A lease that did not complete cleanly must not return its socket to the pool. + The request is on the wire before the failure and its reply is still in flight, so + parking the socket would hand the next caller a reply belonging to someone else. Cancellation counts as well as a raise: asyncio.gather cancels its siblings on the first failure, so this is the routine case rather than an exotic one. """ + # Answers the first request only after it has been abandoned, which is what leaves the + # stale reply in flight. + peer = _Peer(delay_first_reply=0.6) ctx = zmq.asyncio.Context() pool = ZMQSocketPool(ctx, "owner", "put_get_socket") async def poisoned(): - with pool.lease(peer.info): - # Nothing is sent, so no reply is left in flight to confuse the next request; - # the socket is poisoned purely by the lease not completing. + with pool.lease(peer.info) as sock: + await sock.send_multipart([b"doomed"]) + # Fail only once the peer has the request, so its reply really is outstanding. + while not peer.identities: + await asyncio.sleep(0.01) if failure == "exception": raise RuntimeError("handler blew up") await asyncio.sleep(60) # cancelled here, after the lease was handed out - if failure == "exception": - with pytest.raises(RuntimeError): - await poisoned() - else: - task = asyncio.create_task(poisoned()) - await asyncio.sleep(0.1) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - assert await _round_trip(pool, peer.info) == b"reply-to-req" - assert peer.callers == 1, "the discarded socket never reached the peer" - - pool.close() - ctx.destroy(linger=0) + try: + if failure == "exception": + with pytest.raises(RuntimeError): + await poisoned() + else: + task = asyncio.create_task(poisoned()) + while not peer.identities: + await asyncio.sleep(0.01) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + abandoned = set(peer.identities) + # The reply to "doomed" must not surface as the answer to "second". + assert await _round_trip(pool, peer.info, b"second") == b"reply-to-second" + assert set(peer.identities) - abandoned, "the poisoned socket was reused" + finally: + pool.close() + ctx.destroy(linger=0) + peer.stop() def test_sockets_are_not_reused_across_event_loops(peer): diff --git a/transfer_queue/client.py b/transfer_queue/client.py index eafad617..de761c88 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -726,13 +726,12 @@ async def async_get_consumption_status( RuntimeError: If communication fails or controller returns error response Example: - >>> async def main(): - ... # Get consumption status - ... global_index, consumption_status = await client.async_get_consumption_status( - ... task_name="generate_sequences", - ... partition_id="train_0" - ... ) - ... print(f"Global index: {global_index}, Consumption status: {consumption_status}") + >>> # Get consumption status + >>> global_index, consumption_status = await client.async_get_consumption_status( + ... task_name="generate_sequences", + ... partition_id="train_0" + ... ) + >>> print(f"Global index: {global_index}, Consumption status: {consumption_status}") """ try: @@ -774,13 +773,12 @@ async def async_get_production_status( RuntimeError: If communication fails or controller returns error response Example: - >>> async def main(): - ... # Get production status - ... global_index, production_status = await client.async_get_production_status( - ... data_fields=["input_ids", "attention_mask"], - ... partition_id="train_0" - ... ) - ... print(f"Global index: {global_index}, Production status: {production_status}") + >>> # Get production status + >>> global_index, production_status = await client.async_get_production_status( + ... data_fields=["input_ids", "attention_mask"], + ... partition_id="train_0" + ... ) + >>> print(f"Global index: {global_index}, Production status: {production_status}") """ try: response_msg = await self._request_controller( @@ -816,13 +814,12 @@ async def async_check_consumption_status( RuntimeError: If communication fails or controller returns error response Example: - >>> async def main(): - ... # Check if all samples have been consumed - ... is_consumed = await client.async_check_consumption_status( - ... task_name="generate_sequences", - ... partition_id="train_0" - ... ) - ... print(f"All samples consumed: {is_consumed}") + >>> # Check if all samples have been consumed + >>> is_consumed = await client.async_check_consumption_status( + ... task_name="generate_sequences", + ... partition_id="train_0" + ... ) + >>> print(f"All samples consumed: {is_consumed}") """ _, consumption_status = await self.async_get_consumption_status( @@ -853,13 +850,12 @@ async def async_check_production_status( RuntimeError: If communication fails or controller returns error response Example: - >>> async def main(): - ... # Check if all samples are ready for consumption - ... is_ready = await client.async_check_production_status( - ... data_fields=["input_ids", "attention_mask"], - ... partition_id="train_0" - ... ) - ... print(f"All samples ready: {is_ready}") + >>> # Check if all samples are ready for consumption + >>> is_ready = await client.async_check_production_status( + ... data_fields=["input_ids", "attention_mask"], + ... partition_id="train_0" + ... ) + >>> print(f"All samples ready: {is_ready}") """ _, production_status = await self.async_get_production_status( data_fields=data_fields, @@ -894,13 +890,12 @@ async def async_reset_consumption( RuntimeError: If communication fails or controller returns error response Example: - >>> async def main(): - ... # Reset consumption for train task to re-train on same data - ... success = await client.async_reset_consumption( - ... partition_id="train_0", - ... task_name="train" - ... ) - ... print(f"Reset successful: {success}") + >>> # Reset consumption for train task to re-train on same data + >>> success = await client.async_reset_consumption( + ... partition_id="train_0", + ... task_name="train" + ... ) + >>> print(f"Reset successful: {success}") """ body = {"partition_id": partition_id} if task_name is not None: @@ -933,9 +928,8 @@ async def async_get_partition_list( list[str]: List of partition ids managed by the controller Example: - >>> async def main(): - ... partition_ids = await client.get_partition_list() - ... print(f"Available partitions: {partition_ids}") + >>> partition_ids = await client.get_partition_list() + >>> print(f"Available partitions: {partition_ids}") """ try: response_msg = await self._request_controller( diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 537da503..6114753a 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -378,10 +378,8 @@ async def load_checkpoint(self, checkpoint_dir: str) -> None: def close(self) -> None: """Close all ZMQ sockets/contexts and stop the notify loop.""" - # A subclass may reject its config before calling super().__init__(), as the KV - # managers do, leaving nothing here allocated. __del__ calls close() anyway, so - # return rather than burying the constructor's error under an AttributeError. - # notify_pool is the last thing __init__ sets before its first fallible step. + # A subclass may reject its config before super().__init__() runs, and __del__ calls + # close() anyway; return rather than bury that error under an AttributeError. if not hasattr(self, "notify_pool"): return From 365c3ff69b7be749d5a9d04e3843c101473d28f6 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 9 Sep 2026 11:10:34 +0800 Subject: [PATCH 20/28] [refactor] Trim the comments this branch added Five blocks ran past the four-line ceiling and three more repeated what the code or a nearby docstring already says. Comments only in code, none removed outright. Most of the excess was duplication: lease exclusivity and the maxsize semantics are documented on ZMQSocketPool, so restating them at each call site left three places to keep in sync. Two blocks explained a state no supported path reaches -- a bucket orphaned by endpoint remapping, and pool isolation that the field names already carry -- which the simplicity gate asks us not to describe. What is left is the part a reader cannot infer: why the pre-context checks come first, why reuse makes the socket budget track concurrency rather than request count, why a late ACK forces a close, and why owner_id alone cannot make a ZMQ identity unique. Signed-off-by: OutstanderWang --- transfer_queue/client.py | 20 ++++++-------------- transfer_queue/storage/managers/base.py | 11 ++++------- transfer_queue/utils/zmq_utils.py | 12 ++++-------- 3 files changed, 14 insertions(+), 29 deletions(-) diff --git a/transfer_queue/client.py b/transfer_queue/client.py index de761c88..3eaa1aeb 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -47,10 +47,8 @@ # 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 -# Idle sockets kept per (loop, endpoint) bucket for controller RPC, at least 1. A soft cap: -# bursts beyond it still get sockets, so this bounds the steady state rather than the peak. -# Scoped to this one pool; the storage RPC, notify, and metrics pools take ZMQSocketPool's -# own default, as none of them reaches a concurrency worth tuning separately. +# Idle sockets kept per (loop, endpoint) bucket for controller RPC; see ZMQSocketPool for +# the cap's semantics and which pools take its default instead. TQ_CONTROLLER_RPC_POOL_SIZE = int(os.environ.get("TQ_CONTROLLER_RPC_POOL_SIZE", 64)) # Pre-bound decorator for controller socket operations. @@ -100,12 +98,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, with sockets leased from a pool over it rather - # than built per request; a lease is exclusive because ZMQ sockets are not - # thread-safe and replies are matched to requests by arrival order. - # Everything checkable without the context is checked first: the finalizer is not - # armed until __init__ returns, so a raise after allocation leaks the context and its - # native I/O threads. + # Check everything that does not need the context first: the finalizer is not armed + # until __init__ returns, so raising after allocation leaks it and its I/O 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"Client ZMQ I/O thread pool size must be at least 1, got {io_threads}") @@ -158,10 +152,8 @@ def __init__( except BaseException: self.zmq_context.destroy(linger=0) raise - # Sockets are leased from this pool and reused across requests, so the context's - # socket budget above is consumed by the concurrency high-water mark, not by - # request count. Controller RPC only -- the storage backend keeps its own pools, so - # neither scenario can disturb the other's sockets. + # Reused across requests, so the socket budget above is consumed by the concurrency + # high-water mark rather than by request count. self.controller_rpc_pool = ZMQSocketPool( self.zmq_context, client_id, diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 6114753a..34b75026 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -87,9 +87,8 @@ def __init__( # 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 - # Notify traffic gets its own pool. It runs on a dedicated loop, so it could not share - # sockets with another scenario in any case, and keeping it separate means no other - # scenario's sockets are reachable from here. + # Notify runs on a dedicated loop, so it could not share sockets with another + # scenario in any case. self.notify_pool = ZMQSocketPool(self.zmq_context, self.storage_manager_id, "request_handle_socket") self._connect_to_controller() @@ -304,10 +303,8 @@ async def _notify_and_wait(self, request_msg: list) -> None: ) return except Exception as e: - # Notification failure has always been logged rather than raised, so a slow - # controller does not fail the put that triggered it. Close the socket instead - # of reusing it: an ACK may still be in flight, and the next lessee would read - # it as its own reply. The pool discards a closed socket on its next lease. + # Logged rather than raised, so a slow controller does not fail the put. Close + # the socket: a late ACK would otherwise be read as the next lessee's reply. logger.error(f"[{self.storage_manager_id}]: Data status update failed: {type(e).__name__}: {e}") sock.close(linger=0) diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index d6f99337..5c128c39 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -436,16 +436,12 @@ def __init__( self._socket_name = socket_name self._timeout = timeout self._maxsize = maxsize - # Keyed by lease owner (see _lease_owner), then by address. The address, not the peer - # id: a peer restarted under the same id at a new address must not be handed a socket - # still connected to the old one. The bucket it moved off is then never asked for - # again and lingers until close(), which no supported path can reach today -- adding - # runtime endpoint remapping means retiring those buckets too. + # Keyed by lease owner (see _lease_owner), then by address rather than peer id: a peer + # restarted under the same id at a new address must not get a socket wired to the old. self._idle: dict[Any, dict[str, list[zmq.Socket]]] = {} self._lock = threading.Lock() - # A ROUTER silently drops a second peer claiming an identity it already has, so - # identities must not collide between processes: owner_id alone does not suffice - # (client ids are pid-derived, and pids repeat across nodes). + # A ROUTER silently drops a second peer claiming an identity it already has, and + # owner_id alone repeats across nodes because client ids are pid-derived. self._identity_prefix = f"{owner_id}_{uuid4().hex[:8]}" self._counter = itertools.count() From 9f9f847c6aba175b5800a129474023827bb40938 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 9 Sep 2026 11:27:11 +0800 Subject: [PATCH 21/28] [feat] Size every role's socket pool from TQ_SOCKET_POOL_SIZE The knob only reached the client's controller RPC pool, so the storage RPC, notify, and metrics pools were fixed at the constructor default with no way to tune them. Review asked for one name covering all four; TQ_SOCKET_POOL_SIZE replaces TQ_CONTROLLER_RPC_POOL_SIZE, keeping the same default of 64. It lives beside ZMQSocketPool because that is what the cap belongs to, and is resolved in __init__ rather than bound as a signature default: a default read at import freezes whatever the environment held when the module first loaded, which is the failure the per-call socket timeout already had to fix. An explicit maxsize still wins, so a caller can opt out. The cap remains per (owner, address), so one value means different totals for a pool dialling one peer and a pool dialling N. That is unchanged behaviour, and the docstring says so. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 6 +++--- tests/test_zmq_socket_pool.py | 29 +++++++++++++++++++++++++++++ transfer_queue/client.py | 9 +++------ transfer_queue/utils/zmq_utils.py | 29 +++++++++++++---------------- 4 files changed, 48 insertions(+), 25 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index c4d1ca0f..a0441e64 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -197,15 +197,15 @@ def _spy(*args, **kwargs): def test_client_rejects_invalid_socket_pool_size(echo_controller): - """A bad TQ_CONTROLLER_RPC_POOL_SIZE must name the variable, not silently disable reuse. + """A bad TQ_SOCKET_POOL_SIZE must name the variable, not silently disable reuse. Below 1 nothing is ever parked, so every request pays a fresh connect while the client still looks pooled. """ for bad in (-1, 0): - with patch("transfer_queue.client.TQ_CONTROLLER_RPC_POOL_SIZE", bad): + with patch("transfer_queue.client.TQ_SOCKET_POOL_SIZE", bad): with _no_context_left_open() as created: - with pytest.raises(ValueError, match="TQ_CONTROLLER_RPC_POOL_SIZE must be at least 1"): + with pytest.raises(ValueError, match="TQ_SOCKET_POOL_SIZE must be at least 1"): AsyncTransferQueueClient( client_id="client_invalid_socket_pool", controller_info=echo_controller.zmq_server_info, diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py index a23ef1e9..9f56c2a0 100644 --- a/tests/test_zmq_socket_pool.py +++ b/tests/test_zmq_socket_pool.py @@ -33,11 +33,13 @@ import asyncio import threading +from unittest.mock import patch import pytest import zmq import zmq.asyncio +import transfer_queue.utils.zmq_utils as zmq_utils from transfer_queue.utils.enum_utils import Role from transfer_queue.utils.zmq_utils import ZMQServerInfo, ZMQSocketPool @@ -241,3 +243,30 @@ async def test_burst_beyond_pool_size_is_served(peer): pool.close() ctx.destroy(linger=0) + + +def test_pool_size_comes_from_the_env_var_at_construction(): + """Every role's pool takes TQ_SOCKET_POOL_SIZE, resolved per pool rather than at import. + + A default bound in the signature would freeze whatever the environment held when this + module first loaded, which is what makes such a knob look settable but do nothing. + """ + ctx = zmq.Context() + try: + with patch.object(zmq_utils, "TQ_SOCKET_POOL_SIZE", 7): + assert ZMQSocketPool(ctx, "owner", "put_get_socket")._maxsize == 7 + # An explicit argument still wins, which is what lets a caller opt out. + assert ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=3)._maxsize == 3 + finally: + ctx.destroy(linger=0) + + +def test_pool_size_below_one_is_rejected(): + """Below 1 nothing is ever parked, so reuse is off while the pool still looks pooled.""" + ctx = zmq.Context() + try: + with patch.object(zmq_utils, "TQ_SOCKET_POOL_SIZE", 0): + with pytest.raises(ValueError, match="at least 1"): + ZMQSocketPool(ctx, "owner", "put_get_socket") + finally: + ctx.destroy(linger=0) diff --git a/transfer_queue/client.py b/transfer_queue/client.py index 3eaa1aeb..0f7f13f1 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -29,6 +29,7 @@ from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( + TQ_SOCKET_POOL_SIZE, ZMQMessage, ZMQRequestType, ZMQServerInfo, @@ -47,9 +48,6 @@ # 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 -# Idle sockets kept per (loop, endpoint) bucket for controller RPC; see ZMQSocketPool for -# the cap's semantics and which pools take its default instead. -TQ_CONTROLLER_RPC_POOL_SIZE = int(os.environ.get("TQ_CONTROLLER_RPC_POOL_SIZE", 64)) # Pre-bound decorator for controller socket operations. with_controller_socket = with_zmq_socket( @@ -103,10 +101,10 @@ def __init__( 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}") - if TQ_CONTROLLER_RPC_POOL_SIZE < 1: + if TQ_SOCKET_POOL_SIZE < 1: # Name the variable: the pool's own error cannot say which knob supplied the value. raise ValueError( - f"TQ_CONTROLLER_RPC_POOL_SIZE must be at least 1, got {TQ_CONTROLLER_RPC_POOL_SIZE}. " + f"TQ_SOCKET_POOL_SIZE must be at least 1, got {TQ_SOCKET_POOL_SIZE}. " f"The pool always reuses at least one socket per endpoint; it cannot be disabled." ) @@ -158,7 +156,6 @@ def __init__( self.zmq_context, client_id, "request_handle_socket", - maxsize=TQ_CONTROLLER_RPC_POOL_SIZE, ) # Backstop for a client that is never closed, so the context and its I/O threads do diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 5c128c39..c399c51a 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -15,6 +15,7 @@ import asyncio import itertools +import os import socket import threading import time @@ -36,15 +37,9 @@ logger = get_logger(__name__) -# Identity prefixes of the peers allowed to reach a storage unit. The storage proxy drops -# anything else, so an identity built without these prefixes is silently unreachable. -STORAGE_MANAGER_IDENTITY_PREFIX = "TQ_STORAGE_" -METRICS_COLLECTOR_IDENTITY_PREFIX = "metrics_collector_" -STORAGE_CLIENT_IDENTITY_PREFIXES = ( - STORAGE_MANAGER_IDENTITY_PREFIX.encode(), - METRICS_COLLECTOR_IDENTITY_PREFIX.encode(), -) - +# Idle sockets kept per (owner, address) bucket by every ZMQSocketPool. See ZMQSocketPool +# for what the cap does and does not bound. +TQ_SOCKET_POOL_SIZE = int(os.environ.get("TQ_SOCKET_POOL_SIZE", 64)) bytestr: TypeAlias = bytes | bytearray | memoryview @@ -411,7 +406,7 @@ def __init__( socket_name: str, *, timeout: int | None = None, - maxsize: int = 64, + maxsize: int | None = None, ): """ Args: @@ -420,13 +415,15 @@ def __init__( owner_id: Identity prefix for pooled sockets, for readable peer-side logs. socket_name: Port key in ``ZMQServerInfo.ports`` that every lease dials. timeout: Send/recv timeout in seconds applied to every socket, or None for none. - maxsize: Idle sockets kept per bucket, at least 1. A soft cap: a burst beyond it - still gets sockets, and the excess is closed on return rather than made to - wait. Counted per (owner, address), not per pool, so a pool dialling N peers - may hold N*maxsize idle sockets. Only the client's controller RPC pool tunes - this (``TQ_CONTROLLER_RPC_POOL_SIZE``); the storage RPC, notify, and metrics - pools take this default, the latter two holding one socket at a time anyway. + maxsize: Idle sockets kept per bucket, at least 1, defaulting to + ``TQ_SOCKET_POOL_SIZE``. A soft cap: a burst beyond it still gets sockets, + and the excess is closed on return rather than made to wait. Counted per + (owner, address), not per pool, so a pool dialling N peers may hold + N*maxsize idle sockets. """ + # Resolved here rather than in the signature so the env var stays patchable; a + # default bound at import froze whatever the environment held when this module loaded. + maxsize = TQ_SOCKET_POOL_SIZE if maxsize is None else maxsize if maxsize < 1: # Below 1 nothing is ever parked, so every request pays a fresh connect while # still looking pooled. Reject it rather than silently disable reuse. From 96e218aa755ab6ebc7f8316114c376d2d688d578 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 9 Sep 2026 12:19:05 +0800 Subject: [PATCH 22/28] [fix] Default the socket pool to 8 and warn when it can exhaust the context The cap is per (owner, address), so it multiplies by peer count while the context's socket ceiling does not. At 64 a storage manager addressing two thousand units could park 128k idle sockets against a budget of 8192, and the lease that hits the limit fails with EMFILE rather than degrading. Drop the default to 8. One socket per peer already removes the repeated handshake, which is what pooling was for; a larger cap only helps when a single peer sees concurrent requests, and that is the case worth opting into explicitly rather than paying for by default. Warn at construction when TQ_SOCKET_POOL_SIZE times the registered unit count exceeds the context's ZMQ_MAX_SOCKETS, naming both knobs. The storage manager is the first point that knows how many peers it will dial, and the pool cannot tell on its own because it connects lazily. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 45 +++++++++++++++++++ .../managers/simple_storage_manager.py | 23 ++++++++++ transfer_queue/utils/zmq_utils.py | 7 +-- 3 files changed, 72 insertions(+), 3 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index a0441e64..75a8abd1 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -293,6 +293,51 @@ def test_each_scenario_gets_its_own_pool(echo_controller): client.close() +@pytest.mark.parametrize( + "pool_size, units, warns", + [ + (8, 4, False), # 32 idle sockets against a 64-socket budget + (8, 32, True), # 256 would exceed it + ], +) +def test_warns_when_pool_size_times_units_exceeds_the_context(echo_controller, caplog, pool_size, units, warns): + """The per-address cap multiplies by storage-unit count; the context ceiling does not. + + At a few thousand units the product passes ZMQ_MAX_SOCKETS, where opening a socket + fails outright, so the mismatch is worth naming at construction rather than at a lease. + """ + client = AsyncTransferQueueClient( + client_id="client_pool_budget", + controller_info=echo_controller.zmq_server_info, + zmq_max_sockets=64, + ) + zmq_info = { + f"storage_{i}": ZMQServerInfo( + role=Role.STORAGE, + id=f"storage_{i}", + ip="127.0.0.1", + ports={"put_get_socket": 5600 + i}, + ) + for i in range(units) + } + + with ( + patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"), + patch("transfer_queue.storage.managers.simple_storage_manager.TQ_SOCKET_POOL_SIZE", pool_size), + caplog.at_level("WARNING"), + ): + manager = AsyncSimpleStorageManager( + echo_controller.zmq_server_info, + {"zmq_info": zmq_info}, + zmq_context=client.zmq_context, + ) + + assert ("above this context's ZMQ_MAX_SOCKETS" in caplog.text) is warns + + manager.close() + client.close() + + def test_close_after_failed_handshake_still_releases_the_context(echo_controller): """A manager whose base constructor raised must still tear down on close(). diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 6fb5ef1b..3780bc9b 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -32,6 +32,7 @@ from transfer_queue.storage.simple_storage import KEY_NOT_FOUND_MARKER, StorageKeyNotFoundError from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( + TQ_SOCKET_POOL_SIZE, ZMQMessage, ZMQRequestType, ZMQServerInfo, @@ -102,6 +103,28 @@ def __init__( raise ValueError("AsyncSimpleStorageManager requires non-empty 'zmq_info' in config.") self.storage_unit_infos = self._register_servers(server_infos) + self._warn_if_pool_can_exhaust_context(len(self.storage_unit_infos)) + + def _warn_if_pool_can_exhaust_context(self, num_units: int) -> None: + """Warn when the pool may park more sockets than the context can hold. + + The cap is per (owner, address), so it multiplies by storage-unit count while the + context ceiling does not. At a few thousand units the product passes ZMQ_MAX_SOCKETS, + where a lease fails with EMFILE instead of degrading. + """ + try: + budget = self.zmq_context.get(zmq.MAX_SOCKETS) + except zmq.ZMQError: # pragma: no cover - context already terminating + return + worst_case = TQ_SOCKET_POOL_SIZE * num_units + if worst_case > budget: + logger.warning( + f"[{self.storage_manager_id}]: storage RPC pool may hold up to " + f"{TQ_SOCKET_POOL_SIZE} x {num_units} = {worst_case} idle sockets, above this " + f"context's ZMQ_MAX_SOCKETS ({budget}). Concurrent requests can then fail to " + f"open a socket. Lower TQ_SOCKET_POOL_SIZE or raise TQ_CLIENT_ZMQ_MAX_SOCKETS " + f"(with enough file descriptors, see ulimit -n)." + ) def _register_servers(self, server_infos: "ZMQServerInfo | dict[Any, ZMQServerInfo]"): """Register and validate server information. diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index c399c51a..c6176325 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -37,9 +37,10 @@ logger = get_logger(__name__) -# Idle sockets kept per (owner, address) bucket by every ZMQSocketPool. See ZMQSocketPool -# for what the cap does and does not bound. -TQ_SOCKET_POOL_SIZE = int(os.environ.get("TQ_SOCKET_POOL_SIZE", 64)) +# Idle sockets kept per (owner, address) bucket by every ZMQSocketPool. Small because the +# cap multiplies by peer count: one socket per peer already avoids the repeated handshake, +# and a large value only helps when one peer sees concurrent requests. +TQ_SOCKET_POOL_SIZE = int(os.environ.get("TQ_SOCKET_POOL_SIZE", 8)) bytestr: TypeAlias = bytes | bytearray | memoryview From ab3e45c442aa06bbe3ae376646f94ec500150031 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 9 Sep 2026 13:15:36 +0800 Subject: [PATCH 23/28] [fix] Stop the metrics collector accumulating one socket per storage unit Pooling these queries made the collector hold a socket per unit for as long as the process ran. Sockets are bucketed per address and only swept when their owner ends, and the collector is a permanent daemon thread, so nothing ever retired them: a few thousand units meant a few thousand resident sockets. They land on the controller's context, which never raises MAX_SOCKETS from libzmq's default of 1023. Past that a lease fails with EMFILE, so the units past the limit lose their metrics every cycle, and the controller itself cannot open a socket either -- a monitoring change taking out the control plane. TQ_SOCKET_POOL_SIZE does not bound this. It caps sockets per (owner, address) and every bucket here holds exactly one; the count comes from the number of addresses. Close the socket after each query, as the notify path already does. Collection walks every unit once per cycle, so a kept socket is reused only a cycle later while occupying budget for the whole walk. Resident sockets go from one per unit to zero, and the per-cycle handshake is immaterial against a 10s interval. Signed-off-by: OutstanderWang --- tests/test_metrics.py | 14 ++++++++------ transfer_queue/metrics.py | 4 ++++ 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 62e9c25b..a8ac3e2b 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -235,12 +235,12 @@ def test_missing_context_is_reported(self): with pytest.raises(RuntimeError, match="without a ZMQ context"): exporter._get_socket_pool() - def test_collector_reuses_one_socket_per_storage_unit(self): - """The collector runs on a plain thread with no event loop, and must still reuse. + def test_collector_parks_no_socket_between_queries(self): + """A queried unit must leave nothing in the pool. - Its pool keys by thread rather than by loop, and the collect loop is one long-lived - daemon thread, so every cycle should reach a storage unit over the same connection. - Counted from the storage unit, which sees one ZMQ identity per socket dialled. + Collection walks every unit once per cycle, so a parked socket is reused only a + cycle later while holding a slot in the controller context's budget for the whole + walk. At a few thousand units that budget is what runs out first. """ identities: set[bytes] = set() ctx_peer = zmq.Context() @@ -272,7 +272,9 @@ def serve(): exporter = TQMetricsExporter(zmq_context=ctx) for _ in range(3): assert exporter._query_storage_unit(su_info, "storage_0") == {} - assert len(identities) == 1, "each collection cycle opened its own socket" + # A parked socket would be reused, so a fresh identity per query is the + # externally visible proof that nothing was kept. + assert len(identities) == 3, "a queried unit left its socket in the pool" finally: running = False server.join(timeout=2.0) diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index bb475561..656adc54 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -408,6 +408,10 @@ def _query_storage_unit(self, su_info: ZMQServerInfo, su_id: str) -> dict[str, A sock.send_multipart(request_msg.serialize()) response_frames = sock.recv_multipart(copy=False) response_msg = ZMQMessage.deserialize(response_frames) + # Closed rather than parked: collection walks every unit once per cycle, so a + # kept socket is reused only a cycle later while occupying the controller + # context's budget for the whole walk -- one per unit, at any scale. + sock.close(linger=0) if response_msg.request_type == ZMQRequestType.METRICS_RESPONSE: return response_msg.body return None From b6ec5ef6144a2dc945310a7e69b9f8c0b26594ee Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 9 Sep 2026 17:57:34 +0800 Subject: [PATCH 24/28] [perf, fix] Wait for a pooled socket instead of opening one past the cap maxsize bounded only what the pool parked, not what it opened: a burst of N concurrent requests opened N sockets and closed all but maxsize on return. Peak socket count therefore tracked real concurrency, which is what makes the budget unpredictable -- maxsize times peer count was a number nothing enforced, and past ZMQ_MAX_SOCKETS a lease fails with EMFILE. Async callers now use alease, which takes a permit per (owner, address) before leasing. Concurrency past the cap waits for a socket to come back rather than opening another, so sockets in flight follow configuration. It is also faster, since waiting costs less than the handshake it replaces: locally, 300-way concurrency ran 48ms against 71ms, with peak sockets 8 against 300. Permits are keyed like the buckets, so a fan-out across peers is not serialized by a cap meant to bound one peer's concurrency. A permit lives until its body ends, so leasing inside a lease can wait on one the task -- or a sibling mid-cycle -- already holds, and that hangs with no timeout. Any nesting now raises, including across pools: separate semaphores do not break a cycle, only the ordering does. An RPC's reply is what frees its permit, so a second RPC belongs after the first returns; notify already works that way, running once the puts it reports have completed. The synchronous lessee keeps lease(): the metrics collector issues one request at a time, so it never queues and cannot await. Signed-off-by: OutstanderWang --- tests/test_zmq_shared_context.py | 6 +- tests/test_zmq_socket_pool.py | 118 ++++++++++++++++++++++++ transfer_queue/storage/managers/base.py | 2 +- transfer_queue/utils/zmq_utils.py | 69 ++++++++++++-- 4 files changed, 184 insertions(+), 11 deletions(-) diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 75a8abd1..127e4d56 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -140,8 +140,10 @@ def _spy_create(ctx, *args, **kwargs): 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 + # Every socket must come from the SAME context, and it must be the client's. The count + # is bounded by the pool's cap rather than by the call count, since concurrency past it + # waits for a socket instead of opening one. + assert seen_contexts, "no socket was created at all" 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 diff --git a/tests/test_zmq_socket_pool.py b/tests/test_zmq_socket_pool.py index 9f56c2a0..6af2ebcf 100644 --- a/tests/test_zmq_socket_pool.py +++ b/tests/test_zmq_socket_pool.py @@ -245,6 +245,124 @@ async def test_burst_beyond_pool_size_is_served(peer): ctx.destroy(linger=0) +@pytest.mark.asyncio +async def test_alease_caps_sockets_in_flight(peer): + """alease waits for a socket rather than opening one past maxsize. + + This is what makes the context's socket budget a function of configuration: without it + the peak equals real concurrency, which at a few thousand peers exhausts ZMQ_MAX_SOCKETS + and fails a lease with EMFILE. + """ + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=4) + + async def one(i): + async with pool.alease(peer.info) as sock: + await sock.send_multipart([f"c{i}".encode()]) + return (await sock.recv_multipart())[0] + + results = await asyncio.gather(*[one(i) for i in range(40)]) + + assert sorted(results) == sorted(f"reply-to-c{i}".encode() for i in range(40)), "a request was dropped" + assert peer.callers == 4, "concurrency past maxsize opened extra sockets instead of waiting" + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_alease_permits_are_per_address(peer): + """A busy peer must not stall requests to a different one. + + Permits are keyed like the buckets, so a fan-out across N peers is not serialized by a + cap meant to bound one peer's concurrency. + """ + other = _Peer(peer_id="peer_1", tag=b"other-") + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=1) + + async def hold(): + async with pool.alease(peer.info) as sock: + await sock.send_multipart([b"slow"]) + await sock.recv_multipart() + await asyncio.sleep(0.3) # keeps peer's only permit + + async def other_peer(): + async with pool.alease(other.info) as sock: + await sock.send_multipart([b"req"]) + return (await sock.recv_multipart())[0] + + # Would time out if one peer's permit gated the other. + _, reply = await asyncio.wait_for(asyncio.gather(hold(), other_peer()), timeout=5) + assert reply == b"other-req" + + pool.close() + ctx.destroy(linger=0) + other.stop() + + +@pytest.mark.asyncio +async def test_nested_alease_raises_instead_of_hanging(peer): + """A lease inside another lease must fail loudly rather than wait on a held permit. + + A permit is released when its body ends, so nesting can wait on one the same task -- + or a sibling mid-cycle -- already holds. That hangs with no timeout and no error, the + failure mode hardest to diagnose in production. + """ + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=8) + + with pytest.raises(RuntimeError, match="already holding"): + async with pool.alease(peer.info): + async with pool.alease(peer.info): + pass + + # The outer permit is returned, so the pool still works. + async with pool.alease(peer.info) as sock: + await sock.send_multipart([b"after"]) + assert (await sock.recv_multipart())[0] == b"reply-to-after" + + pool.close() + ctx.destroy(linger=0) + + +@pytest.mark.asyncio +async def test_nesting_is_rejected_across_pools(peer): + """Two pools do not make nesting safe: a cycle between them deadlocks just as well. + + Each task holds what the other waits for, and separate semaphores do not break that. + """ + other = _Peer(peer_id="peer_1", tag=b"other-") + ctx = zmq.asyncio.Context() + first = ZMQSocketPool(ctx, "first", "put_get_socket", maxsize=1) + second = ZMQSocketPool(ctx, "second", "put_get_socket", maxsize=1) + + with pytest.raises(RuntimeError, match="already holding"): + async with first.alease(peer.info): + async with second.alease(other.info): + pass + + first.close() + second.close() + ctx.destroy(linger=0) + other.stop() + + +@pytest.mark.asyncio +async def test_consecutive_leases_are_not_nesting(peer): + """Leasing again after the previous body ended is the supported shape.""" + ctx = zmq.asyncio.Context() + pool = ZMQSocketPool(ctx, "owner", "put_get_socket", maxsize=2) + + for i in range(3): + async with pool.alease(peer.info) as sock: + await sock.send_multipart([f"s{i}".encode()]) + assert (await sock.recv_multipart())[0] == f"reply-to-s{i}".encode() + + pool.close() + ctx.destroy(linger=0) + + def test_pool_size_comes_from_the_env_var_at_construction(): """Every role's pool takes TQ_SOCKET_POOL_SIZE, resolved per pool rather than at import. diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 34b75026..6ea17734 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -276,7 +276,7 @@ async def _notify_and_wait(self, request_msg: list) -> None: """Send a data status notification to the controller and block until ACK is received.""" # Acquiring the lease sits outside the handler below: a missing socket name or a dead # context is a configuration/lifecycle fault the caller must see, not a slow ACK. - with self.notify_pool.lease(self.controller_info) as sock: + async with self.notify_pool.alease(self.controller_info) as sock: try: await sock.send_multipart(request_msg) logger.debug( diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index c6176325..2d21fab3 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -19,8 +19,9 @@ import socket import threading import time -from collections.abc import Iterator, Sequence -from contextlib import contextmanager +from collections.abc import AsyncIterator, Iterator, Sequence +from contextlib import asynccontextmanager, contextmanager +from contextvars import ContextVar from dataclasses import dataclass from functools import wraps from typing import Any, Callable, TypeAlias @@ -382,6 +383,11 @@ def _owner_finished(owner: Any) -> bool: return owner.is_closed() +# Addresses whose permit the current task already holds. A ContextVar because asyncio +# copies the context per task, so siblings under one gather do not see each other's. +_held_permits: ContextVar[frozenset[str]] = ContextVar("_held_permits", default=frozenset()) + + class ZMQSocketPool: """Lends connected DEALER sockets for one request scenario, reusing them across calls. @@ -438,6 +444,10 @@ def __init__( # restarted under the same id at a new address must not get a socket wired to the old. self._idle: dict[Any, dict[str, list[zmq.Socket]]] = {} self._lock = threading.Lock() + # One semaphore per (owner, address), so a caller waits for a socket to come back + # instead of opening an extra one. Async only: an async lease can await, while the + # synchronous lessee (metrics) issues one request at a time and never queues. + self._permits: dict[Any, dict[str, asyncio.Semaphore]] = {} # A ROUTER silently drops a second peer claiming an identity it already has, and # owner_id alone repeats across nodes because client ids are pid-derived. self._identity_prefix = f"{owner_id}_{uuid4().hex[:8]}" @@ -450,12 +460,11 @@ def lease(self, peer: ZMQServerInfo) -> Iterator[zmq.Socket]: A plain (non-async) contextmanager on purpose: ``with`` still sees exceptions and ``CancelledError`` raised across ``await``s in its body, so this one definition serves both async and synchronous callers. - """ - port = peer.ports.get(self._socket_name) - if port is None: - raise RuntimeError(f"Socket '{self._socket_name}' not configured for server '{peer.id}'") - address = format_zmq_address(peer.ip, port) + Concurrency beyond ``maxsize`` opens extra sockets here rather than waiting; async + callers that want the cap enforced use ``alease``. + """ + address = self._address(peer) sock = self._take(address) or self._connect(peer, address) try: yield sock @@ -468,6 +477,49 @@ def lease(self, peer: ZMQServerInfo) -> Iterator[zmq.Socket]: else: self._release(address, sock) + @asynccontextmanager + async def alease(self, peer: ZMQServerInfo) -> AsyncIterator[zmq.Socket]: + """Like ``lease``, but waits for a socket instead of opening one past ``maxsize``. + + This bounds sockets in flight at ``maxsize`` per (owner, address), so the context's + socket budget follows configuration rather than peak concurrency. + + A permit is held for the whole body, so a task that leases inside another lease can + wait on a permit it -- or a sibling mid-cycle -- already holds, and that hangs with + no timeout. Any nesting therefore raises, including across pools: an RPC's reply is + what releases its permit, so a second RPC belongs after the first returns. Notify + already works this way, running once the puts it reports have completed. + """ + address = self._address(peer) + held = _held_permits.get() + if held: + raise RuntimeError( + f"Lease on {address} from a task already holding {sorted(held)}. A lease " + f"keeps its permit until its body ends, so nesting one inside another can " + f"wait on a permit that is already held and hang. Complete the outer " + f"request first, then start this one." + ) + token = _held_permits.set(held | {address}) + permit = self._permit(address) + try: + async with permit: + with self.lease(peer) as sock: + yield sock + finally: + _held_permits.reset(token) + + def _address(self, peer: ZMQServerInfo) -> str: + port = peer.ports.get(self._socket_name) + if port is None: + raise RuntimeError(f"Socket '{self._socket_name}' not configured for server '{peer.id}'") + return format_zmq_address(peer.ip, port) + + def _permit(self, address: str) -> asyncio.Semaphore: + """The permit gating this (owner, address), created on first use by that owner.""" + with self._lock: + # Keyed like _idle, since a semaphore belongs to the loop that awaits it. + return self._permits.setdefault(_lease_owner(), {}).setdefault(address, asyncio.Semaphore(self._maxsize)) + def _owner_buckets(self) -> dict[str, list[zmq.Socket]]: """Buckets for the current lease owner, first evicting any owner that has finished. @@ -478,6 +530,7 @@ def _owner_buckets(self) -> dict[str, list[zmq.Socket]]: """ for owner in [o for o in self._idle if _owner_finished(o)]: self._close_all(self._idle.pop(owner)) + self._permits.pop(owner, None) return self._idle.setdefault(_lease_owner(), {}) def _take(self, address: str) -> zmq.Socket | None: @@ -580,7 +633,7 @@ async def wrapper(self, *args, **kwargs): if pool is None: raise RuntimeError("get_pool returned None") - with pool.lease(server_info) as sock: + async with pool.alease(server_info) as sock: kwargs["socket"] = sock return await func(self, *args, **kwargs) From de7a96ff3df0a8c12c656b38b3da586517d4c971 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 9 Sep 2026 19:51:32 +0800 Subject: [PATCH 25/28] [fix] Raise the controller context's socket ceiling above libzmq's default The controller never set MAX_SOCKETS, so its context kept libzmq's default of 1023. Its own two ROUTERs are nowhere near that, but the metrics exporter borrows this context to query storage units, and a controller that cannot open a socket stops answering requests at all -- a monitoring cost taking out the control plane. Set it to 4096, clamped to this build's ZMQ_SOCKET_LIMIT, before the first socket is opened: libzmq applies the ceiling at socket creation. Metrics now closes each query's socket, so today's peak is one; the headroom is for a collector that fans out instead of walking units in sequence. Signed-off-by: OutstanderWang --- tests/test_controller.py | 23 ++++++++++++++++++++++- transfer_queue/controller.py | 8 ++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/tests/test_controller.py b/tests/test_controller.py index 8a735f6a..32c1603a 100644 --- a/tests/test_controller.py +++ b/tests/test_controller.py @@ -21,7 +21,7 @@ import torch import zmq -from transfer_queue.controller import TransferQueueController +from transfer_queue.controller import TQ_CONTROLLER_ZMQ_MAX_SOCKETS, TransferQueueController from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, create_zmq_socket # Set up logging @@ -1465,6 +1465,27 @@ def test_controller_checkpoint_load_nonexistent_file(self, ray_setup, tmp_path): print("✓ load_checkpoint raises on missing file") +class TestControllerSocketBudget: + """The controller's context must hold more sockets than libzmq's default 1023. + + The metrics exporter borrows this context to query storage units, so the budget has to + cover the fleet rather than just the controller's own two ROUTERs. Past the ceiling a + socket cannot be opened at all, which would take out the control plane too. + """ + + def test_context_ceiling_is_raised_above_the_libzmq_default(self, ray_setup): + controller = TransferQueueController.remote() + try: + # Read it from inside the actor: the ceiling is only meaningful on the context + # the controller actually built. + budget = ray.get(controller.__ray_call__.remote(lambda self: self.zmq_context.get(zmq.MAX_SOCKETS))) + finally: + ray.kill(controller) + + assert budget > 1023, "the controller kept libzmq's default socket ceiling" + assert budget == TQ_CONTROLLER_ZMQ_MAX_SOCKETS + + class TestTransferQueueControllerBadRequests: """The request loop must survive requests it cannot decode, does not handle, or fails on. diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index 188c011c..e8f25f84 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -56,6 +56,11 @@ TQ_CONTROLLER_GET_METADATA_TIMEOUT = int(os.environ.get("TQ_CONTROLLER_GET_METADATA_TIMEOUT", 1)) TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL = int(os.environ.get("TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL", 5)) +# Socket ceiling for the controller's context, above libzmq's default of 1023. The two +# ROUTERs are all it binds, but the metrics exporter borrows this context to query storage +# units, so the budget has to cover a large fleet. Raising it needs file descriptors to +# match (``ulimit -n``). +TQ_CONTROLLER_ZMQ_MAX_SOCKETS = int(os.environ.get("TQ_CONTROLLER_ZMQ_MAX_SOCKETS", 4096)) # Sample pre-allocation for StreamingDataLoader compatibility. # By pre-allocating sample indices (typically global_batch_size), consumers can accurately @@ -1725,6 +1730,9 @@ def kv_retrieve_keys( def _init_zmq_socket(self): """Initialize ZMQ sockets for communication.""" self.zmq_context = zmq.Context() + # Before any socket is opened on it: libzmq applies MAX_SOCKETS at socket creation. + socket_limit = self.zmq_context.get(zmq.SOCKET_LIMIT) + self.zmq_context.set(zmq.MAX_SOCKETS, min(TQ_CONTROLLER_ZMQ_MAX_SOCKETS, socket_limit)) self._node_ip = get_node_ip_address() while True: From 5e7cc3ebfa0df98af89ad9492d940e9bd6b01a89 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 9 Sep 2026 19:51:45 +0800 Subject: [PATCH 26/28] [fix] Lower the default socket pool size to 4 With alease the cap also bounds sockets in flight to one peer, so it multiplies by peer count against a fixed context budget. At 8 the worst case for two thousand storage units was 16000 sockets against a budget of 8192 -- reachable, not theoretical: a single peer under 50-way concurrency does saturate its cap. Four keeps that product inside the budget (8004) while still absorbing the overlap a peer actually sees, since one put or get sends a single merged request per unit and concurrency there comes only from operations overlapping. Signed-off-by: OutstanderWang --- transfer_queue/utils/zmq_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 2d21fab3..449a39c9 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -38,10 +38,10 @@ logger = get_logger(__name__) -# Idle sockets kept per (owner, address) bucket by every ZMQSocketPool. Small because the -# cap multiplies by peer count: one socket per peer already avoids the repeated handshake, -# and a large value only helps when one peer sees concurrent requests. -TQ_SOCKET_POOL_SIZE = int(os.environ.get("TQ_SOCKET_POOL_SIZE", 8)) +# Idle sockets kept per (owner, address) bucket by every ZMQSocketPool, and with alease the +# sockets in flight to one peer. Small because it multiplies by peer count: one socket per +# peer already avoids the repeated handshake, and 4 x 2000 units still fits the budget. +TQ_SOCKET_POOL_SIZE = int(os.environ.get("TQ_SOCKET_POOL_SIZE", 4)) bytestr: TypeAlias = bytes | bytearray | memoryview From 4b75fba65202a575bcb6b9778ad0bfa8c9717a92 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 9 Sep 2026 21:38:06 +0800 Subject: [PATCH 27/28] [refactor] Trim the comments added since the last pass One block ran past the four-line ceiling and seven more said something the code or a nearby docstring already says. Comments only. Two were duplication introduced by earlier edits: the permit keying rationale now sits on the field declaration rather than repeated at its accessor, and the per-role isolation argument belongs to ZMQSocketPool's docstring rather than to each decorator that picks a pool. One narrated what the code plainly does -- naming a daemon thread that the thread's own name and target already give. What is left is the part a reader cannot infer: why the pre-context checks come first, why ZMQ_SOCKET_LIMIT forces the last check below allocation, why a late ACK forces a close, why the pool cap is resolved at construction rather than at import, and why owner_id alone cannot make a ZMQ identity unique. Signed-off-by: OutstanderWang --- transfer_queue/controller.py | 8 ++------ transfer_queue/metrics.py | 5 ++--- transfer_queue/storage/managers/base.py | 5 ++--- .../storage/managers/simple_storage_manager.py | 2 -- transfer_queue/utils/zmq_utils.py | 14 +++++--------- 5 files changed, 11 insertions(+), 23 deletions(-) diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py index e8f25f84..b45a85a6 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -56,10 +56,8 @@ TQ_CONTROLLER_GET_METADATA_TIMEOUT = int(os.environ.get("TQ_CONTROLLER_GET_METADATA_TIMEOUT", 1)) TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL = int(os.environ.get("TQ_CONTROLLER_GET_METADATA_CHECK_INTERVAL", 5)) -# Socket ceiling for the controller's context, above libzmq's default of 1023. The two -# ROUTERs are all it binds, but the metrics exporter borrows this context to query storage -# units, so the budget has to cover a large fleet. Raising it needs file descriptors to -# match (``ulimit -n``). +# Above libzmq's default of 1023: the metrics exporter borrows this context to query +# storage units, so the budget covers a fleet rather than the two ROUTERs bound here. TQ_CONTROLLER_ZMQ_MAX_SOCKETS = int(os.environ.get("TQ_CONTROLLER_ZMQ_MAX_SOCKETS", 4096)) # Sample pre-allocation for StreamingDataLoader compatibility. @@ -2348,8 +2346,6 @@ def start_metrics(self, port: int = 0) -> str: # Lend the controller's context rather than let the exporter build a second one. self._metrics = TQMetricsExporter(zmq_context=self.zmq_context) self._metrics_endpoint = self._metrics.start(node_ip=self._node_ip, port=port) - # Launch a daemon thread that periodically pushes controller state - # snapshots to the exporter, keeping them process-isolated. self._metrics_snapshot_thread = Thread( target=self._metrics_snapshot_loop, name="TQMetricsSnapshotThread", diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index 656adc54..492d4ae6 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -408,9 +408,8 @@ def _query_storage_unit(self, su_info: ZMQServerInfo, su_id: str) -> dict[str, A sock.send_multipart(request_msg.serialize()) response_frames = sock.recv_multipart(copy=False) response_msg = ZMQMessage.deserialize(response_frames) - # Closed rather than parked: collection walks every unit once per cycle, so a - # kept socket is reused only a cycle later while occupying the controller - # context's budget for the whole walk -- one per unit, at any scale. + # Closed rather than parked: collection walks every unit once per cycle, so + # a kept socket holds budget for the whole walk to save one handshake. sock.close(linger=0) if response_msg.request_type == ZMQRequestType.METRICS_RESPONSE: return response_msg.body diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 6ea17734..28211355 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -399,9 +399,8 @@ def close(self) -> None: else: logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.") - # This manager always owns its notify pool, even over a borrowed context. Only after - # the notify thread is gone, since Socket.close() is not thread-safe and that thread - # holds the leases; linger=0 means this cannot hang. + # This manager always owns its notify pool, even over a borrowed context. Only once + # the notify thread is gone: Socket.close() is not thread-safe and it holds the leases. if notify_thread_stopped: self.notify_pool.close() diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 3780bc9b..b986701a 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -50,8 +50,6 @@ # Pre-bound decorator for storage-unit socket operations. with_storage_unit_socket = with_zmq_socket( get_peer=lambda self, target: self.storage_unit_infos[target], - # Storage RPC has its own pool, separate from the notify pool on the same context: the - # two dial different peers with different timeouts and could never share a socket. get_pool=lambda self: self.storage_rpc_pool, resolve_target=lambda args, kwargs: kwargs.get("target_storage_unit"), ) diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 449a39c9..3e4dacc5 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -38,9 +38,8 @@ logger = get_logger(__name__) -# Idle sockets kept per (owner, address) bucket by every ZMQSocketPool, and with alease the -# sockets in flight to one peer. Small because it multiplies by peer count: one socket per -# peer already avoids the repeated handshake, and 4 x 2000 units still fits the budget. +# Cap for every pool; see ZMQSocketPool for what it bounds. Small because it multiplies by +# peer count, and one socket per peer already avoids the repeated handshake. TQ_SOCKET_POOL_SIZE = int(os.environ.get("TQ_SOCKET_POOL_SIZE", 4)) bytestr: TypeAlias = bytes | bytearray | memoryview @@ -444,9 +443,8 @@ def __init__( # restarted under the same id at a new address must not get a socket wired to the old. self._idle: dict[Any, dict[str, list[zmq.Socket]]] = {} self._lock = threading.Lock() - # One semaphore per (owner, address), so a caller waits for a socket to come back - # instead of opening an extra one. Async only: an async lease can await, while the - # synchronous lessee (metrics) issues one request at a time and never queues. + # One semaphore per (owner, address), keyed like _idle because a semaphore belongs + # to the loop that awaits it. Used by alease only; see there. self._permits: dict[Any, dict[str, asyncio.Semaphore]] = {} # A ROUTER silently drops a second peer claiming an identity it already has, and # owner_id alone repeats across nodes because client ids are pid-derived. @@ -470,8 +468,7 @@ def lease(self, peer: ZMQServerInfo) -> Iterator[zmq.Socket]: yield sock except BaseException: # Poisoned: the request may already be on the wire, so its reply could still - # arrive and the next lessee would read it as its own. Timeouts and - # cancellation count -- asyncio.gather cancels siblings on the first failure. + # arrive and the next lessee would read it as its own. Cancellation counts too. sock.close(linger=0) raise else: @@ -517,7 +514,6 @@ def _address(self, peer: ZMQServerInfo) -> str: def _permit(self, address: str) -> asyncio.Semaphore: """The permit gating this (owner, address), created on first use by that owner.""" with self._lock: - # Keyed like _idle, since a semaphore belongs to the loop that awaits it. return self._permits.setdefault(_lease_owner(), {}).setdefault(address, asyncio.Semaphore(self._maxsize)) def _owner_buckets(self) -> dict[str, list[zmq.Socket]]: From 56f2e7df97ba018161a3993d6bda57cf6fe4f7ce Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Thu, 10 Sep 2026 14:44:49 +0800 Subject: [PATCH 28/28] [fix] Keep the storage identity allowlist through the rebase Rebasing onto main replayed this branch over the identity filter added in #168, and the pool's changes to the same region dropped the prefixes it filters on. Restore them: the storage proxy rejects any identity lacking one, so without these constants the filter has nothing to match and the module does not import. The pooled sockets already satisfy the filter -- the storage manager pool passes storage_manager_id and the metrics pool now derives its owner id from METRICS_COLLECTOR_IDENTITY_PREFIX rather than repeating the literal, so a change to the prefix reaches both ends. Signed-off-by: OutstanderWang --- transfer_queue/metrics.py | 4 +++- transfer_queue/storage/managers/base.py | 1 + transfer_queue/utils/zmq_utils.py | 9 +++++++++ 3 files changed, 13 insertions(+), 1 deletion(-) diff --git a/transfer_queue/metrics.py b/transfer_queue/metrics.py index 492d4ae6..12914585 100644 --- a/transfer_queue/metrics.py +++ b/transfer_queue/metrics.py @@ -389,7 +389,9 @@ def _get_socket_pool(self) -> ZMQSocketPool: ) self._zmq_socket_pool = ZMQSocketPool( self._zmq_ctx, - "metrics_collector", + # The storage proxy drops identities without this prefix, and the pool + # builds each socket's identity from the owner id. + METRICS_COLLECTOR_IDENTITY_PREFIX.rstrip("_"), "put_get_socket", timeout=TQ_METRICS_STORAGE_TIMEOUT, ) diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 28211355..369f800e 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -37,6 +37,7 @@ from transfer_queue.storage.clients.base import StorageClientFactory from transfer_queue.utils.logging_utils import get_logger from transfer_queue.utils.zmq_utils import ( + STORAGE_MANAGER_IDENTITY_PREFIX, ZMQMessage, ZMQRequestType, ZMQServerInfo, diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index 3e4dacc5..a94cd3ab 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -38,6 +38,15 @@ logger = get_logger(__name__) +# Identity prefixes of the peers allowed to reach a storage unit. The storage proxy drops +# anything else, so an identity built without these prefixes is silently unreachable. +STORAGE_MANAGER_IDENTITY_PREFIX = "TQ_STORAGE_" +METRICS_COLLECTOR_IDENTITY_PREFIX = "metrics_collector_" +STORAGE_CLIENT_IDENTITY_PREFIXES = ( + STORAGE_MANAGER_IDENTITY_PREFIX.encode(), + METRICS_COLLECTOR_IDENTITY_PREFIX.encode(), +) + # Cap for every pool; see ZMQSocketPool for what it bounds. Small because it multiplies by # peer count, and one socket per peer already avoids the repeated handshake. TQ_SOCKET_POOL_SIZE = int(os.environ.get("TQ_SOCKET_POOL_SIZE", 4))