deisa.ray package

Submodules

Module contents

class deisa.ray.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

class deisa.ray.Deisa(feedback_queue_size: int = 1024, *args: Any, **kwargs: Any)[source]

Bases: IDeisa

Entry point that orchestrates analytics callbacks on Ray.

Provides an API for registering sliding window callbacks and executing them as arrays arrive from simulation ranks.

execute_callbacks() → None[source]

Execute the registered simulation callback loop.

Notes

Supports a single registered callback at present. Manages array retrieval from the head actor, windowed array delivery, and garbage collection between iterations.

register(*callback_args: CallbackArgs, exception_handler: ExceptionHandler = <function IDeisa.__default_exception_handler>, when: Literal['AND', 'OR']='AND') → Callable[source]

Decorator that registers a sliding-window analytics callback.

Parameters:

  • *callback_args (CallbackArgs) – Array descriptions the callback should receive.

  • exception_handler (Optional[Callable], optional) – Handler invoked when the user callback raises. Defaults to deisa.core.IDeisa.__default_exception_handler().

  • when (Literal["AND", "OR"], optional) – Governs whether all arrays ("AND") or any array ("OR") must be available before the callback runs. Defaults to "AND".

Returns:

Decorator that registers simulation_callback with the window handler.

Return type:

Callable

register_callback(callback: Callback, *callback_args: CallbackArgs, exception_handler: ExceptionHandler = <function IDeisa.__default_exception_handler>, when: Literal['AND', 'OR']='AND') → Callable[source]

Register the analytics callback and array descriptions.

Parameters:

  • callback (Callable) – Function to run for each iteration; receives arrays as kwargs and timestep.

  • *callback_args (CallbackArgs) – Descriptions of arrays to stream to the callback (with optional sliding windows). Maximum iterations to execute. Default is a large sentinel.

  • exception_handler (Optional[Callable]) – Exception handler to handle any exception thrown by simulation (like division by zero). Defaults to printing the error and moving on.

  • when (Literal['AND', 'OR']) – When callback have multiple arrays, govern when callback should be called. AND: only call callback if ALL required arrays have been shared for a given timestep. OR: call callback if ANY array has been shared for a given timestep.

Returns:

The original callback, allowing decorator-style usage.

Return type:

Callable

set(key: str, value: Any, timestep: int) → None[source]

Publish a feedback value for bridges.

Parameters:

  • key (Hashable) – Identifier for the shared value.

  • value (Any) – Value to store.

  • timestep (Hashable) – Timestep associated with value.

Notes

Timestamped values are stored in a fixed-size queue on the head actor. For a given key, timesteps must be strictly increasing; publishing the same timestep twice or publishing an older timestep raises ValueError. Bridges retrieve them collectively with bridge.get("foo", timestep=t).