Repository: dask/distributed

Issue:
Regression: frequent deadlocks when gather_dep fails to contact peer
`test_worker_metrics.py::test_gather_dep_network_error` has started being heavily flaky recently.
The failure has nothing to do with metrics; a `Worker.gather_dep` method that fails to open a new RPC channel to its peer seems now to be left dangling forever. 

```python
@gen_cluster(
    client=True,
    nthreads=[("", 1)],
    config={"distributed.comm.timeouts.connect": "500ms"},
)
async def test_gather_dep_network_error(c, s, a):
    x = c.submit(inc, 1, key="x")
    await wait(x)
    async with BlockedGatherDep(s.address) as b:
        y = c.submit(inc, x, key="y", workers=[b.address])
        await b.in_gather_dep.wait()   # <-- Before b tries to connect to a
        await a.close()
        b.block_gather_dep.set()  # <-- b will now attempt and gracefully fail RPC call to a.get_data()
        await wait(y)  # <--- Times out here
```

This may be (happy to be proven wrong) a major regression that could deadlock entire clusters. I suggest treating it as a blocker for the next release.

This also shows we are missing a dedicated unit test; this regression should not have been fortuitously caught by a test that is about metrics.

The [test report](https://dask.github.io/distributed/test_short_report.html) seems to strongly correlate it to #7969:

![image](https://github.com/dask/distributed/assets/6213168/baaf3c21-08d2-4d88-9b95-a4a99cfcb881)

CC @graingert @jrbourbeau @fjetter @hendrikmakait 

Discussion:

Patch:
diff --git a/distributed/diagnostics/plugin.py b/distributed/diagnostics/plugin.py
index 93fbf2a7..02ae91a2 100644
--- a/distributed/diagnostics/plugin.py
+++ b/distributed/diagnostics/plugin.py
@@ -123,6 +123,7 @@ class SchedulerPlugin:
         start: TaskStateState,
         finish: TaskStateState,
         *args: Any,
+        stimulus_id: str,
         **kwargs: Any,
     ) -> None:
         """Run whenever a task changes state
@@ -143,6 +144,8 @@ class SchedulerPlugin:
             One of released, waiting, processing, memory, error.
         finish : string
             Final state of the transition.
+        stimulus_id: string
+            ID of stimulus causing the transition.
         *args, **kwargs :
             More options passed when transitioning
             This may include worker ID, compute time, etc.
@@ -164,7 +167,7 @@ class SchedulerPlugin:
         """
 
     def remove_worker(
-        self, scheduler: Scheduler, worker: str
+        self, scheduler: Scheduler, worker: str, *, stimulus_id: str, **kwargs: Any
     ) -> None | Awaitable[None]:
         """Run when a worker leaves the cluster
 
diff --git a/distributed/diagnostics/progress.py b/distributed/diagnostics/progress.py
index 21532fcf..11520d6a 100644
--- a/distributed/diagnostics/progress.py
+++ b/distributed/diagnostics/progress.py
@@ -104,7 +104,9 @@ class Progress(SchedulerPlugin):
         logger.debug("Set up Progress keys")
 
         for k in errors:
-            self.transition(k, None, "erred", exception=True)
+            self.transition(
+                k, None, "erred", stimulus_id="progress-setup", exception=True
+            )
 
     def transition(self, key, start, finish, *args, **kwargs):
         if key in self.keys and start == "processing" and finish == "memory":
@@ -240,7 +242,9 @@ class MultiProgress(Progress):
                 self.keys[k] = set()
 
         for k in errors:
-            self.transition(k, None, "erred", exception=True)
+            self.transition(
+                k, None, "erred", stimulus_id="multiprogress-setup", exception=True
+            )
         logger.debug("Set up Progress keys")
 
     def transition(self, key, start, finish, *args, **kwargs):
diff --git a/distributed/diagnostics/websocket.py b/distributed/diagnostics/websocket.py
index 5e77e1aa..4df5de6d 100644
--- a/distributed/diagnostics/websocket.py
+++ b/distributed/diagnostics/websocket.py
@@ -48,6 +48,8 @@ class WebsocketPlugin(SchedulerPlugin):
             One of released, waiting, processing, memory, error.
         finish : string
             Final state of the transition.
+        stimulus_id: string
+            ID of stimulus causing the transition.
         *args, **kwargs : More options passed when transitioning
             This may include worker ID, compute time, etc.
         """
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 8cbfaa13..109a9e67 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -1978,7 +1978,9 @@ class SchedulerState:
                     self.tasks[ts.key] = ts
                 for plugin in list(self.plugins.values()):
                     try:
-                        plugin.transition(key, start, actual_finish, **kwargs)
+                        plugin.transition(
+                            key, start, actual_finish, stimulus_id=stimulus_id, **kwargs
+                        )
                     except Exception:
                         logger.info("Plugin failed with exception", exc_info=True)
                 if ts.state == "forgotten":
@@ -5069,7 +5071,19 @@ class Scheduler(SchedulerState, ServerNode):
         awaitables = []
         for plugin in list(self.plugins.values()):
             try:
-                result = plugin.remove_worker(scheduler=self, worker=address)
+                try:
+                    result = plugin.remove_worker(
+                        scheduler=self, worker=address, stimulus_id=stimulus_id
+                    )
+                except TypeError:
+                    parameters = inspect.signature(plugin.remove_worker).parameters
+                    if "stimulus_id" not in parameters and not any(
+                        p.kind is p.VAR_KEYWORD for p in parameters.values()
+                    ):
+                        # Deprecated (see add_plugin)
+                        result = plugin.remove_worker(scheduler=self, worker=address)  # type: ignore
+                    else:
+                        raise
                 if inspect.isawaitable(result):
                     awaitables.append(result)
             except Exception as e:
@@ -5724,6 +5738,15 @@ class Scheduler(SchedulerState, ServerNode):
                 category=UserWarning,
             )
 
+        parameters = inspect.signature(plugin.remove_worker).parameters
+        if not any(p.kind is p.VAR_KEYWORD for p in parameters.values()):
+            warnings.warn(
+                "The signature of `SchedulerPlugin.remove_worker` now requires `**kwargs` "
+                "to ensure that plugins remain forward-compatible. Not including "
+                "`**kwargs` in the signature will no longer be supported in future versions.",
+                FutureWarning,
+            )
+
         self.plugins[name] = plugin
 
     def remove_plugin(
@@ -8420,7 +8443,7 @@ class WorkerStatusPlugin(SchedulerPlugin):
         except CommClosedError:
             scheduler.remove_plugin(name=self.name)
 
-    def remove_worker(self, scheduler: Scheduler, worker: str) -> None:
+    def remove_worker(self, scheduler: Scheduler, worker: str, **kwargs: Any) -> None:
         try:
             self.bcomm.send(["remove", worker])
         except CommClosedError:
diff --git a/distributed/shuffle/_scheduler_plugin.py b/distributed/shuffle/_scheduler_plugin.py
index 88029173..911f4f89 100644
--- a/distributed/shuffle/_scheduler_plugin.py
+++ b/distributed/shuffle/_scheduler_plugin.py
@@ -309,7 +309,9 @@ class ShuffleSchedulerPlugin(SchedulerPlugin):
         original_restrictions = ts.annotations.pop("shuffle_original_restrictions")
         self.scheduler.set_restrictions({ts.key: original_restrictions})
 
-    def remove_worker(self, scheduler: Scheduler, worker: str) -> None:
+    def remove_worker(
+        self, scheduler: Scheduler, worker: str, *, stimulus_id: str, **kwargs: Any
+    ) -> None:
         from time import time
 
         stimulus_id = f"shuffle-failed-worker-left-{time()}"
@@ -342,6 +344,7 @@ class ShuffleSchedulerPlugin(SchedulerPlugin):
         start: TaskStateState,
         finish: TaskStateState,
         *args: Any,
+        stimulus_id: str,
         **kwargs: Any,
     ) -> None:
         if finish not in ("released", "forgotten"):
diff --git a/distributed/stealing.py b/distributed/stealing.py
index afc76ef5..45b26101 100644
--- a/distributed/stealing.py
+++ b/distributed/stealing.py
@@ -21,7 +21,13 @@ from distributed.utils import log_errors, recursive_to_dict
 
 if TYPE_CHECKING:
     # Recursive imports
-    from distributed.scheduler import Scheduler, SchedulerState, TaskState, WorkerState
+    from distributed.scheduler import (
+        Scheduler,
+        SchedulerState,
+        TaskState,
+        TaskStateState,
+        WorkerState,
+    )
 
 # Stealing requires multiple network bounces and if successful also task
 # submission which may include code serialization. Therefore, be very
@@ -155,7 +161,7 @@ class WorkStealing(SchedulerPlugin):
     def add_worker(self, scheduler: Any = None, worker: Any = None) -> None:
         self.stealable[worker] = tuple(set() for _ in range(15))
 
-    def remove_worker(self, scheduler: Scheduler, worker: str) -> None:
+    def remove_worker(self, scheduler: Scheduler, worker: str, **kwargs: Any) -> None:
         del self.stealable[worker]
 
     def teardown(self) -> None:
@@ -167,10 +173,8 @@ class WorkStealing(SchedulerPlugin):
     def transition(
         self,
         key: str,
-        start: str,
-        finish: str,
-        compute_start: Any = None,
-        compute_stop: Any = None,
+        start: TaskStateState,
+        finish: TaskStateState,
         *args: Any,
         **kwargs: Any,
     ) -> None:
diff --git a/distributed/worker.py b/distributed/worker.py
index fd4941a2..fe52f796 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -2081,7 +2081,10 @@ class Worker(BaseWorker, ServerNode):
                 stimulus_id=f"gather-dep-success-{time()}",
             )
 
-        except OSError:
+        # Note: CancelledError and asyncio.TimeoutError are rare conditions
+        # that can be raised by the network stack.
+        # See https://github.com/dask/distributed/issues/8006
+        except (OSError, asyncio.CancelledError, asyncio.TimeoutError):
             logger.exception("Worker stream died during communication: %s", worker)
             self.state.log.append(
                 ("gather-dep-failed", worker, to_gather, stimulus_id, time())
@@ -2094,6 +2097,8 @@ class Worker(BaseWorker, ServerNode):
 
         except Exception as e:
             # e.g. data failed to deserialize
+            # FIXME this will deadlock the cluster
+            #       https://github.com/dask/distributed/issues/6705
             logger.exception(e)
             self.state.log.append(
                 ("gather-dep-failed", worker, to_gather, stimulus_id, time())


Test patch:
diff --git a/distributed/diagnostics/tests/test_scheduler_plugin.py b/distributed/diagnostics/tests/test_scheduler_plugin.py
index 92f72310..7ab9e825 100644
--- a/distributed/diagnostics/tests/test_scheduler_plugin.py
+++ b/distributed/diagnostics/tests/test_scheduler_plugin.py
@@ -16,7 +16,8 @@ async def test_simple(c, s, a, b):
             scheduler.add_plugin(self, name="counter")
             self.count = 0
 
-        def transition(self, key, start, finish, *args, **kwargs):
+        def transition(self, key, start, finish, *args, stimulus_id, **kwargs):
+            assert stimulus_id is not None
             if start == "processing" and finish == "memory":
                 self.count += 1
 
@@ -51,8 +52,9 @@ async def test_add_remove_worker(s):
             assert scheduler is s
             events.append(("add_worker", worker))
 
-        def remove_worker(self, worker, scheduler):
+        def remove_worker(self, worker, scheduler, *, stimulus_id, **kwargs):
             assert scheduler is s
+            assert stimulus_id is not None
             events.append(("remove_worker", worker))
 
     plugin = MyPlugin()
@@ -80,6 +82,70 @@ async def test_add_remove_worker(s):
     assert events == []
 
 
+@gen_cluster(nthreads=[])
+async def test_remove_worker_renamed_kwargs_allowed(s):
+    events = []
+
+    class MyPlugin(SchedulerPlugin):
+        name = "MyPlugin"
+
+        def remove_worker(self, worker, scheduler, **kwds):
+            assert scheduler is s
+            events.append(("remove_worker", worker))
+
+    plugin = MyPlugin()
+    s.add_plugin(plugin)
+    assert events == []
+
+    a = Worker(s.address)
+    await a
+    await a.close()
+
+    assert events == [
+        ("remove_worker", a.address),
+    ]
+
+    events[:] = []
+    s.remove_plugin(plugin.name)
+    async with Worker(s.address):
+        pass
+    assert events == []
+
+
+@gen_cluster(nthreads=[])
+async def test_remove_worker_without_kwargs_deprecated(s):
+    events = []
+
+    class DeprecatedPlugin(SchedulerPlugin):
+        name = "DeprecatedPlugin"
+
+        def remove_worker(self, worker, scheduler):
+            assert scheduler is s
+            events.append(("remove_worker", worker))
+
+    plugin = DeprecatedPlugin()
+    with pytest.warns(
+        FutureWarning,
+        match="The signature of `SchedulerPlugin.remove_worker` now requires `\\*\\*kwargs`",
+    ):
+        s.add_plugin(plugin)
+    assert events == []
+
+    a = Worker(s.address)
+    await a
+    await a.close()
+
+    assert events == [
+        ("remove_worker", a.address),
+    ]
+
+    events[:] = []
+    s.remove_plugin(plugin.name)
+    async with Worker(s.address):
+        pass
+    assert events == []
+
+
 @gen_cluster(nthreads=[])
 async def test_async_add_remove_worker(s):
     events = []
@@ -91,7 +157,7 @@ async def test_async_add_remove_worker(s):
             assert scheduler is s
             events.append(("add_worker", worker))
 
-        async def remove_worker(self, worker, scheduler):
+        async def remove_worker(self, worker, scheduler, **kwargs):
             assert scheduler is s
             events.append(("remove_worker", worker))
 
@@ -135,7 +201,7 @@ async def test_async_and_sync_add_remove_worker(s):
             await asyncio.sleep(0)
             events.append((self.name, "add_worker", worker))
 
-        async def remove_worker(self, scheduler, worker):
+        async def remove_worker(self, scheduler, worker, **kwargs):
             assert scheduler is s
             self.in_remove_worker.set()
             await self.block_remove_worker.wait()
@@ -149,7 +215,7 @@ async def test_async_and_sync_add_remove_worker(s):
             assert scheduler is s
             events.append((self.name, "add_worker", worker))
 
-        def remove_worker(self, worker, scheduler):
+        def remove_worker(self, worker, scheduler, **kwargs):
             assert scheduler is s
             events.append((self.name, "remove_worker", worker))
 
@@ -229,7 +295,7 @@ async def test_failing_async_add_remove_worker(s):
             await asyncio.sleep(0)
             raise RuntimeError("Async add_worker failed")
 
-        async def remove_worker(self, scheduler, worker):
+        async def remove_worker(self, scheduler, worker, **kwargs):
             assert scheduler is s
             await asyncio.sleep(0)
             raise RuntimeError("Async remove_worker failed")
@@ -257,7 +323,7 @@ async def test_failing_sync_add_remove_worker(s):
             assert scheduler is s
             raise RuntimeError("Async add_worker failed")
 
-        def remove_worker(self, scheduler, worker):
+        def remove_worker(self, scheduler, worker, **kwargs):
             assert scheduler is s
             raise RuntimeError("Async remove_worker failed")
 
diff --git a/distributed/diagnostics/tests/test_task_stream.py b/distributed/diagnostics/tests/test_task_stream.py
index 3c152a14..0e020695 100644
--- a/distributed/diagnostics/tests/test_task_stream.py
+++ b/distributed/diagnostics/tests/test_task_stream.py
@@ -94,16 +94,20 @@ async def test_no_startstops(c, s, a, b):
     await wait(future)
     assert len(tasks.buffer) == 1
 
-    tasks.transition(future.key, "processing", "erred")
+    tasks.transition(future.key, "processing", "erred", stimulus_id="s1")
     # Transition was not recorded because it didn't contain `startstops`
     assert len(tasks.buffer) == 1
 
-    tasks.transition(future.key, "processing", "erred", startstops=[])
+    tasks.transition(future.key, "processing", "erred", stimulus_id="s2", startstops=[])
     # Transition was not recorded because `startstops` was empty
     assert len(tasks.buffer) == 1
 
     tasks.transition(
-        future.key, "processing", "erred", startstops=[dict(start=time(), stop=time())]
+        future.key,
+        "processing",
+        "erred",
+        stimulus_id="s3",
+        startstops=[dict(start=time(), stop=time())],
     )
     assert len(tasks.buffer) == 2
 
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
index cf902c0b..29b78961 100644
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -3351,6 +3351,78 @@ async def test_gather_dep_no_longer_in_flight_tasks(c, s, a):
         assert not any("missing-dep" in msg for msg in f2_story)
 
 
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_gather_dep_cancelled_error(c, s, a):
+    """Something somewhere in the networking stack raises CancelledError while
+    gather_dep is running
+
+    See Also
+    --------
+    test_get_data_cancelled_error
+    https://github.com/dask/distributed/issues/8006
+    """
+    async with BlockedGetData(s.address) as b:
+        x = c.submit(inc, 1, key="x", workers=[b.address])
+        y = c.submit(inc, x, key="y", workers=[a.address])
+        await b.in_get_data.wait()
+        tasks = {
+            task for task in asyncio.all_tasks() if "gather_dep" in task.get_name()
+        }
+        assert tasks
+        # There should be only one task but cope with finding more just in case a
+        # previous test didn't properly clean up
+        for task in tasks:
+            task.cancel()
+
+        b.block_get_data.set()
+        assert await y == 3
+
+    assert_story(
+        a.state.story("x"),
+        [
+            ("x", "fetch", "flight", "flight", {}),
+            ("x", "flight", "missing", "missing", {}),
+            ("x", "missing", "fetch", "fetch", {}),
+            ("x", "fetch", "flight", "flight", {}),
+            ("x", "flight", "memory", "memory", {"y": "ready"}),
+        ],
+    )
+
+
+@gen_cluster(client=True, nthreads=[("", 1)], timeout=5)
+async def test_get_data_cancelled_error(c, s, a):
+    """Something somewhere in the networking stack raises CancelledError while
+    get_data is running
+
+    See Also
+    --------
+    test_gather_dep_cancelled_error
+    https://github.com/dask/distributed/issues/8006
+    """
+
+    class FlakyInboundRPC(Worker):
+        flake = 0
+
+        def handle_comm(self, comm):
+            if self.flake:
+                self.flake -= 1
+
+                async def write(*args, **kwargs):
+                    raise asyncio.CancelledError()
+
+                comm.write = write
+
+            return super().handle_comm(comm)
+
+    async with FlakyInboundRPC(s.address) as b:
+        x = c.submit(inc, 1, key="x", workers=[b.address])
+        await wait(x)
+        b.flake = 2
+        y = c.submit(inc, x, key="y", workers=[a.address])
+        assert await y == 3
+        assert b.flake == 0
+
+
 @gen_cluster(client=True, nthreads=[("", 1)])
 async def test_Worker__to_dict(c, s, a):
     x = c.submit(inc, 1, key="x")