Feat long live zmq context pool

transfer_queue/client.py

@@ -37,12 +37,21 @@
 logger = get_logger(__name__)
 
 TQ_NUM_THREADS = int(os.environ.get("TQ_NUM_THREADS", 8))
+# Size of the client context's native I/O-thread pool. The context serves every
+# controller RPC (all backends) and, for SimpleStorage, storage-unit requests too,
+# so this knob is client-scoped rather than backend-scoped.
+TQ_CLIENT_ZMQ_IO_THREADS = int(os.environ.get("TQ_CLIENT_ZMQ_IO_THREADS", 8))
+# Per-context socket ceiling. Unset (or empty) means libzmq's default (1023). Because the
+# context is now shared per client instead of created per call, all in-flight sockets draw
+# on a single budget; raise this if a large num_data_storage_units fan-out exhausts it.
+TQ_CLIENT_ZMQ_MAX_SOCKETS = os.environ.get("TQ_CLIENT_ZMQ_MAX_SOCKETS") or None
 
 # Pre-bound decorator for controller socket operations.
 with_controller_socket = with_zmq_socket(
     "request_handle_socket",
     get_identity=lambda self: self.client_id,
     get_peer=lambda self, target: self._controller,
+    get_context=lambda self: self.zmq_context,
 )
 
 
@@ -57,19 +66,54 @@ def __init__(
         self,
         client_id: str,
         controller_info: ZMQServerInfo,
+        zmq_io_threads: int | None = None,
+        zmq_max_sockets: int | None = None,
     ):
         """Initialize the asynchronous TransferQueue client.
 
         Args:
             client_id: Unique identifier for this client instance
             controller_info: Single controller ZMQ server information
+            zmq_io_threads: Fixed size of the client context's native I/O-thread pool.
+                Defaults to ``TQ_CLIENT_ZMQ_IO_THREADS`` (8).
+            zmq_max_sockets: Maximum number of sockets the client context may hold open
+                at once. Defaults to ``TQ_CLIENT_ZMQ_MAX_SOCKETS``, and to libzmq's own
+                default (1023) when that is unset. This ceiling is per-client: a single
+                ``put``/``get`` fans out one socket per storage unit, so deep fan-out
+                combined with high concurrency can approach it. Raising it also requires
+                enough file descriptors (``ulimit -n``).
         """
         if controller_info is None:
             raise ValueError("controller_info cannot be None")
         if not isinstance(controller_info, ZMQServerInfo):
             raise TypeError(f"controller_info must be ZMQServerInfo, got {type(controller_info)}")
         self.client_id = client_id
         self._controller: ZMQServerInfo = controller_info
+        # One long-lived context per client. Its fixed native I/O-thread pool is shared
+        # by all controller RPCs and, for SimpleStorage only, storage-unit requests.
+        # Sockets remain per-request because ZMQ sockets are not thread-safe.
+        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)
+
+        max_sockets = zmq_max_sockets
+        if max_sockets is None and TQ_CLIENT_ZMQ_MAX_SOCKETS is not None:
+            try:
+                max_sockets = int(TQ_CLIENT_ZMQ_MAX_SOCKETS)
+            except ValueError as e:
+                # Name the variable: a bare int() error gives no clue which knob is wrong.
+                raise ValueError(
+                    f"TQ_CLIENT_ZMQ_MAX_SOCKETS must be an integer, got {TQ_CLIENT_ZMQ_MAX_SOCKETS!r}"
+                ) from e
+        if max_sockets is not None:
+            socket_limit = self.zmq_context.get(zmq.SOCKET_LIMIT)
+            if not 1 <= max_sockets <= socket_limit:
+                raise ValueError(
+                    f"Client ZMQ max sockets must be between 1 and this build's "
+                    f"ZMQ_SOCKET_LIMIT ({socket_limit}), got {max_sockets}"
+                )
+            self.zmq_context.set(zmq.MAX_SOCKETS, max_sockets)
         logger.info(f"[{self.client_id}]: Registered Controller server {controller_info.id} at {controller_info.ip}")
 
     def initialize_storage_manager(
@@ -79,6 +123,10 @@ def initialize_storage_manager(
     ):
         """Initialize the storage manager.
 
+        The client's long-lived ZMQ context is offered to every backend uniformly; each
+        registered manager decides whether to borrow it or keep its own, so the client
+        needs no knowledge of specific backend names.
+
         Args:
             manager_type: Type of storage manager to create. Supported types include:
                           AsyncSimpleStorageManager, KVStorageManager (under development), etc.
@@ -88,7 +136,10 @@ def initialize_storage_manager(
 
         """
         self.storage_manager = StorageManagerFactory.create(
-            manager_type, controller_info=self._controller, config=config
+            manager_type,
+            controller_info=self._controller,
+            config=config,
+            zmq_context=self.zmq_context,
         )
 
     # ==================== Basic API ====================
@@ -1087,14 +1138,55 @@ async def async_kv_list(
             raise RuntimeError(f"[{self.client_id}]: Error in kv_list: {str(e)}") from e
 
     def close(self) -> None:
-        """Close the client and cleanup resources including storage manager."""
+        """Close the client and cleanup resources including storage manager.
+
+        The caller must ensure no RPCs are still in flight: this tears down the shared
+        ZMQ context via ``destroy()``, which calls ``Socket.close()`` internally and is
+        **not** thread-safe. Driving this client from your own event loop means closing
+        it only after outstanding tasks have been awaited or cancelled. Subclasses that
+        own a loop thread must join it before delegating here (see
+        :meth:`TransferQueueClient.close`).
+        """
         try:
             if hasattr(self, "storage_manager") and self.storage_manager:
                 if hasattr(self.storage_manager, "close"):
                     self.storage_manager.close()
         except Exception as e:
             logger.warning(f"Error closing storage manager: {e}")
 
+        # Tear down the shared context last. destroy(linger=0) force-closes any socket that
+        # leaked from an interrupted RPC then terminates, so shutdown cannot hang.
+        if not self._can_destroy_zmq_context():
+            logger.warning(
+                f"[{self.client_id}]: Skipping zmq_context.destroy() because a thread owning "
+                f"sockets on it is still alive; destroy() is not thread-safe. Leaking the context."
+            )
+            return
+        try:
+            if hasattr(self, "zmq_context") and self.zmq_context is not None:
+                self.zmq_context.destroy(linger=0)
+        except Exception as e:
+            logger.warning(f"[{self.client_id}]: Error terminating zmq_context: {e}")
+
+    def _can_destroy_zmq_context(self) -> bool:
+        """Whether it is safe to call ``destroy()`` on the shared context.
+
+        This class owns no background thread of its own, so the caller's quiescence
+        contract (documented on :meth:`close`) covers the client side. But a storage
+        manager that *borrowed* this context runs a notify thread the client cannot see,
+        and that thread holds sockets on it -- so ask the manager whether it finished
+        shutting down. Subclasses that run their own loop thread extend this.
+        """
+        manager = getattr(self, "storage_manager", None)
+        if manager is not None:
+            can_destroy = getattr(manager, "can_destroy_zmq_context", None)
+            # Only consult a manager that actually shares this context; one with its own
+            # context has no say in when the client's is destroyed.
+            if callable(can_destroy) and getattr(manager, "zmq_context", None) is self.zmq_context:
+                if not can_destroy():
+                    return False
+        return True
+
     # ==================== Checkpoint API ====================
     @with_controller_socket
     async def async_save_controller_checkpoint(
@@ -1230,16 +1322,25 @@ def __init__(
         self,
         client_id: str,
         controller_info: ZMQServerInfo,
+        zmq_io_threads: int | None = None,
+        zmq_max_sockets: int | None = None,
     ):
         """Initialize the synchronous TransferQueue client.
 
         Args:
             client_id: Unique identifier for this client instance
             controller_info: Single controller ZMQ server information
+            zmq_io_threads: Fixed size of the client context's native I/O-thread pool.
+                Defaults to ``TQ_CLIENT_ZMQ_IO_THREADS`` (8).
+            zmq_max_sockets: Maximum number of sockets the client context may hold open
+                at once. Defaults to ``TQ_CLIENT_ZMQ_MAX_SOCKETS``, and to libzmq's own
+                default (1023) when that is unset.
         """
         super().__init__(
             client_id,
             controller_info,
+            zmq_io_threads,
+            zmq_max_sockets,
         )
 
         # create new event loop in a separate thread
@@ -1773,7 +1874,16 @@ def load_storage_checkpoint(self, checkpoint_dir: str) -> None:
         return self._load_storage_checkpoint(checkpoint_dir)
 
     def close(self) -> None:
-        """Close the client and cleanup resources including event loop and thread."""
+        """Close the client and cleanup resources including event loop and thread.
+
+        The ordering here is load-bearing and must not be rearranged: the background
+        loop thread is stopped and joined *before* delegating to
+        :meth:`AsyncTransferQueueClient.close`, which destroys the shared ZMQ context.
+        ``destroy()`` calls ``Socket.close()`` internally and is not thread-safe, so no
+        other thread may hold sockets on that context when it runs. If the join times
+        out, :meth:`_can_destroy_zmq_context` reports False and the context is leaked
+        instead of destroyed unsafely.
+        """
 
         if hasattr(self, "_loop") and self._loop is not None:
             self._loop.call_soon_threadsafe(self._loop.stop)
@@ -1789,3 +1899,14 @@ def close(self) -> None:
                 logger.warning(f"[{self.client_id}]: Error closing event loop: {e}")
 
         super().close()
+
+    def _can_destroy_zmq_context(self) -> bool:
+        """False while the loop thread that owns sockets on the context is still alive.
+
+        Also defers to the base check, which covers a borrowing storage manager's notify
+        thread -- both threads must be gone before ``destroy()`` is safe.
+        """
+        thread = getattr(self, "_thread", None)
+        if thread is not None and thread.is_alive():
+            return False
+        return super()._can_destroy_zmq_context()

transfer_queue/storage/managers/base.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import asyncio
+import inspect
 import itertools
 import os
 import threading
@@ -63,15 +64,25 @@ class StorageManager(ABC):
     """Base class for storage layer. It defines the interface for data operations and
     generally provides handshake & notification capabilities."""
 
-    def __init__(self, controller_info: ZMQServerInfo, config: DictConfig):
+    def __init__(
+        self,
+        controller_info: ZMQServerInfo,
+        config: DictConfig,
+        zmq_context: zmq.asyncio.Context | None = None,
+    ):
         self.storage_manager_id = f"TQ_STORAGE_{uuid4().hex[:8]}"
         self.config = config
         self.controller_info = controller_info
 
         # Handshake socket is sync (used only during initialization)
         self.controller_handshake_socket: zmq.Socket | None = None
 
-        self.zmq_context = zmq.asyncio.Context()
+        # A manager may borrow a caller-owned context whose fixed native I/O thread pool
+        # is shared by controller and storage-unit request sockets; SimpleStorage does.
+        # A manager that is handed nothing owns the context it creates, and only an owner
+        # tears its context down (see close()).
+        self._owns_zmq_context = zmq_context is None
+        self.zmq_context = zmq.asyncio.Context() if zmq_context is None else zmq_context
         self._connect_to_controller()
 
         # Dedicated asyncio loop for ZMQ notify traffic, isolated from the caller's loop
@@ -384,21 +395,59 @@ def close(self) -> None:
         if hasattr(self, "_notify_loop") and self._notify_loop.is_running():
             self._notify_loop.call_soon_threadsafe(self._notify_loop.stop)
 
+        notify_thread_stopped = True
         if hasattr(self, "_notify_thread") and self._notify_thread is not None:
             self._notify_thread.join(timeout=5.0)
             if self._notify_thread.is_alive():
+                notify_thread_stopped = False
                 logger.warning(f"[{self.storage_manager_id}]: Notify ZMQ thread did not stop within 5 second timeout.")
             else:
                 logger.debug(f"[{self.storage_manager_id}]: Notify ZMQ thread shut down.")
 
-        self.zmq_context.term()
+        # Record the outcome even when the context is borrowed: the owner cannot see this
+        # thread, so a borrower that stayed silent here would let the owner destroy() a
+        # context whose sockets are still in use. See can_destroy_zmq_context().
+        self._notify_thread_stopped = notify_thread_stopped
+
+        if self._owns_zmq_context:
+            # Ordering below is load-bearing: destroy() calls Socket.close() internally,
+            # which is NOT thread-safe, so it must run only after the notify thread that
+            # owns sockets on this context is gone. If that thread outlived its join
+            # timeout, leak the context rather than risk a crash during shutdown -- the
+            # process is terminating anyway, so a leaked context is the cheaper outcome.
+            if notify_thread_stopped:
+                # destroy(linger=0) force-closes any socket still open (e.g. from an interrupted
+                # request or the notify path) then terminates, so shutdown cannot hang on term().
+                self.zmq_context.destroy(linger=0)
+            else:
+                logger.warning(
+                    f"[{self.storage_manager_id}]: Skipping zmq_context.destroy() because the notify "
+                    f"thread is still alive; destroy() is not thread-safe while sockets are in use. "
+                    f"The context will be leaked."
+                )
+
+    def can_destroy_zmq_context(self) -> bool:
+        """Whether an owner may safely ``destroy()`` a context this manager borrowed.
+
+        ``destroy()`` calls ``Socket.close()`` internally and is not thread-safe, so the
+        notify thread must be gone first. Only this manager can see that thread, so an
+        owner of a borrowed context must ask before tearing the context down.
+
+        Checks the thread directly rather than trusting the flag recorded by ``close()``,
+        so this is also correct if called before ``close()`` or if the thread exited late.
+        """
+        thread = getattr(self, "_notify_thread", None)
+        return thread is None or not thread.is_alive()
 
     def __del__(self):
         """Destructor to ensure resources are cleaned up."""
         try:
             self.close()
         except Exception as e:
-            logger.error(f"[{self.storage_manager_id}]: Exception during __del__: {str(e)}")
+            # __init__ may have failed before storage_manager_id was set; reaching for it
+            # here would raise AttributeError and mask the exception we mean to report.
+            manager_id = getattr(self, "storage_manager_id", f"<uninitialized {type(self).__name__}>")
+            logger.error(f"[{manager_id}]: Exception during __del__: {str(e)}")
 
 
 class StorageManagerFactory:
@@ -422,12 +471,57 @@ def decorator(manager_cls: type[StorageManager]):
         return decorator
 
     @classmethod
-    def create(cls, manager_type: str, controller_info: ZMQServerInfo, config: dict[str, Any]) -> StorageManager:
-        """Create and return a StorageManager instance."""
+    def create(
+        cls,
+        manager_type: str,
+        controller_info: ZMQServerInfo,
+        config: dict[str, Any],
+        **kwargs: Any,
+    ) -> StorageManager:
+        """Create and return a StorageManager instance.
+
+        Extra keyword arguments are forwarded to the registered class. The factory
+        deliberately knows nothing about any individual backend: each manager decides for
+        itself what to do with what it receives (e.g. whether to borrow a caller-supplied
+        ``zmq_context`` or keep its own).
+
+        Registration is an extension mechanism, so managers written against the older
+        ``(controller_info, config)`` contract must keep working without being updated in
+        lockstep. Any keyword the constructor does not accept is therefore dropped, with a
+        warning, rather than raising ``TypeError``.
+        """
         assert manager_type in cls._registry, (
             f"Unknown manager_type: {manager_type}. Supported managers include: {list(cls._registry.keys())}"
         )
-        return cls._registry[manager_type](controller_info, config)
+        manager_cls = cls._registry[manager_type]
+        accepted = cls._filter_supported_kwargs(manager_cls, kwargs, manager_type)
+        return manager_cls(controller_info, config, **accepted)
+
+    @staticmethod
+    def _filter_supported_kwargs(
+        manager_cls: type[StorageManager], kwargs: dict[str, Any], manager_type: str
+    ) -> dict[str, Any]:
+        """Drop keywords ``manager_cls.__init__`` cannot accept, warning about each.
+
+        A constructor taking ``**kwargs`` is assumed to accept everything.
+        """
+        if not kwargs:
+            return kwargs
+        try:
+            params = inspect.signature(manager_cls.__init__).parameters
+        except (TypeError, ValueError):
+            # Un-introspectable constructor (e.g. a C extension): pass through unchanged
+            # rather than silently dropping arguments the class may well accept.
+            return kwargs
+        if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()):
+            return kwargs
+        supported = {name: value for name, value in kwargs.items() if name in params}
+        for name in kwargs.keys() - supported.keys():
+            logger.warning(
+                f"{manager_cls.__name__} (registered as '{manager_type}') does not accept "
+                f"'{name}'; ignoring it. Add '{name}' to its __init__ signature to use it."
+            )
+        return supported
 
 
 class KVStorageManager(StorageManager):
@@ -436,9 +530,22 @@ class KVStorageManager(StorageManager):
     It maps structured metadata (BatchMeta) to flat lists of keys and values for efficient KV operations.
     """
 
-    def __init__(self, controller_info: ZMQServerInfo, config: dict[str, Any]):
+    def __init__(
+        self,
+        controller_info: ZMQServerInfo,
+        config: dict[str, Any],
+        zmq_context: zmq.asyncio.Context | None = None,
+    ):
         """
         Initialize the KVStorageManager with configuration.
+
+        Args:
+            controller_info: Controller ZMQ server information.
+            config: Backend configuration; must contain ``client_name``.
+            zmq_context: Accepted for interface uniformity but deliberately ignored.
+                KV backends move bulk data through their own SDKs and use ZMQ only for
+                the controller notify/handshake path, so they keep an independent
+                context rather than drawing on a caller's shared socket budget.
         """
         client_name = config.get("client_name", None)
         if client_name is None: