Fix bug

Author: AdrienVannson

doreisa/_scheduler.py

@@ -3,6 +3,7 @@
 from collections import Counter
 
 import ray
+import ray.actor
 from dask.core import get_dependencies
 
 from doreisa._scheduling_actor import ChunkRef, ScheduledByOtherActor
@@ -123,13 +124,8 @@ def explore(k) -> int:
 
     res_ref = scheduling_actors[partition[key]].get_value.remote(graph_id, key)
 
-    # At this point, get_value will not be called anymore for this graph.
-    ray.get(
-        [
-            actor.clear_graph.options(enable_task_events=False).remote(graph_id)
-            for id, actor in enumerate(scheduling_actors)
-            if partitionned_graphs[id]
-        ]
+    clear_graph.remote(
+        [actor for id, actor in enumerate(scheduling_actors) if partitionned_graphs[id]], res_ref, graph_id
     )
 
     if kwargs.get("ray_persist"):
@@ -144,3 +140,12 @@ def explore(k) -> int:
     if isinstance(keys[0], list):
         return [[res]]
     return [res]
+
+
+@ray.remote(max_retries=0, num_cpus=0)
+def clear_graph(scheduling_actors: list[ray.actor.ActorHandle], res: ray.ObjectRef, graph_id: int) -> None:
+    # Wait until the result is ready
+    ray.wait([res], fetch_local=False)
+
+    # Clear the graph
+    ray.get([actor.clear_graph.remote(graph_id) for actor in scheduling_actors])