proxystore.connectors.zmq

proxystore/connectors/zmq.py

ZeroMQ-based distributed in-memory connector implementation.

Each host runs a ZeroMQServer which stores objects in memory. The first ZeroMQConnector created on a host spawns the server if one is not already running. Objects are put in the server on the same host, and keys contain the address of that server so connectors on other hosts can get objects directly from it.

Warning

The servers do not authenticate clients or encrypt data, so anyone who can reach a server's port can read, write, and evict objects. Only use this connector within trusted networks which are not publicly accessible, such as the interconnect of an HPC cluster.

Note

Only the ZeroMQConnector and its key and exception types are public. The server functions and classes (e.g., ZeroMQServer and spawn_server()) are internal implementation details which may change between releases without notice (see Versioning and Compatibility).

Messages are multipart ZeroMQ messages. A request consists of a request ID frame, a JSON header frame (e.g., {"op": "get", "obj_id": "..."}), and optionally a data frame. The reply echoes the request ID and contains a JSON header frame with the status and optionally a data frame. Data frames are sent and received without copies.

ServerTimeoutError

Bases: Exception

Client timed out waiting for a response from a server.

ZeroMQServerError

Bases: Exception

Server failed to process a request.

ZeroMQKey

Bases: NamedTuple

Key to objects stored in a ZeroMQ server.

Attributes:

• obj_id (str) –

  Unique object ID.

• peer_host (str) –

  Address of the server where the object is stored.

• peer_port (int) –

  Port of the server where the object is stored.

ZeroMQConnector

ZeroMQConnector(
    port: int,
    address: str | None = None,
    interface: str | None = None,
    timeout: float = 5,
    request_timeout: float | None = 60,
    clear: bool = False,
)

ZeroMQ-based distributed in-memory connector.

Warning

The servers do not authenticate clients or encrypt data, so anyone who can reach a server's port can read, write, and evict objects. Only use this connector within trusted networks which are not publicly accessible, such as the interconnect of an HPC cluster.

Warning

Objects are stored in the memory of the server process, and there is no limit to the memory used. Objects are lost when the server exits.

Note

The first connector created on a host spawns a ZeroMQServer in a new process which other connectors on the host will use. Closing the connector does not stop the server by default (see clear), but the server is stopped when the process which spawned it exits.

Example

from proxystore.connectors.zmq import ZeroMQConnector

with ZeroMQConnector(port=5555) as connector:
    key = connector.put(b'value')
    assert connector.get(key) == b'value'

Parameters:

• port (int) –

  Port of the server on this host.

• address (str | None, default: None ) –

  Address of this host that other hosts can connect to. Takes precedence over interface if both are provided.

• interface (str | None, default: None ) –

  Network interface to get the address of this host from (e.g., 'eth0'). Only supported on Linux.

• timeout (float, default: 5 ) –

  Timeout in seconds to connect to or spawn the server on this host.

• request_timeout (float | None, default: 60 ) –

  Timeout in seconds to wait for a response from a server before raising a ServerTimeoutError. When operating on multiple objects, the timeout applies to each response. None waits forever.

• clear (bool, default: False ) –

  Stop the server on this host, if it was spawned by this connector, when close() is called. This will lose all objects stored in the server.

Raises:

• ServerTimeoutError –

  If a server on this host could not be connected to or spawned within timeout seconds.

Source code in proxystore/connectors/zmq.py

def __init__(
    self,
    port: int,
    address: str | None = None,
    interface: str | None = None,
    timeout: float = 5,
    request_timeout: float | None = 60,
    clear: bool = False,
) -> None:
    self._address = address
    self._interface = interface
    self.port = port
    self.timeout = timeout
    self.request_timeout = request_timeout
    self.clear = clear

    if self._address is not None:
        self.address = self._address
    elif self._interface is not None:
        self.address = get_interface_address(self._interface)
    else:
        self.address = socket.gethostbyname(socket.gethostname())
        if ipaddress.ip_address(self.address).is_loopback:
            warnings.warn(
                f'The address of this host ({self.address}) resolved '
                'from its hostname is a loopback address so other hosts '
                'will not be able to get objects from this host. '
                'Specify the address or interface to use instead.',
                stacklevel=2,
            )

    self._kill_hook: Any = None
    self.server: subprocess.Popen[bytes] | None = None
    try:
        # Check if the port is open first because ZeroMQ will silently
        # retry connecting to a port which refuses connections.
        socket.create_connection(
            (self.address, self.port),
            timeout=self.timeout,
        ).close()
        wait_for_server(self.address, self.port, timeout=self.timeout)
    except (OSError, ServerTimeoutError):
        self.server = spawn_server(
            self.address,
            self.port,
            timeout=self.timeout,
        )
    else:
        logger.info('Connected to existing server at %s', self.url)

    if self.server is not None:
        self._kill_hook = _register_kill_hook(self.server)

    self._pool = _SocketPool()

url property

url: str

URL of the server on this host.

close

close(*, clear: bool | None = None) -> bool

Close the connector.

Parameters:

• clear (bool | None, default: None ) –

  Stop the server on this host if it was spawned by this connector. This will lose all objects stored in the server. Overrides the default value of clear provided when the connector was instantiated.

Returns:

• bool –

  If the server was stopped, removing its objects.

Source code in proxystore/connectors/zmq.py

def close(self, *, clear: bool | None = None) -> bool:
    """Close the connector.

    Args:
        clear: Stop the server on this host if it was spawned by
            this connector. This will lose all objects stored in the
            server. Overrides the default value of `clear` provided when
            the connector was instantiated.

    Returns:
        If the server was stopped, removing its objects.
    """
    self._pool.close()

    clear = self.clear if clear is None else clear
    if clear and self.server is not None:
        _kill_server(self.server)
        atexit.unregister(self._kill_hook)
        logger.info(
            'Stopped server at %s (pid=%d)',
            self.url,
            self.server.pid,
        )
        self.server = None
        return True
    return False

config

config() -> dict[str, Any]

Get the connector configuration.

The configuration contains all the information needed to reconstruct the connector object.

Source code in proxystore/connectors/zmq.py

def config(self) -> dict[str, Any]:
    """Get the connector configuration.

    The configuration contains all the information needed to reconstruct
    the connector object.
    """
    return {
        'port': self.port,
        'address': self._address,
        'interface': self._interface,
        'timeout': self.timeout,
        'request_timeout': self.request_timeout,
        'clear': self.clear,
    }

from_config classmethod

from_config(config: dict[str, Any]) -> ZeroMQConnector

Create a new connector instance from a configuration.

Parameters:

• config (dict[str, Any]) –

  Configuration returned by .config().

Source code in proxystore/connectors/zmq.py

@classmethod
def from_config(cls, config: dict[str, Any]) -> ZeroMQConnector:
    """Create a new connector instance from a configuration.

    Args:
        config: Configuration returned by `#!python .config()`.
    """
    return cls(**init_kwargs(cls, config))

evict

evict(key: ZeroMQKey) -> None

Evict the object associated with the key.

Parameters:

• key (ZeroMQKey) –

  Key associated with object to evict.

Source code in proxystore/connectors/zmq.py

def evict(self, key: ZeroMQKey) -> None:
    """Evict the object associated with the key.

    Args:
        key: Key associated with object to evict.
    """
    self._request([('evict', key, None)])

exists

exists(key: ZeroMQKey) -> bool

Check if an object associated with the key exists.

Parameters:

• key (ZeroMQKey) –

  Key potentially associated with stored object.

Returns:

• bool –

  If an object associated with the key exists.

Source code in proxystore/connectors/zmq.py

def exists(self, key: ZeroMQKey) -> bool:
    """Check if an object associated with the key exists.

    Args:
        key: Key potentially associated with stored object.

    Returns:
        If an object associated with the key exists.
    """
    (reply,) = self._request([('exists', key, None)])
    return reply.header['exists']

get

get(key: ZeroMQKey) -> BytesLike | None

Get the serialized object associated with the key.

Parameters:

• key (ZeroMQKey) –

  Key associated with the object to retrieve.

Returns:

• BytesLike | None –

  Serialized object or None if the object does not exist.

Source code in proxystore/connectors/zmq.py

def get(self, key: ZeroMQKey) -> BytesLike | None:
    """Get the serialized object associated with the key.

    Args:
        key: Key associated with the object to retrieve.

    Returns:
        Serialized object or `None` if the object does not exist.
    """
    (reply,) = self._request([('get', key, None)])
    return reply.data

get_batch

get_batch(
    keys: Sequence[ZeroMQKey],
) -> list[BytesLike | None]

Get a batch of serialized objects associated with the keys.

Parameters:

• keys (Sequence[ZeroMQKey]) –

  Sequence of keys associated with objects to retrieve.

Returns:

• list[BytesLike | None] –

  List with same order as keys with the serialized objects or None if the corresponding key does not have an associated object.

Source code in proxystore/connectors/zmq.py

def get_batch(self, keys: Sequence[ZeroMQKey]) -> list[BytesLike | None]:
    """Get a batch of serialized objects associated with the keys.

    Args:
        keys: Sequence of keys associated with objects to retrieve.

    Returns:
        List with same order as `keys` with the serialized objects or \
        `None` if the corresponding key does not have an associated object.
    """
    replies = self._request([('get', key, None) for key in keys])
    return [reply.data for reply in replies]

new_key

new_key(obj: BytesLike | None = None) -> ZeroMQKey

Create a new key.

Parameters:

• obj (BytesLike | None, default: None ) –

  Optional object which the key will be associated with. Ignored in this implementation.

Returns:

• ZeroMQKey –

  Key which can be used to retrieve an object once set() has been called on the key.

Source code in proxystore/connectors/zmq.py

def new_key(self, obj: BytesLike | None = None) -> ZeroMQKey:
    """Create a new key.

    Args:
        obj: Optional object which the key will be associated with.
            Ignored in this implementation.

    Returns:
        Key which can be used to retrieve an object once \
        [`set()`][proxystore.connectors.zmq.ZeroMQConnector.set] \
        has been called on the key.
    """
    return self._new_key()

put

put(obj: BytesLike) -> ZeroMQKey

Put a serialized object in the store.

Parameters:

• obj (BytesLike) –

  Serialized object to put in the store.

Returns:

• ZeroMQKey –

  Key which can be used to retrieve the object.

Source code in proxystore/connectors/zmq.py

def put(self, obj: BytesLike) -> ZeroMQKey:
    """Put a serialized object in the store.

    Args:
        obj: Serialized object to put in the store.

    Returns:
        Key which can be used to retrieve the object.
    """
    key = self._new_key()
    self._request([('put', key, obj)])
    return key

put_batch

put_batch(objs: Sequence[BytesLike]) -> list[ZeroMQKey]

Put a batch of serialized objects in the store.

Parameters:

• objs (Sequence[BytesLike]) –

  Sequence of serialized objects to put in the store.

Returns:

• list[ZeroMQKey] –

  List of keys with the same order as objs which can be used to retrieve the objects.

Source code in proxystore/connectors/zmq.py

def put_batch(self, objs: Sequence[BytesLike]) -> list[ZeroMQKey]:
    """Put a batch of serialized objects in the store.

    Args:
        objs: Sequence of serialized objects to put in the store.

    Returns:
        List of keys with the same order as `objs` which can be used to \
        retrieve the objects.
    """
    keys = [self._new_key() for _ in objs]
    self._request(
        [('put', key, obj) for key, obj in zip(keys, objs, strict=True)],
    )
    return keys

set

set(key: ZeroMQKey, obj: BytesLike) -> None

Set the object associated with a key.

Note

The Connector provides write-once, read-many semantics. Thus, set() should only be called once per key, otherwise unexpected behavior can occur.

Parameters:

• key (ZeroMQKey) –

  Key that the object will be associated with.

• obj (BytesLike) –

  Object to associate with the key.

Source code in proxystore/connectors/zmq.py

def set(self, key: ZeroMQKey, obj: BytesLike) -> None:
    """Set the object associated with a key.

    Note:
        The [`Connector`][proxystore.connectors.protocols.Connector]
        provides write-once, read-many semantics. Thus,
        [`set()`][proxystore.connectors.zmq.ZeroMQConnector.set]
        should only be called once per key, otherwise unexpected behavior
        can occur.

    Args:
        key: Key that the object will be associated with.
        obj: Object to associate with the key.
    """
    self._request([('put', key, obj)])

ZeroMQServer

ZeroMQServer()

In-memory storage and request handling for a server.

Use run_server() to serve requests from clients.

Source code in proxystore/connectors/zmq.py

def __init__(self) -> None:
    self.data: dict[str, BytesLike] = {}

handle

handle(frames: Sequence[Frame]) -> list[Any]

Process a request.

Parameters:

• frames (Sequence[Frame]) –

  Frames of the request, excluding the identity frame added by the router socket.

Returns:

• list[Any] –

  Frames of the reply, excluding the identity frame.

Source code in proxystore/connectors/zmq.py

def handle(self, frames: Sequence[zmq.Frame]) -> list[Any]:
    """Process a request.

    Args:
        frames: Frames of the request, excluding the identity frame
            added by the router socket.

    Returns:
        Frames of the reply, excluding the identity frame.
    """
    request_id = frames[0].bytes if len(frames) > 0 else b''
    try:
        header, data = self._handle(frames)
    except (TypeError, ValueError) as e:
        # JSON and Unicode decode errors are subclasses of ValueError.
        logger.debug('Failed to process request: %r', e)
        header = {'status': 'error', 'error': f'{type(e).__name__}: {e}'}
        data = None

    reply: list[Any] = [request_id, json.dumps(header).encode()]
    if data is not None:
        reply.append(data)
    return reply

run_server

run_server(
    address: str,
    port: int,
    *,
    stop: Event | None = None,
    poll_interval: float = 0.1,
) -> None

Serve requests from clients until stopped.

Parameters:

• address (str) –

  Address to bind to.

• port (int) –

  Port to bind to.

• stop (Event | None, default: None ) –

  Event which stops the server when set. If None, the server runs forever.

• poll_interval (float, default: 0.1 ) –

  Max time in seconds between checking stop.

Raises:

• ZMQError –

  If the server fails to bind to the address and port.

Source code in proxystore/connectors/zmq.py

def run_server(
    address: str,
    port: int,
    *,
    stop: threading.Event | None = None,
    poll_interval: float = 0.1,
) -> None:
    """Serve requests from clients until stopped.

    Args:
        address: Address to bind to.
        port: Port to bind to.
        stop: Event which stops the server when set. If `None`, the server
            runs forever.
        poll_interval: Max time in seconds between checking `stop`.

    Raises:
        zmq.ZMQError: If the server fails to bind to the address and port.
    """
    stop = threading.Event() if stop is None else stop
    server = ZeroMQServer()
    context: zmq.Context[zmq.Socket[bytes]] = zmq.Context()
    sock = context.socket(zmq.ROUTER)
    sock.setsockopt(zmq.LINGER, 0)

    try:
        sock.bind(f'tcp://{address}:{port}')
        logger.info('Server listening on tcp://%s:%d', address, port)
        while not stop.is_set():
            if not sock.poll(int(poll_interval * 1000)):
                continue
            identity, *frames = sock.recv_multipart(copy=False)
            reply = server.handle(frames)
            sock.send_multipart([identity, *reply], copy=False)
    finally:
        context.destroy(linger=0)

start_server

start_server(address: str, port: int) -> None

Run a server until SIGINT or SIGTERM is received.

This is the entry point of the process started by spawn_server().

Parameters:

• address (str) –

  Address to bind to.

• port (int) –

  Port to bind to.

Source code in proxystore/connectors/zmq.py

def start_server(address: str, port: int) -> None:  # pragma: no cover
    """Run a server until SIGINT or SIGTERM is received.

    This is the entry point of the process started by
    [`spawn_server()`][proxystore.connectors.zmq.spawn_server].

    Args:
        address: Address to bind to.
        port: Port to bind to.
    """
    stop = threading.Event()

    def _handler(signum: int, frame: FrameType | None) -> None:
        stop.set()

    signal.signal(signal.SIGINT, _handler)
    signal.signal(signal.SIGTERM, _handler)

    run_server(address, port, stop=stop)

spawn_server

spawn_server(
    address: str, port: int, *, timeout: float = 5
) -> Popen[bytes] | None

Spawn a server in a new process.

If another process spawns a server on the same address and port at the same time, the server spawned by this call will fail to bind and exit, and None is returned because the other server can be used instead.

Parameters:

• address (str) –

  Address the server will bind to.

• port (int) –

  Port the server will bind to.

• timeout (float, default: 5 ) –

  Max time in seconds to wait for the server to start.

Returns:

• Popen[bytes] | None –

  The process running the server or None if the server was spawned by a different process.

Raises:

• ServerTimeoutError –

  If the server does not start within timeout seconds.

Source code in proxystore/connectors/zmq.py

def spawn_server(
    address: str,
    port: int,
    *,
    timeout: float = 5,
) -> subprocess.Popen[bytes] | None:
    """Spawn a server in a new process.

    If another process spawns a server on the same address and port at the
    same time, the server spawned by this call will fail to bind and exit,
    and `None` is returned because the other server can be used instead.

    Args:
        address: Address the server will bind to.
        port: Port the server will bind to.
        timeout: Max time in seconds to wait for the server to start.

    Returns:
        The process running the server or `None` if the server was spawned \
        by a different process.

    Raises:
        ServerTimeoutError: If the server does not start within `timeout`
            seconds.
    """
    process = subprocess.Popen(
        [
            sys.executable,
            '-m',
            'proxystore.connectors.zmq',
            '--address',
            address,
            '--port',
            str(port),
        ],
    )

    try:
        pid = wait_for_server(address, port, timeout=timeout)
    except ServerTimeoutError:
        _kill_server(process)
        raise

    if pid != process.pid:
        # The server spawned by this call will fail to bind because another
        # server is already running. Wait for it to exit to avoid leaving a
        # zombie process.
        _kill_server(process)
        logger.info(
            'Using server at tcp://%s:%d spawned by another process (pid=%d)',
            address,
            port,
            pid,
        )
        return None

    logger.info(
        'Spawned server at tcp://%s:%d (pid=%d)',
        address,
        port,
        process.pid,
    )
    return process

wait_for_server

wait_for_server(
    address: str, port: int, timeout: float = 5
) -> int

Wait until a server responds.

Parameters:

• address (str) –

  Address of the server.

• port (int) –

  Port of the server.

• timeout (float, default: 5 ) –

  Max time in seconds to wait for the server to respond.

Returns:

• int –

  The process ID of the server.

Raises:

• ServerTimeoutError –

  If the server does not respond within timeout seconds.

Source code in proxystore/connectors/zmq.py

def wait_for_server(address: str, port: int, timeout: float = 5) -> int:
    """Wait until a server responds.

    Args:
        address: Address of the server.
        port: Port of the server.
        timeout: Max time in seconds to wait for the server to respond.

    Returns:
        The process ID of the server.

    Raises:
        ServerTimeoutError: If the server does not respond within `timeout`
            seconds.
    """
    deadline = time.monotonic() + timeout
    with zmq.Context() as context, context.socket(zmq.DEALER) as sock:
        sock.setsockopt(zmq.LINGER, 0)
        sock.connect(f'tcp://{address}:{port}')
        request_id = uuid.uuid4().bytes
        _send(sock, request_id, _request_header('ping'), None)

        while (remaining := deadline - time.monotonic()) > 0:
            if sock.poll(int(min(remaining, _PING_INTERVAL) * 1000)):
                reply_id, reply = _recv(sock)
                if reply_id == request_id:  # pragma: no branch
                    return reply.header['pid']

    raise ServerTimeoutError(
        f'Server at tcp://{address}:{port} did not respond within '
        f'{timeout} seconds.',
    )

get_interface_address

get_interface_address(interface: str) -> str

Get the IPv4 address of a network interface.

Warning

This function is only supported on Linux.

Parameters:

• interface (str) –

  Name of the network interface (e.g., 'eth0').

Returns:

• str –

  The IPv4 address of the interface.

Raises:

• NotImplementedError –

  If not on Linux.

• OSError –

  If the interface does not exist or has no IPv4 address.

Source code in proxystore/connectors/zmq.py

def get_interface_address(interface: str) -> str:
    """Get the IPv4 address of a network interface.

    Warning:
        This function is only supported on Linux.

    Args:
        interface: Name of the network interface (e.g., `'eth0'`).

    Returns:
        The IPv4 address of the interface.

    Raises:
        NotImplementedError: If not on Linux.
        OSError: If the interface does not exist or has no IPv4 address.
    """
    if sys.platform.startswith('linux'):  # pragma: linux cover
        import fcntl

        siocgifaddr = 0x8915
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
            request = struct.pack('256s', interface[:15].encode())
            result = fcntl.ioctl(s.fileno(), siocgifaddr, request)
        return socket.inet_ntoa(result[20:24])

    raise NotImplementedError(
        'Getting the address of an interface is only supported on Linux. '
        'Specify the address instead.',
    )