Repository: dask/distributed

Issue:
P2P Be robust to connection timeouts during shuffle_receive
When sending shards, we currently rely on the `distributed.comm.timeouts.connect` to be sufficiently large to be able to establish a connection in case of a cold CommPool.

However, if the remote is struggling with a blocked event loop, this timeout can be too short, particularly as long as https://github.com/dask/distributed/pull/7698 is not fixed.

We could consider retrying the sending on CommErrors. If the remote was dead, the shuffle extension would fail the tasks such that I believe we can be arbitrarily generous with retries during this operation.

https://github.com/dask/distributed/blob/2be7f35ee9fa47fde3341bd261f50f823134c296/distributed/shuffle/_worker_plugin.py#L128-L132


Discussion:
When adding retries, we should make sure to include checking whether the shuffle run has been closed in the retried code.
After merging #7698, P2P shuffle tests have begun to become very flaky. I've tracked this down to a change in the order in which the scheduler is informed about the worker leaving and listeners are being stopped. This now causes the shuffle tasks to fail with a `CommClosedError` before the scheduler is able to restart the shuffle. As a result, the scheduler will think the tasks genuinely experienced a problem and will _not_ restart them. 

As a result, implementing a retry mechanism has become increasingly important. @charlesbluca - let's coordinate on this!
I could see us also work around #7698 by warming up our connection pools before the shuffle commences. Something like

```python

class P2PWorkerPlugin:

    async def warmup_comms(self):
        await asyncio.gather(
            self.rpc(w).ping()
            for w in self.workers
        )
```

Maybe the scheduler initiates this during shuffle init, idk. as long as we have fewer than 512 workers (that's the comm pool), this will likely be a good thing and avoid having any transfer tasks fail due to a connection error. The problem still exists, of course but this way most users will likely not be impacted

@hendrikmakait can you elaborate more about the ordering during shutdown that is leading to this? My intuition says that stopping the listener right away as done in #7698 is the correct way and I wonder if we have to change other actions to make this more consistent or if my intuition is wrong
> When adding retries, we should make sure to include checking whether the shuffle run has been closed in the retried code.

Thought about this briefly again. I'm not entirely sure if this is necessary. I'd expect the remote to either ignore the pushed data or even raise a proper exception if the shuffle was already closed. Is it actually necessary to add more logic on sender side?

Just to be clear, I think we should only retry on `OSError`
> Thought about this briefly again. I'm not entirely sure if this is necessary. I'd expect the remote to either ignore the pushed data or even raise a proper exception if the shuffle was already closed. Is it actually necessary to add more logic on sender side?

I've thought of this as an early stopping mechanism for the retries. If the shuffle has been closed (e.g., because a remote worker dropped), there's no point in retrying until we maxed out whatever retry limitations we've set. That worker won't respond anymore.
> rt them.

The problem is that by stopping the listeners, RPCs start throwing errors. Remotes using RPC calls start receiving these errors (notably P2P's send buffers and barrier) and raise them. This causes P2P tasks to transition to `erred`. For all the scheduler knows, those are legitimate errors, because it has not removed the worker yet and therefore we haven't restarted the shuffle, which would have invalidated all those `erred` tasks coming in from workers. If P2P tasks failed due to some _legitimate_ error (or exceeded suspicious count), we will not restart the P2P shuffle. 

For the scheduler to remove the worker, its batched stream to the worker needs to be closed. The worker won't close the stream before it has
* canceled all asynchronous instructions
* torn down all preloads
* closed all extensions
* torn down all plugins
* torn down the scheduler RPC
* stopped all services

so there's now plenty of time for `task-erred` messages to arrive at the scheduler.

Previously, the worker would have stopped its listeners _after_ it had closed the batched stream to the scheduler. So in _almost all_ scenarios, the scheduler would have removed the worker and caused the shuffle to restart _before_ any `task-erred` messages would start coming in. The notable exception is some networking-issue where it takes a while for the `close-stream` message to arrive at the scheduler.
TL;DR: I think it's generally a good idea to close listeners early and not let them potentially interfere with shutdown. The problem is just that it adds a race condition we haven't covered in P2P restarts so far.

The most holistic way of solving this is handling the errors that would be caused by this on the receiving end in P2P. The easiest approach should be retrying with some generously-configured limits to ensure that there will never be a deadlock. Checking if we can stop retrying because the shuffle has been closed (e.g., due to a remote worker leaving) seem like a worthwhile optimization should not be strictly necessary. Cancelled tasks would just block workers longer than necessary.
> After merging #7698, P2P shuffle tests have begun to become very flaky. I've tracked this down to a change in the order in which the scheduler is informed about the worker leaving and listeners are being stopped. This now causes the shuffle tasks to fail with a `CommClosedError` before the scheduler is able to restart the shuffle. As a result, the scheduler will think the tasks genuinely experienced a problem and will _not_ restart them.

This problem has been solved more holistically (see #8088).

Resolving patch:
diff --git a/distributed/shuffle/_core.py b/distributed/shuffle/_core.py
--- a/distributed/shuffle/_core.py
+++ b/distributed/shuffle/_core.py
@@ -14,7 +14,9 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generic, NewType, TypeVar
 
+import dask.config
 from dask.typing import Key
+from dask.utils import parse_timedelta
 
 from distributed.core import PooledRPCCall
 from distributed.exceptions import Reschedule
@@ -23,6 +25,7 @@
 from distributed.shuffle._disk import DiskShardsBuffer
 from distributed.shuffle._exceptions import ShuffleClosedError
 from distributed.shuffle._limiter import ResourceLimiter
+from distributed.utils_comm import retry
 
 if TYPE_CHECKING:
     # TODO import from typing (requires Python >=3.10)
@@ -105,7 +108,7 @@ async def barrier(self) -> None:
         # up the comm pool on scheduler side
         await self.scheduler.shuffle_barrier(id=self.id, run_id=self.run_id)
 
-    async def send(
+    async def _send(
         self, address: str, shards: list[tuple[_T_partition_id, bytes]]
     ) -> None:
         self.raise_if_closed()
@@ -115,6 +118,23 @@ async def send(
             run_id=self.run_id,
         )
 
+    async def send(
+        self, address: str, shards: list[tuple[_T_partition_id, bytes]]
+    ) -> None:
+        retry_count = dask.config.get("distributed.p2p.comm.retry.count")
+        retry_delay_min = parse_timedelta(
+            dask.config.get("distributed.p2p.comm.retry.delay.min"), default="s"
+        )
+        retry_delay_max = parse_timedelta(
+            dask.config.get("distributed.p2p.comm.retry.delay.max"), default="s"
+        )
+        return await retry(
+            partial(self._send, address, shards),
+            count=retry_count,
+            delay_min=retry_delay_min,
+            delay_max=retry_delay_max,
+        )
+
     async def offload(self, func: Callable[..., _T], *args: Any) -> _T:
         self.raise_if_closed()
         with self.time("cpu"):


Test patch:
diff --git a/distributed/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py
--- a/distributed/shuffle/tests/test_shuffle.py
+++ b/distributed/shuffle/tests/test_shuffle.py
@@ -36,6 +36,7 @@
     Scheduler,
     Worker,
 )
+from distributed.core import ConnectionPool
 from distributed.scheduler import TaskState as SchedulerTaskState
 from distributed.shuffle._arrow import (
     convert_shards,
@@ -2310,3 +2311,77 @@ async def test_unpack_gets_rescheduled_from_non_participating_worker(c, s, a):
         shuffleA.block_barrier.set()
         result = await fut
         dd.assert_eq(result, expected)
+
+
+class FlakyConnectionPool(ConnectionPool):
+    def __init__(self, *args, failing_connects=0, **kwargs):
+        self.attempts = 0
+        self.failed_attempts = 0
+        self.failing_connects = failing_connects
+        super().__init__(*args, **kwargs)
+
+    async def connect(self, *args, **kwargs):
+        self.attempts += 1
+        if self.attempts > self.failing_connects:
+            return await super().connect(*args, **kwargs)
+
+        with dask.config.set({"distributed.comm.timeouts.connect": "0 ms"}):
+            try:
+                _ = await super().connect(*args, **kwargs)
+            except Exception:
+                self.failed_attempts += 1
+                raise
+
+
+@gen_cluster(
+    client=True,
+    config={"distributed.comm.retry.count": 0, "distributed.p2p.comm.retry.count": 0},
+)
+async def test_p2p_flaky_connect_fails_without_retry(c, s, a, b):
+    df = dask.datasets.timeseries(
+        start="2000-01-01",
+        end="2000-01-10",
+        dtypes={"x": float, "y": float},
+        freq="10 s",
+    )
+    x = dd.shuffle.shuffle(df, "x", shuffle="p2p")
+
+    rpc = await FlakyConnectionPool(failing_connects=1)
+
+    with mock.patch.object(a, "rpc", rpc):
+        with raises_with_cause(
+            expected_exception=RuntimeError,
+            match="P2P shuffling.*transfer",
+            expected_cause=OSError,
+            match_cause=None,
+        ):
+            await c.compute(x)
+
+    await check_worker_cleanup(a)
+    await check_worker_cleanup(b)
+    await c.close()
+    await check_scheduler_cleanup(s)
+
+
+@gen_cluster(
+    client=True,
+    config={"distributed.comm.retry.count": 0, "distributed.p2p.comm.retry.count": 1},
+)
+async def test_p2p_flaky_connect_recover_with_retry(c, s, a, b):
+    df = dask.datasets.timeseries(
+        start="2000-01-01",
+        end="2000-01-10",
+        dtypes={"x": float, "y": float},
+        freq="10 s",
+    )
+    x = dd.shuffle.shuffle(df, "x", shuffle="p2p")
+
+    rpc = await FlakyConnectionPool(failing_connects=1)
+
+    with mock.patch.object(a, "rpc", rpc):
+        await c.compute(x)
+    assert rpc.failed_attempts == 1
+
+    await check_worker_cleanup(a)
+    await check_worker_cleanup(b)
+    await check_scheduler_cleanup(s)