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.
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.
Bases: Exception
Client timed out waiting for a response from a server.
Bases: Exception
Server failed to process a request.
Bases: NamedTuple
Key to objects stored in a ZeroMQ server.
Attributes:
-
obj_id
(str)
–
-
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(
port: int,
address: str | None = None,
interface: str | None = None,
timeout: float = 5,
request_timeout: float | None = 60,
)
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, 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.
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,
) -> None:
self._address = address
self._interface = interface
self.port = port
self.timeout = timeout
self.request_timeout = request_timeout
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()
|
property
URL of the server on this host.
close(kill_server: bool = False) -> None
Close the connector.
Parameters:
-
kill_server
(bool, default:
False
)
–
Stop the server on this host if it was spawned by
this connector. This will lose all objects stored in the
server.
Source code in proxystore/connectors/zmq.py
| def close(self, kill_server: bool = False) -> None:
"""Close the connector.
Args:
kill_server: Stop the server on this host if it was spawned by
this connector. This will lose all objects stored in the
server.
"""
self._pool.close()
if kill_server 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
|
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,
}
|
classmethod
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(**config)
|
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)])
|
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 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 a batch of serialized objects associated with the keys.
Parameters:
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]
|
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 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 a batch of serialized objects in the store.
Parameters:
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 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)])
|
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] = {}
|
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(
address: str,
port: int,
*,
stop: Event | None = None,
poll_interval: float = 0.1,
) -> None
Serve requests from clients until stopped.
Parameters:
-
address
(str)
–
-
port
(int)
–
-
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(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:
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 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:
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 until a server responds.
Parameters:
-
address
(str)
–
-
port
(int)
–
-
timeout
(float, default:
5
)
–
Max time in seconds to wait for the server to respond.
Returns:
-
int
–
The process ID of the server.
Raises:
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(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:
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.',
)
|