Optimize graph sending (#27)

* Optimize graph sending

* Fix code quality

* Bump to 0.1.6

* Update lockfile

Author: AdrienVannson

doreisa/_scheduler.py

@@ -5,7 +5,7 @@
 import ray
 from dask.core import get_dependencies
 
-from doreisa._scheduling_actor import ChunkRef
+from doreisa._scheduling_actor import ChunkRef, ScheduledByOtherActor
 
 
 def doreisa_get(dsk, keys, **kwargs):
@@ -32,7 +32,7 @@ def log(message: str, debug_logs_path: str | None) -> None:
 
     # Find a not too bad scheduling strategy
     # Good scheduling in a tree
-    scheduling = {k: -1 for k in dsk.keys()}
+    partition = {k: -1 for k in dsk.keys()}
 
     # def explore(key, v: int):
     #     # Only works for trees for now
@@ -76,32 +76,43 @@ def explore(k) -> int:
         val = dsk[k]
 
         if isinstance(val, ChunkRef):
-            scheduling[k] = val.actor_id
+            partition[k] = val.actor_id
         else:
             res = [explore(dep) for dep in get_dependencies(dsk, k)]
-            scheduling[k] = Counter(res).most_common(1)[0][0]
+            partition[k] = Counter(res).most_common(1)[0][0]
 
-        return scheduling[k]
+        return partition[k]
 
     explore(key)
 
     log("2. Graph partitionning done", debug_logs_path)
 
-    # Pass the scheduling to the scheduling actors
-    dsk_ref, scheduling_ref = ray.put(dsk), ray.put(scheduling)  # noqa: F841
+    partitionned_graphs: dict[int, dict] = {}
 
-    log("3. Graph put in object store", debug_logs_path)
+    for k, v in dsk.items():
+        actor_id = partition[k]
+
+        if actor_id not in partitionned_graphs:
+            partitionned_graphs[actor_id] = {}
+
+        partitionned_graphs[actor_id][k] = v
+
+        for dep in get_dependencies(dsk, k):
+            if partition[dep] != actor_id:
+                partitionned_graphs[actor_id][dep] = ScheduledByOtherActor(partition[dep])
+
+    log("3. Partitionned graphs created", debug_logs_path)
 
     graph_id = random.randint(0, 2**128 - 1)
 
     ray.get(
         [
-            scheduling_actors[i].store_graph.options(enable_task_events=False).remote(graph_id, dsk_ref, scheduling_ref)
-            for i in range(len(scheduling_actors))
+            actor.store_graph.options(enable_task_events=False).remote(graph_id, partitionned_graphs[id])
+            for id, actor in enumerate(scheduling_actors)
         ]
     )
 
-    log("4. Partitionned graph sent", debug_logs_path)
+    log("4. Partitionned graphs sent", debug_logs_path)
 
     ray.get(
         [
@@ -112,7 +123,7 @@ def explore(k) -> int:
 
     log("5. Graph scheduled", debug_logs_path)
 
-    res = ray.get(ray.get(scheduling_actors[scheduling[key]].get_value.remote(graph_id, key)))
+    res = ray.get(ray.get(scheduling_actors[partition[key]].get_value.remote(graph_id, key)))
 
     log("6. End Doreisa scheduler", debug_logs_path)
 

doreisa/_scheduling_actor.py

@@ -6,7 +6,6 @@
 import ray
 import ray.actor
 import ray.util.dask.scheduler
-from dask.core import get_dependencies
 
 from doreisa import Timestep
 
@@ -30,6 +29,15 @@ class ChunkRef:
     _all_chunks: ray.ObjectRef | None = None
 
 
+@dataclass
+class ScheduledByOtherActor:
+    """
+    Represents a task that is scheduled by another actor in the part of the task graph sent to an actor.
+    """
+
+    actor_id: int
+
+
 class GraphInfo:
     """
     Information about graphs and their scheduling.
@@ -124,7 +132,7 @@ def __init__(self, actor_id: int) -> None:
         # For scheduling
         self.new_graph_available = asyncio.Event()
         self.graph_infos: dict[int, GraphInfo] = {}
-        self.partitionned_graphs: dict[int, tuple[dict, dict[str, int]]] = {}
+        self.partitionned_graphs: dict[int, dict] = {}
 
     def ready(self) -> None:
         pass
@@ -188,18 +196,18 @@ async def add_chunk(
         else:
             await self.chunks_ready_event.wait()
 
-    def store_graph(self, graph_id: int, dsk: dict, scheduling: dict[str, int]) -> None:
+    def store_graph(self, graph_id: int, dsk: dict) -> None:
         """
         Store the given graph in the actor until `schedule_graph` is called.
 
         This allows measuring precisely the time it takes to send the graph to all the
         actors. If needed, this will be optimized using an efficient communication
         method.
         """
-        self.partitionned_graphs[graph_id] = (dsk, scheduling)
+        self.partitionned_graphs[graph_id] = dsk
 
     async def schedule_graph(self, graph_id: int):
-        dsk, scheduling = self.partitionned_graphs.pop(graph_id)
+        dsk = self.partitionned_graphs.pop(graph_id)
 
         # Find the scheduling actors
         if not self.scheduling_actors:
@@ -210,22 +218,13 @@ async def schedule_graph(self, graph_id: int):
         self.new_graph_available.set()
         self.new_graph_available.clear()
 
-        local_keys = {k for k in dsk if scheduling[k] == self.actor_id}
-
-        dependency_keys: set[str] = {dep for k in local_keys for dep in get_dependencies(dsk, k)}  # type: ignore[assignment]
-
-        external_keys = dependency_keys - local_keys
-
-        # Filter the dask array
-        dsk = {k: v for k, v in dsk.items() if k in local_keys}
-
-        # Adapt external keys
-        for k in external_keys:
-            actor = self.scheduling_actors[scheduling[k]]
-            dsk[k] = actor.get_value.options(enable_task_events=False).remote(graph_id, k)
-
-        # Replace the false chunks by the real ObjectRefs
         for key, val in dsk.items():
+            # Adapt external keys
+            if isinstance(val, ScheduledByOtherActor):
+                actor = self.scheduling_actors[val.actor_id]
+                dsk[key] = actor.get_value.options(enable_task_events=False).remote(graph_id, key)
+
+            # Replace the false chunks by the real ObjectRefs
             if isinstance(val, ChunkRef):
                 assert val.actor_id == self.actor_id
 
@@ -234,7 +233,7 @@ async def schedule_graph(self, graph_id: int):
                 dsk[key] = pickle.loads(encoded_ref)
 
         # We will need the ObjectRefs of these keys
-        keys_needed = list(local_keys - dependency_keys)
+        keys_needed = list(dsk.keys())
 
         refs = await remote_ray_dask_get.remote(dsk, keys_needed)
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "doreisa"
-version = "0.1.5"
+version = "0.1.6"
 description = ""
 authors = [{ name = "Adrien Vannson", email = "adrien.vannson@protonmail.com" }]
 requires-python = ">=3.12"