Repository: dask/distributed

Issue:
Client/Scheduler gather not robust to busy worker
When I run a lot of tasks on the CNES HPC with a big Dask cluster (512 threads/128 workers), I sometimes have communication errors between the scheduler and the workers. The error is thrown by the module "distributed/utils_comm.py," because the code tries to read the key 'data' which does not exist in the dictionary.

```python
    /home/ad/briolf/anaconda3/lib/python3.8/site-packages/distributed/utils_comm.py in retry_operation
          383         dask.config.get("distributed.comm.retry.delay.max"), default="s"
          384     )
    ----> 385     return await retry(
          386         partial(coro, *args, **kwargs),
          387         count=retry_count,
    
    /home/ad/briolf/anaconda3/lib/python3.8/site-packages/distributed/utils_comm.py in retry
          368                 delay *= 1 + random.random() * jitter_fraction
          369             await asyncio.sleep(delay)
    ----> 370     return await coro()
          371 
          372 
    
    /home/ad/briolf/anaconda3/lib/python3.8/site-packages/distributed/core.py in send_recv_from_rpc
          860             name, comm.name = comm.name, "ConnectionPool." + key
          861             try:
    ----> 862                 result = await send_recv(comm=comm, op=key, **kwargs)
          863             finally:
          864                 self.pool.reuse(self.addr, comm)
    
    /home/ad/briolf/anaconda3/lib/python3.8/site-packages/distributed/core.py in send_recv
          659         if comm.deserialize:
          660             typ, exc, tb = clean_exception(**response)
    ----> 661             raise exc.with_traceback(tb)
          662         else:
          663             raise Exception(response["text"])
    
    /home/ad/briolf/anaconda3/lib/python3.8/site-packages/distributed/core.py in handle_comm
          499                             result = asyncio.ensure_future(result)
          500                             self._ongoing_coroutines.add(result)
    ----> 501                             result = await result
          502                     except (CommClosedError, CancelledError) as e:
          503                         if self.status == Status.running:
    
    /home/ad/briolf/anaconda3/lib/python3.8/site-packages/distributed/scheduler.py in gather
          5052                 who_has[key] = []
          5053 
    ----> 5054         data, missing_keys, missing_workers = await gather_from_workers(
          5055             who_has, rpc=self.rpc, close=False, serializers=serializers
          5056         )
    
    /home/ad/briolf/anaconda3/lib/python3.8/site-packages/distributed/utils_comm.py in gather_from_workers
          86                     missing_workers.add(worker)
          87                 else:
    ----> 88                     response.update(r["data"])
          89         finally:
          90             for r in rpcs.values():
    
    KeyError: 'data'
```
I modified the module code, to display the content of this dictionary when this error is thrown. I saw that in this context, the dictionary just contains the following data: ```dict(status='busy')```.

I cannot put reproducible code because the software used is not public and I did not discover simple lines of code to reproduce this problem. I can eventually do other tests to give you further information.

**Environment**:

- Dask version: 2021.4.0
- Python version: 3.8.8
- Operating System: CentOS Linux release 7.6.1810
- Install method (conda, pip, source): conda, channel conda-forge


Discussion:
github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/distributed/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

       20 files  ±0         20 suites  ±0   11h 20m 33s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "duration of all tests") + 9m 34s
  3 716 tests ±0    3 602 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "passed tests")  - 1     106 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "skipped / disabled tests") ±0  7 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "failed tests") +1  1 [:fire:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "test errors") ±0 
35 943 runs  ±0  34 185 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "passed tests")  - 3  1 750 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "skipped / disabled tests") +2  7 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "failed tests") +1  1 [:fire:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "test errors") ±0 

For more details on these failures and errors, see [this check](https://github.com/dask/distributed/runs/14989618451).

Results for commit a1d5e46b. ± Comparison against base commit 7b213995.

[test-results]:data:application/gzip;base64,H4sIAJXqrmQC/02Myw6DIBBFf8Ww7mKG4WV/pkHBhNRHg7Bq+u9Fo+jynHtzvmwIo1/Zs+HwaNiaQ7rI5WhTWObCAgxRUWVN204a1YmvNff95hTwy73DpziE6zbYMBalq/AxLnE7FRPzvHdlK+jA2hVoZHVHV0s41a278z3bL9MUUgFm0UkvVNd6ItMZ7jmiJuukEpZMS04jOAGc/f4TgNX/FAEAAA==

crusaderky: This is ready for review however I'd rather have #8013 in before so that we don't muddle the waters.
crusaderky: As discussed online and offline:
- moved fix of #7995 to a future PR
- moved handling of CancelledError to match #8013 to a future PR
- refactored implementation without the callback
- removed cosmetic refactorings as much as possible

crusaderky: @hendrikmakait @fjetter 
This is in theory complete and ready for a second round of review.
However, I'm seeing a lot of instability in CI.

These failures I've never seen before; I'll investigate:
- 1 out of 10 runs failed: test_TaskStreamPlugin (distributed.diagnostics.tests.test_task_stream)
- 1 out of 10 runs failed: test_gather_after_failed_worker (distributed.tests.test_failed_workers)
- 1 out of 10 runs failed: test_dont_steal_fast_tasks_compute_time (distributed.tests.test_steal)

These are known offenders but potentially impacted by this PR; I'll investigate:
- 1 out of 10 runs failed: test_gather_then_submit_after_failed_workers (distributed.tests.test_failed_workers)
- 1 out of 10 runs failed: test_tell_workers_when_peers_have_left (distributed.tests.test_scheduler)
- 2 out of 10 runs failed: test_chaos_rechunk (distributed.tests.test_stress)

These are known offenders which shouldn't be impacted:
- 1 out of 8 runs failed: test_closed_input_only_worker_during_transfer (distributed.shuffle.tests.test_shuffle)
- 1 out of 8 runs failed: test_closed_worker_during_transfer (distributed.shuffle.tests.test_shuffle)
- 1 out of 8 runs failed: test_crashed_worker_during_transfer (distributed.shuffle.tests.test_shuffle)
- 2 out of 8 runs failed: test_restarting_during_transfer_raises_killed_worker (distributed.shuffle.tests.test_shuffle)
- 1 out of 10 runs failed: test_file_descriptors_dont_leak[Nanny] (distributed.tests.test_client)


crusaderky: I've run all tests above 10 times on each CI environment, on main and on this PR, and counted the number of failures, and I can see no regression; this is ready for final review and merge.

On a side note, this was an extremely time-consuming activity and we should urgently pen in the time to fix CI.


| test | Failures in main | Failures in this PR |
| --- | --- | --- |
| diagnostics/tests/test_task_stream.py::test_TaskStreamPlugin | 1 | 0 |
| shuffle/tests/test_shuffle.py::test_clean_after_close | 1 | 0 |
| shuffle/tests/test_shuffle.py::test_closed_input_only_worker_during_transfer | 0 | 1 |
| shuffle/tests/test_shuffle.py::test_closed_worker_during_transfer | 29 | 35 |
| shuffle/tests/test_shuffle.py::test_crashed_worker_during_transfer | 6 | 1 |
| shuffle/tests/test_shuffle.py::test_restarting_during_transfer_raises_killed_worker | 38 | 33 |
| tests/test_client.py::test_file_descriptors_dont_leak | 9 | 10 |
| tests/test_failed_workers.py::test_gather_after_failed_worker | 4 | 2 |
| tests/test_failed_workers.py::test_gather_then_submit_after_failed_workers | 9 | 1 |
| tests/test_scheduler.py::test_tell_workers_when_peers_have_left | 1 | 1 |
| tests/test_stress.py::test_chaos_rechunk | 0 | 1 |


 

Patch:
diff --git a/distributed/client.py b/distributed/client.py
index 290feabd..4e424a73 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -2319,7 +2319,7 @@ class Client(SyncMethodMixin):
         result = pack_data(unpacked, merge(data, bad_data))
         return result
 
-    async def _gather_remote(self, direct, local_worker):
+    async def _gather_remote(self, direct: bool, local_worker: bool) -> dict[str, Any]:
         """Perform gather with workers or scheduler
 
         This method exists to limit and batch many concurrent gathers into a
@@ -2333,15 +2333,16 @@ class Client(SyncMethodMixin):
 
             if direct or local_worker:  # gather directly from workers
                 who_has = await retry_operation(self.scheduler.who_has, keys=keys)
-                data2, missing_keys, missing_workers = await gather_from_workers(
-                    who_has, rpc=self.rpc, close=False
+                data, missing_keys, failed_keys, _ = await gather_from_workers(
+                    who_has, rpc=self.rpc
                 )
-                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: dict[str, Any] = {"status": "OK", "data": data}
+                if missing_keys or failed_keys:
+                    response = await retry_operation(
+                        self.scheduler.gather, keys=missing_keys + failed_keys
+                    )
                     if response["status"] == "OK":
-                        response["data"].update(data2)
+                        response["data"].update(data)
 
             else:  # ask scheduler to gather data for us
                 response = await retry_operation(self.scheduler.gather, keys=keys)
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 9e1b16fd..0e84ec84 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -5889,35 +5889,49 @@ class Scheduler(SchedulerState, ServerNode):
         )
         return keys
 
-    async def gather(self, keys, serializers=None):
+    async def gather(
+        self, keys: Collection[str], serializers: list[str] | None = None
+    ) -> dict[str, Any]:
         """Collect data from workers to the scheduler"""
         stimulus_id = f"gather-{time()}"
-        keys = list(keys)
-        who_has = {}
-        for key in keys:
-            ts: TaskState = self.tasks.get(key)
-            if ts is not None:
-                who_has[key] = [ws.address for ws in ts.who_has]
-            else:
-                who_has[key] = []
+        data = {}
+        missing_keys = list(keys)
+        failed_keys: list[str] = []
+        missing_workers: set[str] = set()
+
+        while missing_keys:
+            who_has = {}
+            for key, workers in self.get_who_has(missing_keys).items():
+                valid_workers = set(workers) - missing_workers
+                if valid_workers:
+                    who_has[key] = valid_workers
+                else:
+                    failed_keys.append(key)
 
-        data, missing_keys, missing_workers = await gather_from_workers(
-            who_has, rpc=self.rpc, close=False, serializers=serializers
-        )
-        if not missing_keys:
-            result = {"status": "OK", "data": data}
-        else:
-            missing_states = [
-                (self.tasks[key].state if key in self.tasks else None)
-                for key in missing_keys
-            ]
-            logger.exception(
-                "Couldn't gather keys %s state: %s workers: %s",
+            (
+                new_data,
                 missing_keys,
-                missing_states,
-                missing_workers,
+                new_failed_keys,
+                new_missing_workers,
+            ) = await gather_from_workers(
+                who_has, rpc=self.rpc, serializers=serializers
             )
-            result = {"status": "error", "keys": missing_keys}
+            data.update(new_data)
+            failed_keys += new_failed_keys
+            missing_workers.update(new_missing_workers)
+
+        self.log_event("all", {"action": "gather", "count": len(keys)})
+
+        if not failed_keys:
+            return {"status": "OK", "data": data}
+
+        failed_states = {
+            key: self.tasks[key].state if key in self.tasks else "forgotten"
+            for key in failed_keys
+        }
+        logger.error("Couldn't gather keys: %s", failed_states)
+
+        if missing_workers:
             with log_errors():
                 # Remove suspicious workers from the scheduler and shut them down.
                 await asyncio.gather(
@@ -5928,15 +5942,9 @@ class Scheduler(SchedulerState, ServerNode):
                         for worker in missing_workers
                     )
                 )
-                for key, workers in missing_keys.items():
-                    logger.exception(
-                        "Shut down workers that don't have promised key: %s, %s",
-                        str(workers),
-                        str(key),
-                    )
+                logger.error("Shut down unresponsive workers:: %s", missing_workers)
 
-        self.log_event("all", {"action": "gather", "count": len(keys)})
-        return result
+        return {"status": "error", "keys": list(failed_keys)}
 
     @log_errors
     async def restart(self, client=None, timeout=30, wait_for_workers=True):
diff --git a/distributed/utils_comm.py b/distributed/utils_comm.py
index e78f303b..1ba22e3d 100644
--- a/distributed/utils_comm.py
+++ b/distributed/utils_comm.py
@@ -24,50 +24,59 @@ logger = logging.getLogger(__name__)
 async def gather_from_workers(
     who_has: Mapping[str, Collection[str]],
     rpc: ConnectionPool,
-    close: bool = True,
+    *,
     serializers: list[str] | None = None,
     who: str | None = None,
-) -> tuple[dict[str, object], dict[str, list[str]], list[str]]:
+) -> tuple[dict[str, object], list[str], list[str], list[str]]:
     """Gather data directly from peers
 
     Parameters
     ----------
-    who_has: dict
-        Dict mapping keys to sets of workers that may have that key
-    rpc: callable
+    who_has:
+        mapping from keys to worker addresses
+    rpc:
+        RPC channel to use
 
-    Returns dict mapping key to value
+    Returns
+    -------
+    Tuple:
+
+    - Successfully retrieved: ``{key: value, ...}``
+    - Keys that were not available on any worker: ``[key, ...]``
+    - Keys that raised exception; e.g. failed to deserialize: ``[key, ...]``
+    - Workers that failed to respond: ``[address, ...]``
 
     See Also
     --------
     gather
     _gather
+    Scheduler.get_who_has
     """
     from distributed.worker import get_data_from_worker
 
-    bad_addresses: set[str] = set()
-    missing_workers = set()
-    original_who_has = who_has
-    new_who_has = {k: set(v) for k, v in who_has.items()}
-    results: dict[str, object] = {}
-    all_bad_keys: set[str] = set()
+    to_gather = {k: set(v) for k, v in who_has.items()}
+    data: dict[str, object] = {}
+    failed_keys: list[str] = []
+    missing_workers: set[str] = set()
+    busy_workers: set[str] = set()
 
-    while len(results) + len(all_bad_keys) < len(who_has):
+    while to_gather:
         d = defaultdict(list)
-        rev = dict()
-        bad_keys = set()
-        for key, addresses in new_who_has.items():
-            if key in results:
+        for key, addresses in to_gather.items():
+            addresses -= missing_workers
+            ready_addresses = addresses - busy_workers
+            if ready_addresses:
+                d[random.choice(list(ready_addresses))].append(key)
+
+        if not d:
+            if busy_workers:
+                await asyncio.sleep(0.15)
+                busy_workers.clear()
                 continue
-            try:
-                addr = random.choice(list(addresses - bad_addresses))
-                d[addr].append(key)
-                rev[key] = addr
-            except IndexError:
-                bad_keys.add(key)
-        if bad_keys:
-            all_bad_keys |= bad_keys
-        coroutines = {
+
+            return data, list(to_gather), failed_keys, list(missing_workers)
+
+        tasks = {
             address: asyncio.create_task(
                 retry_operation(
                     partial(
@@ -77,7 +86,6 @@ async def gather_from_workers(
                         address,
                         who=who,
                         serializers=serializers,
-                        max_connections=False,
                     ),
                     operation="get_data_from_worker",
                 ),
@@ -85,28 +93,35 @@ async def gather_from_workers(
             )
             for address, keys in d.items()
         }
-        response: dict[str, object] = {}
-        for worker, c in coroutines.items():
+        for address, task in tasks.items():
             try:
-                r = await c
+                r = await task
             except OSError:
-                missing_workers.add(worker)
-            except ValueError as e:
-                logger.info(
-                    "Got an unexpected error while collecting from workers: %s", e
+                missing_workers.add(address)
+            except Exception:
+                # For example, deserialization error
+                logger.exception(
+                    "Unexpected error while collecting tasks %s from %s",
+                    d[address],
+                    address,
                 )
-                missing_workers.add(worker)
+                for key in d[address]:
+                    failed_keys.append(key)
+                    del to_gather[key]
             else:
-                response.update(r["data"])
-
-        bad_addresses |= {v for k, v in rev.items() if k not in response}
-        results.update(response)
-
-    return (
-        results,
-        {k: list(original_who_has[k]) for k in all_bad_keys},
-        list(missing_workers),
-    )
+                if r["status"] == "busy":
+                    busy_workers.add(address)
+                    continue
+
+                assert r["status"] == "OK"
+                for key in d[address]:
+                    if key in r["data"]:
+                        data[key] = r["data"][key]
+                        del to_gather[key]
+                    else:
+                        to_gather[key].remove(address)
+
+    return data, [], failed_keys, list(missing_workers)
 
 
 class WrappedKey:
diff --git a/distributed/worker.py b/distributed/worker.py
index 27dbc516..0be327a0 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -40,7 +40,6 @@ from typing import (
     TypedDict,
     TypeVar,
     cast,
-    overload,
 )
 
 from tlz import keymap, pluck
@@ -1308,23 +1307,44 @@ class Worker(BaseWorker, ServerNode):
         return list(self.data)
 
     async def gather(self, who_has: dict[str, list[str]]) -> dict[str, Any]:
-        who_has = {
-            k: [coerce_to_address(addr) for addr in v]
-            for k, v in who_has.items()
-            if k not in self.data
-        }
-        result, missing_keys, missing_workers = await gather_from_workers(
-            who_has=who_has, rpc=self.rpc, who=self.address
-        )
-        self.update_data(data=result)
-        if missing_keys:
-            logger.warning(
-                "Could not find data: %s on workers: %s (who_has: %s)",
+        """Endpoint used by Scheduler.rebalance() and Scheduler.replicate()"""
+        missing_keys = [k for k in who_has if k not in self.data]
+        failed_keys = []
+        missing_workers: set[str] = set()
+        stimulus_id = f"gather-{time()}"
+
+        while missing_keys:
+            to_gather = {}
+            for k in missing_keys:
+                workers = set(who_has[k]) - missing_workers
+                if workers:
+                    to_gather[k] = workers
+                else:
+                    failed_keys.append(k)
+            if not to_gather:
+                break
+
+            (
+                data,
                 missing_keys,
-                missing_workers,
-                who_has,
+                new_failed_keys,
+                new_missing_workers,
+            ) = await gather_from_workers(
+                who_has=to_gather, rpc=self.rpc, who=self.address
             )
-            return {"status": "partial-fail", "keys": missing_keys}
+            self.update_data(data, stimulus_id=stimulus_id)
+            del data
+            failed_keys += new_failed_keys
+            missing_workers.update(new_missing_workers)
+
+            if missing_keys:
+                who_has = await retry_operation(
+                    self.scheduler.who_has, keys=missing_keys
+                )
+
+        if failed_keys:
+            logger.error("Could not find data: %s", failed_keys)
+            return {"status": "partial-fail", "keys": list(failed_keys)}
         else:
             return {"status": "OK"}
 
@@ -1731,23 +1751,13 @@ class Worker(BaseWorker, ServerNode):
     async def get_data(
         self,
         comm: Comm,
-        keys: Collection[str] | None = None,
+        keys: Collection[str],
         who: str | None = None,
         serializers: list[str] | None = None,
-        max_connections: int | None = None,
     ) -> GetDataBusy | Literal[Status.dont_reply]:
-        if max_connections is None:
-            max_connections = self.transfer_outgoing_count_limit
-
-        if keys is None:
-            keys = set()
-
+        max_connections = self.transfer_outgoing_count_limit
         # Allow same-host connections more liberally
-        if (
-            max_connections
-            and comm
-            and get_address_host(comm.peer_address) == get_address_host(self.address)
-        ):
+        if get_address_host(comm.peer_address) == get_address_host(self.address):
             max_connections = max_connections * 2
 
         if self.status == Status.paused:
@@ -2869,41 +2879,12 @@ def secede():
     )
 
 
-@overload
-async def get_data_from_worker(
-    rpc: ConnectionPool,
-    keys: Collection[str],
-    worker: str,
-    *,
-    who: str | None = None,
-    max_connections: Literal[False],
-    serializers: list[str] | None = None,
-    deserializers: list[str] | None = None,
-) -> GetDataSuccess:
-    ...
-
-
-@overload
-async def get_data_from_worker(
-    rpc: ConnectionPool,
-    keys: Collection[str],
-    worker: str,
-    *,
-    who: str | None = None,
-    max_connections: bool | int | None = None,
-    serializers: list[str] | None = None,
-    deserializers: list[str] | None = None,
-) -> GetDataBusy | GetDataSuccess:
-    ...
-
-
 async def get_data_from_worker(
     rpc: ConnectionPool,
     keys: Collection[str],
     worker: str,
     *,
     who: str | None = None,
-    max_connections: bool | int | None = None,
     serializers: list[str] | None = None,
     deserializers: list[str] | None = None,
 ) -> GetDataBusy | GetDataSuccess:
@@ -2933,7 +2914,6 @@ async def get_data_from_worker(
             op="get_data",
             keys=keys,
             who=who,
-            max_connections=max_connections,
         )
         try:
             status = response["status"]
diff --git a/distributed/worker_state_machine.py b/distributed/worker_state_machine.py
index 985db35d..56bba0f6 100644
--- a/distributed/worker_state_machine.py
+++ b/distributed/worker_state_machine.py
@@ -3656,9 +3656,11 @@ class BaseWorker(abc.ABC):
             stim = task.result()
         except asyncio.CancelledError:
             # This should exclusively happen in Worker.close()
+            logger.warning(f"Async instruction for {task} ended with CancelledError")
             return
         except BaseException:  # pragma: nocover
-            logger.exception("async instruction handlers should never raise!")
+            # This should never happen
+            logger.exception(f"Unhandled exception in async instruction for {task}")
             raise
 
         # Capture metric events in _transition_to_memory()


Test patch:
diff --git a/distributed/tests/test_asyncprocess.py b/distributed/tests/test_asyncprocess.py
index d3f1b0bd..76cb33f0 100644
--- a/distributed/tests/test_asyncprocess.py
+++ b/distributed/tests/test_asyncprocess.py
@@ -81,7 +81,7 @@ async def test_simple():
     with pytest.raises(asyncio.TimeoutError):
         await proc.join(timeout=0.02)
     dt = time() - t1
-    assert 0.2 >= dt >= 0.001
+    assert 0.001 <= dt <= 0.5
     assert proc.is_alive()
     assert proc.pid is not None
     assert proc.exitcode is None
@@ -97,7 +97,7 @@ async def test_simple():
     t1 = time()
     await proc.join(timeout=30)
     dt = time() - t1
-    assert dt <= 1.0
+    assert dt <= 2.0
     assert not proc.is_alive()
     assert proc.pid is not None
     assert proc.exitcode == 0
@@ -141,7 +141,7 @@ async def test_simple():
         await asyncio.sleep(0.01)
         gc.collect()
         dt = time() - t1
-        assert dt < 2.0
+        assert dt < 5.0
 
 
 @gen_test()
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
index 5a9c5e80..6ae27705 100644
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -74,7 +74,7 @@ from distributed.client import (
 )
 from distributed.cluster_dump import load_cluster_dump
 from distributed.comm import CommClosedError
-from distributed.compatibility import LINUX, WINDOWS
+from distributed.compatibility import LINUX, MACOS, WINDOWS
 from distributed.core import Status, error_message
 from distributed.diagnostics.plugin import WorkerPlugin
 from distributed.metrics import time
@@ -85,6 +85,7 @@ from distributed.utils import get_mp_context, is_valid_xml, open_port, sync, tmp
 from distributed.utils_test import (
     NO_AMM,
     BlockedGatherDep,
+    BlockedGetData,
     TaskStateMetadataPlugin,
     _UnhashableCallable,
     async_poll_for,
@@ -5165,17 +5166,19 @@ def test_quiet_client_close(loop):
             sleep(0.200)  # stop part-way
         sleep(0.1)  # let things settle
 
-        out = logger.getvalue()
-        lines = out.strip().split("\n")
-        assert len(lines) <= 2
-        for line in lines:
-            assert (
-                not line
-                or "heartbeat from unregistered worker" in line
-                or "unaware of this worker" in line
-                or "garbage" in line
-                or set(line) == {"-"}
-            ), line
+    out = logger.getvalue()
+    lines = out.strip().split("\n")
+    unexpected_lines = [
+        line
+        for line in lines
+        if line
+        and "heartbeat from unregistered worker" not in line
+        and "unaware of this worker" not in line
+        and "garbage" not in line
+        and "ended with CancelledError" not in line
+        and set(line) != {"-"}
+    ]
+    assert not unexpected_lines, lines
 
 
 @pytest.mark.slow
@@ -6383,6 +6386,7 @@ async def test_wait_for_workers(c, s, a, b):
 
 
 @pytest.mark.skipif(WINDOWS, reason="num_fds not supported on windows")
+@pytest.mark.skipif(MACOS, reason="dask/distributed#8075")
 @pytest.mark.parametrize(
     "Worker", [Worker, pytest.param(Nanny, marks=[pytest.mark.slow])]
 )
@@ -8441,3 +8445,28 @@ async def test_resolves_future_in_dict(c, s, a, b):
     outer_future = c.submit(identity, {"x": inner_future, "y": 2})
     result = await outer_future
     assert result == {"x": 1, "y": 2}
+
+
+@pytest.mark.parametrize("direct", [False, True])
+@gen_cluster(client=True, nthreads=[("", 1)], config=NO_AMM)
+async def test_gather_race_vs_AMM(c, s, a, direct):
+    """Test race condition:
+    Client.gather() tries to get a key from a worker, but in the meantime the
+    Active Memory Manager has moved it to another worker
+    """
+    async with BlockedGetData(s.address) as b:
+        x = c.submit(inc, 1, key="x", workers=[b.address])
+        fut = asyncio.create_task(c.gather(x, direct=direct))
+        await b.in_get_data.wait()
+
+        # Simulate AMM replicate from b to a, followed by AMM drop on b
+        # Can't use s.request_acquire_replicas as it would get stuck on b.block_get_data
+        a.update_data({"x": 3})
+        a.batched_send({"op": "add-keys", "keys": ["x"]})
+        await async_poll_for(lambda: len(s.tasks["x"].who_has) == 2, timeout=5)
+        s.request_remove_replicas(b.address, ["x"], stimulus_id="remove")
+        await async_poll_for(lambda: "x" not in b.data, timeout=5)
+
+        b.block_get_data.set()
+
+    assert await fut == 3  # It's from a; it would be 2 if it were from b
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
index af476c09..6dfd2f56 100644
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -2912,7 +2912,7 @@ async def test_gather_bad_worker_removed(c, s, a, b):
 
             sched_logger = sched_logger.getvalue()
             client_logger = client_logger.getvalue()
-            assert "Shut down workers that don't have promised key" in sched_logger
+            assert "Shut down unresponsive workers" in sched_logger
 
             assert "Couldn't gather 1 keys, rescheduling" in client_logger
 
diff --git a/distributed/tests/test_utils_comm.py b/distributed/tests/test_utils_comm.py
index 44b5e52e..ee0eaff0 100644
--- a/distributed/tests/test_utils_comm.py
+++ b/distributed/tests/test_utils_comm.py
@@ -1,14 +1,17 @@
 from __future__ import annotations
 
+import asyncio
+import random
 from unittest import mock
 
 import pytest
 
 from dask.optimization import SubgraphCallable
 
+from distributed import wait
 from distributed.compatibility import asyncio_run
 from distributed.config import get_loop_factory
-from distributed.core import ConnectionPool
+from distributed.core import ConnectionPool, Status
 from distributed.utils_comm import (
     WrappedKey,
     gather_from_workers,
@@ -17,7 +20,7 @@ from distributed.utils_comm import (
     subs_multiple,
     unpack_remotedata,
 )
-from distributed.utils_test import BrokenComm, gen_cluster
+from distributed.utils_test import BarrierGetData, BrokenComm, gen_cluster, inc
 
 
 def test_pack_data():
@@ -41,35 +44,121 @@ def test_subs_multiple():
     assert subs_multiple(dsk, data) == {"a": (sum, [1, 2])}
 
 
+@gen_cluster(client=True, nthreads=[("", 1)] * 10)
+async def test_gather_from_workers_missing_replicas(c, s, *workers):
+    """When a key is replicated on multiple workers, but the who_has is slightly
+    obsolete, gather_from_workers, retries fetching from all known holders of a replica
+    until it finds the key
+    """
+    a = random.choice(workers)
+    x = await c.scatter({"x": 1}, workers=a.address)
+    assert len(s.workers) == 10
+    assert len(s.tasks["x"].who_has) == 1
+
+    rpc = await ConnectionPool()
+    data, missing, failed, bad_workers = await gather_from_workers(
+        {"x": [w.address for w in workers]}, rpc=rpc
+    )
+
+    assert data == {"x": 1}
+    assert missing == []
+    assert failed == []
+    assert bad_workers == []
+
+
 @gen_cluster(client=True)
 async def test_gather_from_workers_permissive(c, s, a, b):
+    """gather_from_workers fetches multiple keys, of which some are missing.
+    Test that the available data is returned with a note for missing data.
+    """
     rpc = await ConnectionPool()
     x = await c.scatter({"x": 1}, workers=a.address)
 
-    data, missing, bad_workers = await gather_from_workers(
+    data, missing, failed, bad_workers = await gather_from_workers(
         {"x": [a.address], "y": [b.address]}, rpc=rpc
     )
 
     assert data == {"x": 1}
-    assert list(missing) == ["y"]
+    assert missing == ["y"]
+    assert failed == []
+    assert bad_workers == []
 
 
 class BrokenConnectionPool(ConnectionPool):
-    async def connect(self, *args, **kwargs):
+    async def connect(self, address, *args, **kwargs):
         return BrokenComm()
 
 
 @gen_cluster(client=True)
 async def test_gather_from_workers_permissive_flaky(c, s, a, b):
+    """gather_from_workers fails to connect to a worker"""
     x = await c.scatter({"x": 1}, workers=a.address)
 
     rpc = await BrokenConnectionPool()
-    data, missing, bad_workers = await gather_from_workers({"x": [a.address]}, rpc=rpc)
+    data, missing, failed, bad_workers = await gather_from_workers(
+        {"x": [a.address]}, rpc=rpc
+    )
 
-    assert missing == {"x": [a.address]}
+    assert data == {}
+    assert missing == ["x"]
+    assert failed == []
     assert bad_workers == [a.address]
 
 
+@gen_cluster(
+    client=True,
+    nthreads=[],
+    config={"distributed.worker.memory.pause": False},
+)
+async def test_gather_from_workers_busy(c, s):
+    """gather_from_workers receives a 'busy' response from a worker"""
+    async with BarrierGetData(s.address, barrier_count=2) as w:
+        x = await c.scatter({"x": 1}, workers=[w.address])
+        await wait(x)
+        # Throttle to 1 simultaneous connection
+        w.status = Status.paused
+
+        rpc1 = await ConnectionPool()
+        rpc2 = await ConnectionPool()
+        out1, out2 = await asyncio.gather(
+            gather_from_workers({"x": [w.address]}, rpc=rpc1),
+            gather_from_workers({"x": [w.address]}, rpc=rpc2),
+        )
+        assert w.barrier_count == -1  # w.get_data() has been hit 3 times
+        assert out1 == out2 == ({"x": 1}, [], [], [])
+
+
+@pytest.mark.parametrize("when", ["pickle", "unpickle"])
+@gen_cluster(client=True)
+async def test_gather_from_workers_serialization_error(c, s, a, b, when):
+    """A task fails to (de)serialize. Tasks from other workers are fetched
+    successfully.
+    """
+
+    class BadReduce:
+        def __reduce__(self):
+            if when == "pickle":
+                1 / 0
+            else:
+                return lambda: 1 / 0, ()
+
+    rpc = await ConnectionPool()
+    x = c.submit(BadReduce, key="x", workers=[a.address])
+    y = c.submit(inc, 1, key="y", workers=[a.address])
+    z = c.submit(inc, 2, key="z", workers=[b.address])
+    await wait([x, y, z])
+    data, missing, failed, bad_workers = await gather_from_workers(
+        {"x": [a.address], "y": [a.address], "z": [b.address]}, rpc=rpc
+    )
+
+    assert data == {"z": 3}
+    assert missing == []
+    # x and y were serialized together with a single call to pickle; can't tell which
+    # raised
+    assert failed == ["x", "y"]
+    assert bad_workers == []
+
+
 def test_retry_no_exception(cleanup):
     n_calls = 0
     retval = object()
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
index 29b78961..e99f2db1 100644
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -547,7 +547,7 @@ async def test_gather_missing_keys(c, s, a, b):
     async with rpc(a.address) as aa:
         resp = await aa.gather(who_has={x.key: [b.address], "y": [b.address]})
 
-    assert resp == {"status": "partial-fail", "keys": {"y": (b.address,)}}
+    assert resp == {"status": "partial-fail", "keys": ("y",)}
     assert a.data[x.key] == b.data[x.key] == "x"
 
 
@@ -563,21 +563,31 @@ async def test_gather_missing_workers(c, s, a, b):
     async with rpc(a.address) as aa:
         resp = await aa.gather(who_has={x.key: [b.address], "y": [bad_addr]})
 
-    assert resp == {"status": "partial-fail", "keys": {"y": (bad_addr,)}}
+    assert resp == {"status": "partial-fail", "keys": ("y",)}
     assert a.data[x.key] == b.data[x.key] == "x"
 
 
-@pytest.mark.parametrize("missing_first", [False, True])
-@gen_cluster(client=True, worker_kwargs={"timeout": "100ms"})
-async def test_gather_missing_workers_replicated(c, s, a, b, missing_first):
+@pytest.mark.slow
+@pytest.mark.parametrize("know_real", [False, True, True, True, True])  # Read below
+@gen_cluster(client=True, worker_kwargs={"timeout": "1s"}, config=NO_AMM)
+async def test_gather_missing_workers_replicated(c, s, a, b, know_real):
     """A worker owning a redundant copy of a key is missing.
     The key is successfully gathered from other workers.
+
+    know_real=False
+        gather() will try to connect to the bad address, fail, and then query the
+        scheduler who will respond with the good address. Then gather will successfully
+        retrieve the key from the good address.
+    know_real=True
+        50% of the times, gather() will try to connect to the bad address, fail, and
+        immediately connect to the good address.
+        The other 50% of the times it will directly connect to the good address,
+        hence why this test is repeated.
     """
     assert b.address.startswith("tcp://127.0.0.1:")
     x = await c.scatter("x", workers=[b.address])
     bad_addr = "tcp://127.0.0.1:12345"
-    # Order matters! Test both
-    addrs = [bad_addr, b.address] if missing_first else [b.address, bad_addr]
+    addrs = [bad_addr, b.address] if know_real else [bad_addr]
     async with rpc(a.address) as aa:
         resp = await aa.gather(who_has={x.key: addrs})
     assert resp == {"status": "OK"}
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
index ee2300c7..7976a1f5 100644
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -2212,6 +2212,7 @@ class BlockedGatherDep(Worker):
     See also
     --------
     BlockedGetData
+    BarrierGetData
     BlockedExecute
     """
 
@@ -2233,6 +2234,7 @@ class BlockedGetData(Worker):
 
     See also
     --------
+    BarrierGetData
     BlockedGatherDep
     BlockedExecute
     """
@@ -2282,6 +2284,7 @@ class BlockedExecute(Worker):
     --------
     BlockedGatherDep
     BlockedGetData
+    BarrierGetData
     """
 
     def __init__(self, *args, **kwargs):
@@ -2311,6 +2314,32 @@ class BlockedExecute(Worker):
         return await super()._maybe_deserialize_task(ts)
 
 
+class BarrierGetData(Worker):
+    """Block get_data RPC call until at least barrier_count connections are going on
+    in parallel at the same time
+
+    See also
+    --------
+    BlockedGatherDep
+    BlockedGetData
+    BlockedExecute
+    """
+
+    def __init__(self, *args, barrier_count, **kwargs):
+        # TODO just use asyncio.Barrier (needs Python >=3.11)
+        self.barrier_count = barrier_count
+        self.wait_get_data = asyncio.Event()
+        super().__init__(*args, **kwargs)
+
+    async def get_data(self, comm, *args, **kwargs):
+        self.barrier_count -= 1
+        if self.barrier_count > 0:
+            await self.wait_get_data.wait()
+        else:
+            self.wait_get_data.set()
+        return await super().get_data(comm, *args, **kwargs)
+
+
 @contextmanager
 def freeze_data_fetching(w: Worker, *, jump_start: bool = False) -> Iterator[None]:
     """Prevent any task from transitioning from fetch to flight on the worker while