deisa.ray.bridge module

Bridge between MPI ranks and the Ray-based analytics system.

This module exposes the Bridge class used by simulation ranks to register their data chunks and exchange information with analytics running on top of Ray.

class deisa.ray.bridge.Bridge(comm: ICommunicator, arrays_metadata: Dict[str, Dict], *args: Any, **kwargs: Any)[source]

Bases: IBridge

Bridge between MPI ranks and Ray cluster for distributed array processing.

Each Bridge instance is created by an MPI rank to connect to the Ray cluster and send data chunks. Each Bridge is responsible for managing a chunk of data from the decomposed distributed array.

Parameters:

  • comm (deisa.core.ICommunicator) – Communication backend for the simulation ranks. The bridge ID is derived from comm.Get_rank().

  • arrays_metadata (Mapping[str, Mapping[str, Any]]) – Metadata describing the array layout managed by this bridge.

  • _node_id (str or None, optional) – Node identifier used for testing or custom scheduling. Defaults to None.

  • scheduling_actor_cls (Type, optional) – Class used to materialize the scheduling actor. Defaults to deisa.ray.scheduling_actor.SchedulingActor.

  • _init_retries (int, optional) – Number of attempts to create and ready the node actor. Defaults to 3.

  • object_store_memory_timeout_s (float or None, optional) – Maximum time to wait for Ray object store memory before raising MemoryError. Defaults to 100 seconds. Set to None to wait forever.

node_id

The ID of the node this Bridge is associated with.

Type:

str

Notes

The Bridge automatically initializes Ray if it hasn’t been initialized yet. The scheduling actor is created with a detached lifetime to persist beyond the Bridge initialization. The actor uses node affinity scheduling to ensure it runs on the specified node.

Examples

Create a bridge for a simulation rank that owns one array temperature:

arrays_metadata = {
    "temperature": {
        "global_shape": (40, 40),
        "chunk_shape": (10, 10),
        "chunk_position": (0, 0),
    }
}
bridge = Bridge(
    arrays_metadata=arrays_metadata,
    comm=comm,
)

bridge.send(
    array_name="temperature",
    chunk=np.zeros((10, 10), dtype=np.float64),
    timestep=0
)
close(timestep: int) → None[source]

Close the bridge by signaling analytics that the simulation finished.

Parameters:

timestep (int) – The timestep index corresponding to the sentinel chunk.

get(key: str, timestep: int | None = None, default: Any = None) → list | Any | None[source]

Retrieve feedback from analytics to influence the simulation.

Bridge 0 queries the global head actor directly, then broadcasts the lookup result to every bridge in the communicator.

Parameters:

  • key (str) – The key that is being retrieved from the Analytics.

  • timestep (Optional[int], optional) – Timestep associated with the requested feedback value. When omitted, returns the entire retained queue for key.

  • default (Any, optional) – Value returned when no feedback exists for key and timestep. Defaults to None.

Notes

This remains a collective operation when a communicator is used: all bridges must call get in the same order so the broadcast completes. The retained feedback queue is fixed-size, so old entries may be dropped if analytics publishes more values than the queue can hold. Callback execution is intentionally one timestep behind: analytics processes a timestep only after a later timestep or the close sentinel arrives. As a result, feedback for the final simulated timestep may only be published after close and is not meant to drive another simulation step.

Returns:

The feedback value for timestep, the full retained queue when timestep is omitted, or default when no feedback exists.

Return type:

Any | None

Warning

Feedback timing is asynchronous and not reproducible run to run. The head queue may be populated at slightly different times, and this bridge may read it at slightly different times. Simulation code should decide how to react whenever a signal becomes available, and must not rely on exactly when an analytics event becomes visible for simulation correctness.

send(array_name: str, chunk: ndarray, timestep: int) → None[source]

Make a chunk of data available to the analytics.

This method stores the chunk in Ray’s object store, and sends a reference to the node actor. The method blocks until the data is processed by the node actor.

Parameters:

  • array_name (str) – The name of the array this chunk belongs to.

  • chunk (numpy.ndarray) – The chunk of data to be sent to the analytics.

  • timestep (int) – The timestep index for this chunk of data.

Notes

The chunk is stored in Ray’s object store with the node actor as the owner, ensuring the reference persists even after the simulation script terminates. This method blocks until the node actor has the chunk.

Raises:

ContractError – When the scheduling node detects a contract violation for the provided chunk.

Returns:

Blocks until the node actor processes the chunk.

Return type:

None