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_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/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/tests/test_metrics.py b/tests/test_metrics.py index 15c43882..a8ac3e2b 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -16,12 +16,17 @@ """Unit tests for the Prometheus metrics exporter (transfer_queue.metrics).""" import time -from unittest.mock import MagicMock +from threading import Thread +from unittest.mock import MagicMock, patch import pytest try: + 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): @@ -207,6 +212,77 @@ 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) + with patch("zmq.Context") as minted: + exporter._get_socket_pool() + minted.assert_not_called() + 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() + + def test_collector_parks_no_socket_between_queries(self): + """A queried unit must leave nothing in the pool. + + 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() + 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") == {} + # 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) + ctx.destroy(linger=0) + router.close(linger=0) + ctx_peer.term() + + class TestStorageMetricsCollection: def test_collect_with_no_storage_units(self): """No storage units registered — collect should be a no-op.""" diff --git a/tests/test_zmq_shared_context.py b/tests/test_zmq_shared_context.py index 2bd9a6c7..127e4d56 100644 --- a/tests/test_zmq_shared_context.py +++ b/tests/test_zmq_shared_context.py @@ -17,11 +17,13 @@ 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 +from contextlib import contextmanager from threading import Thread from unittest.mock import patch @@ -138,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 @@ -168,6 +172,83 @@ 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_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_SOCKET_POOL_SIZE", bad): + with _no_context_left_open() as created: + 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, + ) + 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): client = AsyncTransferQueueClient( client_id="client_simple_storage_context", @@ -188,6 +269,105 @@ def test_simple_storage_borrows_client_context(echo_controller): client.close() +def test_each_scenario_gets_its_own_pool(echo_controller): + """Controller RPC, storage RPC and notify must each hold a separate pool. + + 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_scenario_pools", + 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, + ) + + 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" + + manager.close() + 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(). + + 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", @@ -266,8 +446,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, like SimpleStorage does. + + With no client it creates its own context. Either way it builds its own notify pool. + """ class Borrower(StorageManager): def _connect_to_controller(self): @@ -282,7 +465,7 @@ async def get_data(self, *args, **kwargs): async def clear_data(self, *args, **kwargs): return None - return Borrower(None, {}, zmq_context=zmq_context) + 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): @@ -295,7 +478,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 +497,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 +509,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..6af2ebcf --- /dev/null +++ b/tests/test_zmq_socket_pool.py @@ -0,0 +1,390 @@ +# 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'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 +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. + +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. +""" + +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 + + +class _Peer: + """A ROUTER that echoes one reply per request, recording who dialled it. + + ``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-"): + 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_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 + self.running = True + 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) + while self.running: + 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. + 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, self._tag + 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() + + +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] + + +@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. + """ + # 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", "put_get_socket", timeout=1) + try: + with pytest.raises(zmq.error.Again): + await _round_trip(pool, peer.info, b"first") + + # 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) + peer.stop() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["exception", "cancellation"]) +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) 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 + + 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): + """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", "put_get_socket") + + async def lease_twice(tag): + # Twice per loop, so a socket IS reused within a loop -- which is what makes the + # cross-loop comparison meaningful rather than trivially true. + for i in range(2): + assert await _round_trip(pool, peer.info, f"{tag}{i}".encode()) == f"reply-to-{tag}{i}".encode() + + for tag in ("a", "b", "c"): + asyncio.run(lease_twice(tag)) + + assert peer.callers == 3, "one socket per loop, reused within it" + + pool.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. + + 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", "put_get_socket") + 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_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 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_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. + + 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 01e50b85..0f7f13f1 100644 --- a/transfer_queue/client.py +++ b/transfer_queue/client.py @@ -29,9 +29,11 @@ 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, + ZMQSocketPool, with_zmq_socket, ) @@ -49,10 +51,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.controller_rpc_pool, ) @@ -61,6 +61,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__( @@ -91,12 +96,17 @@ 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. + # 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}") - self.zmq_context = zmq.asyncio.Context(io_threads=io_threads) + 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_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." + ) max_sockets = zmq_max_sockets explicitly_requested = max_sockets is not None @@ -109,26 +119,44 @@ 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 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}") 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 + # 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, + "request_handle_socket", + ) # 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 @@ -159,7 +187,8 @@ def initialize_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. + 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: @@ -236,33 +265,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, @@ -302,10 +333,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 @@ -387,30 +419,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: @@ -467,16 +503,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: @@ -678,10 +716,10 @@ async def async_get_consumption_status( Example: >>> # Get consumption status - >>> global_index, consumption_status = asyncio.run(client.async_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}") """ @@ -725,10 +763,10 @@ async def async_get_production_status( Example: >>> # Get production status - >>> global_index, production_status = asyncio.run(client.async_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: @@ -766,10 +804,10 @@ async def async_check_consumption_status( Example: >>> # Check if all samples have been consumed - >>> is_consumed = asyncio.run(client.async_check_consumption_status( + >>> is_consumed = await client.async_check_consumption_status( ... task_name="generate_sequences", ... partition_id="train_0" - ... )) + ... ) >>> print(f"All samples consumed: {is_consumed}") """ @@ -802,10 +840,10 @@ async def async_check_production_status( Example: >>> # Check if all samples are ready for consumption - >>> is_ready = asyncio.run(client.async_check_production_status( + >>> 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( @@ -842,10 +880,10 @@ async def async_reset_consumption( Example: >>> # Reset consumption for train task to re-train on same data - >>> success = asyncio.run(client.async_reset_consumption( + >>> success = await client.async_reset_consumption( ... partition_id="train_0", ... task_name="train" - ... )) + ... ) >>> print(f"Reset successful: {success}") """ body = {"partition_id": partition_id} @@ -879,7 +917,7 @@ 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()) + >>> partition_ids = await client.get_partition_list() >>> print(f"Available partitions: {partition_ids}") """ try: @@ -1053,6 +1091,8 @@ def close(self) -> None: ) return try: + # Close pooled sockets before the context that owns them. + 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/controller.py b/transfer_queue/controller.py index ef71bb4c..b45a85a6 100644 --- a/transfer_queue/controller.py +++ b/transfer_queue/controller.py @@ -56,6 +56,9 @@ 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)) +# 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. # By pre-allocating sample indices (typically global_batch_size), consumers can accurately @@ -1725,6 +1728,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: @@ -2337,10 +2343,9 @@ 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. 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 e90e9c59..12914585 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__) @@ -68,13 +66,22 @@ 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_sockets: dict[str, zmq.Socket] = {} + 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() self._known_consumption_labels: set[tuple[str, str]] = set() @@ -372,49 +379,50 @@ 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: - 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 + def _get_socket_pool(self) -> ZMQSocketPool: + """Return the lazily-created socket pool for storage-unit queries.""" + if self._zmq_socket_pool is None: + 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, + # 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, + ) + 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) 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) + # 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 + 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..369f800e 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -41,6 +41,7 @@ ZMQMessage, ZMQRequestType, ZMQServerInfo, + ZMQSocketPool, create_zmq_socket, ) @@ -87,6 +88,9 @@ 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 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() # Dedicated asyncio loop for ZMQ notify traffic, isolated from the caller's loop @@ -271,56 +275,39 @@ 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. + async with self.notify_pool.alease(self.controller_info) 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: + # 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) @abstractmethod async def put_data( @@ -389,6 +376,11 @@ 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 super().__init__() runs, and __del__ calls + # close() anyway; return rather than bury that error under an AttributeError. + if not hasattr(self, "notify_pool"): + return + if self.controller_handshake_socket: try: if not self.controller_handshake_socket.closed: @@ -408,13 +400,16 @@ 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 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() + 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: - # linger=0 force-closes sockets left by an interrupted request, so this - # cannot hang on term(). self.zmq_context.destroy(linger=0) else: logger.warning( diff --git a/transfer_queue/storage/managers/simple_storage_manager.py b/transfer_queue/storage/managers/simple_storage_manager.py index 88340006..b986701a 100644 --- a/transfer_queue/storage/managers/simple_storage_manager.py +++ b/transfer_queue/storage/managers/simple_storage_manager.py @@ -32,9 +32,11 @@ 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, + ZMQSocketPool, with_zmq_socket, ) @@ -47,14 +49,9 @@ # 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, + 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, ) @@ -80,6 +77,13 @@ def __init__( zmq_context: zmq.asyncio.Context | None = None, ): 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) @@ -97,6 +101,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. @@ -659,4 +685,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. 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() diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py index c4b04ce7..a94cd3ab 100644 --- a/transfer_queue/utils/zmq_utils.py +++ b/transfer_queue/utils/zmq_utils.py @@ -13,9 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio +import itertools +import os import socket +import threading import time -from collections.abc import Sequence +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 @@ -41,6 +47,9 @@ 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)) bytestr: TypeAlias = bytes | bytearray | memoryview @@ -360,47 +369,263 @@ 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() + + +# 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. + + 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: + 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, + socket_name: str, + *, + timeout: int | None = None, + maxsize: int | None = None, + ): + """ + 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, 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. + 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 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() + # 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. + self._identity_prefix = f"{owner_id}_{uuid4().hex[:8]}" + self._counter = itertools.count() + + @contextmanager + 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. + + 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 + 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. Cancellation counts too. + sock.close(linger=0) + raise + 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: + 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. + + 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)) + self._permits.pop(owner, None) + return self._idle.setdefault(_lease_owner(), {}) + + def _take(self, address: str) -> zmq.Socket | None: + """Pop a live idle socket for *address*, discarding any found closed.""" + with self._lock: + bucket = self._owner_buckets().get(address) + while bucket: + sock = bucket.pop() + if not sock.closed: + return sock + return None + + 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. + if not sock.closed: + 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) -> 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 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 + # 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. 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[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: + 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. + """Create a reusable async decorator that injects a pooled request socket. - Lifecycle: get owner's shared context -> create/connect socket -> inject -> close 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 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. + The socket name and timeout come from the pool, which serves one request scenario. 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 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): @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 +634,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) + async with pool.alease(server_info) 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