maxwelllink.sockets.aggregated module

Two-layer socket aggregation for MaxwellLink.

This module adds an opt-in transport layer on top of the existing SocketHub implementation without modifying the original hub logic.

The new design introduces two roles:

• AggregatedSocketHub: an EM-side hub that keeps the public hub API expected by MaxwellLink solvers, but aggregates multiple molecule requests into one upstream connection per HPC node.

• LocalSocketHubBridge: a node-local bridge process/thread that talks to AggregatedSocketHub upstream while reusing an ordinary downstream SocketHub to fan out work to multiple existing Python/socket-only drivers.

This preserves existing SocketHub behavior while enabling a two-layer communication topology:

EM solver -> AggregatedSocketHub ==TCP==> LocalSocketHubBridge
          -> local SocketHub ==TCP/UNIX==> many molecular drivers

class maxwelllink.sockets.aggregated.AggregatedBridge

Bases: object

Convenience handle for one hub-owned local bridge.

Instances of this class are returned by AggregatedSocketHub.add_bridge(). They provide a light wrapper around LocalSocketHubBridge so existing input scripts only need to:

1. create bridge handles from the hub,

2. attach molecules to a handle via append(), and

3. launch downstream drivers against address.

Parameters:

• hub (AggregatedSocketHub) – Owning hub that created this handle.

• group_id (str) – Aggregate group identifier this handle manages.

• bridge (LocalSocketHubBridge) – The underlying node-local bridge this handle wraps.

hub

The owning hub.

Type:

AggregatedSocketHub

group_id

The aggregate group identifier this handle manages.

Type:

str

__init__(*, hub, group_id, bridge)

Parameters:

• hub (AggregatedSocketHub)

• group_id (str)

• bridge (LocalSocketHubBridge)

property address: str

Address string downstream UNIX-socket drivers should use.

Returns:

The configured UNIX-socket address for local drivers.

Return type:

str

Raises:

RuntimeError – If this bridge was not configured with a UNIX socket.

append(molecules)

Attach one molecule or an iterable of molecules to this bridge group.

The helper only mutates molecule.init_payload["aggregate_group"] and therefore works with existing mxl.Molecule / SocketMolecule objects without changing solver-side logic.

Parameters:

molecules (molecule or iterable of molecules) – One molecule or an iterable of molecules to assign to this group.

Raises:

• TypeError – If an item does not expose an init_payload attribute.

• ValueError – If a molecule belongs to a different hub or is already assigned to a different aggregate group.

Return type:

None

property local_endpoint: dict

Return the downstream endpoint mapping for driver launch code.

Returns:

A copy of the downstream endpoint mapping, with either a "unixsocket" key or "host"/"port" keys.

Return type:

dict

start()

Start the underlying local bridge thread.

Returns:

The (daemon) thread running the bridge loop.

Return type:

threading.Thread

stop(wait=2.0)

Stop the underlying local bridge.

Parameters:

wait (float, default: 2.0) – Maximum time (seconds) to wait for the bridge thread to join.

Return type:

None

property unixsocket: str | None

Configured UNIX-socket driver address, if any.

Returns:

The configured UNIX-socket address, or None when the bridge uses TCP downstream.

Return type:

str or None

property unixsocket_path: str | None

Resolved filesystem path for the local UNIX socket.

Returns:

The resolved socket path, or None when no UNIX socket is used.

Return type:

str or None

class maxwelllink.sockets.aggregated.AggregatedSocketHub

Bases: SocketHub

EM-side hub that aggregates multiple molecule requests into one bridge link.

This class keeps the same public methods used by MaxwellLink solvers (register_molecule_return_id, wait_until_bound, all_bound, step_barrier) while mapping many molecule IDs onto a smaller number of bridge connections.

Molecules are assigned to a bridge group through init_payload["aggregate_group"]. All molecules sharing the same group are sent together to one LocalSocketHubBridge.

Parameters:

• host (str or None, optional) – Interface to bind the upstream TCP server to. None, "", or "0.0.0.0" bind all interfaces; bridges then connect back over 127.0.0.1.

• port (int or None, default: 31415) – TCP port for the upstream server.

• timeout (float, default: 60000.0) – Default operation timeout (seconds) used for binding and stepping.

• latency (float, default: 0.01) – Polling interval (seconds) for the bind/step loops.

__init__(host=None, port=31415, timeout=60000.0, latency=0.01)

Initialize the socket hub.

Parameters:

• host (str or None, default: None) – Host address for AF_INET sockets. Ignored when using a UNIX socket.

• port (int or None, default: 31415) – TCP port for AF_INET sockets. Ignored for UNIX sockets.

• unixsocket (str or None, default: None) – Path (or name under /tmp/socketmxl_*) for a UNIX domain socket. When provided, host and port are ignored.

• timeout (float, default: 60000.0) – Socket timeout (seconds) for client operations.

• latency (float, default: 0.01) – Polling sleep (seconds) between hub sweeps; can be very small for local runs.

add_bridge(local_unixsocket)

Create, start, and return one hub-owned local UNIX-socket bridge.

This is the convenience entry point intended for minimal edits when migrating an existing single-layer SocketHub script to the new two-layer transport.

Parameters:

local_unixsocket (str) – Non-empty downstream UNIX-socket address local drivers connect to.

Returns:

A started handle wrapping the new node-local bridge.

Return type:

AggregatedBridge

Raises:

ValueError – If local_unixsocket is empty or already owned by another bridge on this hub.

all_bound(molecule_ids, require_init=True)

Check if all given molecule IDs are bound (and optionally initialized).

Parameters:

• molecule_ids (iterable of int) – Molecule IDs to check.

• require_init (bool, default: True) – Also require that clients completed INIT.

Returns:

True if all are bound (and initialized if requested), else False.

Return type:

bool

graceful_shutdown(reason=None, wait=2.0)

Politely ask all connected drivers to exit and wait briefly for BYE.

Parameters:

• reason (str or None, optional) – Optional reason to log for shutdown.

• wait (float, default: 2.0) – Seconds to wait for clean replies.

init_remote_bridges(molecules, *, molecules_per_bridge, unix_prefix='bridge_', save_file='aggregation.json')

Partition molecules across remote bridge groups and save a manifest.

This helper does not start any bridge threads locally. Instead it assigns molecule.init_payload["aggregate_group"] for each molecule and writes one JSON manifest that bridge-node scripts can consume via run_bridge_node().

Parameters:

• molecules (molecule or iterable of molecules) – Molecules to distribute across remote bridges.

• molecules_per_bridge (int) – Maximum number of molecules assigned to one bridge.

• unix_prefix (str, default: "bridge_") – Prefix used to generate downstream UNIX socket names f"{unix_prefix}{idx}".

• save_file (str, default: "aggregation.json") – Path where the bridge manifest should be written.

Returns:

The generated bridge specifications in order.

Return type:

list[RemoteBridgeSpec]

Raises:

ValueError – If no molecules are supplied or molecules_per_bridge is not a positive integer.

Notes

This method records the generated specs on self.remote_bridges and the full manifest on self.remote_bridge_info as a side effect, but does not start any bridge threads.

register_molecule(molecule_id)

Reserve a slot for a given molecule ID (client may connect later).

Parameters:

molecule_id (int) – Molecule ID to register.

Raises:

ValueError – If the molecule ID is already registered.

Return type:

None

register_molecule_return_id()

Reserve a slot for a molecule and return an auto-assigned ID.

Returns:

The assigned unique molecule ID.

Return type:

int

remote_bridge_info: dict | None

remote_bridges: list[RemoteBridgeSpec]

step_barrier(requests, timeout=None)

Dispatch all requested fields group-by-group and collect grouped replies.

Parameters:

• requests (dict[int, dict]) – Mapping from molecule ID to a request dict with keys: - "efield_au" : array-like (3,) field vector in a.u. - "init" : dict, optional INIT payload for a first bind.

• timeout (float, optional) – Maximum time (seconds) to wait for every group to reply. Defaults to the hub’s timeout setting.

Returns:

Mapping molid -> {"amp": ndarray(3,), "extra": bytes}, matching the SocketHub.step_barrier contract. Returns {} when paused, when a molecule is not yet bound, or on a mid-step disconnect or timeout.

Return type:

dict[int, dict]

Raises:

RuntimeError – If a bridge replies with the wrong set of molecule ids.

Notes

A single pending group is served by a direct blocking receive; multiple groups are awaited on the aggregate selector so whichever bridge becomes readable first is collected next.

stop()

Stop the aggregate hub and clean up bridge groups coherently.

The base SocketHub.stop() assumes one client per molecule, which is not true here. This override shuts down each bridge once and clears all molecule bindings associated with that bridge.

wait_until_bound(init_payloads, require_init=True, timeout=None)

Wait until all requested molecules are served by initialized bridges.

Molecules are grouped through init_payload["aggregate_group"] and each group must be backed by exactly one connected bridge.

Parameters:

• init_payloads (dict[int, dict]) – Mapping from molecule ID to INIT payload to use on bind.

• require_init (bool, default: True) – Also require that each backing bridge completed its AGGINIT handshake.

• timeout (float or None, optional) – Maximum time to wait (seconds). Uses the hub default if None.

Returns:

True if every requested molecule became bound (and, when require_init is set, initialized) within the time limit, else False.

Return type:

bool

class maxwelllink.sockets.aggregated.LocalSocketHubBridge

Bases: object

Bridge process/thread that fans out aggregate requests to a local SocketHub.

Upstream:

one TCP connection to AggregatedSocketHub

Downstream:

one ordinary SocketHub using either TCP or UNIX sockets, connected to many existing MaxwellLink socket drivers.

Parameters:

• group_id (str) – Non-empty aggregate group identifier this bridge serves.

• upstream_host (str) – Host of the upstream AggregatedSocketHub.

• upstream_port (int) – TCP port of the upstream hub.

• timeout (float, default: 60.0) – Operation timeout (seconds) for both the upstream link and the downstream local hub.

• latency (float, default: 0.01) – Polling interval (seconds) propagated to the downstream local hub.

• local_host (str, default: "127.0.0.1") – Downstream bind host, used only when local_unixsocket is None.

• local_port (int or None, optional) – Downstream TCP port. Ignored when a UNIX socket is used.

• local_unixsocket (str or None, optional) – Downstream UNIX-socket address. When both this and local_port are None, a sanitized name derived from group_id is used.

Raises:

ValueError – If group_id is empty.

__init__(*, group_id, upstream_host, upstream_port, timeout=60.0, latency=0.01, local_host='127.0.0.1', local_port=None, local_unixsocket=None)

Parameters:

• group_id (str)

• upstream_host (str)

• upstream_port (int)

• timeout (float)

• latency (float)

• local_host (str)

• local_port (int | None)

• local_unixsocket (str | None)

property local_endpoint: dict

Return the downstream socket endpoint local drivers should connect to.

Returns:

{"unixsocket": <name>} when a UNIX socket is configured, otherwise {"host": <host>, "port": <port>}.

Return type:

dict

run()

Run the bridge loop until the hub sends STOP or disconnects.

Raises:

RuntimeError – If the upstream hub sends an unrecognized aggregate header.

Return type:

None

Notes

Connects upstream (with retry), sends HELLO, then services AGGINIT, AGGSTEP, and STOP frames in a loop. Upstream transport errors end the loop quietly; the downstream local hub is always stopped on exit.

start()

Start the bridge loop in a daemon thread and return the thread handle.

Returns:

The running daemon thread. If a thread is already alive it is returned unchanged rather than starting a second one.

Return type:

threading.Thread

stop(wait=2.0)

Stop the bridge loop and close the downstream local hub.

Parameters:

wait (float, default: 2.0) – Maximum time (seconds) to wait for the bridge thread to join after signalling it to stop.

Return type:

None

class maxwelllink.sockets.aggregated.RemoteBridgeSpec

Bases: object

One remote aggregate bridge entry produced by init_remote_bridges.

idx

Zero-based bridge index used by run_bridge_node().

Type:

int

group_id

Aggregate group identifier transmitted upstream.

Type:

str

unixsocket

Downstream UNIX-socket address local drivers should connect to.

Type:

str

n_molecules

Number of molecules assigned to this bridge.

Type:

int

__init__(idx, group_id, unixsocket, n_molecules)

Parameters:

• idx (int)

• group_id (str)

• unixsocket (str)

• n_molecules (int)

Return type:

None

classmethod from_dict(payload)

Build one bridge specification from JSON-decoded manifest data.

Parameters:

payload (Mapping) – Mapping carrying idx, group_id, unixsocket, and n_molecules entries, as written by to_dict().

Returns:

The reconstructed, type-coerced specification.

Return type:

RemoteBridgeSpec

Raises:

KeyError – If a required field is missing from payload.

group_id: str

idx: int

n_molecules: int

to_dict()

Return a JSON-serializable bridge specification mapping.

Returns:

Mapping with the idx, group_id, unixsocket, and n_molecules fields coerced to plain JSON types.

Return type:

dict

unixsocket: str

maxwelllink.sockets.aggregated.mxl_bridge_main(argv=None)

CLI entry point for running one aggregate bridge from a manifest.

Parameters:

argv (list[str] or None, optional) – Argument vector to parse. None (the default) parses sys.argv[1:].

Returns:

Process exit code; 0 on a clean shutdown.

Return type:

int

Examples

mxl_bridge --info aggregation.json --idx 0

maxwelllink.sockets.aggregated.run_bridge_node(info='aggregation.json', *, idx=0)

Run one bridge node from a manifest written by init_remote_bridges.

Parameters:

• info (str or path-like, default: "aggregation.json") – JSON manifest written by AggregatedSocketHub.init_remote_bridges().

• idx (int, default: 0) – Zero-based bridge index identifying which bridge entry in info this node should start.

Raises:

IndexError – If no bridge entry in info has the requested idx.

Return type:

None

Notes

The call blocks until the bridge thread exits or a KeyboardInterrupt is received, after which the bridge is stopped on a best-effort basis.