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 57m 31s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "duration of all tests") + 3m 56s
  3 942 tests +  3    3 824 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "passed tests") +  6     118 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "skipped / disabled tests") +1  0 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "failed tests")  - 4 
49 580 runs  +27  47 163 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "passed tests") +27  2 417 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "skipped / disabled tests") +4  0 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.11.0/README.md#the-symbols "failed tests")  - 4 

Results for commit 4f016af3. ± Comparison against base commit a923909a.

[test-results]:data:application/gzip;base64,H4sIADmcWmUC/03Myw7CIBCF4VdpWLuYgSnT+jIGoSTEXgyFlfHdpSptl/93kvMSPozDKq6N5Esj1hzSUS5Hk8IylyYFLRYqa9p21ZOseVuztZt1kg57hGcxxG4nb8JYCHYYYlziX2Ket1/q265m/SVGrXb7/UpCrnT6/fb51i7TFFIJQR5QG6+UU3Bn1tZZUMgeBknaQgvcW+PYiPcHA5vnFBQBAAA=

crusaderky: @hendrikmakait all comments have been addressed

Patch:
diff --git a/distributed/shuffle/_core.py b/distributed/shuffle/_core.py
index 134f310d..f04079a8 100644
--- a/distributed/shuffle/_core.py
+++ b/distributed/shuffle/_core.py
@@ -24,6 +24,7 @@ from dask.utils import parse_timedelta
 
 from distributed.core import PooledRPCCall
 from distributed.exceptions import Reschedule
+from distributed.metrics import context_meter, thread_time
 from distributed.protocol import to_serialize
 from distributed.shuffle._comms import CommShardsBuffer
 from distributed.shuffle._disk import DiskShardsBuffer
@@ -287,7 +288,11 @@ 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}")
-        shards = self._shard_partition(data, partition_id)
+        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
 
@@ -305,7 +310,11 @@ class ShuffleRun(Generic[_T_partition_id, _T_partition_type]):
         if not self.transferred:
             raise RuntimeError("`get_output_partition` called before barrier task")
         sync(self._loop, self.flush_receive)
-        return self._get_output_partition(partition_id, key, **kwargs)
+        with (
+            context_meter.meter("p2p-get-output-noncpu"),
+            context_meter.meter("p2p-get-output-cpu", func=thread_time),
+        ):
+            return self._get_output_partition(partition_id, key, **kwargs)
 
     @abc.abstractmethod
     def _get_output_partition(
diff --git a/distributed/shuffle/_disk.py b/distributed/shuffle/_disk.py
index 038a6b98..6017949f 100644
--- a/distributed/shuffle/_disk.py
+++ b/distributed/shuffle/_disk.py
@@ -10,6 +10,7 @@ from typing import Any
 
 from toolz import concat
 
+from distributed.metrics import context_meter
 from distributed.shuffle._buffer import ShardsBuffer
 from distributed.shuffle._limiter import ResourceLimiter
 from distributed.shuffle._pickle import pickle_bytelist
@@ -179,11 +180,18 @@ class DiskShardsBuffer(ShardsBuffer):
             raise RuntimeError("Tried to read from file before done.")
 
         try:
-            with self.time("read"):
-                with self._directory_lock.read():
-                    if self._closed:
-                        raise RuntimeError("Already closed")
-                    data, size = self._read((self.directory / str(id)).resolve())
+            with self._directory_lock.read():
+                if self._closed:
+                    raise RuntimeError("Already closed")
+                fname = (self.directory / str(id)).resolve()
+                # Note: don't add `with context_meter.meter("p2p-disk-read"):` to
+                # 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)
+                context_meter.digest_metric("p2p-disk-read", 1, "count")
+                context_meter.digest_metric("p2p-disk-read", size, "bytes")
         except FileNotFoundError:
             raise KeyError(id)
 
diff --git a/distributed/shuffle/_limiter.py b/distributed/shuffle/_limiter.py
index 35b30267..a55682d1 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 time
+from distributed.metrics import context_meter, time
 
 _T = TypeVar("_T", int, None)
 
@@ -26,6 +26,7 @@ class ResourceLimiter(Generic[_T]):
     """
 
     limit: _T
+    metrics_label: str | None
     time_blocked_total: float
     time_blocked_avg: float
 
@@ -33,8 +34,9 @@ class ResourceLimiter(Generic[_T]):
     _condition: asyncio.Condition
     _waiters: int
 
-    def __init__(self, limit: _T):
+    def __init__(self, limit: _T, metrics_label: str | None = None):
         self.limit = limit
+        self.metrics_label = metrics_label
         self._acquired = 0
         self._condition = asyncio.Condition()
         self._waiters = 0
@@ -63,19 +65,24 @@ class ResourceLimiter(Generic[_T]):
 
     async def wait_for_available(self) -> None:
         """Block until the counter drops below limit"""
+        if self.limit is not None:
+            # Include non-blocking calls in the calculation of the average
+            # seconds / count
+            context_meter.digest_metric(self.metrics_label, 1, "count")
+        if not self.full:
+            self.time_blocked_avg *= 0.9
+            return
+
         start = time()
-        duration = 0.0
-        try:
-            if not self.full:
-                return
+        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
-        finally:
-            self.time_blocked_total += duration
-            self.time_blocked_avg = self.time_blocked_avg * 0.9 + duration * 0.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"""
diff --git a/distributed/shuffle/_rechunk.py b/distributed/shuffle/_rechunk.py
index 65b036a7..f621525c 100644
--- a/distributed/shuffle/_rechunk.py
+++ b/distributed/shuffle/_rechunk.py
@@ -114,6 +114,7 @@ from dask.highlevelgraph import HighLevelGraph, MaterializedLayer
 from dask.typing import Key
 
 from distributed.core import PooledRPCCall
+from distributed.metrics import context_meter
 from distributed.shuffle._core import (
     NDIndex,
     ShuffleId,
@@ -395,10 +396,13 @@ class ArrayRechunkRun(ShuffleRun[NDIndex, "np.ndarray"]):
 
     def _shard_partition(
         self, data: np.ndarray, partition_id: NDIndex
-    ) -> dict[str, tuple[NDIndex, Any]]:
+    ) -> dict[str, tuple[NDIndex, list[tuple[NDIndex, tuple[NDIndex, np.ndarray]]]]]:
         out: dict[str, list[tuple[NDIndex, tuple[NDIndex, np.ndarray]]]] = defaultdict(
             list
         )
+        shards_size = 0
+        shards_count = 0
+
         ndsplits = product(*(axis[i] for axis, i in zip(self.split_axes, partition_id)))
 
         for ndsplit in ndsplits:
@@ -410,9 +414,15 @@ class ArrayRechunkRun(ShuffleRun[NDIndex, "np.ndarray"]):
             if shard.base is not None:
                 shard = shard.copy()
 
+            shards_size += shard.nbytes
+            shards_count += 1
+
             out[self.worker_for[chunk_index]].append(
                 (chunk_index, (shard_index, shard))
             )
+
+        context_meter.digest_metric("p2p-shards", shards_size, "bytes")
+        context_meter.digest_metric("p2p-shards", shards_count, "count")
         return {k: (partition_id, v) for k, v in out.items()}
 
     def _get_output_partition(
diff --git a/distributed/shuffle/_shuffle.py b/distributed/shuffle/_shuffle.py
index 798697f2..7ea53253 100644
--- a/distributed/shuffle/_shuffle.py
+++ b/distributed/shuffle/_shuffle.py
@@ -27,6 +27,7 @@ from dask.typing import Key
 
 from distributed.core import PooledRPCCall
 from distributed.exceptions import Reschedule
+from distributed.metrics import context_meter
 from distributed.shuffle._arrow import (
     check_dtype_support,
     check_minimal_arrow_version,
@@ -498,6 +499,11 @@ class DataFrameShuffleRun(ShuffleRun[int, "pd.DataFrame"]):
             self.worker_for,
         )
         out = {k: (partition_id, serialize_table(t)) for k, t in out.items()}
+
+        nbytes = sum(len(b) for _, b in out.values())
+        context_meter.digest_metric("p2p-shards", nbytes, "bytes")
+        context_meter.digest_metric("p2p-shards", len(out), "count")
+
         return out
 
     def _get_output_partition(
diff --git a/distributed/shuffle/_worker_plugin.py b/distributed/shuffle/_worker_plugin.py
index 39b1c126..eddf00d0 100644
--- a/distributed/shuffle/_worker_plugin.py
+++ b/distributed/shuffle/_worker_plugin.py
@@ -276,8 +276,12 @@ class ShuffleWorkerPlugin(WorkerPlugin):
         # Initialize
         self.worker = worker
         self.shuffle_runs = _ShuffleRunManager(self)
-        self.memory_limiter_comms = ResourceLimiter(parse_bytes("100 MiB"))
-        self.memory_limiter_disk = ResourceLimiter(parse_bytes("1 GiB"))
+        self.memory_limiter_comms = ResourceLimiter(
+            parse_bytes("100 MiB"), metrics_label="p2p-comms-limiter"
+        )
+        self.memory_limiter_disk = ResourceLimiter(
+            parse_bytes("1 GiB"), metrics_label="p2p-disk-limiter"
+        )
         self.closed = False
         self._executor = ThreadPoolExecutor(self.worker.state.nthreads)
 


Test patch:
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
index e3a1e702..47a58ea2 100644
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -281,7 +281,7 @@ jobs:
         if: >
           always() &&
           (steps.run_tests.outcome == 'success' || steps.run_tests.outcome == 'failure')
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: ${{ env.TEST_ID }}
           path: reports
@@ -291,7 +291,7 @@ jobs:
         if: >
           always() &&
           (steps.run_tests.outcome == 'success' || steps.run_tests.outcome == 'failure')
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: ${{ env.TEST_ID }}_cluster_dumps
           path: test_cluster_dump
@@ -305,7 +305,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Upload
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: Event File
           path: ${{ github.event_path }}
diff --git a/distributed/shuffle/tests/test_metrics.py b/distributed/shuffle/tests/test_metrics.py
new file mode 100644
index 00000000..905e8f13
--- /dev/null
+++ b/distributed/shuffle/tests/test_metrics.py
@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import pytest
+
+import dask.datasets
+
+from distributed import Scheduler
+from distributed.utils_test import gen_cluster
+
+da = pytest.importorskip("dask.array")
+
+
+def assert_metrics(s: Scheduler, *keys: tuple[str, ...]) -> None:
+    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
+
+
+@gen_cluster(client=True, config={"optimization.fuse.active": False})
+async def test_rechunk(c, s, a, b):
+    x = da.random.random((10, 10), chunks=(-1, 1))
+    x = x.rechunk((1, -1), method="p2p")
+    await c.compute(x)
+    await a.heartbeat()
+    await b.heartbeat()
+
+    assert_metrics(
+        s,
+        ("execute", "rechunk-transfer", "p2p-shard-partition-cpu", "seconds"),
+        ("execute", "rechunk-transfer", "p2p-shard-partition-noncpu", "seconds"),
+        ("execute", "rechunk-transfer", "p2p-shards", "bytes"),
+        ("execute", "rechunk-transfer", "p2p-shards", "count"),
+        ("execute", "rechunk-transfer", "p2p-comms-limiter", "count"),
+        ("execute", "rechunk", "p2p-disk-read", "bytes"),
+        ("execute", "rechunk", "p2p-disk-read", "count"),
+        ("execute", "rechunk", "p2p-get-output-cpu", "seconds"),
+        ("execute", "rechunk", "p2p-get-output-noncpu", "seconds"),
+    )
+
+
+@gen_cluster(client=True)
+async def test_dataframe(c, s, a, b):
+    """Metrics are *almost* agnostic in dataframe shuffle vs. array rechunk.
+    The only exception is the 'p2p-shards' metric, which is implemented separately.
+    """
+    dd = pytest.importorskip("dask.dataframe")
+
+    df = dask.datasets.timeseries(
+        start="2000-01-01",
+        end="2000-01-10",
+        dtypes={"x": float, "y": float},
+        freq="10 s",
+    )
+    shuffled = dd.shuffle.shuffle(df, "x", shuffle="p2p", npartitions=20)
+    await c.compute(shuffled)
+    await a.heartbeat()
+    await b.heartbeat()
+
+    assert_metrics(
+        s,
+        ("execute", "shuffle-transfer", "p2p-shard-partition-cpu", "seconds"),
+        ("execute", "shuffle-transfer", "p2p-shard-partition-noncpu", "seconds"),
+        ("execute", "shuffle-transfer", "p2p-shards", "bytes"),
+        ("execute", "shuffle-transfer", "p2p-shards", "count"),
+        ("execute", "shuffle-transfer", "p2p-comms-limiter", "count"),
+        ("execute", "shuffle_p2p", "p2p-disk-read", "bytes"),
+        ("execute", "shuffle_p2p", "p2p-disk-read", "count"),
+        ("execute", "shuffle_p2p", "p2p-get-output-cpu", "seconds"),
+        ("execute", "shuffle_p2p", "p2p-get-output-noncpu", "seconds"),
+    )