Repository: dask/distributed

Issue:
Use metering for P2P shuffling instrumentation
The P2P extensions (primarily the buffers) are implementing their own version for diagnostics, e.g. here https://github.com/dask/distributed/blob/429ef8cc682da36017a28e061d1c773066470fe3/distributed/shuffle/_buffer.py#L267-L272

I think P2P diagnostics would greatly benefit from the metering / fine grained metrics and we should consider replacing (os extending) the custom ctx managers there with the metrics.meter (mind that there is also a special dashboard for P2P shuffling that is using this custom diagnostics. I don't want to break this dashboard).

cc @crusaderky @hendrikmakait 

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

       27 files  ±  0         27 suites  ±0   12h 3m 28s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "duration of all tests") + 9m 53s
  3 942 tests +  3    3 821 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "passed tests") +  3     118 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "skipped / disabled tests") +1  3 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "failed tests")  - 1 
49 580 runs  +27  47 160 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "passed tests") +24  2 417 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "skipped / disabled tests") +4  3 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "failed tests")  - 1 

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

Results for commit 82ac08a6. ± Comparison against base commit a923909a.

[test-results]:data:application/gzip;base64,H4sIAPGnWmUC/02Myw6DIBBFf8Ww7oLH6EB/xiCPhNRHg7Ay/fdCG9TlOXfmHMSH2e3k2XF8dGTPIV1kc9QpbGthEEBlUWVNdRcKeMNxz8ZUJzm73Cu8i2Ps/Bq9DnM9O4WLcYvF0GJiXmsXVC8bti4gGy7373Jg2NSt++N71mzLElIBIrk2VOpBMaEEWuo5Dgiid1RSa9EDSkf1NJHPF90e+zcUAQAA

crusaderky: @hendrikmakait I've reinstated all metrics that are visible from the dashboard, as discussed. This is ready for review again.

Patch:
diff --git a/distributed/actor.py b/distributed/actor.py
index 88b7421e..1fdbf5da 100644
--- a/distributed/actor.py
+++ b/distributed/actor.py
@@ -246,11 +246,11 @@ class BaseActorFuture(abc.ABC, Awaitable[_T]):
 
     @abc.abstractmethod
     def result(self, timeout: str | timedelta | float | None = None) -> _T:
-        ...  # pragma: nocover
+        ...
 
     @abc.abstractmethod
     def done(self) -> bool:
-        ...  # pragma: nocover
+        ...
 
     def __repr__(self) -> Literal["<ActorFuture>"]:
         return "<ActorFuture>"
diff --git a/distributed/dashboard/components/scheduler.py b/distributed/dashboard/components/scheduler.py
index b7a41515..f4c91d7d 100644
--- a/distributed/dashboard/components/scheduler.py
+++ b/distributed/dashboard/components/scheduler.py
@@ -4471,12 +4471,8 @@ class Shuffling(DashboardComponent):
                 data[f"{prefix}_memory"].append(d[prefix]["memory"])
                 data[f"{prefix}_memory_limit"].append(memory_limit)
                 data[f"{prefix}_buckets"].append(d[prefix]["buckets"])
-                data[f"{prefix}_avg_duration"].append(
-                    d[prefix]["diagnostics"].get("avg_duration", 0)
-                )
-                data[f"{prefix}_avg_size"].append(
-                    d[prefix]["diagnostics"].get("avg_size", 0)
-                )
+                data[f"{prefix}_avg_duration"].append(d[prefix]["avg_duration"])
+                data[f"{prefix}_avg_size"].append(d[prefix]["avg_size"])
                 data[f"{prefix}_read"].append(d[prefix]["read"])
                 data[f"{prefix}_written"].append(d[prefix]["written"])
                 if self.scheduler.workers[worker].last_seen < now - 5:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index e7325b7a..b173a478 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -2823,7 +2823,7 @@ class SchedulerState:
                 # It's ok to forget a task that nobody needs
                 pass
             else:
-                raise AssertionError("Unreachable", ts)  # pragma: nocover
+                raise AssertionError("Unreachable", str(ts))  # pragma: nocover
 
         if ts.actor:
             for ws in ts.who_has or ():
diff --git a/distributed/shuffle/_buffer.py b/distributed/shuffle/_buffer.py
index 961f56cd..60cc0a86 100644
--- a/distributed/shuffle/_buffer.py
+++ b/distributed/shuffle/_buffer.py
@@ -1,10 +1,10 @@
 from __future__ import annotations
 
+import abc
 import asyncio
-import contextlib
 import logging
 from collections import defaultdict
-from collections.abc import Iterator, Sized
+from collections.abc import Sized
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 from distributed.metrics import context_meter, time
@@ -49,13 +49,14 @@ class ShardsBuffer(Generic[ShardType]):
     sizes_detail: defaultdict[str, list[int]]
     concurrency_limit: int
     memory_limiter: ResourceLimiter
-    diagnostics: dict[str, float]
     max_message_size: int
 
     bytes_total: int
     bytes_memory: int
     bytes_written: int
     bytes_read: int
+    avg_size: float
+    avg_duration: float
 
     _accepts_input: bool
     _inputs_done: bool
@@ -78,7 +79,6 @@ class ShardsBuffer(Generic[ShardType]):
         self.concurrency_limit = concurrency_limit
         self._inputs_done = False
         self.memory_limiter = memory_limiter
-        self.diagnostics: dict[str, float] = defaultdict(float)
         self._tasks = [
             asyncio.create_task(self._background_task())
             for _ in range(concurrency_limit)
@@ -91,6 +91,8 @@ class ShardsBuffer(Generic[ShardType]):
         self.bytes_memory = 0
         self.bytes_written = 0
         self.bytes_read = 0
+        self.avg_size = 0.0
+        self.avg_duration = 0.0
 
     def heartbeat(self) -> dict[str, Any]:
         return {
@@ -99,7 +101,8 @@ class ShardsBuffer(Generic[ShardType]):
             "buckets": len(self.shards),
             "written": self.bytes_written,
             "read": self.bytes_read,
-            "diagnostics": self.diagnostics,
+            "avg_size": self.avg_size,
+            "avg_duration": self.avg_duration,
             "memory_limit": self.memory_limiter.limit,
         }
 
@@ -113,12 +116,8 @@ class ShardsBuffer(Generic[ShardType]):
             self.bytes_written += size
 
             stop = time()
-            self.diagnostics["avg_size"] = (
-                0.98 * self.diagnostics["avg_size"] + 0.02 * size
-            )
-            self.diagnostics["avg_duration"] = 0.98 * self.diagnostics[
-                "avg_duration"
-            ] + 0.02 * (stop - start)
+            self.avg_size = 0.98 * self.avg_size + 0.02 * size
+            self.avg_duration = 0.98 * self.avg_duration + 0.02 * (stop - start)
         except Exception as e:
             self._exception = e
             self._inputs_done = True
@@ -126,11 +125,12 @@ class ShardsBuffer(Generic[ShardType]):
             await self.memory_limiter.decrease(size)
             self.bytes_memory -= size
 
+    @abc.abstractmethod
     async def _process(self, id: str, shards: list[ShardType]) -> None:
-        raise NotImplementedError()
+        ...
 
     def read(self, id: str) -> ShardType:
-        raise NotImplementedError()
+        raise NotImplementedError()  # pragma: nocover
 
     @property
     def empty(self) -> bool:
@@ -262,10 +262,3 @@ class ShardsBuffer(Generic[ShardType]):
 
     async def __aexit__(self, exc: Any, typ: Any, traceback: Any) -> None:
         await self.close()
-
-    @contextlib.contextmanager
-    def time(self, name: str) -> Iterator[None]:
-        start = time()
-        yield
-        stop = time()
-        self.diagnostics[name] += stop - start
diff --git a/distributed/shuffle/_comms.py b/distributed/shuffle/_comms.py
index a9776be5..50094afd 100644
--- a/distributed/shuffle/_comms.py
+++ b/distributed/shuffle/_comms.py
@@ -68,5 +68,5 @@ class CommShardsBuffer(ShardsBuffer):
     async def _process(self, address: str, shards: list[tuple[Any, Any]]) -> None:
         """Send one message off to a neighboring worker"""
         # Consider boosting total_size a bit here to account for duplication
-        with self.time("send"), context_meter.meter("send"):
+        with context_meter.meter("send"):
             await self.send(address, shards)
diff --git a/distributed/shuffle/_core.py b/distributed/shuffle/_core.py
index 810bf1cf..dfb4e885 100644
--- a/distributed/shuffle/_core.py
+++ b/distributed/shuffle/_core.py
@@ -6,7 +6,6 @@ import contextlib
 import itertools
 import pickle
 import time
-from collections import defaultdict
 from collections.abc import Callable, Generator, Hashable, Iterable, Iterator, Sequence
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
@@ -64,7 +63,6 @@ class ShuffleRun(Generic[_T_partition_id, _T_partition_type]):
     closed: bool
     _disk_buffer: DiskShardsBuffer | MemoryShardsBuffer
     _comm_buffer: CommShardsBuffer
-    diagnostics: dict[str, float]
     received: set[_T_partition_id]
     total_recvd: int
     start_time: float
@@ -124,7 +122,6 @@ class ShuffleRun(Generic[_T_partition_id, _T_partition_type]):
         # TODO: reduce number of connections to number of workers
         # MultiComm.max_connections = min(10, n_workers)
 
-        self.diagnostics = defaultdict(float)
         self.transferred = False
         self.received = set()
         self.total_recvd = 0
@@ -182,13 +179,6 @@ class ShuffleRun(Generic[_T_partition_id, _T_partition_type]):
         with context_meter.add_callback(callback, allow_offload="background" in where):
             yield
 
-    @contextlib.contextmanager
-    def time(self, name: str) -> Iterator[None]:
-        start = time.time()
-        yield
-        stop = time.time()
-        self.diagnostics[name] += stop - start
-
     async def barrier(self, run_ids: Sequence[int]) -> int:
         self.raise_if_closed()
         consistent = all(run_id == self.run_id for run_id in run_ids)
@@ -233,7 +223,7 @@ class ShuffleRun(Generic[_T_partition_id, _T_partition_type]):
         self, func: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs
     ) -> _T:
         self.raise_if_closed()
-        with self.time("cpu"), context_meter.meter("offload"):
+        with context_meter.meter("offload"):
             return await run_in_executor_with_context(
                 self.executor, func, *args, **kwargs
             )
@@ -244,7 +234,6 @@ class ShuffleRun(Generic[_T_partition_id, _T_partition_type]):
         return {
             "disk": self._disk_buffer.heartbeat(),
             "comm": comm_heartbeat,
-            "diagnostics": self.diagnostics,
             "start": self.start_time,
         }
 
@@ -434,10 +423,10 @@ class ShuffleSpec(abc.ABC, Generic[_T_partition_id]):
     id: ShuffleId
     disk: bool
 
-    @abc.abstractproperty
+    @property
+    @abc.abstractmethod
     def output_partitions(self) -> Generator[_T_partition_id, None, None]:
         """Output partitions"""
-        raise NotImplementedError
 
     @abc.abstractmethod
     def pick_worker(self, partition: _T_partition_id, workers: Sequence[str]) -> str:
diff --git a/distributed/shuffle/_disk.py b/distributed/shuffle/_disk.py
index 0c0b30ce..86a678a8 100644
--- a/distributed/shuffle/_disk.py
+++ b/distributed/shuffle/_disk.py
@@ -167,7 +167,6 @@ class DiskShardsBuffer(ShardsBuffer):
 
         with (
             self._directory_lock.read(),
-            self.time("write"),
             context_meter.meter("disk-write"),
             serialize_meter_ctx,
         ):
@@ -198,8 +197,7 @@ class DiskShardsBuffer(ShardsBuffer):
                 # measure seconds here, as it would shadow "p2p-get-output-cpu" and
                 # "p2p-get-output-noncpu". Also, for rechunk it would not measure
                 # the whole disk access, as _read returns memory-mapped buffers.
-                with self.time("read"):
-                    data, size = self._read(fname)
+                data, size = self._read(fname)
                 context_meter.digest_metric("p2p-disk-read", 1, "count")
                 context_meter.digest_metric("p2p-disk-read", size, "bytes")
         except FileNotFoundError:
diff --git a/distributed/shuffle/_limiter.py b/distributed/shuffle/_limiter.py
index a55682d1..c144de86 100644
--- a/distributed/shuffle/_limiter.py
+++ b/distributed/shuffle/_limiter.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 import asyncio
 from typing import Generic, TypeVar
 
-from distributed.metrics import context_meter, time
+from distributed.metrics import context_meter
 
 _T = TypeVar("_T", int, None)
 
@@ -27,8 +27,6 @@ class ResourceLimiter(Generic[_T]):
 
     limit: _T
     metrics_label: str | None
-    time_blocked_total: float
-    time_blocked_avg: float
 
     _acquired: int
     _condition: asyncio.Condition
@@ -40,8 +38,6 @@ class ResourceLimiter(Generic[_T]):
         self._acquired = 0
         self._condition = asyncio.Condition()
         self._waiters = 0
-        self.time_blocked_total = 0.0
-        self.time_blocked_avg = 0.0
 
     def __repr__(self) -> str:
         return f"<ResourceLimiter limit: {self.limit} available: {self.available}>"
@@ -70,20 +66,14 @@ class ResourceLimiter(Generic[_T]):
             # seconds / count
             context_meter.digest_metric(self.metrics_label, 1, "count")
         if not self.full:
-            self.time_blocked_avg *= 0.9
             return
 
-        start = time()
         with context_meter.meter(self.metrics_label):
             async with self._condition:
                 self._waiters += 1
                 await self._condition.wait_for(lambda: not self.full)
                 self._waiters -= 1
 
-        duration = time() - start
-        self.time_blocked_total += duration
-        self.time_blocked_avg = self.time_blocked_avg * 0.9 + duration * 0.1
-
     def increase(self, value: int) -> None:
         """Increase the internal counter by value"""
         self._acquired += value
diff --git a/distributed/shuffle/_memory.py b/distributed/shuffle/_memory.py
index dc9a94f4..92e4b5ee 100644
--- a/distributed/shuffle/_memory.py
+++ b/distributed/shuffle/_memory.py
@@ -23,22 +23,20 @@ class MemoryShardsBuffer(ShardsBuffer):
     async def _process(self, id: str, shards: list[Any]) -> None:
         # TODO: This can be greatly simplified, there's no need for
         # background threads at all.
-        with self.time("write"):
-            self._shards[id].extend(shards)
+        self._shards[id].extend(shards)
 
     def read(self, id: str) -> Any:
         self.raise_on_exception()
         if not self._inputs_done:
             raise RuntimeError("Tried to read from file before done.")
 
-        with self.time("read"):
-            shards = self._shards.pop(id)  # Raises KeyError
-            self.bytes_read += sum(map(sizeof, shards))
-            # Don't keep the serialized and the deserialized shards
-            # in memory at the same time
-            data = []
-            while shards:
-                shard = shards.pop()
-                data.append(self._deserialize(shard))
+        shards = self._shards.pop(id)  # Raises KeyError
+        self.bytes_read += sum(map(sizeof, shards))
+        # Don't keep the serialized and the deserialized shards
+        # in memory at the same time
+        data = []
+        while shards:
+            shard = shards.pop()
+            data.append(self._deserialize(shard))
 
         return data
diff --git a/distributed/shuffle/_rechunk.py b/distributed/shuffle/_rechunk.py
index e50d888a..6bb6e3f0 100644
--- a/distributed/shuffle/_rechunk.py
+++ b/distributed/shuffle/_rechunk.py
@@ -106,6 +106,7 @@ from itertools import product
 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):
@@ -442,8 +703,7 @@ class ArrayRechunkRun(ShuffleRun[NDIndex, "np.ndarray"]):
         data = self._read_from_disk(partition_id)
         # Copy the memory-mapped buffers from disk into memory.
         # This is where we'll spend most time.
-        with self._disk_buffer.time("read"):
-            return convert_chunk(data)
+        return convert_chunk(data)
 
     def deserialize(self, buffer: Any) -> Any:
         return buffer
diff --git a/distributed/shuffle/_scheduler_plugin.py b/distributed/shuffle/_scheduler_plugin.py
index 2fb605d0..09d97fff 100644
--- 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 @@ class ShuffleSchedulerPlugin(SchedulerPlugin):
         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 @@ class ShuffleSchedulerPlugin(SchedulerPlugin):
         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.
diff --git a/pyproject.toml b/pyproject.toml
index f3b38536..3dc62d79 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -243,6 +243,8 @@ exclude_lines = [
     "TYPE_CHECKING",
     "except ImportError",
     "@overload",
+    "@abstractmethod",
+    "@abc.abstractmethod"
 ]
 
 [tool.coverage.html]


Test patch:
diff --git a/distributed/http/scheduler/tests/test_scheduler_http.py b/distributed/http/scheduler/tests/test_scheduler_http.py
index 1718d589..d51ddf33 100644
--- a/distributed/http/scheduler/tests/test_scheduler_http.py
+++ b/distributed/http/scheduler/tests/test_scheduler_http.py
@@ -98,7 +98,13 @@ async def test_prefix(c, s, a, b):
             assert is_valid_xml(body)
 
 
-@gen_cluster(client=True, clean_kwargs={"threads": False})
+@gen_cluster(
+    client=True,
+    clean_kwargs={"threads": False},
+    config={
+        "distributed.admin.system-monitor.gil.enabled": True,
+    },
+)
 async def test_prometheus(c, s, a, b):
     pytest.importorskip("prometheus_client")
 
diff --git a/distributed/http/worker/tests/test_worker_http.py b/distributed/http/worker/tests/test_worker_http.py
index 4d8b12dd..36213589 100644
--- a/distributed/http/worker/tests/test_worker_http.py
+++ b/distributed/http/worker/tests/test_worker_http.py
@@ -15,7 +15,13 @@ from distributed.utils_test import (
 )
 
 
-@gen_cluster(client=True, nthreads=[("127.0.0.1", 1)])
+@gen_cluster(
+    client=True,
+    nthreads=[("127.0.0.1", 1)],
+    config={
+        "distributed.admin.system-monitor.gil.enabled": True,
+    },
+)
 async def test_prometheus(c, s, a):
     pytest.importorskip("prometheus_client")
 
diff --git a/distributed/shuffle/tests/test_buffer.py b/distributed/shuffle/tests/test_buffer.py
index a3ad4b07..c0782d54 100644
--- a/distributed/shuffle/tests/test_buffer.py
+++ b/distributed/shuffle/tests/test_buffer.py
@@ -11,7 +11,7 @@ from dask.utils import parse_bytes
 from distributed.shuffle._buffer import ShardsBuffer
 from distributed.shuffle._limiter import ResourceLimiter
 from distributed.utils import wait_for
-from distributed.utils_test import gen_test
+from distributed.utils_test import captured_context_meter, gen_test
 
 
 def gen_bytes(percentage: float, limit: int) -> bytes:
@@ -51,54 +51,64 @@ limit = parse_bytes("10.0 MiB")
 async def test_memory_limit(big_payloads):
     small_payload = {"small": gen_bytes(0.1, limit)}
 
-    limiter = ResourceLimiter(limit)
+    limiter = ResourceLimiter(limit, "my_label")
 
     async with BufferTest(
         memory_limiter=limiter,
         concurrency_limit=2,
     ) as buf:
-        # It's OK to write nothing
-        await buf.write({})
-
-        many_small = [asyncio.create_task(buf.write(small_payload)) for _ in range(9)]
-        buf.allow_process.set()
-        many_small = asyncio.gather(*many_small)
-        # Puts that do not breach the limit do not block
-        await many_small
-        assert buf.memory_limiter.time_blocked_total == 0
-        buf.allow_process.clear()
-
-        many_small = [asyncio.create_task(buf.write(small_payload)) for _ in range(11)]
-        assert buf.memory_limiter
-        while buf.memory_limiter.available:
-            await asyncio.sleep(0.1)
-
-        new_put = asyncio.create_task(buf.write(small_payload))
-        with pytest.raises(asyncio.TimeoutError):
-            await wait_for(asyncio.shield(new_put), 0.1)
-        buf.allow_process.set()
-        many_small = asyncio.gather(*many_small)
-        await new_put
-
-        while not buf.memory_limiter.empty:
-            await asyncio.sleep(0.1)
-        buf.allow_process.clear()
-        big_tasks = [
-            asyncio.create_task(buf.write(big_payload)) for big_payload in big_payloads
-        ]
-        small = asyncio.create_task(buf.write(small_payload))
-        with pytest.raises(asyncio.TimeoutError):
-            await wait_for(asyncio.shield(asyncio.gather(*big_tasks)), 0.1)
-        with pytest.raises(asyncio.TimeoutError):
-            await wait_for(asyncio.shield(small), 0.1)
-        # Puts only return once we're below memory limit
-        buf.allow_process.set()
-        await asyncio.gather(*big_tasks)
-        await small
-        # Once the big write is through, we can write without blocking again
-        before = buf.memory_limiter.time_blocked_total
-        await buf.write(small_payload)
-        assert before == buf.memory_limiter.time_blocked_total
+        with captured_context_meter() as metrics:
+            # It's OK to write nothing
+            await buf.write({})
+            assert buf.avg_size == 0
+
+            many_small = [
+                asyncio.create_task(buf.write(small_payload)) for _ in range(9)
+            ]
+            buf.allow_process.set()
+            many_small = asyncio.gather(*many_small)
+            # Puts that do not breach the limit do not block
+            await many_small
+            assert metrics.keys() == {("my_label", "count")}  # non-blocking only
+            buf.allow_process.clear()
+
+            many_small = [
+                asyncio.create_task(buf.write(small_payload)) for _ in range(11)
+            ]
+            assert buf.memory_limiter
+            while buf.memory_limiter.available:
+                await asyncio.sleep(0.1)
+
+            new_put = asyncio.create_task(buf.write(small_payload))
+            with pytest.raises(asyncio.TimeoutError):
+                await wait_for(asyncio.shield(new_put), 0.1)
+            buf.allow_process.set()
+            await asyncio.gather(new_put, *many_small)
+
+            while not buf.memory_limiter.empty:
+                await asyncio.sleep(0.1)
+            buf.allow_process.clear()
+            big_tasks = [
+                asyncio.create_task(buf.write(big_payload))
+                for big_payload in big_payloads
+            ]
+            small = asyncio.create_task(buf.write(small_payload))
+            with pytest.raises(asyncio.TimeoutError):
+                await wait_for(asyncio.shield(asyncio.gather(*big_tasks)), 0.1)
+            with pytest.raises(asyncio.TimeoutError):
+                await wait_for(asyncio.shield(small), 0.1)
+            # Puts only return once we're below memory limit
+            buf.allow_process.set()
+            await asyncio.gather(*big_tasks)
+            await small
+            # Once the big write is through, we can write without blocking again
+            before_count = metrics["my_label", "count"]
+            before_seconds = metrics["my_label", "seconds"]
+            assert before_seconds > 0
+            await buf.write(small_payload)
+            assert metrics["my_label", "count"] == before_count + 1
+            assert metrics["my_label", "seconds"] == before_seconds
+            assert buf.avg_size > 0
 
 
 class BufferShardsBroken(ShardsBuffer):
diff --git a/distributed/shuffle/tests/test_limiter.py b/distributed/shuffle/tests/test_limiter.py
index 46253423..e57bccec 100644
--- a/distributed/shuffle/tests/test_limiter.py
+++ b/distributed/shuffle/tests/test_limiter.py
@@ -7,7 +7,7 @@ import pytest
 from distributed.metrics import time
 from distributed.shuffle._limiter import ResourceLimiter
 from distributed.utils import wait_for
-from distributed.utils_test import gen_test
+from distributed.utils_test import captured_context_meter, gen_test
 
 
 @gen_test()
@@ -61,23 +61,23 @@ async def test_limiter_basic():
 @gen_test()
 async def test_unlimited_limiter():
     res = ResourceLimiter(None)
+    with captured_context_meter() as metrics:
+        assert res.empty
+        assert res.available is None
+        assert not res.full
 
-    assert res.empty
-    assert res.available is None
-    assert not res.full
+        res.increase(3)
+        assert not res.empty
+        assert res.available is None
+        assert not res.full
 
-    res.increase(3)
-    assert not res.empty
-    assert res.available is None
-    assert not res.full
-
-    res.increase(2**40)
-    assert not res.empty
-    assert res.available is None
-    assert not res.full
+        res.increase(2**40)
+        assert not res.empty
+        assert res.available is None
+        assert not res.full
 
-    await res.wait_for_available()
-    assert res.time_blocked_total == 0
+        await res.wait_for_available()
+        assert not metrics  # Did not block
 
 
 @gen_test()
@@ -102,34 +102,32 @@ async def test_limiter_concurrent_decrease_releases_waiter():
 
 @gen_test()
 async def test_limiter_statistics():
-    res = ResourceLimiter(1)
+    with captured_context_meter() as metrics:
+        res = ResourceLimiter(1, "foo")
+        await res.wait_for_available()
+        assert metrics == {("foo", "count"): 1}
 
-    assert res.time_blocked_avg == 0.0
-    assert res.time_blocked_total == 0.0
+        res.increase(1)
+        start = time()
 
-    await res.wait_for_available()
-
-    assert res.time_blocked_avg == 0.0
-    assert res.time_blocked_total == 0.0
-
-    res.increase(1)
-    start = time()
-    blocked_wait = asyncio.create_task(res.wait_for_available())
-
-    await asyncio.sleep(0.05)
-
-    assert not blocked_wait.done()
+        blocked_wait = asyncio.create_task(res.wait_for_available())
+        await asyncio.sleep(0.1)
+        assert not blocked_wait.done()
+        await res.decrease(1)
+        await blocked_wait
 
-    await res.decrease(1)
+        stop = time()
+        t_before = metrics["foo", "seconds"]
+        assert 0 < t_before < stop - start
+        assert metrics["foo", "count"] == 2
 
-    await blocked_wait
-    stop = time()
-    assert stop - start >= res.time_blocked_total > 0.0
-    assert res.time_blocked_total > res.time_blocked_avg
+        await res.wait_for_available()
+        assert metrics == {("foo", "count"): 3, ("foo", "seconds"): t_before}
 
-    before_total = res.time_blocked_total
-    before_avg = res.time_blocked_avg
 
-    await res.wait_for_available()
-    assert before_total == res.time_blocked_total
-    assert before_avg > res.time_blocked_avg
+@gen_test()
+async def test_limiter_no_limit_no_statistics():
+    with captured_context_meter() as metrics:
+        res = ResourceLimiter(None)
+        await res.wait_for_available()
+        assert not metrics
diff --git a/distributed/shuffle/tests/test_rechunk.py b/distributed/shuffle/tests/test_rechunk.py
index b7552695..4744dfa6 100644
--- a/distributed/shuffle/tests/test_rechunk.py
+++ b/distributed/shuffle/tests/test_rechunk.py
@@ -19,6 +19,8 @@ import dask
 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 @@ from distributed.shuffle._rechunk import (
     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 @@ async def test_rechunk_in_memory_shards_dont_share_buffer(c, s, a, b):
 
 
 @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]
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
index 37df9ba5..9c167d04 100644
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -1814,6 +1814,7 @@ def config_for_cluster_tests(**extra_config):
             "distributed.scheduler.validate": True,
             "distributed.worker.validate": True,
             "distributed.worker.profile.enabled": False,
+            "distributed.admin.system-monitor.gil.enabled": False,
         },
         **extra_config,
     ):