Properly forward cancellation reason (#9028)

Author: fjetter

distributed/client.py

@@ -1839,11 +1839,11 @@ def _handle_lost_data(self, key=None):
         if state is not None:
             state.lose()
 
-    def _handle_cancelled_keys(self, keys):
+    def _handle_cancelled_keys(self, keys, reason=None, msg=None):
         for key in keys:
             state = self.futures.get(key)
             if state is not None:
-                state.cancel()
+                state.cancel(reason=reason, msg=msg)
 
     def _handle_retried_key(self, key=None):
         state = self.futures.get(key)
@@ -2796,7 +2796,15 @@ def scatter(
     async def _cancel(self, futures, reason=None, msg=None, force=False):
         # FIXME: This method is asynchronous since interacting with the FutureState below requires an event loop.
         keys = list({f.key for f in futures_of(futures)})
-        self._send_to_scheduler({"op": "cancel-keys", "keys": keys, "force": force})
+        self._send_to_scheduler(
+            {
+                "op": "cancel-keys",
+                "keys": keys,
+                "force": force,
+                "reason": reason,
+                "msg": msg,
+            }
+        )
         for k in keys:
             st = self.futures.pop(k, None)
             if st is not None:
@@ -2823,7 +2831,12 @@ def cancel(self, futures, asynchronous=None, force=False, reason=None, msg=None)
             Message that will be attached to the cancelled future
         """
         return self.sync(
-            self._cancel, futures, asynchronous=asynchronous, force=force, msg=msg
+            self._cancel,
+            futures,
+            asynchronous=asynchronous,
+            force=force,
+            msg=msg,
+            reason=reason,
         )
 
     async def _retry(self, futures):

distributed/scheduler.py

@@ -4865,7 +4865,14 @@ async def update_graph(
             lost_keys = self._find_lost_dependencies(dsk, dependencies, keys)
 
             if lost_keys:
-                self.report({"op": "cancelled-keys", "keys": lost_keys}, client=client)
+                self.report(
+                    {
+                        "op": "cancelled-keys",
+                        "keys": lost_keys,
+                        "reason": "lost dependencies",
+                    },
+                    client=client,
+                )
                 self.client_releases_keys(
                     keys=lost_keys, client=client, stimulus_id=stimulus_id
                 )
@@ -5572,7 +5579,7 @@ async def remove_worker_from_events() -> None:
         return "OK"
 
     def stimulus_cancel(
-        self, keys: Collection[Key], client: str, force: bool = False
+        self, keys: Collection[Key], client: str, force: bool, reason: str, msg: str
     ) -> None:
         """Stop execution on a list of keys"""
         logger.info("Client %s requests to cancel %d keys", client, len(keys))
@@ -5591,7 +5598,11 @@ def stimulus_cancel(
             if force or ts.who_wants == {cs}:  # no one else wants this key
                 if ts.dependents:
                     self.stimulus_cancel(
-                        [dts.key for dts in ts.dependents], client, force=force
+                        [dts.key for dts in ts.dependents],
+                        client,
+                        force=force,
+                        reason=reason,
+                        msg=msg,
                     )
                 logger.info("Scheduler cancels key %s.  Force=%s", key, force)
                 cancelled_keys.append(key)
@@ -5603,7 +5614,14 @@ def stimulus_cancel(
                 client=cs.client_key,
                 stimulus_id=f"cancel-key-{time()}",
             )
-        self.report({"op": "cancelled-keys", "keys": cancelled_keys})
+        self.report(
+            {
+                "op": "cancelled-keys",
+                "keys": cancelled_keys,
+                "reason": reason,
+                "msg": msg,
+            }
+        )
 
     def client_desires_keys(self, keys: Collection[Key], client: str) -> None:
         cs = self.clients.get(client)
@@ -8948,7 +8966,7 @@ def request_remove_replicas(
 
 def _task_to_report_msg(ts: TaskState) -> dict[str, Any] | None:
     if ts.state == "forgotten":
-        return {"op": "cancelled-keys", "keys": [ts.key]}
+        return {"op": "cancelled-keys", "keys": [ts.key], "reason": "already forgotten"}
     elif ts.state == "memory":
         return {"op": "key-in-memory", "key": ts.key}
     elif ts.state == "erred":

distributed/tests/test_client.py

@@ -2641,28 +2641,23 @@ async def test_futures_of_cancelled_raises(c, s, a, b):
         await asyncio.sleep(0.01)
     await c.cancel([x], reason="testreason")
 
-    # Note: The scheduler currently doesn't remember the reason but rather
-    # forgets the task immediately. The reason is currently. only raised if the
-    # client checks on it. Therefore, we expect an unknown reason and definitely
-    # not a scheduler disconnected which would otherwise indicate a bug, e.g. an
-    # AssertionError during transitioning.
-    with pytest.raises(CancelledError, match="(reason: unknown|testreason)"):
+    with pytest.raises(CancelledError, match="reason: testreason"):
         await x
     while x.key in s.tasks:
         await asyncio.sleep(0.01)
 
-    with pytest.raises(CancelledError, match="(reason: unknown|testreason)"):
+    with pytest.raises(CancelledError, match="reason: lost dependencies"):
         get_obj = c.get({"x": (inc, x), "y": (inc, 2)}, ["x", "y"], sync=False)
         gather_obj = c.gather(get_obj)
         await gather_obj
 
-    with pytest.raises(CancelledError, match="(reason: unknown|testreason)"):
+    with pytest.raises(CancelledError, match="reason: lost dependencies"):
         await c.submit(inc, x)
 
-    with pytest.raises(CancelledError, match="(reason: unknown|testreason)"):
+    with pytest.raises(CancelledError, match="reason: lost dependencies"):
         await c.submit(add, 1, y=x)
 
-    with pytest.raises(CancelledError, match="(reason: unknown|testreason)"):
+    with pytest.raises(CancelledError, match="reason: lost dependencies"):
         await c.gather(c.map(add, [1], y=x))