Add stimulus_id to SchedulerPlugin.update_graph hook

distributed/diagnostics/plugin.py

@@ -101,6 +101,7 @@ def update_graph(
         annotations: dict[str, dict[Key, Any]],
         priority: dict[Key, tuple[int | float, ...]],
         dependencies: dict[Key, set[Key]],
+        stimulus_id: str,
         **kwargs: Any,
     ) -> None:
         """Run when a new graph / tasks enter the scheduler
@@ -129,6 +130,8 @@ def update_graph(
                 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:
                 It is recommended to allow plugins to accept more parameters to
                 ensure future compatibility.

distributed/diagnostics/tests/test_scheduler_plugin.py

@@ -493,6 +493,7 @@ def update_graph(  # type: ignore
             annotations,
             priority,
             dependencies,
+            stimulus_id,
             **kwargs,
         ) -> None:
             assert scheduler is s
@@ -505,6 +506,7 @@ def update_graph(  # type: ignore
             assert len(priority) == 1
             assert isinstance(priority["foo"], tuple)
             assert dependencies == {"foo": set()}
+            assert stimulus_id is not None
             self.success = True
 
     plugin = UpdateGraph()
@@ -533,6 +535,7 @@ def update_graph(  # type: ignore
             annotations,
             priority,
             dependencies,
+            stimulus_id,
             **kwargs,
         ) -> None:
             assert scheduler is s
@@ -553,6 +556,7 @@ def update_graph(  # type: ignore
                 assert k in dependencies
             assert dependencies["f1"] == set()
             assert dependencies["sum"] == {"f1", "f3"}
+            assert stimulus_id is not None
 
             self.success = True
 

distributed/scheduler.py

@@ -4780,6 +4780,7 @@ def _create_taskstate_from_graph(
                     dependencies=dependencies,
                     annotations=dict(annotations_for_plugin),
                     priority=priority,
+                    stimulus_id=stimulus_id,
                 )
             except Exception as e:
                 logger.exception(e)
@@ -4846,6 +4847,7 @@ async def update_graph(
         stimulus_id: str | None = None,
     ) -> None:
         start = time()
+        stimulus_id = stimulus_id or f"update-graph-{start}"
         self._active_graph_updates += 1
         try:
             logger.debug("Received new graph. Deserializing...")
@@ -4923,7 +4925,7 @@ async def update_graph(
                 # objects. This should be removed
                 global_annotations=annotations,
                 start=start,
-                stimulus_id=stimulus_id or f"update-graph-{start}",
+                stimulus_id=stimulus_id,
             )
             task_state_created = time()
             metrics.update(