deisa.ray.types module

class deisa.ray.types.ArrayPerTimestep[source]

Bases: object

Internal class tracking chunks for a specific array and timestep.

Tracks the per-timestep chunks owned by a single scheduling actor. Each instance is keyed by timestep inside a PartialArray.

chunks_ready_event

Triggered when all chunks for this timestep owned by the node actor have arrived and been forwarded to the head actor.

Type:

asyncio.Event

local_chunks

Mapping of bridge_id to the double ObjectRef for that chunk. Once forwarded to the head actor the value is replaced with pickled bytes to free memory.

Type:

AsyncDict[int, ray.ObjectRef | bytes]

dtype

Dtype shared by all chunks received for this array timestep.

Type:

np.dtype or None

class deisa.ray.types.ChunkRef(ref: ObjectRef, actorid: int, array_name: str, timestep: Hashable, bridge_id: int)[source]

Bases: object

Represents a chunk of an array in a Dask task graph.

This class is used as a placeholder in Dask task graphs to represent a chunk of data. The task corresponding to this object must be scheduled by the actor who has the actual data. This class is used since Dask tends to inline simple tuples, which would prevent proper scheduling.

Parameters:

  • ref (ray.ObjectRef) – Ray ObjectRef that eventually points to the chunk data. This is a ref of a ref produced by the patched Dask scheduler.

  • actorid (int) – The ID of the scheduling actor that owns this chunk.

  • array_name (str) – The real name of the array, without the timestep suffix.

  • timestep (Timestep) – The timestep this chunk belongs to.

  • bridge_id (int) – Identifier of the bridge that produced this chunk.

Notes

This class is used to prevent Dask from inlining simple tuples in the task graph, which would break the scheduling mechanism. The behavior may change in newer versions of Dask.

class deisa.ray.types.DaskArrayData(name)[source]

Bases: object

Information about a Dask array being built.

Tracks metadata and per-timestep state for a Dask array assembled from chunks sent by scheduling actors.

Parameters:

name (str) – Array name registered with the head actor.

name

Array name without timestep suffix.

Type:

str

fully_defined

Set when every chunk owner has been registered.

Type:

asyncio.Event

nb_chunks_per_dim

Number of chunks per dimension. Set when first chunk owner is registered.

Type:

tuple[int, …] or None

nb_chunks

Total number of chunks in the array. Set when first chunk owner is registered.

Type:

int or None

chunks_size

For each dimension, the size of chunks in that dimension. None values indicate unknown chunk sizes.

Type:

list[list[int | None]] or None

dtype_by_timestep

NumPy dtype for each timestep, set when chunk data arrives.

Type:

dict[Timestep, np.dtype]

position_to_node_actorID

Mapping from chunk position to the scheduling actor responsible for that chunk.

Type:

dict[tuple[int, …], int]

position_to_bridgeID

Mapping from chunk position to the producing bridge ID.

Type:

dict[tuple, int]

nb_scheduling_actors

Number of unique scheduling actors owning chunks of this array. Set when all chunk owners are known.

Type:

int or None

chunk_refs

For each timestep, the list of per-actor references that keep chunk payloads alive in the object store.

Type:

dict[Timestep, list[ray.ObjectRef]]

pos_to_ref_by_timestep

For each timestep, the (position, ref) pairs provided by scheduling actors. Used when distributed scheduling is disabled.

Type:

defaultdict

add_chunk_ref(chunk_ref: ObjectRef, timestep: Hashable, pos_to_ref: dict[tuple, ObjectRef]) → bool[source]

Add a reference sent by a scheduling actor.

Parameters:

  • chunk_ref (ray.ObjectRef) – Ray object reference to a chunk sent by a scheduling actor.

  • timestep (Timestep) – The timestep this chunk belongs to.

  • pos_to_ref (dict[tuple, ray.ObjectRef]) – Mapping of chunk position to the (double) Ray ObjectRef provided by the scheduling actor for this timestep.

Returns:

True if all chunks for this timestep are ready (i.e., all scheduling actors have sent their chunks), False otherwise.

Return type:

bool

Notes

This method adds the chunk reference to the list for the given timestep. It returns True only when all chunk owners are known and all have sent their chunks for this timestep.

get_full_array(timestep: Hashable, *, distributing_scheduling_enabled: bool) → Array[source]

Return the full Dask array for a given timestep.

Parameters:

timestep (Timestep) – The timestep for which the full array should be returned.

Returns:

A Dask array representing the full decomposed array for the given timestep. The array is constructed from chunk references stored in Ray’s object store.

Return type:

da.Array

Raises:

AssertionError – If not all chunk owners have been registered, or if the number of chunks is inconsistent.

Notes

When distributed scheduling is enabled the graph uses ChunkRef placeholders that keep data owner information. Otherwise the concrete chunk payloads are inlined. Chunk reference lists are deleted after embedding in the graph to avoid leaking memory.

update_dtype(timestep: Hashable, dtype: dtype) → None[source]

Record or validate the dtype for an array timestep.

Parameters:

  • timestep (Timestep) – Timestep whose dtype is being recorded.

  • dtype (numpy.dtype) – Dtype reported for chunks in this timestep.

Raises:

AssertionError – If a different dtype was already registered for timestep.

update_meta(nb_chunks_per_dim: tuple[int, ...], position: tuple[int, ...], size: tuple[int, ...], node_actor_id: int, bridge_id: int) → None[source]

Register a scheduling actor as the owner of a chunk at a specific position.

This method records which scheduling actor is responsible for a chunk and updates the array metadata. If this is the first chunk registered, it initializes the array dimensions.

Parameters:

  • nb_chunks_per_dim (tuple[int, ...]) – Number of chunks per dimension in the array decomposition.

  • position (tuple[int, ...]) – The position of the chunk in the array decomposition.

  • size (tuple[int, ...]) – The size of the chunk along each dimension.

  • node_actor_id (int) – Scheduling actor that owns this chunk.

  • bridge_id (int) – Bridge identifier that produced the chunk (used for lookups).

Raises:

AssertionError – If the chunk position is out of bounds, or if subsequent chunks have inconsistent dimensions or sizes compared to the first chunk.

class deisa.ray.types.DeisaArray(dask, name, chunks, t: int, dtype=None, meta=None, shape=None)[source]

Bases: Array

property timestep: int

Return the simulation timestep associated with this array.

Returns:

Timestep provided when the array was created.

Return type:

int

to_hdf5(fname: str, dataset: str) → None[source]

Save data to a HDF5 file (using HDF5 VDS).

Parameters:

  • fname (str) – The name of the final file where the data will be stored.

  • dataset (str) – The name of the dataset in the hdf5 where the data will be stored.

Notes

This method creates files for each chunk, then links the files using HDF5 VDS. The chunk file is named as .{filename}-{chunk_position}.h5.

to_zarr(url, component=None, storage_options=None, region=None, compute=True, return_stored=False, mode='a') → None[source]

Save this array using the Zarr storage format.

Parameters:

  • url (str or pathlib.Path) – Directory, store, or path-like target passed to dask.array.to_zarr().

  • component (str or None, optional) – Component path inside the Zarr store.

  • storage_options (dict or None, optional) – Storage options forwarded to Dask for fsspec-backed targets.

  • region (tuple[slice, ...] or None, optional) – Region to write within an existing array.

  • compute (bool, optional) – Whether to compute immediately. Defaults to True.

  • return_stored (bool, optional) – Whether Dask should return stored chunks. Defaults to False.

  • mode (str, optional) – Zarr write mode. Defaults to "a".

Returns:

Return value from dask.array.to_zarr().

Return type:

Any

Notes

This method persists the array first, resolves local path-like targets to absolute paths, and then delegates to dask.array.to_zarr().

class deisa.ray.types.GraphInfo[source]

Bases: object

Information about graphs and their scheduling.

Tracks scheduling status and produced references for a Dask task graph scheduled by a SchedulingActor.

scheduled_event

Event set once the graph has been submitted to Ray.

Type:

asyncio.Event

refs

Mapping from task key to the double Ray ObjectRef returned by the patched Dask-on-Ray scheduler.

Type:

dict[str, ray.ObjectRef]

class deisa.ray.types.PartialArray[source]

Bases: object

Internal class tracking metadata and chunks for an array.

Maintains metadata for a single array on one scheduling actor, including the set of locally owned chunks and per-timestep chunk collections.

ready_event

Set once all local bridges have provided metadata for the array and the head actor has been notified.

Type:

asyncio.Event

chunks_contained_meta

Metadata tuples (bridge_id, chunk_position, chunk_size) for the chunks owned by this actor.

Type:

set[tuple[int, tuple[int, …], tuple[int, …]]]

bid_to_pos

Maps bridge_id to the corresponding chunk position for quick lookup when forwarding payloads.

Type:

dict[int, tuple]

per_timestep_arrays

Per-timestep structures that hold chunk references as they arrive.

Type:

AsyncDict[Timestep, ArrayPerTimestep]

class deisa.ray.types.ScheduledByOtherActor(actor_id: ActorID)[source]

Bases: object

Represents a task that is scheduled by another actor.

This class is used as a placeholder in Dask task graphs to indicate that a task should be scheduled by a different actor. When a task graph is sent to an actor, tasks marked with this class will be delegated to the specified actor.

Parameters:

actor_id (int) – The ID of the scheduling actor that should schedule this task.

Notes

This is used to handle cross-actor task dependencies in distributed Dask computations where different parts of the task graph are handled by different scheduling actors.

class deisa.ray.types.Window(name: str, size: int = 1)[source]

Bases: object

Description of an array with optional windowing support.

Parameters:

  • name (str) – Name of the array passed to a registered analytics callback.

  • size (int, optional) – Number of latest timesteps to pass for this array. Defaults to 1.

Examples

Request the latest five temperature arrays in a callback:

deisa.register_callback(callback, Window("temperature", size=5))
deisa.ray.types.chunk_fname(fname: str, dataset: str, block_id: tuple[int, ...] = ())[source]

Create the hidden HDF5 filename for one stored chunk.

Parameters:

  • fname (str) – Final HDF5 VDS file path.

  • dataset (str) – Dataset name inside the final HDF5 file.

  • block_id (tuple[int, ...], optional) – Chunk grid position. Defaults to ().

Returns:

Hidden chunk file path next to fname.

Return type:

pathlib.Path

deisa.ray.types.create_vds(fname: str, dataset: str, chunk_shape: tuple[int, ...], data_shape: tuple[int, ...], nb_chunks_per_dim: tuple[int, ...], data_dtype: dtype) → None[source]

Create an HDF5 virtual dataset aggregating all chunk files.

Parameters:

  • fname (str) – Final HDF5 VDS file path.

  • dataset (str) – Dataset name to create in fname.

  • chunk_shape (tuple[int, ...]) – Shape of each chunk.

  • data_shape (tuple[int, ...]) – Shape of the full logical dataset.

  • nb_chunks_per_dim (tuple[int, ...]) – Number of chunks along each dimension.

  • data_dtype (numpy.dtype) – NumPy dtype of the dataset.

Returns:

Creates or updates fname with the virtual dataset.

Return type:

None

deisa.ray.types.save_chunk(chunk: ndarray, fname: str, dataset: str, block_id: tuple[int, ...] | None = None) → None[source]

Save one chunk to an individual HDF5 file.

Parameters:

  • chunk (numpy.ndarray) – Chunk payload to write.

  • fname (str) – Final HDF5 VDS file path used to derive the chunk filename.

  • dataset (str) – Dataset name to create in the chunk file.

  • block_id (tuple[int, ...] or None, optional) – Chunk grid position used in the chunk filename.

Returns:

Writes the chunk file for later VDS assembly.

Return type:

None

deisa.ray.types.to_hdf5(fname: str, sources: dict[str, DeisaArray]) → None[source]

Save data to a HDF5 file (using HDF5 VDS).

Parameters:

  • fname (str) – Final HDF5 file path.

  • sources (dict[str, DeisaArray]) – Mapping from dataset names in fname to arrays to persist.

Returns:

Creates chunk files and the virtual datasets in fname.

Return type:

None

Notes

This method creates files for each chunk, then links the files using HDF5 VDS. The chunk file is named as .{filename}-{dataset}-{chunk_position}.h5. This may cause errors with datasets with special characters in the name.