[perf, fix] Reuse a long-lived ZMQ context instead of creating one per call by OutstanderWang · Pull Request #145 · Ascend/TransferQueue

OutstanderWang · 2026-07-31T06:45:36Z

[perf,fix] Reuse a long-lived ZMQ context pool instead of creating one per call

Motivation

Every client→controller RPC and every SimpleStorage manager→storage-unit request went through the with_zmq_socket decorator, which created a brand-new zmq.asyncio.Context per call, opened a socket on it, and then term()-ed the context in the finally block. A ZMQ context owns a native I/O-thread pool and a set of internal signaler file descriptors; standing one up and tearing it down on every request is both wasteful and, under concurrency, unsafe. This surfaced as intermittent Bad file descriptor errors, SIGABRT crashes, and occasional hangs inside ctx.term() on the hot data path.

The goal of this change is to make each client own one long-lived context that is created once, reused across all of its RPCs, and terminated exactly once at shutdown — while keeping sockets per-request (ZMQ sockets are not thread-safe).

What needs to be improved

Per-call context churn. with_zmq_socket did context = zmq.asyncio.Context() on entry and context.term() in finally for every decorated call. Repeatedly allocating/destroying contexts recreates libzmq's internal signaler pipe FDs; under concurrent in-flight calls these FDs can collide, yielding Bad file descriptor / SIGABRT.
term() on the event loop can hang. context.term() blocks until every socket on the context is closed and all pending messages are handled. Running it in the request finally block could stall the asyncio loop if a socket lingered from an interrupted RPC.
No shared I/O-thread pool. Because each call had its own throwaway context, there was no way to size or share the native I/O-thread pool that actually moves bytes; every request paid context-startup cost.

What this changes

with_zmq_socket reuses an owner-provided context. The decorator now takes a required get_context(self) callable and does context = get_context(self) instead of constructing one. It creates and closes only the per-call DEALER socket (sock.close(linger=0)); it never creates or terminates a context. The docstring documents the invariant: contexts are thread-safe and event-loop-agnostic, so one shared context is safe even when decorated methods run on different loops/threads, as long as each socket is created and fully used within a single awaited call.
Client owns one long-lived context. AsyncTransferQueueClient.__init__ creates self.zmq_context = zmq.asyncio.Context(io_threads=…) once, sized by a new simple_storage_zmq_io_threads arg / TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS env var (default 8, validated ≥ 1). with_controller_socket binds get_context=lambda self: self.zmq_context, so all backends reuse it on the controller-RPC path.
SimpleStorage borrows the client context. initialize_storage_manager passes zmq_context=self.zmq_context to the factory only for manager_type == "SimpleStorage". AsyncSimpleStorageManager's with_storage_unit_socket binds get_context=lambda self: self.zmq_context, so both the notify path and per-call storage-unit request sockets share the client's fixed I/O-thread pool.
Ownership-aware teardown, no double-free. StorageManager records self._owns_zmq_context = zmq_context is None. A manager that created its own context tears it down with zmq_context.destroy(linger=0); a manager that borrowed the client's context does not terminate it. The client terminates its own context once in close() via destroy(linger=0) (force-closes any leaked socket so shutdown can't hang).
Other backends are unaffected. Mooncake / Yuanrong / RayStore managers receive zmq_context=None, so they keep their own independent long-lived context for the controller notify/handshake path and never touch the shared pool.

Scope

The knob simple_storage_zmq_io_threads / TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS sizes the client context's native I/O-thread pool. Note this context also serves the client→controller path for every backend, not only SimpleStorage.
KV backends (Mooncake/Yuanrong/RayStore) move their bulk data through their own SDKs, not through with_zmq_socket, so this change only affects their controller-side ZMQ traffic (via the shared decorator fix), not their data plane.

Tests

New tests/test_zmq_shared_context.py (7 cases):
- shared context is reused across concurrent calls,
- client context has a fixed I/O-thread pool,
- client rejects an invalid (<1) pool size,
- SimpleStorage borrows the client context,
- SimpleStorage does not destroy the borrowed context,
- other backends do not borrow the client context,
- close() destroys the context exactly once.
Full suite run locally on a fresh single-node Ray: 547 passed, 10 skipped, 8 errors — the 10 skips (GDR/GPU, Mooncake-CUDA) and 8 errors (Yuanrong SDK absent) are all pre-existing environment gaps in files untouched by this branch; 0 failures. The complete e2e lifecycle suite (test_core_consistency, cross-shard, production-status, reset, clear, dynamic-shape, memory-safety) passes.

ascend-robot · 2026-07-31T06:45:46Z

CLA Signature Pass

OutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍

0oshowero0 · 2026-07-31T08:29:49Z

CC @ji-huazhong

0oshowero0 · 2026-07-31T08:30:27Z

Please run the following pre-commit:

# install pre-commit
pip install pre-commit

# run the following command in your repo folder, then fix the check before committing your code
pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always

0oshowero0 · 2026-07-31T08:31:54Z

+        if self._owns_zmq_context:
+            # destroy(linger=0) force-closes any socket still open (e.g. from an interrupted
+            # request or the notify path) then terminates, so shutdown cannot hang on term().
+            self.zmq_context.destroy(linger=0)


destroy it's not thread-safe. TransferQueueClient has a dedicate thread for running async loop. Will this brings new potential problems?

Thank you for pointing out the issue. I have fixed it by only destroy the socket when it is verfied to exist.

0oshowero0 · 2026-07-31T08:34:54Z

 logger = get_logger(__name__)

 TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8))
+TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS = int(os.environ.get("TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS", 8))


If this is the config that only affects the zmq request in simple_storage_manager.py, we should not put it in client.py.

It seems that this config will also affect all the storage managers, we may need to rename the env var

It is only for the cliend. I have renamed it as TQ_CLIENT_ZMQ_IO_THREADS temporarily. Pls ckeck it and let me know if you have some better ideas about it.

0oshowero0 · 2026-07-31T08:40:41Z

+                "SimpleStorage ZMQ I/O thread pool size must be at least 1, "
+                f"got {io_threads}"
+            )
+        self.zmq_context = zmq.asyncio.Context(io_threads=io_threads)


The shared zmq context may limit the number of socket budgets. For example, on macOS, a single context can support 1023 sockets. This is a behavior change that we need to be careful. Maybe we need to expose the MAX_SOCKET setting in env var?

Now we have a shared context. Do we further need a socket pool?

I have added the TQ_CLIENT_ZMQ_MAX_SOCKETS, which is 1023 by default.

It is a good idea to further employ the socket pool, which can benifit the small payload a lot and save the port resources. Shall we consider to impelement it in another new PR.

Maybe another PR :)

0oshowero0 · 2026-07-31T08:44:53Z

+        if manager_type == "SimpleStorage":
+            create_kwargs["zmq_context"] = self.zmq_context


This may break the Factory registry mechanism. We can just pass the kwargs in create and let each registered StorageManager to decide their behaviour

I've fixed it by passing the kwargs in create as you suggested. But I do not know all these backends well, pls help check the all the backends and give me the feedbacks, thanks.

0oshowero0 · 2026-07-31T08:46:46Z

+    assert client.zmq_context.closed
+
+
+def test_other_backends_do_not_borrow_client_context(echo_controller):


Do we need this?

I have deleted it

0oshowero0 · 2026-07-31T08:49:02Z

+        self._owns_zmq_context = zmq_context is None
+        self.zmq_context = zmq_context or zmq.asyncio.Context()


Suggest align the judgement by using None check

It's aligned now.

ascend-robot · 2026-07-31T14:30:06Z

CLA Signature Guide

@OutstanderWang , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
f6cede3f style: apply ruff-format to clie...	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

…call The with_zmq_socket decorator created a brand-new zmq.asyncio.Context() per RPC call and context.term()'d it in the finally block. Under the high-concurrency agent-loop store path this churned libzmq's signaler file descriptors and crashed the worker (signaler.cpp Bad file descriptor -> SIGABRT), and could also hang on the blocking term() (aggravated by sock.close(linger=-1)). Fix: the decorator now reuses the owner's long-lived context via a required get_context callable, and only creates/closes the DEALER socket per call. The context is created once per owner and terminated once at close(). Contexts are thread-safe and event-loop-agnostic, so a single shared context is safe across loops/threads; each socket stays per-call on one loop. - zmq_utils.with_zmq_socket: add required get_context; drop per-call Context()/term(); change sock.close(linger=-1) -> linger=0. - client.AsyncTransferQueueClient: own a shared self.zmq_context; destroy(linger=0) in close(). - simple_storage_manager: feed the base StorageManager's self.zmq_context via get_context. - base.StorageManager.close(): term() -> destroy(linger=0) so a leaked socket cannot hang shutdown. - tests: add test_zmq_shared_context.py asserting concurrent RPCs reuse one context and it is closed exactly once. Microbenchmark: ~7.6x faster socket setup/teardown (~210us saved per call). Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

This reverts commit b54c4093d71348b3c82572ff18c3a35a0bb9dd2b. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

This reverts commit a8bfbd81c68226f0c679ce3673c467866d470881. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

- Guard zmq_context.destroy() against live threads. destroy() calls Socket.close() internally and is not thread-safe, but both TransferQueueClient.close() and StorageManager.close() joined their threads with a warn-only timeout and then destroyed anyway. Add a _can_destroy_zmq_context() veto that leaks the context with a loud warning instead, and document the load-bearing close() ordering. - Rename the I/O-thread knob to reflect its real scope: the context serves every backend's controller RPCs, not just SimpleStorage. TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS -> TQ_CLIENT_ZMQ_IO_THREADS, kwarg simple_storage_zmq_io_threads -> zmq_io_threads. - Expose zmq_max_sockets / TQ_CLIENT_ZMQ_MAX_SOCKETS. Sharing one context per client means all in-flight sockets share libzmq's 1023-per-context budget, where previously each call had a private one. Opt-in, validated against the build's ZMQ_SOCKET_LIMIT. - Keep StorageManagerFactory.create backend-agnostic: forward **kwargs and let each registered manager decide, instead of hard-coding "SimpleStorage" in the factory and the client. KV managers accept zmq_context and deliberately keep their own. - Align the ownership check on an explicit `is None` test. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

- Guard zmq_context.destroy() against live threads. destroy() calls Socket.close() internally and is not thread-safe, but both TransferQueueClient.close() and StorageManager.close() joined their threads with a warn-only timeout and then destroyed anyway. Add a _can_destroy_zmq_context() veto that leaks the context with a loud warning instead, and document the load-bearing close() ordering. - Rename the I/O-thread knob to reflect its real scope: the context serves every backend's controller RPCs, not just SimpleStorage. TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS -> TQ_CLIENT_ZMQ_IO_THREADS, kwarg simple_storage_zmq_io_threads -> zmq_io_threads. - Expose zmq_max_sockets / TQ_CLIENT_ZMQ_MAX_SOCKETS. Sharing one context per client means all in-flight sockets share libzmq's 1023-per-context budget, where previously each call had a private one. Opt-in, validated against the build's ZMQ_SOCKET_LIMIT. - Keep StorageManagerFactory.create backend-agnostic: forward **kwargs and let each registered manager decide, instead of hard-coding "SimpleStorage" in the factory and the client. KV managers accept zmq_context and deliberately keep their own. - Align the ownership check on an explicit `is None` test. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

ascend-robot · 2026-08-03T09:36:28Z

CLA Signature Pass

OutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍

- Guard zmq_context.destroy() against live threads. destroy() calls Socket.close() internally and is not thread-safe, but both TransferQueueClient.close() and StorageManager.close() joined their threads with a warn-only timeout and then destroyed anyway. Add a _can_destroy_zmq_context() veto that leaks the context with a loud warning instead, and document the load-bearing close() ordering. - Rename the I/O-thread knob to reflect its real scope: the context serves every backend's controller RPCs, not just SimpleStorage. TQ_SIMPLE_STORAGE_ZMQ_IO_THREADS -> TQ_CLIENT_ZMQ_IO_THREADS, kwarg simple_storage_zmq_io_threads -> zmq_io_threads. - Expose zmq_max_sockets / TQ_CLIENT_ZMQ_MAX_SOCKETS. Sharing one context per client means all in-flight sockets share libzmq's 1023-per-context budget, where previously each call had a private one. Opt-in, validated against the build's ZMQ_SOCKET_LIMIT. - Keep StorageManagerFactory.create backend-agnostic: forward **kwargs and let each registered manager decide, instead of hard-coding "SimpleStorage" in the factory and the client. KV managers accept zmq_context and deliberately keep their own. - Align the ownership check on an explicit `is None` test. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

ascend-robot · 2026-08-03T09:39:37Z

CLA Signature Pass

OutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍

The socket-ceiling knob had tests only for the kwarg, never the env var -- which is the deployment-facing path, since the kwarg needs a code edit. Exercising it end to end surfaced two bugs: - An empty value (TQ_CLIENT_ZMQ_MAX_SOCKETS=, the usual shell idiom for clearing a variable) hit int('') and crashed the client instead of falling back to libzmq's default. Treat empty as unset. - A non-numeric value raised a bare int() error naming no variable, which is hard to trace in a worker log. Name the variable. Also add a regression test that drives the real StorageManagerFactory with an independently-registered third-party manager. The existing factory tests patch create() out, so nothing executed the **kwargs forwarding; verified the new test fails if the old `if manager_type == "SimpleStorage"` special-case is reintroduced, while the rest of the suite stays green. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

ascend-robot · 2026-08-03T12:20:08Z

CLA Signature Pass

OutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Both were reachable and both are now covered by tests verified to fail when the fix is reverted. 1. A borrowing manager's stuck notify thread could not veto destroy(). StorageManager.close() detected the failed join but only acted on it inside `if self._owns_zmq_context`, so a manager that borrowed the client's context warned and returned silently. The client's veto checked only its own loop thread, so it went on to destroy a context whose sockets that thread might still hold -- the documented non-thread-safe Socket.close() hazard. The manager now records the outcome unconditionally and exposes can_destroy_zmq_context(); the client consults it, but only for a manager that actually shares the context, and TransferQueueClient now combines that with its own loop-thread check instead of replacing it. 2. Externally registered managers on the old (controller_info, config) contract raised TypeError, because the client passes zmq_context unconditionally. Registration is an extension mechanism, so managers are not required to update in lockstep: the factory now drops keywords a constructor cannot accept, warning with the class and parameter name so the drop is never silent. Constructors taking **kwargs, and un-introspectable ones, pass through unchanged. Also fix __del__ reaching for storage_manager_id on a half-constructed object, which raised AttributeError and masked the real constructor error -- found while reproducing issue 2. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

ascend-robot · 2026-08-03T12:56:07Z

CLA Signature Pass

OutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍

0oshowero0 · 2026-08-04T10:52:30Z

+# Per-context socket ceiling. Unset (or empty) means libzmq's default (1023). Because the
+# context is now shared per client instead of created per call, all in-flight sockets draw
+# on a single budget; raise this if a large num_data_storage_units fan-out exhausts it.
+TQ_CLIENT_ZMQ_MAX_SOCKETS = os.environ.get("TQ_CLIENT_ZMQ_MAX_SOCKETS") or None


I think we should increase the max sockets by default. In large-scale training we might hit the 1024 limit. That's mainly because the storage manager who borrows the context might have many sockets.

In the large-scale training, I think distributed agentloop workers and train_engine DP worker has limited samples to process (concurrently). For train_engine DP worker, the worst case is that all samples in different storage units, but its local batch size cannot be large due to GPU memory limitation.

Therefore, I have set it as 8192 as the defensive strategy. In practice, if there are 8192 sockets running, they occupy about 0.5*8192=4096 GB buffer.

0oshowero0 · 2026-08-04T10:56:46Z

+        # Record the outcome even when the context is borrowed: the owner cannot see this
+        # thread, so a borrower that stayed silent here would let the owner destroy() a
+        # context whose sockets are still in use. See can_destroy_zmq_context().
+        self._notify_thread_stopped = notify_thread_stopped


Why we need this? There are no other places read this flag

I have deleted the dead code.

0oshowero0 · 2026-08-04T10:59:14Z

 )


 class AsyncTransferQueueClient:


We may need a weakref.finalize for the client class. Users may forget to deliberately call .close()

https://github.com/Ascend/TransferQueue/blob/main/transfer_queue/storage/managers/base.py#L449

I have added it, pls refer to https://github.com/Ascend/TransferQueue/pull/145/changes#diff-94ebe743c62dfc7a3a16db0137ec631c8ef9f6d4637f838b508de67a6f5b58bcR136

0oshowero0 · 2026-08-04T11:02:24Z

+def test_factory_is_backend_agnostic(echo_controller):
+    """The client offers its context to every backend uniformly, naming none of them."""
+    client = AsyncTransferQueueClient(
+        client_id="client_other_storage_context",
+        controller_info=echo_controller.zmq_server_info,
+    )
+    config = {"client_name": "unused"}
+
+    with patch("transfer_queue.client.StorageManagerFactory.create") as create_manager:
+        client.initialize_storage_manager("OtherStorage", config)
+
+    create_manager.assert_called_once_with(
+        "OtherStorage",
+        controller_info=echo_controller.zmq_server_info,
+        config=config,
+        zmq_context=client.zmq_context,
+    )
+
+    client.close()
+
+
+def test_kv_backends_keep_own_context(echo_controller):
+    """KV managers accept the shared context but deliberately keep an independent one.
+
+    They move bulk data through their own SDKs and use ZMQ only for the controller
+    notify/handshake path, so they must not draw on the client's socket budget.
+    """
+    client = AsyncTransferQueueClient(
+        client_id="client_kv_own_context",
+        controller_info=echo_controller.zmq_server_info,
+    )
+
+    with (
+        patch("transfer_queue.storage.managers.base.StorageManager._connect_to_controller"),
+        patch("transfer_queue.storage.managers.base.StorageClientFactory.create"),
+    ):
+        manager = KVStorageManager(
+            echo_controller.zmq_server_info,
+            {"client_name": "unused"},
+            zmq_context=client.zmq_context,
+        )
+
+    assert manager.zmq_context is not client.zmq_context
+    assert manager._owns_zmq_context
+
+    manager.close()
+    assert not client.zmq_context.closed
+
+    client.close()
+
+
+def test_factory_forwards_kwargs_to_registered_manager(echo_controller):
+    """The real factory forwards **kwargs verbatim, knowing no backend by name.
+
+    The other factory test patches ``create`` out, so this one exercises the real
+    dispatch: a third-party manager registered from outside this package must receive
+    ``zmq_context`` without the factory special-casing its name. A regression to the old
+    ``if manager_type == "SimpleStorage"`` branch would silently drop the kwarg here, and
+    a manager whose signature drifts would raise TypeError -- neither of which mypy can
+    catch through ``**kwargs: Any``.
+    """
+    received = {}
+
+    @StorageManagerFactory.register("THIRD_PARTY_PROBE")
+    class ThirdPartyManager(StorageManager):
+        def __init__(self, controller_info, config, zmq_context=None):
+            received["zmq_context"] = zmq_context
+            received["config"] = config
+            super().__init__(controller_info, config, zmq_context=zmq_context)
+
+        def _connect_to_controller(self):
+            pass
+
+        def _do_handshake_with_controller(self):
+            pass
+
+        async def put_data(self, *args, **kwargs):
+            return None
+
+        async def get_data(self, *args, **kwargs):
+            return None
+
+        async def clear_data(self, *args, **kwargs):
+            return None
+
+        async def notify_data_update(self, *args, **kwargs):
+            return None
+
+    try:
+        client = AsyncTransferQueueClient(
+            client_id="client_third_party_factory",
+            controller_info=echo_controller.zmq_server_info,
+        )
+        config = {"marker": "forwarded"}
+
+        client.initialize_storage_manager("THIRD_PARTY_PROBE", config)
+
+        # The kwarg survived dispatch through the unpatched factory...
+        assert received["zmq_context"] is client.zmq_context
+        assert received["config"] == config
+        # ...and a manager that opts in genuinely borrows rather than re-creating.
+        assert client.storage_manager.zmq_context is client.zmq_context
+        assert not client.storage_manager._owns_zmq_context
+
+        # A borrower must not tear down a context it does not own.
+        client.storage_manager.close()
+        assert not client.zmq_context.closed
+
+        client.close()
+        assert client.zmq_context.closed
+    finally:
+        StorageManagerFactory._registry.pop("THIRD_PARTY_PROBE", None)
+


We may not need these test

I have deleted the redundant tests.

0oshowero0 · 2026-08-04T11:04:02Z

+def test_factory_tolerates_legacy_manager_signature(echo_controller):
+    """A manager on the old (controller_info, config) contract must still construct.
+
+    Registration is an extension mechanism, so third-party managers are not required to
+    add ``zmq_context`` in lockstep. The factory drops keywords a constructor cannot
+    accept instead of raising TypeError.
+    """


We can remove this test

0oshowero0 · 2026-08-04T11:04:40Z

+def test_client_applies_max_sockets(echo_controller):
+    """The per-context socket ceiling is configurable, since it is now shared per client."""
+    client = AsyncTransferQueueClient(
+        client_id="client_max_sockets",
+        controller_info=echo_controller.zmq_server_info,
+        zmq_max_sockets=2048,
+    )
+
+    assert client.zmq_context.get(zmq.MAX_SOCKETS) == 2048
+
+    client.close()
+
+
+def test_client_rejects_max_sockets_above_build_limit(echo_controller):
+    """Values above this libzmq build's ZMQ_SOCKET_LIMIT are rejected up front."""
+    probe = zmq.Context()
+    socket_limit = probe.get(zmq.SOCKET_LIMIT)
+    probe.term()
+
+    with pytest.raises(ValueError, match="ZMQ_SOCKET_LIMIT"):
+        AsyncTransferQueueClient(
+            client_id="client_max_sockets_too_big",
+            controller_info=echo_controller.zmq_server_info,
+            zmq_max_sockets=socket_limit + 1,
+        )
+
+
+def test_max_sockets_from_env_var(echo_controller):
+    """TQ_CLIENT_ZMQ_MAX_SOCKETS configures the ceiling without touching call sites.
+
+    The env var is the deployment-facing knob (the kwarg requires editing code), so it
+    needs its own coverage. Patched at the module constant because it is read at import.
+    """
+    with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", "4096"):
+        client = AsyncTransferQueueClient(
+            client_id="client_max_sockets_env",
+            controller_info=echo_controller.zmq_server_info,
+        )
+
+    assert client.zmq_context.get(zmq.MAX_SOCKETS) == 4096
+    client.close()
+
+
+def test_explicit_max_sockets_overrides_env_var(echo_controller):
+    """An explicit kwarg wins over the env var, matching the io_threads precedence."""
+    with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", "4096"):
+        client = AsyncTransferQueueClient(
+            client_id="client_max_sockets_precedence",
+            controller_info=echo_controller.zmq_server_info,
+            zmq_max_sockets=2048,
+        )
+
+    assert client.zmq_context.get(zmq.MAX_SOCKETS) == 2048
+    client.close()
+
+
+def test_unset_max_sockets_leaves_libzmq_default(echo_controller):
+    """Opt-in only: with nothing configured, libzmq's own default must be untouched."""
+    probe = zmq.Context()
+    default = probe.get(zmq.MAX_SOCKETS)
+    probe.term()
+
+    with patch("transfer_queue.client.TQ_CLIENT_ZMQ_MAX_SOCKETS", None):
+        client = AsyncTransferQueueClient(
+            client_id="client_max_sockets_unset",
+            controller_info=echo_controller.zmq_server_info,
+        )
+
+    assert client.zmq_context.get(zmq.MAX_SOCKETS) == default
+    client.close()


We can remove these tests

these tests have been removed.

Raise the client ZMQ socket ceiling by default, drop a dead flag, and add a finalizer backstop for clients that are never closed. - max sockets: default to DEFAULT_CLIENT_ZMQ_MAX_SOCKETS (8192) instead of inheriting libzmq's 1023, which large-scale training can exhaust since a borrowing storage manager draws on the same per-client budget. Validation now splits on provenance: an explicitly requested value (kwarg or env var) still raises when out of range, while the built-in default clamps to the build's ZMQ_SOCKET_LIMIT rather than making the client fail to construct on a build whose limit is lower. - StorageManager.close(): remove the write to self._notify_thread_stopped. No code read it -- can_destroy_zmq_context() checks the thread liveness directly so it stays correct when called before close(). - AsyncTransferQueueClient: register a weakref.finalize to destroy the context when close() is never called, so the context and its I/O threads do not leak for the process lifetime. close() detaches it, including on the path that deliberately leaks the context, so an unsafe destroy() cannot run later. Note TransferQueueClient's loop thread references the client, deferring collection to interpreter exit; close() remains the supported path. - tests: drop the factory/backend-agnosticism, max-sockets config, and legacy signature cases per review. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

Raise the client ZMQ socket ceiling by default, drop a dead flag, and add a finalizer backstop for clients that are never closed. - max sockets: default to DEFAULT_CLIENT_ZMQ_MAX_SOCKETS (8192) instead of inheriting libzmq's 1023, which large-scale training can exhaust since a borrowing storage manager draws on the same per-client budget. Validation now splits on provenance: an explicitly requested value (kwarg or env var) still raises when out of range, while the built-in default clamps to the build's ZMQ_SOCKET_LIMIT rather than making the client fail to construct on a build whose limit is lower. - StorageManager.close(): remove the write to self._notify_thread_stopped. No code read it -- can_destroy_zmq_context() checks the thread liveness directly so it stays correct when called before close(). - AsyncTransferQueueClient: register a weakref.finalize to destroy the context when close() is never called, so the context and its I/O threads do not leak for the process lifetime. close() detaches it, including on the path that deliberately leaks the context, so an unsafe destroy() cannot run later. Note TransferQueueClient's loop thread references the client, deferring collection to interpreter exit; close() remains the supported path. - tests: drop the factory/backend-agnosticism, max-sockets config, and legacy signature cases per review. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

ascend-robot · 2026-08-04T12:55:33Z

CLA Signature Guide

@OutstanderWang , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
cff5aeaa fix: address second round of PR ...	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

Inline comment blocks in this PR ran 4-6 lines while the surrounding files use 1-2, and several docstrings had grown to full paragraphs. Condense both to the essential why, keeping every load-bearing caveat (destroy() thread-safety, per-call context churn, ownership rules). No logic changes. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

ascend-robot · 2026-08-04T13:18:48Z

CLA Signature Guide

@OutstanderWang , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
cff5aeaa fix: address second round of PR ...	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

…hange - TransferQueueClient.__init__ docstring still documented libzmq's 1023 as the zmq_max_sockets default after DEFAULT_CLIENT_ZMQ_MAX_SOCKETS (8192) landed. - _release_zmq_context: drop the `context is not None` check. The finalizer is armed only after the context is constructed, so it can never be None. - test Borrower stub: drop the _do_handshake_with_controller and notify_data_update overrides. Neither is abstract, and the handshake is reachable only from _connect_to_controller, which the stub already overrides to pass. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

ascend-robot · 2026-08-04T13:28:01Z

CLA Signature Guide

@OutstanderWang , thanks for your pull request.

The following commit(s) are not associated with a signed Contributor License Agreement (CLA).

Commit	Reason
cff5aeaa fix: address second round of PR ...	the email used in the commit is not linked to a signed CLA! please verify that it matches the email you used when signing the CLA.

To sign CLA, click here.

To check if your email is configured correctly, refer to the FAQs.

Once you've signed the CLA or updating your email, please comment /check-cla to revalidate CLA status.

Raise the client ZMQ socket ceiling by default, drop a dead flag, and add a finalizer backstop for clients that are never closed. - max sockets: default to DEFAULT_CLIENT_ZMQ_MAX_SOCKETS (8192) instead of inheriting libzmq's 1023, which large-scale training can exhaust since a borrowing storage manager draws on the same per-client budget. Validation now splits on provenance: an explicitly requested value (kwarg or env var) still raises when out of range, while the built-in default clamps to the build's ZMQ_SOCKET_LIMIT rather than making the client fail to construct on a build whose limit is lower. - StorageManager.close(): remove the write to self._notify_thread_stopped. No code read it -- can_destroy_zmq_context() checks the thread liveness directly so it stays correct when called before close(). - AsyncTransferQueueClient: register a weakref.finalize to destroy the context when close() is never called, so the context and its I/O threads do not leak for the process lifetime. close() detaches it, including on the path that deliberately leaks the context, so an unsafe destroy() cannot run later. Note TransferQueueClient's loop thread references the client, deferring collection to interpreter exit; close() remains the supported path. - tests: drop the factory/backend-agnosticism, max-sockets config, and legacy signature cases per review. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

ascend-robot · 2026-08-04T14:33:39Z

CLA Signature Pass

OutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍

Raise the client ZMQ socket ceiling by default, drop a dead flag, and add a finalizer backstop for clients that are never closed. - max sockets: default to DEFAULT_CLIENT_ZMQ_MAX_SOCKETS (8192) instead of inheriting libzmq's 1023, which large-scale training can exhaust since a borrowing storage manager draws on the same per-client budget. Validation now splits on provenance: an explicitly requested value (kwarg or env var) still raises when out of range, while the built-in default clamps to the build's ZMQ_SOCKET_LIMIT rather than making the client fail to construct on a build whose limit is lower. - StorageManager.close(): remove the write to self._notify_thread_stopped. No code read it -- can_destroy_zmq_context() checks the thread liveness directly so it stays correct when called before close(). - AsyncTransferQueueClient: register a weakref.finalize to destroy the context when close() is never called, so the context and its I/O threads do not leak for the process lifetime. close() detaches it, including on the path that deliberately leaks the context, so an unsafe destroy() cannot run later. Note TransferQueueClient's loop thread references the client, deferring collection to interpreter exit; close() remains the supported path. - tests: drop the factory/backend-agnosticism, max-sockets config, and legacy signature cases per review.

ascend-robot · 2026-08-04T14:36:31Z

CLA Signature Pass

OutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍

ascend-robot · 2026-08-04T14:39:42Z

CLA Signature Pass

OutstanderWang, thanks for your pull request. All authors of the commits have signed the CLA. 👍

ascend-robot added the ascend-cla/yes label Jul 31, 2026

0oshowero0 changed the title ~~Feat long live zmq context pool~~ [perf, fix] Reuse a long-lived ZMQ context pool instead of creating one per call Jul 31, 2026

0oshowero0 requested a review from ji-huazhong July 31, 2026 08:30

0oshowero0 reviewed Jul 31, 2026

View reviewed changes

0oshowero0 changed the title ~~[perf, fix] Reuse a long-lived ZMQ context pool instead of creating one per call~~ [perf, fix] Reuse a long-lived ZMQ context instead of creating one per call Jul 31, 2026

0oshowero0 reviewed Jul 31, 2026

View reviewed changes

ascend-robot added ascend-cla/no and removed ascend-cla/yes labels Jul 31, 2026

OutstanderWang added 8 commits August 3, 2026 16:34

feat: share fixed ZMQ context pool across client requests

a47c3c5

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

fix: limit shared ZMQ context to SimpleStorage

12ce1b0

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

Revert "fix: limit shared ZMQ context to SimpleStorage"

b8907a5

This reverts commit b54c4093d71348b3c82572ff18c3a35a0bb9dd2b. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

Revert "feat: share fixed ZMQ context pool across client requests"

c82191c

This reverts commit a8bfbd81c68226f0c679ce3673c467866d470881. Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

feat: share fixed ZMQ context pool with SimpleStorage

2039e29

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

style: apply ruff-format to client.py

5894383

Signed-off-by: OutstanderWang <wangweiyanster@gmail.com>

OutstanderWang force-pushed the feat_long_live_zmq_context_pool branch from f6cede3 to 6b87e31 Compare August 3, 2026 09:36

ascend-robot added ascend-cla/yes and removed ascend-cla/no labels Aug 3, 2026

0oshowero0 reviewed Aug 4, 2026

View reviewed changes

ascend-robot added ascend-cla/no and removed ascend-cla/yes labels Aug 4, 2026

OutstanderWang force-pushed the feat_long_live_zmq_context_pool branch from 01e6a0f to 33a67ca Compare August 4, 2026 14:33

ascend-robot added ascend-cla/yes and removed ascend-cla/no labels Aug 4, 2026

OutstanderWang force-pushed the feat_long_live_zmq_context_pool branch from 33a67ca to 7152756 Compare August 4, 2026 14:36

OutstanderWang force-pushed the feat_long_live_zmq_context_pool branch from 7152756 to 309c0aa Compare August 4, 2026 14:39

		if manager_type == "SimpleStorage":
		create_kwargs["zmq_context"] = self.zmq_context

		assert client.zmq_context.closed


		def test_other_backends_do_not_borrow_client_context(echo_controller):

		self._owns_zmq_context = zmq_context is None
		self.zmq_context = zmq_context or zmq.asyncio.Context()

Uh oh!

Conversation

OutstanderWang commented Jul 31, 2026

[perf,fix] Reuse a long-lived ZMQ context pool instead of creating one per call

Motivation

What needs to be improved

What this changes

Scope

Tests

Uh oh!

ascend-robot commented Jul 31, 2026

CLA Signature Pass

Uh oh!

0oshowero0 commented Jul 31, 2026

Uh oh!

0oshowero0 commented Jul 31, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

0oshowero0 Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

OutstanderWang Aug 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ascend-robot commented Jul 31, 2026

CLA Signature Guide

Uh oh!

ascend-robot commented Aug 3, 2026

CLA Signature Pass

Uh oh!

ascend-robot commented Aug 3, 2026

CLA Signature Pass

Uh oh!

ascend-robot commented Aug 3, 2026

CLA Signature Pass

Uh oh!

ascend-robot commented Aug 3, 2026

CLA Signature Pass

Uh oh!

Choose a reason for hiding this comment

Uh oh!

OutstanderWang Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

0oshowero0 Jul 31, 2026 •

edited

Loading

OutstanderWang Aug 3, 2026 •

edited

Loading

OutstanderWang Aug 4, 2026 •

edited

Loading

OutstanderWang Aug 4, 2026 •

edited

Loading