Repository: dask/distributed

Issue:
P2P Rechunking graph prohibits culling and partial release
The P2P rechunk graph is similar to the shuffle graph one large globally connected graph. However, many rechunk applications only operate on slices of the data and is often not connecting all in- and output partitions but only planes in the dataset.

Having only slices/planes connected by the graph is beneficial since it allows us to cull the graph, process downstream releaser earlier and release disk spaces sooner.

A real life example showing this effect is using the public ERA5 dataset (see code below) but a simpler variant can be constructed using artificial data.

<details>
<summary>ERA5 example</summary>

```python
import itertools

import coiled
import dask
import distributed
import flox.xarray
import fsspec
import numpy as np
import pandas as pd
import xarray as xr
fs = fsspec.filesystem("s3")
prefix = "s3://era5-pds/zarr"
variables = [
    store.split("/")[-1] for store in fs.glob("era5-pds/zarr/1979/01/data/*.zarr")
]
years = [path.split("/")[-1] for path in fs.glob("era5-pds/zarr/*")]
months = [f"{m:02d}" for m in range(1, 13)]
last_months = [path.split("/")[-1] for path in fs.glob(f"era5-pds/zarr/{years[-1]}/*")]
var = ("precipitation_amount_1hour_Accumulation.zarr",)
var = (variables[0],)
all_stores = [
    "/".join(t)
    for t in itertools.chain(
        itertools.product((prefix,), years[:-1], months, ("data",), var),
        itertools.product((prefix,), years[-2:], last_months, ("data",), var),
    )
]

def preprocess(ds):
    """Edit the dataset so it combines nicely."""
    (time_dim,) = [dim for dim in ds.dims if "time" in dim]
    ds = ds.rename({time_dim: "time"})
    bounds_var = [var for var in ds.data_vars if "bounds" in var]
    ds = ds.drop_vars(bounds_var)
    return ds

ds = xr.open_mfdataset(
    all_stores,
    engine="zarr",
    combine="nested",
    concat_dim="time",
    preprocess=preprocess,
    join="override",
    parallel=True,
)

import dask
dask.config.set({"array.rechunk.method": "p2p"})

def rechunk_to_frequency(ds, **kwargs):
    newchunks = {}
    for dim, freq in kwargs.items():
        newchunks[dim] = tuple(
            ds[dim]
            .copy(data=np.ones(ds[dim].shape, dtype=np.int64))
            .to_dataframe()
            .resample(freq)
            .sum()[dim]
            .values
        )
    return newchunks

newchunks = rechunk_to_frequency(ds, time="A")
rechunked = ds.chunk(**newchunks, lat=75, lon=75)

# Run any kind of reduction on this. This reduces per year
# so every reduction step should only load 1/45 of the data
result = flox.xarray.xarray_reduce(
    rechunked,
    rechunked.time.dt.year,
    func="quantile",
    skipna=False,
    q=0.9,
    method="blockwise",
)

# Even though we're selecting only a thin slices, this is loading the entire data
# since culling is not possible. The entire data also has to be stored to disk temporarily
# such that the cluster has to be really large for this to work even though the actual
# slices should only be a couple GBs
result.isel(year=slice(2)).compute()
```

</details>

A simpler, toy example showing the same principle is just concatenating a single dimension.
```python

import dask.array as da
# This is a simple example where we're just compacting the first dimension
# The ERA5 example simultaneously also splits another dimension but this boils
# down to the same principle
arr = da.random.random((10, 45), chunks=(1, 1))
rechunked = arr.rechunk((-1, 1))
dask.optimize(rechunked[:, :, :2])[0]
```

```python
import dask
import dask.array as da

arr = da.random.random((10, 45), chunks=(1, 1))
for method in ["p2p", "tasks"]:
    dask.config.set({"array.rechunk.method": method})
    rechunked = arr.rechunk((-1, 1))
    
    len_with_slicing = len(dask.optimize(rechunked[:, :2])[0].__dask_graph__())
    len_without_slicing = len(dask.optimize(rechunked)[0].__dask_graph__())
    print(f"Tasks after slicing\t{method=}\t{len_with_slicing/len_without_slicing*100:2.0f}%")
```

```python
Tasks after slicing	method='p2p'	96%
Tasks after slicing	method='tasks'	 5%
```

For P2P to do the same, we should instead of generating one large P2P layer generate many small P2P layers such that we only connect the relevant slices.


Discussion:
_In theory_, detecting such split points should be pretty easy to do by walking through the input and output chunks of each individual axis and checking where the chunking matches up. If an axis contains unknown sizes, it must be unchanged, so that's even the best case for us.
to be clear, there are also slightly less trivial cases, e.g.

```python
original = (1, 50, 1)
new = (5, 10, -1)
```

which is...
- a simple concat on the first dimension
- a split in the second dimension
- a full shuffle on the third

but I still would expect this to break apart into many decoupled P2Ps

Resolving patch:
diff --git a/distributed/shuffle/_rechunk.py b/distributed/shuffle/_rechunk.py
--- a/distributed/shuffle/_rechunk.py
+++ b/distributed/shuffle/_rechunk.py
@@ -106,6 +106,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, NamedTuple
 
+import toolz
 from tornado.ioloop import IOLoop
 
 import dask
@@ -166,45 +167,305 @@ def rechunk_unpack(
         )
 
 
-def rechunk_p2p(x: da.Array, chunks: ChunkedAxes) -> da.Array:
-    import numpy as np
+class _Partial(NamedTuple):
+    """Information used to perform a partial rechunk along a single axis"""
+
+    #: Slice of the old chunks along this axis that belong to the partial
+    old: slice
+    #: Slice of the new chunks along this axis that belong to the partial
+    new: slice
+    #: Index of the first value of the left-most old chunk along this axis
+    #: to include in this partial. Everything left to this index belongs to
+    #: the previous partial.
+    left_start: int
+    #: Index of the first value of the right-most old chunk along this axis
+    #: to exclude from this partial.
+    #: This corresponds to `left_start` of the subsequent partial.
+    right_stop: int
+
+
+class _NDPartial(NamedTuple):
+    """Information used to perform a partial rechunk along all axes"""
+
+    #: n-dimensional slice of the old chunks along each axis that belong to the partial
+    old: NDSlice
+    #: n-dimensional slice of the new chunks along each axis that belong to the partial
+    new: NDSlice
+    #: Indices of the first value of the left-most old chunk along each axis
+    #: to include in this partial. Everything left to this index belongs to
+    #: the previous partial.
+    left_starts: NDIndex
+    #: Indices of the first value of the right-most old chunk along each axis
+    #: to exclude from this partial.
+    #: This corresponds to `left_start` of the subsequent partial.
+    right_stops: NDIndex
+
 
+def rechunk_p2p(x: da.Array, chunks: ChunkedAxes) -> da.Array:
     import dask.array as da
 
     if x.size == 0:
         # Special case for empty array, as the algorithm below does not behave correctly
         return da.empty(x.shape, chunks=chunks, dtype=x.dtype)
 
-    dsk: dict = {}
+    dsk = {}
     token = tokenize(x, chunks)
-    _barrier_key = barrier_key(ShuffleId(token))
-    name = f"rechunk-transfer-{token}"
+    name = f"rechunk-p2p-{token}"
+    for ndpartial in _split_partials(x, chunks):
+        if all(slc.stop == slc.start + 1 for slc in ndpartial.new):
+            # Single output chunk
+            dsk.update(partial_concatenate(x, chunks, ndpartial, name))
+        else:
+            dsk.update(partial_rechunk(x, chunks, ndpartial, name))
+    layer = MaterializedLayer(dsk)
+    graph = HighLevelGraph.from_collections(name, layer, dependencies=[x])
+    arr = da.Array(graph, name, chunks, meta=x)
+    return arr
+
+
+def _split_partials(
+    x: da.Array, chunks: ChunkedAxes
+) -> Generator[_NDPartial, None, None]:
+    """Split the rechunking into partials that can be performed separately"""
+    partials_per_axis = _split_partials_per_axis(x, chunks)
+    for partial_per_axis in product(*partials_per_axis):
+        old, new, left_starts, right_stops = zip(*partial_per_axis)
+        yield _NDPartial(old, new, left_starts, right_stops)
+
+
+def _split_partials_per_axis(
+    x: da.Array, chunks: ChunkedAxes
+) -> tuple[tuple[_Partial, ...], ...]:
+    """Split the rechunking into partials that can be performed separately
+    on each axis"""
+    from dask.array.rechunk import old_to_new
+
+    chunked_shape = tuple(len(axis) for axis in chunks)
+    _old_to_new = old_to_new(x.chunks, chunks)
+
+    sliced_axes = _partial_slices(_old_to_new, chunked_shape)
+
+    partial_axes = []
+    for axis_index, slices in enumerate(sliced_axes):
+        partials = []
+        for slice_ in slices:
+            last_old_chunk: int
+            first_old_chunk, first_old_slice = _old_to_new[axis_index][slice_.start][0]
+            last_old_chunk, last_old_slice = _old_to_new[axis_index][slice_.stop - 1][
+                -1
+            ]
+            partials.append(
+                _Partial(
+                    old=slice(first_old_chunk, last_old_chunk + 1),
+                    new=slice_,
+                    left_start=first_old_slice.start,
+                    right_stop=last_old_slice.stop,
+                )
+            )
+        partial_axes.append(tuple(partials))
+    return tuple(partial_axes)
+
+
+def _partial_slices(
+    old_to_new: list[list[list[tuple[int, slice]]]], chunked_shape: NDIndex
+) -> tuple[tuple[slice, ...], ...]:
+    """Compute the slices of the new chunks that can be computed separately"""
+    sliced_axes = []
+
+    for axis_index, old_to_new_axis in enumerate(old_to_new):
+        # Two consecutive output chunks A and B belong to the same partial rechunk
+        # if B is fully included in the right-most input chunk of A, i.e.,
+        # separating A and B would not allow us to cull more input tasks.
+
+        # Index of the last input chunk of this partial rechunk
+        last_old_chunk: int | None = None
+        partial_splits = [0]
+        recipe: list[tuple[int, slice]]
+        for new_chunk_index, recipe in enumerate(old_to_new_axis):
+            if len(recipe) == 0:
+                continue
+            current_last_old_chunk, old_slice = recipe[-1]
+            if last_old_chunk is None:
+                last_old_chunk = current_last_old_chunk
+            elif last_old_chunk != current_last_old_chunk:
+                partial_splits.append(new_chunk_index)
+                last_old_chunk = current_last_old_chunk
+        partial_splits.append(chunked_shape[axis_index])
+        sliced_axes.append(
+            tuple(slice(a, b) for a, b in toolz.sliding_window(2, partial_splits))
+        )
+    return tuple(sliced_axes)
+
+
+def _partial_ndindex(ndslice: NDSlice) -> np.ndindex:
+    import numpy as np
+
+    return np.ndindex(tuple(slice.stop - slice.start for slice in ndslice))
+
+
+def _global_index(partial_index: NDIndex, partial_offset: NDIndex) -> NDIndex:
+    return tuple(index + offset for index, offset in zip(partial_index, partial_offset))
+
+
+def partial_concatenate(
+    x: da.Array,
+    chunks: ChunkedAxes,
+    ndpartial: _NDPartial,
+    name: str,
+) -> dict[Key, Any]:
+    import numpy as np
+
+    from dask.array.chunk import getitem
+    from dask.array.core import concatenate3
+
+    dsk: dict[Key, Any] = {}
+
+    partial_token = tokenize(x, chunks, ndpartial.new)
+    slice_name = f"rechunk-slice-{partial_token}"
+    old_offset = tuple(slice_.start for slice_ in ndpartial.old)
+
+    shape = tuple(slice_.stop - slice_.start for slice_ in ndpartial.old)
+    rec_cat_arg = np.empty(shape, dtype="O")
+
+    partial_old = _compute_partial_old_chunks(ndpartial, x.chunks)
+
+    for old_partial_index in _partial_ndindex(ndpartial.old):
+        old_global_index = _global_index(old_partial_index, old_offset)
+        # TODO: Precompute slicing to avoid duplicate work
+        ndslice = ndslice_for(
+            old_partial_index, partial_old, ndpartial.left_starts, ndpartial.right_stops
+        )
+
+        original_shape = tuple(
+            axis[index] for index, axis in zip(old_global_index, x.chunks)
+        )
+        if _slicing_is_necessary(ndslice, original_shape):
+            rec_cat_arg[old_partial_index] = (slice_name,) + old_global_index
+            dsk[(slice_name,) + old_global_index] = (
+                getitem,
+                (x.name,) + old_global_index,
+                ndslice,
+            )
+        else:
+            rec_cat_arg[old_partial_index] = (x.name,) + old_global_index
+    global_index = tuple(int(slice_.start) for slice_ in ndpartial.new)
+    dsk[(name,) + global_index] = (
+        concatenate3,
+        rec_cat_arg.tolist(),
+    )
+    return dsk
+
+
+def _compute_partial_old_chunks(
+    partial: _NDPartial, chunks: ChunkedAxes
+) -> ChunkedAxes:
+    _partial_old = []
+    for axis_index in range(len(partial.old)):
+        c = list(chunks[axis_index][partial.old[axis_index]])
+        c[0] = c[0] - partial.left_starts[axis_index]
+        if (stop := partial.right_stops[axis_index]) is not None:
+            c[-1] = stop
+        _partial_old.append(tuple(c))
+    return tuple(_partial_old)
+
+
+def _slicing_is_necessary(slice: NDSlice, shape: tuple[int | None, ...]) -> bool:
+    """Return True if applying the slice alters the shape, False otherwise."""
+    return not all(
+        slc.start == 0 and (size is None and slc.stop is None or slc.stop == size)
+        for slc, size in zip(slice, shape)
+    )
+
+
+def partial_rechunk(
+    x: da.Array,
+    chunks: ChunkedAxes,
+    ndpartial: _NDPartial,
+    name: str,
+) -> dict[Key, Any]:
+    from dask.array.chunk import getitem
+
+    dsk: dict[Key, Any] = {}
+
+    old_partial_offset = tuple(slice_.start for slice_ in ndpartial.old)
+
+    partial_token = tokenize(x, chunks, ndpartial.new)
+    _barrier_key = barrier_key(ShuffleId(partial_token))
+    slice_name = f"rechunk-slice-{partial_token}"
+    transfer_name = f"rechunk-transfer-{partial_token}"
     disk: bool = dask.config.get("distributed.p2p.disk")
+
+    ndim = len(x.shape)
+
+    partial_old = _compute_partial_old_chunks(ndpartial, x.chunks)
+    partial_new: ChunkedAxes = tuple(
+        chunks[axis_index][ndpartial.new[axis_index]] for axis_index in range(ndim)
+    )
+
     transfer_keys = []
-    for index in np.ndindex(tuple(len(dim) for dim in x.chunks)):
-        transfer_keys.append((name,) + index)
-        dsk[(name,) + index] = (
-            rechunk_transfer,
-            (x.name,) + index,
-            token,
-            index,
-            chunks,
-            x.chunks,
-            disk,
+    for partial_index in _partial_ndindex(ndpartial.old):
+        ndslice = ndslice_for(
+            partial_index, partial_old, ndpartial.left_starts, ndpartial.right_stops
         )
 
-    dsk[_barrier_key] = (shuffle_barrier, token, transfer_keys)
+        global_index = _global_index(partial_index, old_partial_offset)
 
-    name = f"rechunk-p2p-{token}"
+        original_shape = tuple(
+            axis[index] for index, axis in zip(global_index, x.chunks)
+        )
+        if _slicing_is_necessary(ndslice, original_shape):
+            input_task = (slice_name,) + global_index
+            dsk[(slice_name,) + global_index] = (
+                getitem,
+                (x.name,) + global_index,
+                ndslice,
+            )
+        else:
+            input_task = (x.name,) + global_index
 
-    for index in np.ndindex(tuple(len(dim) for dim in chunks)):
-        dsk[(name,) + index] = (rechunk_unpack, token, index, _barrier_key)
+        transfer_keys.append((transfer_name,) + global_index)
+        dsk[(transfer_name,) + global_index] = (
+            rechunk_transfer,
+            input_task,
+            partial_token,
+            partial_index,
+            partial_new,
+            partial_old,
+            disk,
+        )
 
-    with dask.annotate(shuffle=lambda key: key[1:]):
-        layer = MaterializedLayer(dsk)
-        graph = HighLevelGraph.from_collections(name, layer, dependencies=[x])
+    dsk[_barrier_key] = (shuffle_barrier, partial_token, transfer_keys)
 
-        return da.Array(graph, name, chunks, meta=x)
+    new_partial_offset = tuple(axis.start for axis in ndpartial.new)
+    for partial_index in _partial_ndindex(ndpartial.new):
+        global_index = _global_index(partial_index, new_partial_offset)
+        dsk[(name,) + global_index] = (
+            rechunk_unpack,
+            partial_token,
+            partial_index,
+            _barrier_key,
+        )
+    return dsk
+
+
+def ndslice_for(
+    partial_index: NDIndex,
+    chunks: ChunkedAxes,
+    left_starts: NDIndex,
+    right_stops: NDIndex,
+) -> NDSlice:
+    slices = []
+    shape = tuple(len(axis) for axis in chunks)
+    for axis_index, chunked_axis in enumerate(chunks):
+        chunk_index = partial_index[axis_index]
+        start = left_starts[axis_index] if chunk_index == 0 else 0
+        stop = (
+            right_stops[axis_index]
+            if chunk_index == shape[axis_index] - 1
+            else chunked_axis[chunk_index] + start
+        )
+        slices.append(slice(start, stop))
+    return tuple(slices)
 
 
 class Split(NamedTuple):
diff --git a/distributed/shuffle/_scheduler_plugin.py b/distributed/shuffle/_scheduler_plugin.py
--- a/distributed/shuffle/_scheduler_plugin.py
+++ b/distributed/shuffle/_scheduler_plugin.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextlib
+import itertools
 import logging
 from collections import defaultdict
 from typing import TYPE_CHECKING, Any
@@ -49,6 +50,7 @@ class ShuffleSchedulerPlugin(SchedulerPlugin):
     heartbeats: defaultdict[ShuffleId, dict]
     _shuffles: defaultdict[ShuffleId, set[SchedulerShuffleState]]
     _archived_by_stimulus: defaultdict[str, set[SchedulerShuffleState]]
+    _shift_counter: itertools.count[int]
 
     def __init__(self, scheduler: Scheduler):
         self.scheduler = scheduler
@@ -65,6 +67,7 @@ def __init__(self, scheduler: Scheduler):
         self.scheduler.add_plugin(self, name="shuffle")
         self._shuffles = defaultdict(set)
         self._archived_by_stimulus = defaultdict(set)
+        self._shift_counter = itertools.count()
 
     async def start(self, scheduler: Scheduler) -> None:
         worker_plugin = ShuffleWorkerPlugin()
@@ -203,6 +206,11 @@ def _calculate_worker_for(self, spec: ShuffleSpec) -> dict[Any, str]:
         else:
             workers = list(self.scheduler.workers)
 
+        # Ensure homogeneous cluster utilization when there are multiple small,
+        # independent shuffles going on at the same time, e.g. due to partial rechunking
+        shift_by = next(self._shift_counter) % len(workers)
+        workers = workers[shift_by:] + workers[:shift_by]
+
         # Check if this shuffle shares output tasks with a different shuffle that has
         # already been initialized and needs to be taken into account when
         # mapping output partitions to workers.


Test patch:
diff --git a/distributed/shuffle/tests/test_rechunk.py b/distributed/shuffle/tests/test_rechunk.py
--- a/distributed/shuffle/tests/test_rechunk.py
+++ b/distributed/shuffle/tests/test_rechunk.py
@@ -19,6 +19,8 @@
 from dask.array.core import concatenate3
 from dask.array.rechunk import normalize_chunks, rechunk
 from dask.array.utils import assert_eq
+from dask.base import flatten
+from dask.utils import key_split
 
 from distributed import Event
 from distributed.protocol.utils_test import get_host_array
@@ -31,7 +33,7 @@
     split_axes,
 )
 from distributed.shuffle.tests.utils import AbstractShuffleTestPool
-from distributed.utils_test import gen_cluster, gen_test, raises_with_cause
+from distributed.utils_test import gen_cluster, gen_test
 
 NUMPY_GE_124 = parse_version(np.__version__) >= parse_version("1.24")
 
@@ -204,6 +206,62 @@ async def test_rechunk_configuration(c, s, *ws, config_value, keyword):
     assert np.all(await c.compute(x2) == a)
 
 
+@gen_cluster(client=True)
+async def test_cull_p2p_rechunk_independent_partitions(c, s, *ws):
+    a = np.random.default_rng().uniform(0, 1, 1000).reshape((10, 10, 10))
+    x = da.from_array(a, chunks=(1, 5, 1))
+    new = (5, 1, -1)
+    rechunked = rechunk(x, chunks=new, method="p2p")
+    (dsk,) = dask.optimize(rechunked)
+    culled = rechunked[:5, :2]
+    (dsk_culled,) = dask.optimize(culled)
+
+    # The culled graph requires only 1/2 of the input tasks
+    n_inputs = len(
+        [1 for key in dsk.dask.get_all_dependencies() if key[0].startswith("array-")]
+    )
+    n_culled_inputs = len(
+        [
+            1
+            for key in dsk_culled.dask.get_all_dependencies()
+            if key[0].startswith("array-")
+        ]
+    )
+    assert n_culled_inputs == n_inputs / 4
+    # The culled graph should also have less than 1/4 the tasks
+    assert len(dsk_culled.dask) < len(dsk.dask) / 4
+
+    assert np.all(await c.compute(culled) == a[:5, :2])
+
+
+@gen_cluster(client=True)
+async def test_cull_p2p_rechunk_overlapping_partitions(c, s, *ws):
+    a = np.random.default_rng().uniform(0, 1, 500).reshape((10, 10, 5))
+    x = da.from_array(a, chunks=(1, 5, 1))
+    new = (5, 3, -1)
+    rechunked = rechunk(x, chunks=new, method="p2p")
+    (dsk,) = dask.optimize(rechunked)
+    culled = rechunked[:5, :2]
+    (dsk_culled,) = dask.optimize(culled)
+
+    # The culled graph requires only 1/4 of the input tasks
+    n_inputs = len(
+        [1 for key in dsk.dask.get_all_dependencies() if key[0].startswith("array-")]
+    )
+    n_culled_inputs = len(
+        [
+            1
+            for key in dsk_culled.dask.get_all_dependencies()
+            if key[0].startswith("array-")
+        ]
+    )
+    assert n_culled_inputs == n_inputs / 4
+    # The culled graph should also have less than 1/4 the tasks
+    assert len(dsk_culled.dask) < len(dsk.dask) / 4
+
+    assert np.all(await c.compute(culled) == a[:5, :2])
+
+
 @pytest.mark.parametrize("disk", [True, False])
 @gen_cluster(client=True)
 async def test_rechunk_2d(c, s, *ws, disk):
@@ -234,12 +292,7 @@ async def test_rechunk_4d(c, s, *ws, disk):
     old = ((5, 5),) * 4
     a = np.random.default_rng().uniform(0, 1, 10000).reshape((10,) * 4)
     x = da.from_array(a, chunks=old)
-    new = (
-        (10,),
-        (10,),
-        (10,),
-        (8, 2),
-    )  # This has been altered to return >1 output partition
+    new = ((10,),) * 4
     with dask.config.set({"distributed.p2p.disk": disk}):
         x2 = rechunk(x, chunks=new, method="p2p")
     assert x2.chunks == new
@@ -247,27 +300,6 @@ async def test_rechunk_4d(c, s, *ws, disk):
     assert np.all(await c.compute(x2) == a)
 
 
-@gen_cluster(client=True)
-async def test_rechunk_with_single_output_chunk_raises(c, s, *ws):
-    """See distributed#7816
-
-    See Also
-    --------
-    dask.array.tests.test_rechunk.test_rechunk_4d
-    """
-    old = ((5, 5),) * 4
-    a = np.random.default_rng().uniform(0, 1, 10000).reshape((10,) * 4)
-    x = da.from_array(a, chunks=old)
-    new = ((10,),) * 4
-    x2 = rechunk(x, chunks=new, method="p2p")
-    assert x2.chunks == new
-    # FIXME: distributed#7816
-    with raises_with_cause(
-        RuntimeError, "failed during transfer", RuntimeError, "Barrier task"
-    ):
-        await c.compute(x2)
-
-
 @gen_cluster(client=True)
 async def test_rechunk_expand(c, s, *ws):
     """
@@ -774,7 +806,6 @@ async def test_rechunk_empty_chunks(c, s, *ws):
     assert_eq(await c.compute(x), await c.compute(y))
 
 
-@pytest.mark.skip(reason="FIXME: We should avoid P2P in this case")
 @gen_cluster(client=True)
 async def test_rechunk_avoid_needless_chunking(c, s, *ws):
     x = da.ones(16, chunks=2)
@@ -1206,7 +1237,7 @@ def blocked(chunk, in_map, block_map):
 
 
 @pytest.mark.parametrize("nworkers", [1, 2, 41, 50])
-def test_worker_for_homogeneous_distribution(nworkers):
+def test_pick_worker_homogeneous_distribution(nworkers):
     old = ((1, 2, 3, 4), (5,) * 6)
     new = ((5, 5), (12, 18))
     workers = [str(i) for i in range(nworkers)]
@@ -1218,3 +1249,29 @@ def test_worker_for_homogeneous_distribution(nworkers):
     assert sum(count.values()) > 0
     assert sum(count.values()) == len(list(spec.output_partitions))
     assert abs(max(count.values()) - min(count.values())) <= 1
+
+
+@pytest.mark.slow
+@gen_cluster(
+    client=True,
+    nthreads=[("", 1)] * 5,
+    # Future-proof: disable auto-rebalancing
+    config={"distributed.scheduler.active-memory-manager.start": False},
+)
+async def test_partial_rechunk_homogeneous_distribution(c, s, *workers):
+    da = pytest.importorskip("dask.array")
+
+    # This rechunk operation can be split into 10 independent shuffles with 4 output
+    # chunks each. This is less than the number of workers, so we are at risk of
+    # choosing the same 4 output workers in each separate shuffle.
+    arr = da.random.random((40, 40), chunks=(4, 4))
+    arr = arr.rechunk((1, -1), method="p2p")
+    # Test that we are, in fact, triggering partial rechunks
+    assert sum(key_split(k) == "shuffle-barrier" for k in arr.__dask_graph__()) == 10
+
+    arr = await c.persist(arr)
+    # Don't count input and intermediate keys that have not been released yet
+    out_keys = set(flatten(arr.__dask_keys__()))
+    nchunks = [len(w.data.keys() & out_keys) for w in workers]
+    # There are 40 output chunks and 5 workers. Expect exactly 8 chunks per worker.
+    assert nchunks == [8, 8, 8, 8, 8]