Repository: dask/distributed

Issue:
Instrumentation of SpillBuffer
The only way we currently have to observe disk access is `startstops` we measure whenever we load/store data.

However, with our `SpillBuffer` we have the possibility to introduce many instrumentation hooks to get much better insights into what's going on

For instance
- Total number of keys written / read
- Number of currently managed files
- Average and total size of data written to disk per key
- Total duration spent writing / reading
- Average / max time spent reading / writing per key
- Number of evicts (memory.spill)
- Number of writes without evict because buffer is full right away
- ...

These metrics should not be tracked on `TaskState` level.


Discussion:
> Number of writes without evict because buffer is full right away

Not sure if I understand this one. The only case where a key doesn't end up in fast when you write to the SpillBuffer is when it's individually larger than the target threshold.
> Not sure if I understand this one. The only case where a key doesn't end up in fast when you write to the SpillBuffer is when it's individually larger than the target threshold.

I thought there was some logic that would put a key into slow if this key would push us over the limit. If that's not the case, ignore this, I don't think data shards larger than the limit are a very common problem we need to build instrumentation for
> I thought there was some logic that would put a key into slow if this key would push us over the limit.

No, if inserting a key pushes us over the limit, the least recently used keys are pushed out. The latest inserted one is on top of the LRU pipe and is the only one guaranteed to be in fast at the end of the insertion.
The thing I was hoping to differentiate with this comment is the spilling that happens as a result of setitem vs the spilling that happens when the memory_manager is evicting.
If it's not the same key, that's fine but if there is an evict coupled to the setitem, that might be an interesting datapoint

Resolving patch:
diff --git a/distributed/http/worker/prometheus/core.py b/distributed/http/worker/prometheus/core.py
--- a/distributed/http/worker/prometheus/core.py
+++ b/distributed/http/worker/prometheus/core.py
@@ -1,11 +1,12 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Iterator
 from time import time
 from typing import ClassVar
 
 import prometheus_client
-from prometheus_client.core import CounterMetricFamily, GaugeMetricFamily
+from prometheus_client.core import CounterMetricFamily, GaugeMetricFamily, Metric
 
 from distributed.http.prometheus import PrometheusCollector
 from distributed.http.utils import RequestHandler
@@ -30,8 +31,7 @@ def __init__(self, server: Worker):
                 "Digest-based metrics require crick to be installed."
             )
 
-    def collect(self):
-
+    def collect(self) -> Iterator[Metric]:
         ws = self.server.state
 
         tasks = GaugeMetricFamily(
@@ -40,7 +40,12 @@ def collect(self):
             labels=["state"],
         )
         for k, n in ws.task_counts.items():
-            tasks.add_metric([k], n)
+            if k == "memory" and hasattr(self.server.data, "slow"):
+                n_spilled = len(self.server.data.slow)
+                tasks.add_metric(["memory"], n - n_spilled)
+                tasks.add_metric(["disk"], n_spilled)
+            else:
+                tasks.add_metric([k], n)
         yield tasks
 
         yield GaugeMetricFamily(
@@ -59,13 +64,14 @@ def collect(self):
         )
 
         yield GaugeMetricFamily(
-            self.build_name("latency_seconds"),
+            self.build_name("latency"),
             "Latency of worker connection",
+            unit="seconds",
             value=self.server.latency,
         )
 
         try:
-            spilled_memory, spilled_disk = self.server.data.spilled_total
+            spilled_memory, spilled_disk = self.server.data.spilled_total  # type: ignore
         except AttributeError:
             spilled_memory, spilled_disk = 0, 0  # spilling is disabled
         process_memory = self.server.monitor.get_process_memory()
@@ -130,28 +136,8 @@ def collect(self):
             value=self.server.transfer_outgoing_count_total,
         )
 
-        # all metrics using digests require crick to be installed
-        # the following metrics will export NaN, if the corresponding digests are None
-        if self.crick_available:
-            yield GaugeMetricFamily(
-                self.build_name("tick_duration_median_seconds"),
-                "Median tick duration at worker",
-                value=self.server.digests["tick-duration"].components[1].quantile(50),
-            )
-
-            yield GaugeMetricFamily(
-                self.build_name("task_duration_median_seconds"),
-                "Median task runtime at worker",
-                value=self.server.digests["task-duration"].components[1].quantile(50),
-            )
-
-            yield GaugeMetricFamily(
-                self.build_name("transfer_bandwidth_median_bytes"),
-                "Bandwidth for transfer at worker",
-                value=self.server.digests["transfer-bandwidth"]
-                .components[1]
-                .quantile(50),
-            )
+        yield from self.collect_crick()
+        yield from self.collect_spillbuffer()
 
         now = time()
         max_tick_duration = max(
@@ -159,13 +145,14 @@ def collect(self):
             now - self.server._last_tick,
         )
         yield GaugeMetricFamily(
-            self.build_name("tick_duration_maximum_seconds"),
+            self.build_name("tick_duration_maximum"),
             "Maximum tick duration observed since Prometheus last scraped metrics",
+            unit="seconds",
             value=max_tick_duration,
         )
 
         yield CounterMetricFamily(
-            self.build_name("tick_count_total"),
+            self.build_name("tick_count"),
             "Total number of ticks observed since the server started",
             value=self.server._tick_counter,
         )
@@ -201,9 +188,96 @@ def collect(self):
 
         yield evloop_blocked_total
         yield evloop_blocked_max
-
         self.server.digests_max.clear()
 
+    def collect_crick(self) -> Iterator[Metric]:
+        # All metrics using digests require crick to be installed.
+        # The following metrics will export NaN, if the corresponding digests are None
+        if not self.crick_available:
+            return
+
+        yield GaugeMetricFamily(
+            self.build_name("tick_duration_median"),
+            "Median tick duration at worker",
+            unit="seconds",
+            value=self.server.digests["tick-duration"].components[1].quantile(50),
+        )
+
+        yield GaugeMetricFamily(
+            self.build_name("task_duration_median"),
+            "Median task runtime at worker",
+            unit="seconds",
+            value=self.server.digests["task-duration"].components[1].quantile(50),
+        )
+
+        yield GaugeMetricFamily(
+            self.build_name("transfer_bandwidth_median"),
+            "Bandwidth for transfer at worker",
+            unit="bytes",
+            value=self.server.digests["transfer-bandwidth"].components[1].quantile(50),
+        )
+
+    def collect_spillbuffer(self) -> Iterator[Metric]:
+        """SpillBuffer-specific metrics.
+
+        Additionally, you can obtain derived metrics as follows:
+
+        cache hit ratios:
+          by keys  = spill_count.memory_read / (spill_count.memory_read + spill_count.disk_read)
+          by bytes = spill_bytes.memory_read / (spill_bytes.memory_read + spill_bytes.disk_read)
+
+        mean times per key:
+          pickle   = spill_time.pickle     / spill_count.disk_write
+          write    = spill_time.disk_write / spill_count.disk_write
+          unpickle = spill_time.unpickle   / spill_count.disk_read
+          read     = spill_time.disk_read  / spill_count.disk_read
+
+        mean bytes per key:
+          write    = spill_bytes.disk_write / spill_count.disk_write
+          read     = spill_bytes.disk_read  / spill_count.disk_read
+
+        mean bytes per second:
+          write    = spill_bytes.disk_write / spill_time.disk_write
+          read     = spill_bytes.disk_read  / spill_time.disk_read
+        """
+        try:
+            get_metrics = self.server.data.get_metrics  # type: ignore
+        except AttributeError:
+            return  # spilling is disabled
+        metrics = get_metrics()
+
+        total_bytes = CounterMetricFamily(
+            self.build_name("spill_bytes"),
+            "Total size of memory and disk accesses caused by managed data "
+            "since the latest worker restart",
+            labels=["activity"],
+        )
+        # Note: memory_read is used to calculate cache hit ratios (see docstring)
+        for k in ("memory_read", "disk_read", "disk_write"):
+            total_bytes.add_metric([k], metrics[f"{k}_bytes_total"])
+        yield total_bytes
+
+        total_counts = CounterMetricFamily(
+            self.build_name("spill_count"),
+            "Total number of memory and disk accesses caused by managed data "
+            "since the latest worker restart",
+            labels=["activity"],
+        )
+        # Note: memory_read is used to calculate cache hit ratios (see docstring)
+        for k in ("memory_read", "disk_read", "disk_write"):
+            total_counts.add_metric([k], metrics[f"{k}_count_total"])
+        yield total_counts
+
+        total_times = CounterMetricFamily(
+            self.build_name("spill_time"),
+            "Total time spent spilling/unspilling since the latest worker restart",
+            unit="seconds",
+            labels=["activity"],
+        )
+        for k in ("pickle", "disk_write", "disk_read", "unpickle"):
+            total_times.add_metric([k], metrics[f"{k}_time_total"])
+        yield total_times
+
 
 class PrometheusHandler(RequestHandler):
     _collector: ClassVar[WorkerMetricCollector | None] = None
diff --git a/distributed/spill.py b/distributed/spill.py
--- a/distributed/spill.py
+++ b/distributed/spill.py
@@ -1,10 +1,11 @@
 from __future__ import annotations
 
 import logging
-import time
 from collections.abc import Iterator, Mapping, MutableMapping, Sized
 from contextlib import contextmanager
+from dataclasses import dataclass
 from functools import partial
+from time import perf_counter
 from typing import Any, Literal, NamedTuple, Protocol, cast
 
 from packaging.version import parse as parse_version
@@ -64,6 +65,49 @@ def evict(self) -> int:
         ...  # pragma: nocover
 
 
+@dataclass
+class FastMetrics:
+    """Cumulative metrics for SpillBuffer.fast since the latest worker restart"""
+
+    read_count_total: int = 0
+    read_bytes_total: int = 0
+
+    def log_read(self, key_bytes: int) -> None:
+        self.read_count_total += 1
+        self.read_bytes_total += key_bytes
+
+
+@dataclass
+class SlowMetrics:
+    """Cumulative metrics for SpillBuffer.slow since the latest worker restart"""
+
+    read_count_total: int = 0
+    read_bytes_total: int = 0
+    read_time_total: float = 0
+    write_count_total: int = 0
+    write_bytes_total: int = 0
+    write_time_total: float = 0
+    pickle_time_total: float = 0
+    unpickle_time_total: float = 0
+
+    def log_write(
+        self,
+        key_bytes: int,
+        pickle_time: float,
+        write_time: float,
+    ) -> None:
+        self.write_count_total += 1
+        self.write_bytes_total += key_bytes
+        self.pickle_time_total += pickle_time
+        self.write_time_total += write_time
+
+    def log_read(self, key_bytes: int, read_time: float, unpickle_time: float) -> None:
+        self.read_count_total += 1
+        self.read_bytes_total += key_bytes
+        self.read_time_total += read_time
+        self.unpickle_time_total += unpickle_time
+
+
 # zict.Buffer[str, Any] requires zict >= 2.2.0
 class SpillBuffer(zict.Buffer):
     """MutableMapping that automatically spills out dask key/value pairs to disk when
@@ -85,6 +129,7 @@ class SpillBuffer(zict.Buffer):
     last_logged: float
     min_log_interval: float
     logged_pickle_errors: set[str]
+    fast_metrics: FastMetrics
 
     def __init__(
         self,
@@ -108,6 +153,8 @@ def __init__(
         self.min_log_interval = min_log_interval
         self.logged_pickle_errors = set()  # keys logged with pickle error
 
+        self.fast_metrics = FastMetrics()
+
     @contextmanager
     def handle_errors(self, key: str | None) -> Iterator[None]:
         try:
@@ -118,7 +165,7 @@ def handle_errors(self, key: str | None) -> Iterator[None]:
             (key_e,) = e.args
             assert key_e in self.fast
             assert key_e not in self.slow
-            now = time.time()
+            now = perf_counter()
             if now - self.last_logged >= self.min_log_interval:
                 logger.warning(
                     "Spill file on disk reached capacity; keeping data in memory"
@@ -127,7 +174,7 @@ def handle_errors(self, key: str | None) -> Iterator[None]:
             raise HandledError()
         except OSError:
             # Typically, this is a disk full error
-            now = time.time()
+            now = perf_counter()
             if now - self.last_logged >= self.min_log_interval:
                 logger.error(
                     "Spill to disk failed; keeping data in memory", exc_info=True
@@ -212,6 +259,15 @@ def evict(self) -> int:
         except HandledError:
             return -1
 
+    def __getitem__(self, key: str) -> Any:
+        if key in self.fast:
+            # Note: don't log from self.fast.__getitem__, because that's called every
+            # time a key is evicted, and we don't want to count those events here.
+            nbytes = cast(int, self.fast.weights[key])
+            self.fast_metrics.log_read(nbytes)
+
+        return super().__getitem__(key)
+
     def __delitem__(self, key: str) -> None:
         super().__delitem__(key)
         self.logged_pickle_errors.discard(key)
@@ -230,6 +286,11 @@ def disk(self) -> Mapping[str, Any]:
         """
         return self.slow
 
+    @property
+    def _slow_uncached(self) -> Slow:
+        slow = cast(zict.Cache, self.slow).data if has_zict_220 else self.slow
+        return cast(Slow, slow)
+
     @property
     def spilled_total(self) -> SpilledSize:
         """Number of bytes spilled to disk. Tuple of
@@ -240,8 +301,46 @@ def spilled_total(self) -> SpilledSize:
         The two may differ substantially, e.g. if sizeof() is inaccurate or in case of
         compression.
         """
-        slow = cast(zict.Cache, self.slow).data if has_zict_220 else self.slow
-        return cast(Slow, slow).total_weight
+        return self._slow_uncached.total_weight
+
+    def get_metrics(self) -> dict[str, float]:
+        """Metrics to be exported to Prometheus or to be parsed directly.
+
+        From these you may generate derived metrics:
+
+        cache hit ratio:
+          by keys  = memory_read_count_total / (memory_read_count_total + disk_read_count_total)
+          by bytes = memory_read_bytes_total / (memory_read_bytes_total + disk_read_bytes_total)
+
+        mean times per key:
+          pickle   = pickle_time_total     / disk_write_count_total
+          write    = disk_write_time_total / disk_write_count_total
+          unpickle = unpickle_time_total   / disk_read_count_total
+          read     = disk_read_time_total  / disk_read_count_total
+
+        mean bytes per key:
+          write    = disk_write_bytes_total / disk_write_count_total
+          read     = disk_read_bytes_total  / disk_read_count_total
+
+        mean bytes per second:
+          write    = disk_write_bytes_total / disk_write_time_total
+          read     = disk_read_bytes_total  / disk_read_time_total
+        """
+        fm = self.fast_metrics
+        sm = self._slow_uncached.metrics
+
+        out = {
+            "memory_count": len(self.fast),
+            "memory_bytes": self.fast.total_weight,
+            "disk_count": len(self.slow),
+            "disk_bytes": self._slow_uncached.total_weight.disk,
+        }
+
+        for k, v in fm.__dict__.items():
+            out[f"memory_{k}"] = v
+        for k, v in sm.__dict__.items():
+            out[k if "pickle" in k else f"disk_{k}"] = v
+        return out
 
 
 def _in_memory_weight(key: str, value: Any) -> int:
@@ -266,6 +365,7 @@ class Slow(zict.Func):
     max_weight: int | Literal[False]
     weight_by_key: dict[str, SpilledSize]
     total_weight: SpilledSize
+    metrics: SlowMetrics
 
     def __init__(self, spill_directory: str, max_weight: int | Literal[False] = False):
         super().__init__(
@@ -276,8 +376,27 @@ def __init__(self, spill_directory: str, max_weight: int | Literal[False] = Fals
         self.max_weight = max_weight
         self.weight_by_key = {}
         self.total_weight = SpilledSize(0, 0)
+        self.metrics = SlowMetrics()
+
+    def __getitem__(self, key: str) -> Any:
+        t0 = perf_counter()
+        pickled = self.d[key]
+        assert isinstance(pickled, bytes)
+        t1 = perf_counter()
+        out = self.load(pickled)  # type: ignore
+        t2 = perf_counter()
+
+        # For the sake of simplicity, we're not metering failure use cases.
+        self.metrics.log_read(
+            key_bytes=len(pickled),
+            read_time=t1 - t0,
+            unpickle_time=t2 - t1,
+        )
+
+        return out
 
     def __setitem__(self, key: str, value: Any) -> None:
+        t0 = perf_counter()
         try:
             # FIXME https://github.com/python/mypy/issues/708
             pickled = self.dump(value)  # type: ignore
@@ -291,6 +410,7 @@ def __setitem__(self, key: str, value: Any) -> None:
             frame.nbytes if isinstance(frame, memoryview) else len(frame)
             for frame in pickled
         )
+        t1 = perf_counter()
 
         if has_zict_210:
             # Thanks to Buffer.__setitem__, we never update existing
@@ -313,10 +433,19 @@ def __setitem__(self, key: str, value: Any) -> None:
         # This may raise OSError, which is caught by SpillBuffer above.
         self.d[key] = pickled
 
+        t2 = perf_counter()
+
         weight = SpilledSize(safe_sizeof(value), pickled_size)
         self.weight_by_key[key] = weight
         self.total_weight += weight
 
+        # For the sake of simplicity, we're not metering failure use cases.
+        self.metrics.log_write(
+            key_bytes=pickled_size,
+            pickle_time=t1 - t0,
+            write_time=t2 - t1,
+        )
+
     def __delitem__(self, key: str) -> None:
         super().__delitem__(key)
         self.total_weight -= self.weight_by_key.pop(key)


Test patch:
diff --git a/distributed/http/worker/tests/test_worker_http.py b/distributed/http/worker/tests/test_worker_http.py
--- a/distributed/http/worker/tests/test_worker_http.py
+++ b/distributed/http/worker/tests/test_worker_http.py
@@ -30,10 +30,13 @@ async def test_prometheus(c, s, a):
         "dask_worker_event_loop_blocked_time_seconds_total",
         "dask_worker_latency_seconds",
         "dask_worker_memory_bytes",
+        "dask_worker_spill_bytes_total",
+        "dask_worker_spill_count_total",
+        "dask_worker_spill_time_seconds_total",
         "dask_worker_tasks",
+        "dask_worker_threads",
         "dask_worker_tick_count_total",
         "dask_worker_tick_duration_maximum_seconds",
-        "dask_worker_threads",
         "dask_worker_transfer_incoming_bytes",
         "dask_worker_transfer_incoming_count",
         "dask_worker_transfer_incoming_count_total",
@@ -85,29 +88,32 @@ async def test_metrics_when_prometheus_client_not_installed(
 async def test_prometheus_collect_task_states(c, s, a):
     pytest.importorskip("prometheus_client")
 
-    async def fetch_state_metrics():
+    async def assert_metrics(**kwargs):
+        expect = {
+            "constrained": 0,
+            "executing": 0,
+            "fetch": 0,
+            "flight": 0,
+            "long-running": 0,
+            "memory": 0,
+            "disk": 0,
+            "missing": 0,
+            "other": 0,
+            "ready": 0,
+            "waiting": 0,
+        }
+        expect.update(kwargs)
+
         families = await fetch_metrics(a.http_server.port, prefix="dask_worker_")
-        active_metrics = {
+        actual = {
             sample.labels["state"]: sample.value
             for sample in families["dask_worker_tasks"].samples
         }
-        return active_metrics
 
-    assert not a.state.tasks
-    active_metrics = await fetch_state_metrics()
-    assert active_metrics == {
-        "constrained": 0.0,
-        "executing": 0.0,
-        "fetch": 0.0,
-        "flight": 0.0,
-        "long-running": 0.0,
-        "memory": 0.0,
-        "missing": 0.0,
-        "other": 0.0,
-        "ready": 0.0,
-        "waiting": 0.0,
-    }
+        assert actual == expect
 
+    assert not a.state.tasks
+    await assert_metrics()
     ev = Event()
 
     # submit a task which should show up in the prometheus scraping
@@ -115,41 +121,21 @@ async def fetch_state_metrics():
     while not a.state.executing:
         await asyncio.sleep(0.001)
 
-    active_metrics = await fetch_state_metrics()
-    assert active_metrics == {
-        "constrained": 0.0,
-        "executing": 1.0,
-        "fetch": 0.0,
-        "flight": 0.0,
-        "long-running": 0.0,
-        "memory": 0.0,
-        "missing": 0.0,
-        "other": 0.0,
-        "ready": 0.0,
-        "waiting": 0.0,
-    }
+    await assert_metrics(executing=1)
 
     await ev.set()
     await c.gather(future)
 
+    await assert_metrics(memory=1)
+    a.data.evict()
+    await assert_metrics(disk=1)
+
     future.release()
 
     while future.key in a.state.tasks:
         await asyncio.sleep(0.001)
 
-    active_metrics = await fetch_state_metrics()
-    assert active_metrics == {
-        "constrained": 0.0,
-        "executing": 0.0,
-        "fetch": 0.0,
-        "flight": 0.0,
-        "long-running": 0.0,
-        "memory": 0.0,
-        "missing": 0.0,
-        "other": 0.0,
-        "ready": 0.0,
-        "waiting": 0.0,
-    }
+    await assert_metrics()
 
 
 @gen_cluster(client=True)
diff --git a/distributed/tests/test_spill.py b/distributed/tests/test_spill.py
--- a/distributed/tests/test_spill.py
+++ b/distributed/tests/test_spill.py
@@ -384,3 +384,52 @@ def test_weakref_cache(tmp_path, cls, expect_cached, size):
 
     # Test that we update the weakref cache on getitem
     assert (buf["x"] is x2) == expect_cached
+
+
+def test_metrics(tmp_path):
+    buf = SpillBuffer(str(tmp_path), target=1000)
+    assert buf.get_metrics() == {
+        "disk_bytes": 0,
+        "disk_count": 0,
+        "disk_read_bytes_total": 0,
+        "disk_read_count_total": 0,
+        "disk_read_time_total": 0,
+        "disk_write_bytes_total": 0,
+        "disk_write_count_total": 0,
+        "disk_write_time_total": 0,
+        "memory_bytes": 0,
+        "memory_count": 0,
+        "memory_read_bytes_total": 0,
+        "memory_read_count_total": 0,
+        "pickle_time_total": 0,
+        "unpickle_time_total": 0,
+    }
+
+    a, b, c = "a" * 200, "b" * 150, "c" * 100
+
+    buf["a"] = a
+    buf["b"] = b
+    del buf["b"]
+
+    metrics = buf.get_metrics()
+    assert 200 < metrics["memory_bytes"] < 350
+    assert metrics["memory_count"] == 1
+    for k, v in metrics.items():
+        if "disk" in k or "pickle" in k:
+            assert v == 0, f"{k}={v}"
+
+    assert metrics["memory_read_bytes_total"] == 0
+    assert metrics["memory_read_count_total"] == 0
+    _ = buf["a"]  # Cache hit
+    metrics = buf.get_metrics()
+    assert metrics["memory_read_bytes_total"] > 0
+    assert metrics["memory_read_count_total"] == 1
+
+    buf.evict()
+    buf["c"] = c
+    buf.evict()
+    _ = buf["c"]  # Unspill / cache miss
+
+    metrics = buf.get_metrics()
+    for k, v in metrics.items():
+        assert v > 0, f"{k}={v}"