Migrate dependencies dependents sets from TaskState to Key

distributed/diagnostics/graph_layout.py

@@ -46,17 +46,15 @@ def __init__(self, scheduler):
                 priority=priority,
             )
 
-    def update_graph(
-        self, scheduler, *, dependencies=None, priority=None, tasks=None, **kwargs
-    ):
+    def update_graph(self, scheduler, *, priority=None, tasks=None, **kwargs):
         stack = sorted(
             tasks, key=lambda k: TupleComparable(priority.get(k, 0)), reverse=True
         )
         while stack:
             key = stack.pop()
             if key in self.x or key not in scheduler.tasks:
                 continue
-            deps = dependencies.get(key, ())
+            deps = [ts.key for ts in scheduler.tasks[key].dependencies]
             if deps:
                 if not all(dep in self.y for dep in deps):
                     stack.append(key)

distributed/diagnostics/plugin.py

@@ -100,7 +100,6 @@ def update_graph(
         tasks: list[Key],
         annotations: dict[str, dict[Key, Any]],
         priority: dict[Key, tuple[int | float, ...]],
-        dependencies: dict[Key, set[Key]],
         stimulus_id: str,
         **kwargs: Any,
     ) -> None:
@@ -128,8 +127,6 @@ def update_graph(
                     }
             priority:
                 Task calculated priorities as assigned to the tasks.
-            dependencies:
-                A mapping that maps a key to its dependencies.
             stimulus_id:
                 ID of the stimulus causing the graph update
             **kwargs:

distributed/diagnostics/tests/test_scheduler_plugin.py

@@ -492,7 +492,6 @@ def update_graph(  # type: ignore
             tasks,
             annotations,
             priority,
-            dependencies,
             stimulus_id,
             **kwargs,
         ) -> None:
@@ -505,7 +504,6 @@ def update_graph(  # type: ignore
             assert annotations == {}
             assert len(priority) == 1
             assert isinstance(priority["foo"], tuple)
-            assert dependencies == {"foo": set()}
             assert stimulus_id is not None
             self.success = True
 
@@ -534,7 +532,6 @@ def update_graph(  # type: ignore
             tasks,
             annotations,
             priority,
-            dependencies,
             stimulus_id,
             **kwargs,
         ) -> None:
@@ -552,10 +549,6 @@ def update_graph(  # type: ignore
             }
             assert len(priority) == len(tasks), priority
             assert priority["f2"][0] == -13
-            for k in keys:
-                assert k in dependencies
-            assert dependencies["f1"] == set()
-            assert dependencies["sum"] == {"f1", "f3"}
             assert stimulus_id is not None
 
             self.success = True