Defer external resources until process initialization

README.md

@@ -69,7 +69,7 @@ Optional integrations for different cloud providers can be installed using `plug
 
 Support for parallelisation and hyperparameter optimisation can be installed using `plugboard[ray]`.
 
-Additional optional extras: `plugboard[llm]` for LLM components, `plugboard[redis]` for Redis-based connectors, and `plugboard[websockets]` for WebSocket I/O.
+Additional optional extras: `plugboard[llm]` for LLM components, `plugboard[redis]` for Redis-based connectors, `plugboard[omq]` for the pyomq backend for ZMQ connectors, and `plugboard[websockets]` for WebSocket I/O.
 
 ## ⚡ Quickstart with AI
 

docs/examples/tutorials/running-in-parallel.md

@@ -51,7 +51,7 @@ With some small changes we can make the same model run in parallel on Ray. First
 !!! info
     [`Channel`][plugboard.connector.Channel] objects are used by Plugboard to handle the communication between components. So far we have used [`AsyncioChannel`][plugboard.connector.AsyncioChannel], which is the best option for simple models that don't require parallelisation.
 
-    Plugboard provides different channel classes for use in parallel environments: [`RayChannel`][plugboard.connector.RayChannel] is suitable for single and multi-host Ray environments. [`ZMQChannel`][plugboard.connector.ZMQChannel] is faster, but currently only works on a single host.
+    Plugboard provides different channel classes for use in parallel environments: [`RayChannel`][plugboard.connector.RayChannel] is suitable for single and multi-host Ray environments. [`ZMQChannel`][plugboard.connector.ZMQChannel] is faster, but currently only works on a single host. Set `PLUGBOARD_ZMQ_BACKEND=pyomq` to use the optional pyomq backend instead of PyZMQ.
 
 ```python
 --8<-- "examples/tutorials/004_using_ray/hello_ray.py:ray"

plugboard/_zmq/backend.py

@@ -0,0 +1,47 @@
+"""Selects the ZeroMQ Python backend."""
+
+from __future__ import annotations
+
+import os
+import typing as _t
+
+
+ZMQ_BACKEND_ENV = "PLUGBOARD_ZMQ_BACKEND"
+ZMQ_BACKEND_PYZMQ = "pyzmq"
+ZMQ_BACKEND_PYOMQ = "pyomq"
+ZMQ_BACKENDS = frozenset({ZMQ_BACKEND_PYZMQ, ZMQ_BACKEND_PYOMQ})
+
+
+class ZMQBackendImportError(ImportError):
+    """Raised when the selected ZeroMQ backend cannot be imported."""
+
+
+def _backend_name() -> str:
+    backend = os.environ.get(ZMQ_BACKEND_ENV, ZMQ_BACKEND_PYZMQ).strip().lower()
+    if not backend:
+        return ZMQ_BACKEND_PYZMQ
+    if backend not in ZMQ_BACKENDS:
+        choices = ", ".join(sorted(ZMQ_BACKENDS))
+        raise ValueError(
+            f"Unsupported ZMQ backend {backend!r}. Set {ZMQ_BACKEND_ENV} to one of: {choices}."
+        )
+    return backend
+
+
+def _load_backend() -> tuple[str, _t.Any, _t.Any]:
+    backend = _backend_name()
+    try:
+        if backend == ZMQ_BACKEND_PYOMQ:
+            import pyomq as zmq
+            import pyomq.asyncio as zmq_asyncio
+        else:
+            import zmq
+            import zmq.asyncio as zmq_asyncio
+    except ImportError as e:
+        raise ZMQBackendImportError(
+            f"Failed to import {backend!r} ZMQ backend selected by {ZMQ_BACKEND_ENV}."
+        ) from e
+    return backend, zmq, zmq_asyncio
+
+
+zmq_backend, zmq, zmq_asyncio = _load_backend()

plugboard/_zmq/zmq_proxy.py

@@ -7,8 +7,8 @@
 import typing as _t
 
 from pydantic import BaseModel, Field, ValidationError
-import zmq
-import zmq.asyncio
+
+from plugboard._zmq.backend import zmq, zmq_asyncio
 
 
 try:
@@ -23,8 +23,8 @@
 def create_socket(
     socket_type: int,
     socket_opts: zmq_sockopts_t,
-    ctx: _t.Optional[zmq.asyncio.Context] = None,
-) -> zmq.asyncio.Socket:
+    ctx: _t.Optional[zmq_asyncio.Context] = None,
+) -> zmq_asyncio.Socket:
     """Creates a ZeroMQ socket with the given type and options.
 
     Args:
@@ -35,7 +35,7 @@ def create_socket(
     Returns:
         The created ZMQ socket.
     """
-    _ctx = ctx or zmq.asyncio.Context.instance()
+    _ctx = ctx or zmq_asyncio.Context.instance()
     socket = _ctx.socket(socket_type)
     for opt, value in socket_opts:
         socket.setsockopt(opt, value)
@@ -184,7 +184,7 @@ def _connect_socket_req_socket(self) -> None:
         """Connects the REQ socket to the REP socket in the subprocess."""
         if self._socket_rep_port is None:
             raise RuntimeError("ZMQ proxy socket REP port not set.")
-        self._socket_req_socket: zmq.asyncio.Socket = create_socket(zmq.REQ, [])
+        self._socket_req_socket: zmq_asyncio.Socket = create_socket(zmq.REQ, [])
         socket_rep_socket_address: str = f"{self._zmq_address}:{self._socket_rep_port}"
         self._socket_req_socket.connect(socket_rep_socket_address)
         self._socket_req_lock: asyncio.Lock = asyncio.Lock()
@@ -205,8 +205,8 @@ async def add_push_socket(self, topic: str, maxsize: int = 2000) -> str:
 
     async def _run(self) -> None:
         """Async multiprocessing entrypoint to run ZMQ proxy."""
-        self._push_poller: zmq.asyncio.Poller = zmq.asyncio.Poller()
-        self._push_sockets: dict[str, tuple[str, zmq.asyncio.Socket]] = {}
+        self._push_poller: zmq_asyncio.Poller = zmq_asyncio.Poller()
+        self._push_sockets: dict[str, tuple[str, zmq_asyncio.Socket]] = {}
 
         self._create_proxy_sockets()
 
@@ -285,13 +285,17 @@ def _create_push_socket(self, topic: str, maxsize: int, reuse: bool = True) -> s
     async def _poll_push_sockets(self) -> None:
         """Polls push sockets for messages and sends them to the proxy."""
         while True:
+            # Some backends return immediately when polling an empty socket set.
+            if not self._push_poller.sockets:
+                await asyncio.sleep(1)
+                continue
             # Set a timeout of 1 second to allow for new push sockets to be added
             events = dict(await self._push_poller.poll(timeout=1000))
             async with asyncio.TaskGroup() as tg:
                 for socket in events:
                     tg.create_task(self._handle_push_socket(socket))
 
-    async def _handle_push_socket(self, socket: zmq.asyncio.Socket) -> None:
+    async def _handle_push_socket(self, socket: zmq_asyncio.Socket) -> None:
         msg = await socket.recv_multipart()
         topic = msg[0].decode("utf8")
         _, push_socket = self._push_sockets[topic]

plugboard/connector/connector.py

@@ -17,6 +17,10 @@ class Connector(ABC, ExportMixin):
     def __init__(self, spec: ConnectorSpec, *args: _t.Any, **kwargs: _t.Any) -> None:
         self.spec: ConnectorSpec = spec
 
+    async def init(self) -> None:
+        """Acquire resources required by this connector."""
+        pass
+
     @abstractmethod
     async def connect_send(self) -> Channel:
         """Returns a `Channel` for sending messages."""

plugboard/connector/ray_channel.py

@@ -1,5 +1,6 @@
 """Provides `RayChannel` for use in cluster compute environments."""
 
+import asyncio
 import typing as _t
 
 from plugboard.connector.asyncio_channel import AsyncioChannel
@@ -34,8 +35,18 @@ def __init__(  # noqa: D417
         """
         default_options = {"num_cpus": 0}
         actor_options = actor_options or {}
-        actor_options = {**default_options, **actor_options}
-        self._actor = ray.remote(**actor_options)(_AsyncioChannelActor).remote(**kwargs)
+        self._actor_options = {**default_options, **actor_options}
+        self._channel_kwargs = kwargs
+        self._actor: _t.Any = None
+        self._init_lock = asyncio.Lock()
+
+    async def init(self) -> None:
+        """Create the channel actor when execution starts."""
+        async with self._init_lock:
+            if self._actor is None:
+                self._actor = ray.remote(**self._actor_options)(_AsyncioChannelActor).remote(
+                    **self._channel_kwargs
+                )
 
     @property
     def maxsize(self) -> int:
@@ -73,10 +84,16 @@ def __init__(self, *args: _t.Any, **kwargs: _t.Any) -> None:
             raise ValueError("RayConnector only supports `PIPELINE` type connections.")
         self._channel = RayChannel()
 
+    async def init(self) -> None:
+        """Create the remote channel actor."""
+        await self._channel.init()
+
     async def connect_send(self) -> RayChannel:
         """Returns a `RayChannel` for sending messages."""
+        await self.init()
         return self._channel
 
     async def connect_recv(self) -> RayChannel:
         """Returns a `RayChannel` for receiving messages."""
+        await self.init()
         return self._channel