Stimulus ids scheduler

distributed/client.py

@@ -4276,6 +4276,31 @@ def collections_to_dsk(collections, *args, **kwargs):
         """Convert many collections into a single dask graph, after optimization"""
         return collections_to_dsk(collections, *args, **kwargs)
 
+    async def _story(self, keys=(), on_error="raise"):
+        assert on_error in ("raise", "ignore")
+
+        try:
+            flat_stories = await self.scheduler.get_story(keys=keys)
+            flat_stories = [("scheduler", *msg) for msg in flat_stories]
+        except Exception:
+            if on_error == "raise":
+                raise
+            elif on_error == "ignore":
+                flat_stories = []
+            else:
+                raise ValueError(f"on_error not in {'raise', 'ignore'}")
+
+        responses = await self.scheduler.broadcast(
+            msg={"op": "get_story", "keys": keys}, on_error=on_error
+        )
+        for worker, stories in responses.items():
+            flat_stories.extend((worker, *msg) for msg in stories)
+        return flat_stories
+
+    def story(self, *keys_or_stimulus_ids, on_error="raise"):
+        """Returns a cluster-wide story for the given keys or simtulus_id's"""
+        return self.sync(self._story, keys=keys_or_stimulus_ids, on_error=on_error)
+
     def get_task_stream(
         self,
         start=None,

distributed/cluster_dump.py

@@ -10,9 +10,9 @@
 import fsspec
 import msgpack
 
+from distributed._stories import scheduler_story as _scheduler_story
+from distributed._stories import worker_story as _worker_story
 from distributed.compatibility import to_thread
-from distributed.stories import scheduler_story as _scheduler_story
-from distributed.stories import worker_story as _worker_story
 
 DEFAULT_CLUSTER_DUMP_FORMAT: Literal["msgpack" | "yaml"] = "msgpack"
 DEFAULT_CLUSTER_DUMP_EXCLUDE: Collection[str] = ("run_spec",)

distributed/http/templates/task.html

@@ -118,16 +118,18 @@ <h3 class="title is-5"> Transition Log </h3>
                   <th> Key </th>
                   <th> Start </th>
                   <th> Finish </th>
+                  <th> Stimulus ID </th>
                   <th> Recommended Key </th>
                   <th> Recommended Action </th>
               </thead>
 
-              {% for key, start, finish, recommendations, transition_time in scheduler.story(Task) %}
+              {% for key, start, finish, recommendations, stimulus_id, transition_time in scheduler.story(Task) %}
               <tr>
                   <td> {{ fromtimestamp(transition_time) }} </td>
                   <td> <a href="{{ url_escape(key) }}.html">{{key}}</a> </td>
                   <td> {{ start }} </td>
                   <td> {{ finish }} </td>
+                  <td> {{ stimulus_id }} </td>
                   <td> </td>
                   <td> </td>
               </tr>
@@ -137,6 +139,7 @@ <h3 class="title is-5"> Transition Log </h3>
                   <td> </td>
                   <td> </td>
                   <td> </td>
+                  <td> </td>
                   <td> <a href="{{ url_escape(key2) }}.html">{{key2}}</a> </td>
                   <td> {{ rec }} </td>
               </tr>

distributed/nanny.py

@@ -289,7 +289,10 @@ async def _unregister(self, timeout=10):
         allowed_errors = (TimeoutError, CommClosedError, EnvironmentError, RPCClosed)
         with suppress(allowed_errors):
             await asyncio.wait_for(
-                self.scheduler.unregister(address=self.worker_address), timeout
+                self.scheduler.unregister(
+                    address=self.worker_address, stimulus_id=f"nanny-close-{time()}"
+                ),
+                timeout,
             )
 
     @property