Repository: dask/distributed

Issue:
Tell workers when their peers have left (so they don't hang fetching data from them)
When a worker breaks its connection to the scheduler, or goes too long without heartbeating, the scheduler removes it and reschedules its tasks to run elsewhere. However, other works may continue trying to fetch data from it far beyond that point. This can lead to something that feels like a deadlock, but isn't quite (if you wait long enough, it will resolve), where workers are waiting for their timeouts to expire requesting data from peers that no longer exist.

This is significantly exacerbated by increasing connection timeouts or retries. For example, if you set `distributed.comm.retry.count = 5` (default is 0) and `distributed.comm.timeouts.connect = 60s` (default is 30s), you might experience a 5-minute deadlock as a worker keeps trying to connect to its dead peer, until it finally gives up and asks the scheduler for a new peer to try. If you set `distributed.comm.retry.count = 5` and `distributed.comm.timeouts.tcp = 5m` (default is 30s), and a peer worker is almost out of memory and freezes up but doesn't close its TCP connection (https://github.com/dask/distributed/issues/6110 https://github.com/dask/distributed/issues/6208 https://github.com/dask/distributed/issues/6177), you could have a 25-minute deadlock.

Rather than relying solely on connection timeouts on the worker side, the scheduler should probably inform workers when one of their peers has left. The scheduler is a good source of truth, because workers already have to heartbeat to it, and maintain a long-lived connection (the batched stream).

But naively broadcasting a `remove-worker` message to every worker, every time any worker leaves, probably wouldn't scale well to large clusters. Ideally, we would only inform the workers that are fetching data from the departed worker. We could determine this from `set(wts.processing_on for ts in removed_ws.has_what for wts in ts.waiters if wts.processing_on)` (but something more specific, since not every worker needs to hear the updates about every task).

In that message, we should also include the new `who_has` for the task, with the dead worker removed. (And ensure this actually cancels the `gather_dep` on workers, and doesn't deadlock them if the task has already transitioned to `cancelled` in the meantime.)

cc @crusaderky @fjetter


Discussion:
I've talked about this in various other issues and called this a "circuit breaker" pattern. 

I think the "remove-worker" message (I would probably use a slightly different name) should then abort or close all open connections to that worker. This would allow all code waiting for a connection to get an OSError similar to what would happen if the timeout is actually hit. 

That should be a straight forward implementation and we could use a feature toggle to role this behavior out carefully
> Ideally, we would only inform the workers that are fetching data from the departed worker. 

This is not that simple for pending AMM `acquire-replicas` requests, as the scheduler today doesn't track them. Although in this case I think it would be ok to just state on the scheduler.TaskState "this worker *may* have a pending acquire-replicas going on" and send the remove-worker just in case.
for this we want to be able to cause all checked out comm use for a particular address (the removed worker) to OSError.
ConnectionPool.remove(downed_worker) does this, but we also want connection attempts (tasks in self._connecting) to OSError would it make sense for ConnectionPool.remove to also cancel those?

additionally as the problem here is exacerbated by the retry_operation in get_data_from_worker https://github.com/dask/distributed/blob/d74f50060cfbf1891654e9301447c958910582f4/distributed/worker.py#L2872 would it be better to move the retry logic into `_handle_gather_dep_network_failure` so that instead of assuming the network error is always fatal treat a network error as a busy worker, so that retries are round-robin rather than just hammering the same worker
I think we can safely remove the retry for get_data_from_worker. I believe the worker willhandle this gracefully. It will remove the worker from its internal bookkeeping and, if necessary, asks the scheudler for help. 

>  but we also want connection attempts (tasks in self._connecting) to OSError would it make sense for ConnectionPool.remove to also cancel those?

Yes, that makes sense. I suggest to do this in a dedicated PR. Should be easily isolated from the other changes necessary for this issue.
There are three tasks that can be worked on separately

- The scheduler tells all workers that a worker left. The worker then adjusts internal state, e.g. who_has dicts
- The ConnectionPool removes all available comms. This includes cancellation of currently running connection attempts
- Retries `workers.from get_data_from_workers` is moved up the stack to `utils_comms.gather_from_workers`

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -4884,7 +4884,7 @@ async def remove_worker(
             self.bandwidth_workers.pop((address, w), None)
             self.bandwidth_workers.pop((w, address), None)
 
-        async def remove_worker_from_events():
+        async def remove_worker_from_events() -> None:
             # If the worker isn't registered anymore after the delay, remove from events
             if address not in self.workers and address in self.events:
                 del self.events[address]
@@ -4898,6 +4898,16 @@ async def remove_worker_from_events():
         )
         logger.debug("Removed worker %s", ws)
 
+        for w in self.workers:
+            self.worker_send(
+                w,
+                {
+                    "op": "remove-worker",
+                    "worker": address,
+                    "stimulus_id": stimulus_id,
+                },
+            )
+
         return "OK"
 
     def stimulus_cancel(
@@ -5300,7 +5310,9 @@ async def remove_client_from_events():
                 cleanup_delay, remove_client_from_events
             )
 
-    def send_task_to_worker(self, worker, ts: TaskState, duration: float = -1):
+    def send_task_to_worker(
+        self, worker: str, ts: TaskState, duration: float = -1
+    ) -> None:
         """Send a single computational task to a worker"""
         try:
             msg: dict = self._task_to_msg(ts, duration)
@@ -6815,7 +6827,9 @@ async def _track_retire_worker(
         logger.info("Retired worker %s", ws.address)
         return ws.address, ws.identity()
 
-    def add_keys(self, worker=None, keys=(), stimulus_id=None):
+    def add_keys(
+        self, worker=None, keys=(), stimulus_id=None
+    ) -> Literal["OK", "not found"]:
         """
         Learn that a worker has certain keys
 
@@ -6827,7 +6841,7 @@ def add_keys(self, worker=None, keys=(), stimulus_id=None):
         ws: WorkerState = self.workers[worker]
         redundant_replicas = []
         for key in keys:
-            ts: TaskState = self.tasks.get(key)
+            ts = self.tasks.get(key)
             if ts is not None and ts.state == "memory":
                 if ws not in ts.who_has:
                     self.add_replica(ts, ws)
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -138,6 +138,7 @@
     PauseEvent,
     RefreshWhoHasEvent,
     RemoveReplicasEvent,
+    RemoveWorkerEvent,
     RescheduleEvent,
     RetryBusyWorkerEvent,
     SecedeEvent,
@@ -766,6 +767,7 @@ def __init__(
             "steal-request": self._handle_remote_stimulus(StealRequestEvent),
             "refresh-who-has": self._handle_remote_stimulus(RefreshWhoHasEvent),
             "worker-status-change": self.handle_worker_status_change,
+            "remove-worker": self._handle_remove_worker,
         }
 
         ServerNode.__init__(
@@ -2616,6 +2618,10 @@ def get_current_task(self) -> str:
         """
         return self.active_threads[threading.get_ident()]
 
+    def _handle_remove_worker(self, worker: str, stimulus_id: str) -> None:
+        self.rpc.remove(worker)
+        self.handle_stimulus(RemoveWorkerEvent(worker=worker, stimulus_id=stimulus_id))
+
     def validate_state(self) -> None:
         try:
             self.state.validate_state()
diff --git a/distributed/worker_state_machine.py b/distributed/worker_state_machine.py
--- a/distributed/worker_state_machine.py
+++ b/distributed/worker_state_machine.py
@@ -743,6 +743,12 @@ def from_exception(
         )
 
 
+@dataclass
+class RemoveWorkerEvent(StateMachineEvent):
+    worker: str
+    __slots__ = ("worker",)
+
+
 @dataclass
 class ComputeTaskEvent(StateMachineEvent):
     key: str
@@ -2945,7 +2951,8 @@ def _handle_gather_dep_success(self, ev: GatherDepSuccessEvent) -> RecsInstrs:
                     assert ts.state != "fetch"
                     assert ts not in self.data_needed[ev.worker]
                 ts.who_has.discard(ev.worker)
-                self.has_what[ev.worker].discard(ts.key)
+                if ev.worker in self.has_what:
+                    self.has_what[ev.worker].discard(ts.key)
                 recommendations[ts] = "fetch"
 
         return recommendations, []
@@ -3008,7 +3015,7 @@ def _handle_gather_dep_network_failure(
                 # This can override a recommendation from the previous for loop
                 recommendations[ts] = "missing"
 
-        for key in self.has_what.pop(ev.worker):
+        for key in self.has_what.pop(ev.worker, ()):
             ts = self.tasks[key]
             ts.who_has.remove(ev.worker)
 
@@ -3032,6 +3039,23 @@ def _handle_gather_dep_failure(self, ev: GatherDepFailureEvent) -> RecsInstrs:
 
         return recommendations, []
 
+    @_handle_event.register
+    def _handle_remove_worker(self, ev: RemoveWorkerEvent) -> RecsInstrs:
+        recommendations: Recs = {}
+        for ts in self.data_needed.pop(ev.worker, ()):
+            if self.validate:
+                assert ts.state == "fetch"
+                assert ev.worker in ts.who_has
+            if ts.who_has == {ev.worker}:
+                # This can override a recommendation from the previous for loop
+                recommendations[ts] = "missing"
+
+        for key in self.has_what.pop(ev.worker, ()):
+            ts = self.tasks[key]
+            ts.who_has.remove(ev.worker)
+
+        return recommendations, []
+
     @_handle_event.register
     def _handle_secede(self, ev: SecedeEvent) -> RecsInstrs:
         ts = self.tasks.get(ev.key)


Test patch:
diff --git a/distributed/tests/test_cancelled_state.py b/distributed/tests/test_cancelled_state.py
--- a/distributed/tests/test_cancelled_state.py
+++ b/distributed/tests/test_cancelled_state.py
@@ -13,6 +13,7 @@
     BlockedGetData,
     _LockedCommPool,
     assert_story,
+    async_poll_for,
     gen_cluster,
     inc,
     lock_inc,
@@ -265,19 +266,8 @@ def test_flight_cancelled_error(ws):
 
 @gen_cluster(client=True, nthreads=[("", 1)])
 async def test_in_flight_lost_after_resumed(c, s, b):
-    lock_executing = Lock()
-
-    def block_execution(lock):
-        lock.acquire()
-        return 1
-
     async with BlockedGetData(s.address) as a:
-        fut1 = c.submit(
-            block_execution,
-            lock_executing,
-            workers=[a.address],
-            key="fut1",
-        )
+        fut1 = c.submit(inc, 1, workers=[a.address], key="fut1")
         # Ensure fut1 is in memory but block any further execution afterwards to
         # ensure we control when the recomputation happens
         await wait(fut1)
@@ -286,7 +276,6 @@ def block_execution(lock):
         # This ensures that B already fetches the task, i.e. after this the task
         # is guaranteed to be in flight
         await a.in_get_data.wait()
-        assert fut1.key in b.state.tasks
         assert b.state.tasks[fut1.key].state == "flight"
 
         s.set_restrictions({fut1.key: [a.address, b.address]})
@@ -294,30 +283,12 @@ def block_execution(lock):
         # to be recomputed on B
         await s.remove_worker(a.address, stimulus_id="foo", close=False, safe=True)
 
-        while not b.state.tasks[fut1.key].state == "resumed":
-            await asyncio.sleep(0.01)
+        await wait_for_state(fut1.key, "resumed", b, interval=0)
 
         fut1.release()
         fut2.release()
 
-        while not b.state.tasks[fut1.key].state == "cancelled":
-            await asyncio.sleep(0.01)
-
-        a.block_get_data.set()
-        while b.state.tasks:
-            await asyncio.sleep(0.01)
-
-    assert_story(
-        b.state.story(fut1.key),
-        expect=[
-            # The initial free-keys is rejected
-            ("free-keys", (fut1.key,)),
-            (fut1.key, "resumed", "released", "cancelled", {}),
-            # After gather_dep receives the data, the task is forgotten
-            (fut1.key, "cancelled", "memory", "released", {fut1.key: "forgotten"}),
-            (fut1.key, "released", "forgotten", "forgotten", {}),
-        ],
-    )
+        await async_poll_for(lambda: not b.state.tasks, timeout=5)
 
 
 @gen_cluster(client=True)
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -4366,3 +4366,32 @@ async def test_submit_dependency_of_erred_task(c, s, a, b):
     y = c.submit(inc, x, key="y")
     with pytest.raises(ZeroDivisionError):
         await y
+
+
+@pytest.mark.skipif(
+    condition=sys.version_info < (3, 9),
+    reason="flaky on python 3.8",
+)  # seems related to https://github.com/python/cpython/issues/86296
+@gen_cluster(client=True)
+async def test_tell_workers_when_peers_have_left(c, s, a, b):
+    f = (await c.scatter({"f": 1}, workers=[a.address, b.address], broadcast=True))["f"]
+
+    workers = {a.address: a, b.address: b}
+    connect_timeout = parse_timedelta(
+        dask.config.get("distributed.comm.timeouts.connect"), default="seconds"
+    )
+
+    class BrokenGatherDep(Worker):
+        async def gather_dep(self, worker, *args, **kwargs):
+            w = workers.pop(worker, None)
+            if w is not None and workers:
+                w.listener.stop()
+                s.stream_comms[worker].abort()
+
+            return await super().gather_dep(worker, *args, **kwargs)
+
+    async with BrokenGatherDep(s.address, nthreads=1) as w3:
+        start = time()
+        g = await c.submit(inc, f, key="g", workers=[w3.address])
+        # fails over to the second worker in less than the connect timeout
+        assert time() < start + connect_timeout
diff --git a/distributed/tests/test_worker_state_machine.py b/distributed/tests/test_worker_state_machine.py
--- a/distributed/tests/test_worker_state_machine.py
+++ b/distributed/tests/test_worker_state_machine.py
@@ -38,7 +38,9 @@
     ExecuteSuccessEvent,
     FreeKeysEvent,
     GatherDep,
+    GatherDepBusyEvent,
     GatherDepFailureEvent,
+    GatherDepNetworkFailureEvent,
     GatherDepSuccessEvent,
     Instruction,
     InvalidTaskState,
@@ -48,8 +50,12 @@
     RefreshWhoHasEvent,
     ReleaseWorkerDataMsg,
     RemoveReplicasEvent,
+    RemoveWorkerEvent,
+    RequestRefreshWhoHasMsg,
     RescheduleEvent,
     RescheduleMsg,
+    RetryBusyWorkerEvent,
+    RetryBusyWorkerLater,
     SecedeEvent,
     SerializedTask,
     StateMachineEvent,
@@ -1761,3 +1767,107 @@ def assert_time(actual, expect):
     assert_time(elapsed["x", "cancelled"], 0.15)
     assert_time(elapsed["y", "waiting"], 0.1)
     assert_time(elapsed["z", "flight"], 0.25)
+
+
+@pytest.mark.parametrize("outcome", ["fail", "success", "no-key", "busy"])
+def test_remove_worker_while_in_flight(ws, outcome):
+    ws2 = "127.0.0.1:2"
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy("y", who_has={"x": [ws2]}, stimulus_id="s1"),
+        RemoveWorkerEvent(worker=ws2, stimulus_id="s2"),
+    )
+    ts = ws.tasks["x"]
+    assert ts.state == "flight"
+    assert not ws.data_needed
+    assert not ws.tasks["y"].who_has
+    assert not ws.has_what
+
+    if outcome == "fail":
+        instructions = ws.handle_stimulus(
+            GatherDepNetworkFailureEvent(worker=ws2, total_nbytes=1, stimulus_id="s3")
+        )
+        assert not instructions
+        assert ts.state == "missing"
+    elif outcome == "success":
+        instructions = ws.handle_stimulus(
+            GatherDepSuccessEvent(
+                worker=ws2, total_nbytes=1, data={"x": None}, stimulus_id="s3"
+            )
+        )
+        assert instructions == [
+            AddKeysMsg(stimulus_id="s3", keys=["x"]),
+            Execute(stimulus_id="s3", key="y"),
+        ]
+        assert ts.state == "memory"
+    elif outcome == "no-key":
+        instructions = ws.handle_stimulus(
+            GatherDepSuccessEvent(worker=ws2, total_nbytes=1, data={}, stimulus_id="s3")
+        )
+        assert not instructions
+        assert ts.state == "missing"
+    else:
+        assert outcome == "busy"
+        instructions = ws.handle_stimulus(
+            GatherDepBusyEvent(worker=ws2, total_nbytes=1, stimulus_id="s3")
+        )
+        assert instructions == [
+            RetryBusyWorkerLater(worker=ws2, stimulus_id="s3"),
+            RequestRefreshWhoHasMsg(keys=["x"], stimulus_id="s3"),
+        ]
+        assert ts.state == "missing"
+        instructions = ws.handle_stimulus(
+            RetryBusyWorkerEvent(worker=ws2, stimulus_id="s4")
+        )
+        assert not instructions
+        assert ts.state == "missing"
+
+    assert not ts.who_has
+    assert not ws.has_what
+
+
+def test_remove_worker_while_in_flight_unused_peer(ws):
+    ws2 = "127.0.0.1:2"
+    ws3 = "127.0.0.1:3"
+    ws.handle_stimulus(
+        AcquireReplicasEvent(who_has={"x": [ws2]}, nbytes={"x": 1}, stimulus_id="s1"),
+        AcquireReplicasEvent(
+            who_has={"y": [ws2, ws3]}, nbytes={"y": 1}, stimulus_id="s2"
+        ),
+    )
+
+    ts = ws.tasks["y"]
+    assert ts.state == "flight"
+    assert ts.who_has == {ws2, ws3}
+    assert ts.coming_from == ws3
+
+    ws.handle_stimulus(RemoveWorkerEvent(worker=ws2, stimulus_id="s3"))
+    assert ts.state == "flight"
+    assert ts.who_has == {ws3}
+    assert ts.coming_from == ws3
+
+
+def test_remove_worker_while_in_fetch(ws):
+    ws2 = "127.0.0.1:2"
+    ws3 = "127.0.0.1:3"
+    ws.handle_stimulus(
+        AcquireReplicasEvent(
+            who_has={"x": [ws2], "y": [ws3]}, nbytes={"x": 1, "y": 1}, stimulus_id="s1"
+        ),
+        AcquireReplicasEvent(
+            who_has={"w": [ws2], "z": [ws2, ws3]},
+            nbytes={"w": 1, "z": 1},
+            stimulus_id="s2",
+        ),
+        RemoveWorkerEvent(worker=ws2, stimulus_id="s3"),
+    )
+    assert ws.tasks["w"].state == "missing"
+    assert ws.tasks["z"].state == "fetch"
+    assert not ws.tasks["w"].who_has
+    assert ws.tasks["z"].who_has == {ws3}
+    assert {k: list(v) for k, v in ws.data_needed.items()} == {ws3: [ws.tasks["z"]]}
+    assert ws.has_what == {ws3: {"y", "z"}}
+
+
+def test_remove_worker_unknown(ws):
+    ws2 = "127.0.0.1:2"
+    ws.handle_stimulus(RemoveWorkerEvent(worker=ws2, stimulus_id="s3"))