Repository: dask/distributed

Issue:
P2P blows up memory
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**Describe the issue**:

There are 2 different issues that @hendrikmakait and I identified:

- Doing a repartition before the shuffle loads all data into memory
- After shuffling is finished, the memory is double the amount we would expected, there seem to be unreleased references. The MCVE has around 80GB of unmanaged memory

Potentially, the tasks during unpack are scheduled a bit too fast as well which results in excessive memory usage, but this should be mitigated by the other 2 points.

**Minimal Complete Verifiable Example**:

```python
import dask
import dask.dataframe as dd
from distributed import wait

dask.config.set({"dataframe.convert-string": True})  # use PyArrow strings by default

df = dd.read_parquet(
    "s3://coiled-datasets/uber-lyft-tlc/",
    storage_options={"anon": True},
)
df = df.shuffle("PULocationID").persist()
```

**Anything else we need to know?**:

**Environment**:

- Dask version: main
- Python version: 3.11
- Operating System: mac
- Install method (conda, pip, source): condo


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._

       13 files   -        1         13 suites   - 1   6h 15m 18s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "duration of all tests") - 41m 28s
  3 719 tests +       1    3 609 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "passed tests") +       3     108 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "skipped / disabled tests") ±  0  2 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "failed tests")  - 2 
23 594 runs   - 1 312  22 449 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "passed tests")  - 1 275  1 143 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "skipped / disabled tests")  - 35  2 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.8.0/README.md#the-symbols "failed tests")  - 2 

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

Results for commit 625b9795. ± Comparison against base commit d9a3457d.

[test-results]:data:application/gzip;base64,H4sIADdnuGQC/02Myw6DIBQFf8Ww7gIu7/6MwSskpD4ahJXpvxcatS5nTs7sJMTJb+TZMf7oyFZi/tNYkstxXSoDSGaqqmtuO9fMnthvBbE5RW/uFd+tQ69XH1ycWuoSPqU1VUOrSWVpXeDSigPPLoAQ9nJHlwl+qlv3x/csrvMccwWiQA5WW6mDGhAkmsEF49AjQ6uVAuNHOqLT5PMFYcBAbhQBAAA=


Patch:
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/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py
index 45688c61..664b31d0 100644
--- a/distributed/shuffle/tests/test_shuffle.py
+++ b/distributed/shuffle/tests/test_shuffle.py
@@ -1853,3 +1853,17 @@ async def test_handle_null_partitions_p2p_shuffling(c, s, *workers):
     await c.close()
     await asyncio.gather(*[check_worker_cleanup(w) for w in workers])
     await check_scheduler_cleanup(s)
+
+
+@gen_cluster(client=True)
+async def test_set_index_p2p(c, s, *workers):
+    df = pd.DataFrame({"a": [1, 2, 3, 4, 5, 6, 7, 8], "b": 1})
+    ddf = dd.from_pandas(df, npartitions=3)
+    ddf = ddf.set_index("a", shuffle="p2p", divisions=(1, 3, 8))
+    assert ddf.npartitions == 2
+    result = await c.compute(ddf)
+    dd.assert_eq(result, df.set_index("a"))
+
+    await c.close()
+    await asyncio.gather(*[check_worker_cleanup(w) for w in workers])
+    await check_scheduler_cleanup(s)
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")