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   11h 58m 4s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "duration of all tests") + 4m 29s
  3 942 tests +  3    3 819 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "passed tests") +  1     118 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "skipped / disabled tests") +1  5 [: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 158 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "passed tests") +22  2 417 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "skipped / disabled tests") +4  5 [: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/18832121217).

Results for commit 6a06bc01. ± Comparison against base commit a923909a.

[test-results]:data:application/gzip;base64,H4sIAHejWmUC/03MyQrDIBSF4VcJrrtQr2NfpjhEkGYoRlel796bFpMs/+/AeZOUp3Ej94Hr20C2lutZsRVX87pgC6BGIOFa9x2s4D0fWwthN8Psac/8QmPMHJRcnpDkAWMpa0GhKKUt+6+w0vTsv0IzaQ77/3LBdKfL76+vt2Gd51wxiHJU+UDZCFE7gMi9tRK4FxpAJu+1VSHZqMjnC0ryupUUAQAA


Patch:
diff --git a/distributed/dashboard/components/scheduler.py b/distributed/dashboard/components/scheduler.py
index aacd4b21..b7a41515 100644
--- a/distributed/dashboard/components/scheduler.py
+++ b/distributed/dashboard/components/scheduler.py
@@ -3614,21 +3614,23 @@ class FinePerformanceMetrics(DashboardComponent):
         )
         if spans_ext and self.span_tag_selector.value:
             span = spans_ext.merge_by_tags(*self.span_tag_selector.value)
-            execute_metrics = span.cumulative_worker_metrics
+            metrics = span.cumulative_worker_metrics
         elif spans_ext and spans_ext.spans:
             # Calculate idle time
             span = spans_ext.merge_all()
-            execute_metrics = span.cumulative_worker_metrics
+            metrics = span.cumulative_worker_metrics
         else:
             # Spans extension is not loaded
-            execute_metrics = {
+            metrics = {
                 k: v
                 for k, v in self.scheduler.cumulative_worker_metrics.items()
-                if isinstance(k, tuple) and k[0] == "execute"
+                if isinstance(k, tuple)
+                and len(k) == 4  # Skip get-data, gather-deps, and memory-monitor
             }
 
-        for (context, function, activity, unit), v in execute_metrics.items():
-            assert context == "execute"
+        for (context, function, activity, unit), v in metrics.items():
+            if context != "execute":
+                continue  # TODO visualize 'shuffle' metrics
             assert isinstance(function, str)
             assert isinstance(unit, str)
             assert self.unit_selector.value
diff --git a/distributed/shuffle/_core.py b/distributed/shuffle/_core.py
index 0110a5cc..810bf1cf 100644
--- a/distributed/shuffle/_core.py
+++ b/distributed/shuffle/_core.py
@@ -7,7 +7,7 @@ import itertools
 import pickle
 import time
 from collections import defaultdict
-from collections.abc import Callable, Generator, Iterable, Iterator, Sequence
+from collections.abc import Callable, Generator, Hashable, Iterable, Iterator, Sequence
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
 from enum import Enum
@@ -55,9 +55,11 @@ _T = TypeVar("_T")
 class ShuffleRun(Generic[_T_partition_id, _T_partition_type]):
     id: ShuffleId
     run_id: int
+    span_id: str | None
     local_address: str
     executor: ThreadPoolExecutor
     rpc: Callable[[str], PooledRPCCall]
+    digest_metric: Callable[[Hashable, float], None]
     scheduler: PooledRPCCall
     closed: bool
     _disk_buffer: DiskShardsBuffer | MemoryShardsBuffer
@@ -78,10 +80,12 @@ class ShuffleRun(Generic[_T_partition_id, _T_partition_type]):
         self,
         id: ShuffleId,
         run_id: int,
+        span_id: str | None,
         local_address: str,
         directory: str,
         executor: ThreadPoolExecutor,
         rpc: Callable[[str], PooledRPCCall],
+        digest_metric: Callable[[Hashable, float], None],
         scheduler: PooledRPCCall,
         memory_limiter_disk: ResourceLimiter,
         memory_limiter_comms: ResourceLimiter,
@@ -90,23 +94,33 @@ class ShuffleRun(Generic[_T_partition_id, _T_partition_type]):
     ):
         self.id = id
         self.run_id = run_id
+        self.span_id = span_id
         self.local_address = local_address
         self.executor = executor
         self.rpc = rpc
+        self.digest_metric = digest_metric
         self.scheduler = scheduler
         self.closed = False
-        if disk:
-            self._disk_buffer = DiskShardsBuffer(
-                directory=directory,
-                read=self.read,
-                memory_limiter=memory_limiter_disk,
-            )
-        else:
-            self._disk_buffer = MemoryShardsBuffer(deserialize=self.deserialize)
 
-        self._comm_buffer = CommShardsBuffer(
-            send=self.send, memory_limiter=memory_limiter_comms
-        )
+        # Initialize buffers and start background tasks
+        # Don't log metrics issued by the background tasks onto the dask task that
+        # spawned this object
+        with context_meter.clear_callbacks():
+            with self._capture_metrics("background-disk"):
+                if disk:
+                    self._disk_buffer = DiskShardsBuffer(
+                        directory=directory,
+                        read=self.read,
+                        memory_limiter=memory_limiter_disk,
+                    )
+                else:
+                    self._disk_buffer = MemoryShardsBuffer(deserialize=self.deserialize)
+
+            with self._capture_metrics("background-comms"):
+                self._comm_buffer = CommShardsBuffer(
+                    send=self.send, memory_limiter=memory_limiter_comms
+                )
+
         # TODO: reduce number of connections to number of workers
         # MultiComm.max_connections = min(10, n_workers)
 
@@ -136,6 +150,38 @@ class ShuffleRun(Generic[_T_partition_id, _T_partition_type]):
     def __hash__(self) -> int:
         return self.run_id
 
+    @contextlib.contextmanager
+    def _capture_metrics(self, where: str) -> Iterator[None]:
+        """Capture context_meter metrics as
+
+            {('p2p', <span id>, 'foreground|background...', label, unit): value}
+
+        **Note 1:** When the metric is not logged by a background task
+        (where='foreground'), this produces a duplicated metric under
+
+            {('execute', <span id>, <task prefix>, label, unit): value}
+
+        This is by design so that one can have a holistic view of the whole shuffle
+        process.
+
+        **Note 2:** We're immediately writing to Worker.digests.
+        We don't temporarily store metrics under ShuffleRun as we would lose those
+        recorded between the heartbeat and when the ShuffleRun object is deleted at the
+        end of a run.
+        """
+
+        def callback(label: Hashable, value: float, unit: str) -> None:
+            if not isinstance(label, tuple):
+                label = (label,)
+            if isinstance(label[0], str) and label[0].startswith("p2p-"):
+                label = (label[0][len("p2p-") :], *label[1:])
+            name = ("p2p", self.span_id, where, *label, unit)
+
+            self.digest_metric(name, value)
+
+        with context_meter.add_callback(callback, allow_offload="background" in where):
+            yield
+
     @contextlib.contextmanager
     def time(self, name: str) -> Iterator[None]:
         start = time.time()
@@ -288,12 +334,14 @@ class ShuffleRun(Generic[_T_partition_id, _T_partition_type]):
         self.raise_if_closed()
         if self.transferred:
             raise RuntimeError(f"Cannot add more partitions to {self}")
-        with (
-            context_meter.meter("p2p-shard-partition-noncpu"),
-            context_meter.meter("p2p-shard-partition-cpu", func=thread_time),
-        ):
-            shards = self._shard_partition(data, partition_id)
-        sync(self._loop, self._write_to_comm, shards)
+        # Log metrics both in the "execute" and in the "p2p" contexts
+        with self._capture_metrics("foreground"):
+            with (
+                context_meter.meter("p2p-shard-partition-noncpu"),
+                context_meter.meter("p2p-shard-partition-cpu", func=thread_time),
+            ):
+                shards = self._shard_partition(data, partition_id)
+            sync(self._loop, self._write_to_comm, shards)
         return self.run_id
 
     @abc.abstractmethod
@@ -311,6 +359,8 @@ class ShuffleRun(Generic[_T_partition_id, _T_partition_type]):
             raise RuntimeError("`get_output_partition` called before barrier task")
         sync(self._loop, self.flush_receive)
         with (
+            # Log metrics both in the "execute" and in the "p2p" contexts
+            self._capture_metrics("foreground"),
             context_meter.meter("p2p-get-output-noncpu"),
             context_meter.meter("p2p-get-output-cpu", func=thread_time),
         ):
@@ -372,6 +422,7 @@ class ShuffleRunSpec(Generic[_T_partition_id]):
     run_id: int = field(init=False, default_factory=partial(next, itertools.count(1)))
     spec: ShuffleSpec
     worker_for: dict[_T_partition_id, str]
+    span_id: str | None
 
     @property
     def id(self) -> ShuffleId:
@@ -395,9 +446,10 @@ class ShuffleSpec(abc.ABC, Generic[_T_partition_id]):
     def create_new_run(
         self,
         worker_for: dict[_T_partition_id, str],
+        span_id: str | None,
     ) -> SchedulerShuffleState:
         return SchedulerShuffleState(
-            run_spec=ShuffleRunSpec(spec=self, worker_for=worker_for),
+            run_spec=ShuffleRunSpec(spec=self, worker_for=worker_for, span_id=span_id),
             participating_workers=set(worker_for.values()),
         )
 
@@ -405,6 +457,7 @@ class ShuffleSpec(abc.ABC, Generic[_T_partition_id]):
     def create_run_on_worker(
         self,
         run_id: int,
+        span_id: str | None,
         worker_for: dict[_T_partition_id, str],
         plugin: ShuffleWorkerPlugin,
     ) -> ShuffleRun:
diff --git a/distributed/shuffle/_rechunk.py b/distributed/shuffle/_rechunk.py
index f621525c..e50d888a 100644
--- a/distributed/shuffle/_rechunk.py
+++ b/distributed/shuffle/_rechunk.py
@@ -99,7 +99,7 @@ from __future__ import annotations
 import mmap
 import os
 from collections import defaultdict
-from collections.abc import Callable, Generator, Sequence
+from collections.abc import Callable, Generator, Hashable, Sequence
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from itertools import product
@@ -312,6 +312,8 @@ class ArrayRechunkRun(ShuffleRun[NDIndex, "np.ndarray"]):
         A unique `ShuffleID` this belongs to.
     run_id:
         A unique identifier of the specific execution of the shuffle this belongs to.
+    span_id:
+        Span identifier; see :doc:`spans`
     local_address:
         The local address this Shuffle can be contacted by using `rpc`.
     directory:
@@ -321,8 +323,11 @@ class ArrayRechunkRun(ShuffleRun[NDIndex, "np.ndarray"]):
     rpc:
         A callable returning a PooledRPCCall to contact other Shuffle instances.
         Typically a ConnectionPool.
+    digest_metric:
+        A callable to ingest a performance metric.
+        Typically Server.digest_metric.
     scheduler:
-        A PooledRPCCall to to contact the scheduler.
+        A PooledRPCCall to contact the scheduler.
     memory_limiter_disk:
     memory_limiter_comm:
         A ``ResourceLimiter`` limiting the total amount of memory used in either
@@ -336,10 +341,12 @@ class ArrayRechunkRun(ShuffleRun[NDIndex, "np.ndarray"]):
         new: ChunkedAxes,
         id: ShuffleId,
         run_id: int,
+        span_id: str | None,
         local_address: str,
         directory: str,
         executor: ThreadPoolExecutor,
         rpc: Callable[[str], PooledRPCCall],
+        digest_metric: Callable[[Hashable, float], None],
         scheduler: PooledRPCCall,
         memory_limiter_disk: ResourceLimiter,
         memory_limiter_comms: ResourceLimiter,
@@ -349,10 +356,12 @@ class ArrayRechunkRun(ShuffleRun[NDIndex, "np.ndarray"]):
         super().__init__(
             id=id,
             run_id=run_id,
+            span_id=span_id,
             local_address=local_address,
             directory=directory,
             executor=executor,
             rpc=rpc,
+            digest_metric=digest_metric,
             scheduler=scheduler,
             memory_limiter_comms=memory_limiter_comms,
             memory_limiter_disk=memory_limiter_disk,
@@ -485,6 +494,7 @@ class ArrayRechunkSpec(ShuffleSpec[NDIndex]):
     def create_run_on_worker(
         self,
         run_id: int,
+        span_id: str | None,
         worker_for: dict[NDIndex, str],
         plugin: ShuffleWorkerPlugin,
     ) -> ShuffleRun:
@@ -494,6 +504,7 @@ class ArrayRechunkSpec(ShuffleSpec[NDIndex]):
             new=self.new,
             id=self.id,
             run_id=run_id,
+            span_id=span_id,
             directory=os.path.join(
                 plugin.worker.local_directory,
                 f"shuffle-{self.id}-{run_id}",
@@ -501,6 +512,7 @@ class ArrayRechunkSpec(ShuffleSpec[NDIndex]):
             executor=plugin._executor,
             local_address=plugin.worker.address,
             rpc=plugin.worker.rpc,
+            digest_metric=plugin.worker.digest_metric,
             scheduler=plugin.worker.scheduler,
             memory_limiter_disk=plugin.memory_limiter_disk,
             memory_limiter_comms=plugin.memory_limiter_comms,
diff --git a/distributed/shuffle/_scheduler_plugin.py b/distributed/shuffle/_scheduler_plugin.py
index 065ea0a0..2fb605d0 100644
--- a/distributed/shuffle/_scheduler_plugin.py
+++ b/distributed/shuffle/_scheduler_plugin.py
@@ -146,7 +146,9 @@ class ShuffleSchedulerPlugin(SchedulerPlugin):
             self._raise_if_task_not_processing(key)
             worker_for = self._calculate_worker_for(spec)
             self._ensure_output_tasks_are_non_rootish(spec)
-            state = spec.create_new_run(worker_for)
+            state = spec.create_new_run(
+                worker_for=worker_for, span_id=self.scheduler.tasks[key].group.span_id
+            )
             self.active_shuffles[spec.id] = state
             self._shuffles[spec.id].add(state)
             state.participating_workers.add(worker)
diff --git a/distributed/shuffle/_shuffle.py b/distributed/shuffle/_shuffle.py
index 7ea53253..7d024733 100644
--- a/distributed/shuffle/_shuffle.py
+++ b/distributed/shuffle/_shuffle.py
@@ -7,6 +7,7 @@ from collections.abc import (
     Callable,
     Collection,
     Generator,
+    Hashable,
     Iterable,
     Iterator,
     Sequence,
@@ -397,6 +398,8 @@ class DataFrameShuffleRun(ShuffleRun[int, "pd.DataFrame"]):
         A unique `ShuffleID` this belongs to.
     run_id:
         A unique identifier of the specific execution of the shuffle this belongs to.
+    span_id:
+        Span identifier; see :doc:`spans`
     local_address:
         The local address this Shuffle can be contacted by using `rpc`.
     directory:
@@ -406,8 +409,11 @@ class DataFrameShuffleRun(ShuffleRun[int, "pd.DataFrame"]):
     rpc:
         A callable returning a PooledRPCCall to contact other Shuffle instances.
         Typically a ConnectionPool.
+    digest_metric:
+        A callable to ingest a performance metric.
+        Typically Server.digest_metric.
     scheduler:
-        A PooledRPCCall to to contact the scheduler.
+        A PooledRPCCall to contact the scheduler.
     memory_limiter_disk:
     memory_limiter_comm:
         A ``ResourceLimiter`` limiting the total amount of memory used in either
@@ -426,10 +432,12 @@ class DataFrameShuffleRun(ShuffleRun[int, "pd.DataFrame"]):
         meta: pd.DataFrame,
         id: ShuffleId,
         run_id: int,
+        span_id: str | None,
         local_address: str,
         directory: str,
         executor: ThreadPoolExecutor,
         rpc: Callable[[str], PooledRPCCall],
+        digest_metric: Callable[[Hashable, float], None],
         scheduler: PooledRPCCall,
         memory_limiter_disk: ResourceLimiter,
         memory_limiter_comms: ResourceLimiter,
@@ -441,10 +449,12 @@ class DataFrameShuffleRun(ShuffleRun[int, "pd.DataFrame"]):
         super().__init__(
             id=id,
             run_id=run_id,
+            span_id=span_id,
             local_address=local_address,
             directory=directory,
             executor=executor,
             rpc=rpc,
+            digest_metric=digest_metric,
             scheduler=scheduler,
             memory_limiter_comms=memory_limiter_comms,
             memory_limiter_disk=memory_limiter_disk,
@@ -543,7 +553,11 @@ class DataFrameShuffleSpec(ShuffleSpec[int]):
         return _get_worker_for_range_sharding(self.npartitions, partition, workers)
 
     def create_run_on_worker(
-        self, run_id: int, worker_for: dict[int, str], plugin: ShuffleWorkerPlugin
+        self,
+        run_id: int,
+        span_id: str | None,
+        worker_for: dict[int, str],
+        plugin: ShuffleWorkerPlugin,
     ) -> ShuffleRun:
         return DataFrameShuffleRun(
             column=self.column,
@@ -551,6 +565,7 @@ class DataFrameShuffleSpec(ShuffleSpec[int]):
             worker_for=worker_for,
             id=self.id,
             run_id=run_id,
+            span_id=span_id,
             directory=os.path.join(
                 plugin.worker.local_directory,
                 f"shuffle-{self.id}-{run_id}",
@@ -558,6 +573,7 @@ class DataFrameShuffleSpec(ShuffleSpec[int]):
             executor=plugin._executor,
             local_address=plugin.worker.address,
             rpc=plugin.worker.rpc,
+            digest_metric=plugin.worker.digest_metric,
             scheduler=plugin.worker.scheduler,
             memory_limiter_disk=plugin.memory_limiter_disk
             if self.disk
diff --git a/distributed/shuffle/_worker_plugin.py b/distributed/shuffle/_worker_plugin.py
index eddf00d0..c32fca4b 100644
--- a/distributed/shuffle/_worker_plugin.py
+++ b/distributed/shuffle/_worker_plugin.py
@@ -241,7 +241,10 @@ class _ShuffleRunManager:
                 f" expected run_id > {stale_run_id}"
             )
         shuffle_run = result.spec.create_run_on_worker(
-            result.run_id, result.worker_for, self._plugin
+            run_id=result.run_id,
+            worker_for=result.worker_for,
+            plugin=self._plugin,
+            span_id=result.span_id,
         )
         self._active_runs[shuffle_id] = shuffle_run
         self._runs.add(shuffle_run)
diff --git a/distributed/spans.py b/distributed/spans.py
index ef90b073..37d5ad01 100644
--- a/distributed/spans.py
+++ b/distributed/spans.py
@@ -25,6 +25,9 @@ if TYPE_CHECKING:
     from distributed.scheduler import TaskGroup, TaskStateState
 
 
+CONTEXTS_WITH_SPAN_ID = ("execute", "p2p")
+
+
 @contextmanager
 def span(*tags: str) -> Iterator[str]:
     """Tag group of tasks to be part of a certain group, called a span.
@@ -316,8 +319,10 @@ class Span:
 
         At the moment of writing, all keys are
         ``("execute", <task prefix>, <activity>, <unit>)``
-        but more may be added in the future with a different format; please test for
-        ``k[0] == "execute"``.
+        or
+        ``("p2p", <where>, <activity>, <unit>)``
+        but more may be added in the future with a different format; please test e.g.
+        for ``k[0] == "execute"``.
         """
         out = sum_mappings(
             child._cumulative_worker_metrics for child in self.traverse_spans()
@@ -634,7 +639,7 @@ class SpansWorkerExtension:
         self.digests_total_since_heartbeat = {
             k: v
             for k, v in self.worker.digests_total_since_heartbeat.items()
-            if isinstance(k, tuple) and k[0] == "execute"
+            if isinstance(k, tuple) and k[0] in CONTEXTS_WITH_SPAN_ID
         }
 
     def heartbeat(self) -> dict[tuple[Hashable, ...], float]:
diff --git a/distributed/worker.py b/distributed/worker.py
index d87dfe57..55dd5a77 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -92,7 +92,7 @@ from distributed.protocol.serialize import _is_dumpable
 from distributed.pubsub import PubSubWorkerExtension
 from distributed.security import Security
 from distributed.sizeof import safe_sizeof as sizeof
-from distributed.spans import SpansWorkerExtension
+from distributed.spans import CONTEXTS_WITH_SPAN_ID, SpansWorkerExtension
 from distributed.threadpoolexecutor import ThreadPoolExecutor
 from distributed.threadpoolexecutor import secede as tpe_secede
 from distributed.utils import (
@@ -1041,7 +1041,7 @@ class Worker(BaseWorker, ServerNode):
         # Don't cast int metrics to float
         digests: defaultdict[Hashable, float] = defaultdict(int)
         for k, v in self.digests_total_since_heartbeat.items():
-            if isinstance(k, tuple) and k[0] == "execute":
+            if isinstance(k, tuple) and k[0] in CONTEXTS_WITH_SPAN_ID:
                 k = k[:1] + k[2:]
             digests[k] += v
 


Test patch:
diff --git a/distributed/dashboard/tests/test_scheduler_bokeh.py b/distributed/dashboard/tests/test_scheduler_bokeh.py
index 3faf4657..df9736b0 100644
--- a/distributed/dashboard/tests/test_scheduler_bokeh.py
+++ b/distributed/dashboard/tests/test_scheduler_bokeh.py
@@ -450,6 +450,27 @@ async def test_FinePerformanceMetrics_no_spans(c, s, a, b):
     assert sorted(cl.visible_activities) == ["memory-read"]
 
 
+@gen_cluster(client=True)
+async def test_FinePerformanceMetrics_shuffle(c, s, a, b):
+    da = pytest.importorskip("dask.array")
+
+    x = da.random.random((4, 4), chunks=(1, -1))
+    x = x.rechunk((-1, 1), method="p2p")
+    x = x.sum()
+    await c.compute(x)
+    await a.heartbeat()
+    await b.heartbeat()
+
+    cl = FinePerformanceMetrics(s)
+    cl.update()
+    # execute metrics from shuffle
+    assert "shuffle-barrier" in cl.visible_functions
+    assert "p2p-shard-partition-cpu" in cl.visible_activities
+    # shuffle metrics have been filtered out
+    assert "background-comms" not in cl.visible_functions
+    assert "p2p-partition-cpu" not in cl.visible_activities
+
+
 @gen_cluster(client=True)
 async def test_ClusterMemory(c, s, a, b):
     cl = ClusterMemory(s)
diff --git a/distributed/shuffle/tests/test_metrics.py b/distributed/shuffle/tests/test_metrics.py
index 905e8f13..eec82693 100644
--- a/distributed/shuffle/tests/test_metrics.py
+++ b/distributed/shuffle/tests/test_metrics.py
@@ -11,13 +11,21 @@ da = pytest.importorskip("dask.array")
 
 
 def assert_metrics(s: Scheduler, *keys: tuple[str, ...]) -> None:
+    metrics = dict(s.cumulative_worker_metrics)
+    span = s.extensions["spans"].spans_search_by_name["default",][0]
+    span_metrics = dict(span.cumulative_worker_metrics)
+
     for key in keys:
-        assert key in s.cumulative_worker_metrics
-        value = s.cumulative_worker_metrics[key]
-        if key[-1] == "seconds":
-            assert value >= 0
-        else:  # count or bytes
-            assert value > 0
+        value = metrics.pop(key)
+        assert span_metrics.pop(key) == pytest.approx(value)
+        if key[0] == "execute":
+            key2 = ("p2p", "foreground", key[2].replace("p2p-", ""), key[3])
+            assert metrics.pop(key2) == pytest.approx(value)
+            assert span_metrics.pop(key2) == pytest.approx(value)
+
+    # Check that the test doesn't omit any metrics
+    assert [key for key in metrics if key[0] == "p2p"] == []
+    assert [key for key in span_metrics if key[0] == "p2p"] == []
 
 
 @gen_cluster(client=True, config={"optimization.fuse.active": False})
@@ -39,6 +47,21 @@ async def test_rechunk(c, s, a, b):
         ("execute", "rechunk", "p2p-disk-read", "count"),
         ("execute", "rechunk", "p2p-get-output-cpu", "seconds"),
         ("execute", "rechunk", "p2p-get-output-noncpu", "seconds"),
+        ("p2p", "background-comms", "compress", "seconds"),
+        ("p2p", "background-comms", "idle", "seconds"),
+        ("p2p", "background-comms", "process", "bytes"),
+        ("p2p", "background-comms", "process", "count"),
+        ("p2p", "background-comms", "process", "seconds"),
+        ("p2p", "background-comms", "send", "seconds"),
+        ("p2p", "background-comms", "serialize", "seconds"),
+        ("p2p", "background-disk", "disk-write", "bytes"),
+        ("p2p", "background-disk", "disk-write", "count"),
+        ("p2p", "background-disk", "disk-write", "seconds"),
+        ("p2p", "background-disk", "idle", "seconds"),
+        ("p2p", "background-disk", "process", "bytes"),
+        ("p2p", "background-disk", "process", "count"),
+        ("p2p", "background-disk", "process", "seconds"),
+        ("p2p", "background-disk", "serialize", "seconds"),
     )
 
 
@@ -71,4 +94,18 @@ async def test_dataframe(c, s, a, b):
         ("execute", "shuffle_p2p", "p2p-disk-read", "count"),
         ("execute", "shuffle_p2p", "p2p-get-output-cpu", "seconds"),
         ("execute", "shuffle_p2p", "p2p-get-output-noncpu", "seconds"),
+        ("p2p", "background-comms", "compress", "seconds"),
+        ("p2p", "background-comms", "idle", "seconds"),
+        ("p2p", "background-comms", "process", "bytes"),
+        ("p2p", "background-comms", "process", "count"),
+        ("p2p", "background-comms", "process", "seconds"),
+        ("p2p", "background-comms", "send", "seconds"),
+        ("p2p", "background-comms", "serialize", "seconds"),
+        ("p2p", "background-disk", "disk-write", "bytes"),
+        ("p2p", "background-disk", "disk-write", "count"),
+        ("p2p", "background-disk", "disk-write", "seconds"),
+        ("p2p", "background-disk", "idle", "seconds"),
+        ("p2p", "background-disk", "process", "bytes"),
+        ("p2p", "background-disk", "process", "count"),
+        ("p2p", "background-disk", "process", "seconds"),
     )
diff --git a/distributed/shuffle/tests/test_rechunk.py b/distributed/shuffle/tests/test_rechunk.py
index 6dfc6e32..b7552695 100644
--- a/distributed/shuffle/tests/test_rechunk.py
+++ b/distributed/shuffle/tests/test_rechunk.py
@@ -68,9 +68,11 @@ class ArrayRechunkTestPool(AbstractShuffleTestPool):
             directory=directory / name,
             id=ShuffleId(name),
             run_id=next(AbstractShuffleTestPool._shuffle_run_id_iterator),
+            span_id=None,
             local_address=name,
             executor=self._executor,
             rpc=self,
+            digest_metric=lambda name, value: None,
             scheduler=self,
             memory_limiter_disk=ResourceLimiter(10000000),
             memory_limiter_comms=ResourceLimiter(10000000),
diff --git a/distributed/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py
index c34a85cd..b9864d62 100644
--- a/distributed/shuffle/tests/test_shuffle.py
+++ b/distributed/shuffle/tests/test_shuffle.py
@@ -1578,9 +1578,11 @@ class DataFrameShuffleTestPool(AbstractShuffleTestPool):
             directory=directory / name,
             id=ShuffleId(name),
             run_id=next(AbstractShuffleTestPool._shuffle_run_id_iterator),
+            span_id=None,
             local_address=name,
             executor=self._executor,
             rpc=self,
+            digest_metric=lambda name, value: None,
             scheduler=self,
             memory_limiter_disk=ResourceLimiter(10000000),
             memory_limiter_comms=ResourceLimiter(10000000),