deisa.ray.scheduling_actor module

class deisa.ray.scheduling_actor.NodeActorBase(actor_id: int, arrays_metadata: Dict[str, Dict] = {})[source]

Bases: object

Actor responsible for gathering chunks and exchanging data with analytics.

Each node actor is associated with a specific node and is responsible for:

  • Collecting chunks of arrays sent by simulation nodes (via Bridge)

  • Registering its owned chunks with the head node

  • Providing a small key/value channel (set, get, delete) for feedback between analytics and simulation.

The SchedulingActor subclass adds graph-scheduling behaviour on top of this base functionality.

Parameters:

actor_id (int) – Unique identifier for this node actor, typically derived from the node ID.

actor_id

The unique identifier for this node actor.

Type:

int

actor_handle

Handle to this actor instance.

Type:

RayActorHandle

head

Handle to the head node (HeadNodeActor).

Type:

RayActorHandle

partial_arrays

Per-array containers that capture metadata and chunk references owned by this node actor.

Type:

AsyncDict[str, PartialArray]

feedback

Dictionary for storing feedback values shared between analytics and simulation.

Type:

dict

async add_chunk(bridge_id: int, array_name: str, chunk_ref: list[ObjectRef], dtype, timestep: int, *, _ray_trace_ctx=None) → None[source]

Add a chunk of data to this node actor.

This method is called by Bridge instances to send chunks of arrays to this node actor. When all chunks from a node are received, the actor forwards a position->double-ref mapping to the head node.

Parameters:

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

  • array_name (str) – Name of the array receiving the chunk.

  • chunk_ref (list[ray.ObjectRef]) – Single-element list containing the Ray ObjectRef to the chunk data. The extra list level is kept for Dask compatibility.

  • dtype (np.dtype) – NumPy dtype read from the chunk before it was stored in Ray.

  • timestep (int) – Timestep index the chunk belongs to.

Return type:

None

Raises:

ContractError – If the array has not been registered via register_chunk_meta() before chunks are added.

Notes

This method manages chunk collection and coordination:

  1. Check array is expected

  2. Stores the ref in the per-timestep structure.

  3. When all local chunks have arrived, builds a {chunk_position: double_ref} mapping and sends it to the head actor via HeadNodeActor.chunks_ready().

  4. Pickles stored refs to drop in-memory handles and free memory.

  5. Signals or waits on the per-timestep event so callers block until the node’s share of chunks for the timestep is complete.

async finalize_registration(*, _ray_trace_ctx=None) → None[source]

Publish this node actor’s registered chunk ownership to the head actor.

Returns:

The method is idempotent; subsequent calls return immediately after successful finalization.

Return type:

None

Raises:

AssertionError – If the number of registered chunk metadata entries does not match the local chunk count collected from bridges.

Notes

The method waits until analytics has signaled readiness before registering partial arrays. A lock ensures concurrent calls from multiple bridges sharing the same actor perform registration once.

ready(*, _ray_trace_ctx=None) → None[source]

Check if the node actor is ready.

Returns:

Always returns None. This method serves as a readiness check for the actor.

Return type:

None

Notes

This method can be called to verify that the actor has been successfully initialized and is ready to receive requests. It is used by get_ready_actor_with_retry to ensure the actor is operational before returning its handle.

register_chunk_meta(bridge_id: int, array_name: str, chunk_shape, global_shape, chunk_position, *, _ray_trace_ctx=None) → None[source]

Register metadata for a chunk owned by this node actor.

Parameters:

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

  • array_name (str) – Name of the array the chunk belongs to.

  • chunk_shape (tuple[int, ...]) – Shape of this bridge’s chunk.

  • global_shape (tuple[int, ...]) – Full shape of the distributed array.

  • chunk_position (tuple[int, ...]) – Position of the chunk in the global chunk grid.

Raises:

AssertionError – If the derived global chunk grid for array_name is inconsistent with earlier registrations.