Stimulus id contextvars explicit handlers

distributed/client.py

@@ -209,6 +209,7 @@ def __init__(self, key, client=None, inform=True, state=None):
                     "op": "client-desires-keys",
                     "keys": [stringify(key)],
                     "client": self.client.id,
+                    "stimulus_id": f"client-desires-keys-{time()}",
                 }
             )
 
@@ -472,6 +473,7 @@ def __setstate__(self, state):
                 "tasks": {},
                 "keys": [stringify(self.key)],
                 "client": c.id,
+                "stimulus_id": f"stimulus-id-{time()}",
             }
         )
 
@@ -1265,6 +1267,7 @@ async def _ensure_connected(self, timeout=None):
                     "client": self.id,
                     "reply": False,
                     "versions": version_module.get_versions(),
+                    "stimulus_id": f"client-ensure-connected-{time()}",
                 }
             )
         except Exception:
@@ -1371,17 +1374,22 @@ def _dec_ref(self, key):
             self.refcount[key] -= 1
             if self.refcount[key] == 0:
                 del self.refcount[key]
-                self._release_key(key)
+                self._release_key(key, f"client-release-key-{time()}")
 
-    def _release_key(self, key):
+    def _release_key(self, key, stimulus_id: str):
         """Release key from distributed memory"""
         logger.debug("Release key %s", key)
         st = self.futures.pop(key, None)
         if st is not None:
             st.cancel()
         if self.status != "closed":
             self._send_to_scheduler(
-                {"op": "client-releases-keys", "keys": [key], "client": self.id}
+                {
+                    "op": "client-releases-keys",
+                    "keys": [key],
+                    "client": self.id,
+                    "stimulus_id": stimulus_id,
+                }
             )
 
     async def _handle_report(self):
@@ -1506,7 +1514,9 @@ async def _close(self, fast=False):
                 and self.scheduler_comm.comm
                 and not self.scheduler_comm.comm.closed()
             ):
-                self._send_to_scheduler({"op": "close-client"})
+                self._send_to_scheduler(
+                    {"op": "close-client", "stimulus_id": f"client-close-{time()}"}
+                )
                 self._send_to_scheduler({"op": "close-stream"})
 
             current_task = asyncio.current_task()
@@ -1527,8 +1537,10 @@ async def _close(self, fast=False):
             ):
                 await self.scheduler_comm.close()
 
+            stimulus_id = f"client-close-{time()}"
+
             for key in list(self.futures):
-                self._release_key(key=key)
+                self._release_key(key=key, stimulus_id=stimulus_id)
 
             if self._start_arg is None:
                 with suppress(AttributeError):
@@ -2110,12 +2122,20 @@ async def _gather_remote(self, direct, local_worker):
                 response = {"status": "OK", "data": data2}
                 if missing_keys:
                     keys2 = [key for key in keys if key not in data2]
-                    response = await retry_operation(self.scheduler.gather, keys=keys2)
+                    response = await retry_operation(
+                        self.scheduler.gather,
+                        keys=keys2,
+                        stimulus_id=f"client-gather-remote-{time()}",
+                    )
                     if response["status"] == "OK":
                         response["data"].update(data2)
 
             else:  # ask scheduler to gather data for us
-                response = await retry_operation(self.scheduler.gather, keys=keys)
+                response = await retry_operation(
+                    self.scheduler.gather,
+                    keys=keys,
+                    stimulus_id=f"client-gather-remote-{time()}",
+                )
 
         return response
 
@@ -2201,6 +2221,8 @@ async def _scatter(
             d = await self._scatter(keymap(stringify, data), workers, broadcast)
             return {k: d[stringify(k)] for k in data}
 
+        stimulus_id = f"client-scatter-{time()}"
+
         if isinstance(data, type(range(0))):
             data = list(data)
         input_type = type(data)
@@ -2242,6 +2264,7 @@ async def _scatter(
                 who_has={key: [local_worker.address] for key in data},
                 nbytes=valmap(sizeof, data),
                 client=self.id,
+                stimulus_id=stimulus_id,
             )
 
         else:
@@ -2264,7 +2287,10 @@ async def _scatter(
                 )
 
                 await self.scheduler.update_data(
-                    who_has=who_has, nbytes=nbytes, client=self.id
+                    who_has=who_has,
+                    nbytes=nbytes,
+                    client=self.id,
+                    stimulus_id=stimulus_id,
                 )
             else:
                 await self.scheduler.scatter(
@@ -2273,6 +2299,7 @@ async def _scatter(
                     client=self.id,
                     broadcast=broadcast,
                     timeout=timeout,
+                    stimulus_id=stimulus_id,
                 )
 
         out = {k: Future(k, self, inform=False) for k in data}
@@ -2396,7 +2423,12 @@ def scatter(
 
     async def _cancel(self, futures, force=False):
         keys = list({stringify(f.key) for f in futures_of(futures)})
-        await self.scheduler.cancel(keys=keys, client=self.id, force=force)
+        await self.scheduler.cancel(
+            keys=keys,
+            client=self.id,
+            force=force,
+            stimulus_id=f"client-cancel-{time()}",
+        )
         for k in keys:
             st = self.futures.pop(k, None)
             if st is not None:
@@ -2423,7 +2455,9 @@ def cancel(self, futures, asynchronous=None, force=False):
 
     async def _retry(self, futures):
         keys = list({stringify(f.key) for f in futures_of(futures)})
-        response = await self.scheduler.retry(keys=keys, client=self.id)
+        response = await self.scheduler.retry(
+            keys=keys, client=self.id, stimulus_id=f"client-retry-{time()}"
+        )
         for key in response:
             st = self.futures[key]
             st.retry()
@@ -2922,6 +2956,7 @@ def _graph_to_futures(
                     "fifo_timeout": fifo_timeout,
                     "actors": actors,
                     "code": self._get_computation_code(),
+                    "stimulus_id": f"client-update-graph-hlg-{time()}",
                 }
             )
             return futures
@@ -3347,7 +3382,13 @@ async def _restart(self, timeout=no_default):
         if timeout is not None:
             timeout = parse_timedelta(timeout, "s")
 
-        self._send_to_scheduler({"op": "restart", "timeout": timeout})
+        self._send_to_scheduler(
+            {
+                "op": "restart",
+                "timeout": timeout,
+                "stimulus_id": f"client-restart-{time()}",
+            }
+        )
         self._restart_event = asyncio.Event()
         try:
             await asyncio.wait_for(self._restart_event.wait(), timeout)
@@ -3424,7 +3465,9 @@ async def _rebalance(self, futures=None, workers=None):
             keys = list({stringify(f.key) for f in self.futures_of(futures)})
         else:
             keys = None
-        result = await self.scheduler.rebalance(keys=keys, workers=workers)
+        result = await self.scheduler.rebalance(
+            keys=keys, workers=workers, stimulus_id=f"client-rebalance-{time()}"
+        )
         if result["status"] == "partial-fail":
             raise KeyError(f"Could not rebalance keys: {result['keys']}")
         assert result["status"] == "OK", result
@@ -3459,7 +3502,11 @@ async def _replicate(self, futures, n=None, workers=None, branching_factor=2):
         await _wait(futures)
         keys = {stringify(f.key) for f in futures}
         await self.scheduler.replicate(
-            keys=list(keys), n=n, workers=workers, branching_factor=branching_factor
+            keys=list(keys),
+            n=n,
+            workers=workers,
+            branching_factor=branching_factor,
+            stimulus_id=f"client-replicate-{time()}",
         )
 
     def replicate(self, futures, n=None, workers=None, branching_factor=2, **kwargs):
@@ -4177,6 +4224,7 @@ def retire_workers(
             self.scheduler.retire_workers,
             workers=workers,
             close_workers=close_workers,
+            stimulus_id=f"client-retire-workers-{time()}",
             **kwargs,
         )
 
@@ -5138,6 +5186,7 @@ def fire_and_forget(obj):
                 "op": "client-desires-keys",
                 "keys": [stringify(future.key)],
                 "client": "fire-and-forget",
+                "stimulus_id": f"client-fire-and-forget-{time()}",
             }
         )
 

distributed/core.py

@@ -13,6 +13,7 @@
 from collections import defaultdict
 from collections.abc import Container
 from contextlib import suppress
+from contextvars import copy_context
 from enum import Enum
 from functools import partial
 from typing import Callable, ClassVar, TypedDict, TypeVar
@@ -619,7 +620,9 @@ async def handle_stream(self, comm, extra=None, every_cycle=()):
                                 break
                             handler = self.stream_handlers[op]
                             if is_coroutine_function(handler):
-                                self.loop.add_callback(handler, **merge(extra, msg))
+                                self.loop.add_callback(
+                                    copy_context().run, handler, **merge(extra, msg)
+                                )
                                 await gen.sleep(0)
                             else:
                                 handler(**merge(extra, msg))
@@ -629,7 +632,7 @@ async def handle_stream(self, comm, extra=None, every_cycle=()):
 
                 for func in every_cycle:
                     if is_coroutine_function(func):
-                        self.loop.add_callback(func)
+                        self.loop.add_callback(copy_context().run, func)
                     else:
                         func()
 

distributed/deploy/adaptive.py

@@ -7,6 +7,7 @@
 from dask.utils import parse_timedelta
 
 from distributed.deploy.adaptive_core import AdaptiveCore
+from distributed.metrics import time
 from distributed.protocol import pickle
 from distributed.utils import log_errors
 
@@ -193,6 +194,7 @@ async def scale_down(self, workers):
                 names=workers,
                 remove=True,
                 close_workers=True,
+                stimulus_id=f"scale-down-{time()}",
             )
 
             # close workers more forcefully

distributed/deploy/spec.py

@@ -19,6 +19,7 @@
 from distributed.core import CommClosedError, Status, rpc
 from distributed.deploy.adaptive import Adaptive
 from distributed.deploy.cluster import Cluster
+from distributed.metrics import time
 from distributed.scheduler import Scheduler
 from distributed.security import Security
 from distributed.utils import NoOpAwaitable, TimeoutError, import_term, silence_logging
@@ -318,7 +319,10 @@ async def _correct_state_internal(self):
             to_close = set(self.workers) - set(self.worker_spec)
             if to_close:
                 if self.scheduler.status == Status.running:
-                    await self.scheduler_comm.retire_workers(workers=list(to_close))
+                    await self.scheduler_comm.retire_workers(
+                        workers=list(to_close),
+                        stimulus_id=f"spec-cluster-correct-internal-state-{time()}",
+                    )
                 tasks = [
                     asyncio.create_task(self.workers[w].close())
                     for w in to_close

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

@@ -290,7 +290,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"close-nanny-{time()}"
+                ),
+                timeout,
             )
 
     @property