Repository: dask/distributed

Issue:
Fine performance metrics: apportion to Computations
- Part of https://github.com/dask/distributed/issues/7665

Computation objects are commonly used by third-party Scheduler plugins (e.g. Coiled Analytics) to visualize data.
They can be crudely defined as everything that happened to the cluster between two moments when it was idle.
As there should be only one active at any given time, it makes perfect sense to apportion the metrics to the Computation objects, in addition to the Scheduler's global ever-increasing counter.

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

       20 files  ±  0         20 suites  ±0   12h 1m 33s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "duration of all tests") - 20m 23s
  3 672 tests +  6    3 555 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "passed tests") ±  0     108 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "skipped / disabled tests") ±0    9 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "failed tests") +  6 
35 510 runs  +60  33 722 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "passed tests") +39  1 766 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "skipped / disabled tests") +2  22 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v2.7.0/README.md#the-symbols "failed tests") +19 

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

Results for commit 4092bd5b. ± Comparison against base commit 129b7cb7.

<details>
  <summary>This pull request <b>removes</b> 1 and <b>adds</b> 7 tests. <i>Note that renamed tests count towards both.</i></summary>

```
distributed.tests.test_scheduler ‑ test_cumulative_worker_metrics
```

```
distributed.tests.test_spans ‑ test_default_span
distributed.tests.test_spans ‑ test_no_tags
distributed.tests.test_spans ‑ test_repeat_span
distributed.tests.test_spans ‑ test_worker_metrics
distributed.tests.test_worker_metrics ‑ test_no_spans_extension
distributed.tests.test_worker_metrics ‑ test_reschedule
distributed.tests.test_worker_metrics ‑ test_send_metrics_to_scheduler
```
</details>

[test-results]:data:application/gzip;base64,H4sIANiJf2QC/03Nyw6DIBCF4VcxrLsYBmGkL9OgYkLqpeGyavruHVPFLv9vyOEtpjD7JO4Nwq0RqYR81Viiy2FbuVuFVjHxNe93ZQjPfKQyDLtprS97hhebhK7S5MLMZCv4GLfIsv8Vy5p+G/LMuqsIsdqxS8acdOxeb/53h21ZQuYQLVjsR913aMh4IgfKTjSS0R6k1z1K8EBSic8X76pDwBUBAAA=

crusaderky: A test is failing; investigating
```
  File "/Users/runner/work/distributed/distributed/distributed/spans.py", line 423, in heartbeat
    _, span_id, prefix, activity, unit = k
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: too many values to unpack (expected 5)
```
crusaderky: This is ready for final review. @milesgranger could you re-review distributed/dashboard/components/scheduler.py? I just did major changes to it.

Patch:
diff --git a/distributed/dashboard/components/scheduler.py b/distributed/dashboard/components/scheduler.py
index 23f863b6..238f1acf 100644
--- a/distributed/dashboard/components/scheduler.py
+++ b/distributed/dashboard/components/scheduler.py
@@ -7,6 +7,7 @@ import os
 from collections import OrderedDict, defaultdict
 from collections.abc import Iterable
 from datetime import datetime, timedelta
+from itertools import chain
 from numbers import Number
 from typing import Any, TypeVar
 
@@ -3402,7 +3403,7 @@ class FinePerformanceMetrics(DashboardComponent):
 
         self.function_selector = MultiChoice(value=[], options=[])
         self.function_selector.placeholder = "Select specific functions"
-        self.unit_selector = Select(title="Unit selection", options=[])
+        self.unit_selector = Select(title="Unit selection", options=["seconds"])
         self.unit_selector.on_change("value", handle_selector_chng)
         self.unit_selected = "seconds"
         self.task_exec_by_activity_chart = None
@@ -3418,76 +3419,97 @@ class FinePerformanceMetrics(DashboardComponent):
         formatters = {"bytes": format_bytes, "seconds": format_time}
         return formatters.get(unit, str)(val)
 
+    def get_metrics(
+        self,
+    ) -> tuple[list[tuple[str, str, str, float]], list[tuple[str, str, float]]]:
+        """Pre-process fine performance metrics
+
+        Returns
+        -------
+        - 'execute' metrics: [(function, activity, unit, value), ...]
+        - 'get_data' metrics: [(activity, unit, value), ...]
+        """
+        execute: defaultdict[tuple[str, str, str], float] = defaultdict(float)
+        get_data = []
+
+        for k, v in self.scheduler.cumulative_worker_metrics.items():
+            if not isinstance(k, tuple):
+                continue
+            context, *other, activity, unit = k
+            assert isinstance(unit, str)
+
+            if context == "execute":
+                span_id, function = other
+                assert isinstance(function, str)
+                # Custom metrics can provide any hashable as the label
+                # Squash all span_ids together
+                # TODO offer a filter by span_id, like we already do by function
+                execute[function, str(activity), unit] += v
+            elif context == "get-data":
+                assert not other
+                assert isinstance(activity, str)
+                get_data.append((activity, unit, v))
+
+            # Ignore memory-monitor and gather-dep metrics
+
+        return (
+            [
+                (function, activity, unit, v)
+                for (function, activity, unit), v in sorted(execute.items())
+            ],
+            sorted(get_data),
+        )
+
     @without_property_validation
     @log_errors
     def update(self):
-        items = sorted(
-            # Custom metrics can be any hashable
-            (tuple(map(str, k)), v)
-            for k, v in self.scheduler.cumulative_worker_metrics.items()
-            if isinstance(k, tuple)
-        )
+        execute_metrics, get_data_metrics = self.get_metrics()
 
         activities = {
-            activity
-            for (*_, activity, _), _ in items
-            if activity not in self.task_activities
+            activity for *_, activity, _, _ in chain(execute_metrics, get_data_metrics)
         }
+        activities.difference_update(self.task_activities)
         if activities:
             self.substantial_change = True
             self.task_activities.extend(activities)
 
-        functions = {
-            func
-            for (ctx, func, *_), _ in items
-            if ctx == "execute" and func not in self.task_exec_data["functions"]
-        }
+        units = {unit for *_, unit, _ in chain(execute_metrics, get_data_metrics)}
+        units.difference_update(self.unit_selector.options)
+        self.unit_selector.options.extend(units)
+
+        functions = {function for function, *_ in execute_metrics}
+        functions.difference_update(self.task_exec_data["functions"])
         if functions:
             self.substantial_change = True
             self.task_exec_data["timestamp"].extend(datetime.now() for _ in functions)
             self.function_selector.options.extend(functions)
             self.task_exec_data["functions"].extend(functions)
 
-        for (context, *parts), val in items:
-            if context == "get-data":
-                activity, unit = parts
-
-                if unit not in self.unit_selector.options:
-                    # note append doesn't work here
-                    self.unit_selector.options += [unit]
-
-                if activity not in self.senddata["activity"]:
-                    self.substantial_change = True
-                    self.senddata["activity"].append(activity)
-
-                idx = self.senddata["activity"].index(activity)
-                while idx >= len(self.senddata[f"{activity}_{unit}"]):
-                    self.senddata[f"{activity}_{unit}"].append(0)
-                    self.senddata[f"{activity}_{unit}_text"].append("")
-                self.senddata[f"{activity}_{unit}_text"][idx] = self.format(unit, val)
-                self.senddata[f"{activity}_{unit}"][idx] = val
-
-            elif context == "execute":
-                prefix, activity, unit = parts
-
-                if unit not in self.unit_selector.options:
-                    # note append doesn't work here
-                    self.unit_selector.options += [unit]
-
-                idx = self.task_exec_data["functions"].index(prefix)
-
-                # Some function/activity combos missing, so need to keep columns aligned
-                for op in self.task_activities:
-                    while len(self.task_exec_data[f"{op}_{unit}"]) != len(
-                        self.task_exec_data["functions"]
-                    ):
-                        self.task_exec_data[f"{op}_{unit}"].append(0)
-                        self.task_exec_data[f"{op}_{unit}_text"].append("")
-
-                self.task_exec_data[f"{activity}_{unit}"][idx] = val
-                self.task_exec_data[f"{activity}_{unit}_text"][idx] = self.format(
-                    unit, val
-                )
+        for activity, unit, val in get_data_metrics:
+            if activity not in self.senddata["activity"]:
+                self.substantial_change = True
+                self.senddata["activity"].append(activity)
+
+            idx = self.senddata["activity"].index(activity)
+            while idx >= len(self.senddata[f"{activity}_{unit}"]):
+                self.senddata[f"{activity}_{unit}"].append(0)
+                self.senddata[f"{activity}_{unit}_text"].append("")
+            self.senddata[f"{activity}_{unit}_text"][idx] = self.format(unit, val)
+            self.senddata[f"{activity}_{unit}"][idx] = val
+
+        for prefix, activity, unit, val in execute_metrics:
+            idx = self.task_exec_data["functions"].index(prefix)
+
+            # Some function/activity combos missing, so need to keep columns aligned
+            for op in self.task_activities:
+                while len(self.task_exec_data[f"{op}_{unit}"]) != len(
+                    self.task_exec_data["functions"]
+                ):
+                    self.task_exec_data[f"{op}_{unit}"].append(0)
+                    self.task_exec_data[f"{op}_{unit}_text"].append("")
+
+            self.task_exec_data[f"{activity}_{unit}"][idx] = val
+            self.task_exec_data[f"{activity}_{unit}_text"][idx] = self.format(unit, val)
 
         data = self.task_exec_data.copy()
         # If user has manually selected function(s) then we are only showing them.
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 79da9b89..3cad2f1c 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -99,7 +99,7 @@ from distributed.recreate_tasks import ReplayTaskScheduler
 from distributed.security import Security
 from distributed.semaphore import SemaphoreExtension
 from distributed.shuffle import ShuffleSchedulerExtension
-from distributed.spans import SpansExtension
+from distributed.spans import SpansSchedulerExtension
 from distributed.stealing import WorkStealing
 from distributed.utils import (
     All,
@@ -171,7 +171,7 @@ DEFAULT_EXTENSIONS = {
     "amm": ActiveMemoryManagerExtension,
     "memory_sampler": MemorySamplerExtension,
     "shuffle": ShuffleSchedulerExtension,
-    "spans": SpansExtension,
+    "spans": SpansSchedulerExtension,
     "stealing": WorkStealing,
 }
 
@@ -4448,7 +4448,7 @@ class Scheduler(SchedulerState, ServerNode):
                     recommendations[ts.key] = "erred"
                     break
 
-        spans_ext: SpansExtension | None = self.extensions.get("spans")
+        spans_ext: SpansSchedulerExtension | None = self.extensions.get("spans")
         if spans_ext:
             # new_tasks does not necessarily contain all runnable tasks;
             # _generate_taskstates is not the only thing that calls new_task(). A
diff --git a/distributed/spans.py b/distributed/spans.py
index 5f0524ec..a6b8d8cc 100644
--- a/distributed/spans.py
+++ b/distributed/spans.py
@@ -12,8 +12,8 @@ import dask.config
 from distributed.metrics import time
 
 if TYPE_CHECKING:
-    from distributed import Scheduler
-    from distributed.scheduler import TaskGroup, TaskState, TaskStateState
+    from distributed import Scheduler, Worker
+    from distributed.scheduler import TaskGroup, TaskState, TaskStateState, WorkerState
 
 
 @contextmanager
@@ -124,6 +124,8 @@ class Span:
     #: stop
     enqueued: float
 
+    _cumulative_worker_metrics: defaultdict[tuple[str, ...], float]
+
     # Support for weakrefs to a class with __slots__
     __weakref__: Any
 
@@ -136,6 +138,7 @@ class Span:
         self.enqueued = time()
         self.children = []
         self.groups = set()
+        self._cumulative_worker_metrics = defaultdict(float)
 
     def __repr__(self) -> str:
         return f"Span<name={self.name}, id={self.id}>"
@@ -262,8 +265,25 @@ class Span:
         """
         return sum(tg.nbytes_total for tg in self.traverse_groups())
 
+    @property
+    def cumulative_worker_metrics(self) -> defaultdict[tuple[str, ...], float]:
+        """Replica of Worker.digests_total and Scheduler.cumulative_worker_metrics, but
+        only for the metrics that can be attributed to the current span tree.
+        The span_id has been removed from the key.
+
+        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"``.
+        """
+        out: defaultdict[tuple[str, ...], float] = defaultdict(float)
+        for child in self.traverse_spans():
+            for k, v in child._cumulative_worker_metrics.items():
+                out[k] += v
+        return out
+
 
-class SpansExtension:
+class SpansSchedulerExtension:
     """Scheduler extension for spans support"""
 
     #: All Span objects by id
@@ -358,3 +378,49 @@ class SpansExtension:
             self.root_spans.append(span)
 
         return span
+
+    def heartbeat(
+        self, ws: WorkerState, data: dict[str, dict[tuple[str, ...], float]]
+    ) -> None:
+        """Triggered by SpansWorkerExtension.heartbeat().
+
+        Populate :meth:`Span.cumulative_worker_metrics` with data from the worker.
+
+        See also
+        --------
+        SpansWorkerExtension.heartbeat
+        Span.cumulative_worker_metrics
+        """
+        for span_id, metrics in data.items():
+            span = self.spans[span_id]
+            for k, v in metrics.items():
+                span._cumulative_worker_metrics[k] += v
+
+
+class SpansWorkerExtension:
+    """Worker extension for spans support"""
+
+    worker: Worker
+
+    def __init__(self, worker: Worker):
+        self.worker = worker
+
+    def heartbeat(self) -> dict[str, dict[tuple[str, ...], float]]:
+        """Apportion the metrics that do have a span to the Spans on the scheduler
+
+        Returns
+        -------
+        ``{span_id: {("execute", prefix, activity, unit): value}}``
+
+        See also
+        --------
+        SpansSchedulerExtension.heartbeat
+        Span.cumulative_worker_metrics
+        """
+        out: defaultdict[str, dict[tuple[str, ...], float]] = defaultdict(dict)
+        for k, v in self.worker.digests_total_since_heartbeat.items():
+            if isinstance(k, tuple) and k[0] == "execute":
+                _, span_id, prefix, activity, unit = k
+                assert span_id is not None
+                out[span_id]["execute", prefix, activity, unit] = v
+        return dict(out)
diff --git a/distributed/worker.py b/distributed/worker.py
index 56efc375..6adeb083 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -93,6 +93,7 @@ from distributed.pubsub import PubSubWorkerExtension
 from distributed.security import Security
 from distributed.shuffle import ShuffleWorkerExtension
 from distributed.sizeof import safe_sizeof as sizeof
+from distributed.spans import SpansWorkerExtension
 from distributed.threadpoolexecutor import ThreadPoolExecutor
 from distributed.threadpoolexecutor import secede as tpe_secede
 from distributed.utils import (
@@ -172,6 +173,7 @@ LOG_PDB = dask.config.get("distributed.admin.pdb-on-err")
 DEFAULT_EXTENSIONS: dict[str, type] = {
     "pubsub": PubSubWorkerExtension,
     "shuffle": ShuffleWorkerExtension,
+    "spans": SpansWorkerExtension,
 }
 
 DEFAULT_METRICS: dict[str, Callable[[Worker], Any]] = {}
@@ -1056,8 +1058,6 @@ class Worker(BaseWorker, ServerNode):
             event_loop_interval=self._tick_interval_observed,
         )
 
-        self.digests_total_since_heartbeat.clear()
-
         monitor_recent = self.monitor.recent()
         # Convert {foo.bar: 123} to {foo: {bar: 123}}
         for k, v in monitor_recent.items():
@@ -1251,6 +1251,8 @@ class Worker(BaseWorker, ServerNode):
                     if hasattr(extension, "heartbeat")
                 },
             )
+            self.digests_total_since_heartbeat.clear()
+
             end = time()
             middle = (start + end) / 2
 
diff --git a/distributed/worker_state_machine.py b/distributed/worker_state_machine.py
index 656da321..7afde7c1 100644
--- a/distributed/worker_state_machine.py
+++ b/distributed/worker_state_machine.py
@@ -3604,12 +3604,13 @@ class BaseWorker(abc.ABC):
         self.state = state
         self._async_instructions = set()
 
-    def _start_async_instruction(
+    def _start_async_instruction(  # type: ignore[valid-type]
         self,
         task_name: str,
         func: Callable[P, Awaitable[StateMachineEvent]],
         /,
         *args: P.args,
+        span_id: str | None = None,
         **kwargs: P.kwargs,
     ) -> None:
         """Execute an asynchronous instruction inside an asyncio task"""
@@ -3622,12 +3623,15 @@ class BaseWorker(abc.ABC):
 
         task = asyncio.create_task(wrapper(), name=task_name)
         self._async_instructions.add(task)
-        task.add_done_callback(partial(self._finish_async_instruction, ledger=ledger))
+        task.add_done_callback(
+            partial(self._finish_async_instruction, ledger=ledger, span_id=span_id)
+        )
 
     def _finish_async_instruction(
         self,
         task: asyncio.Task[StateMachineEvent],
         ledger: DelayedMetricsLedger,
+        span_id: str | None,
     ) -> None:
         """The asynchronous instruction just completed; process the returned
         stimulus.
@@ -3647,16 +3651,20 @@ class BaseWorker(abc.ABC):
             # Capture metric events in _transition_to_memory()
             self.handle_stimulus(stim)
 
-        self._finalize_metrics(stim, ledger)
+        self._finalize_metrics(stim, ledger, span_id)
 
     def _finalize_metrics(
         self,
         stim: StateMachineEvent,
         ledger: DelayedMetricsLedger,
+        span_id: str | None = None,
     ) -> None:
-        activity: tuple[str, ...]
+        activity: tuple[str | None, ...]
         coarse_time: str | Literal[False]
 
+        if not isinstance(stim, ExecuteDoneEvent):
+            assert span_id is None
+
         if isinstance(stim, GatherDepSuccessEvent):
             activity = ("gather-dep",)
             for key in stim.data:
@@ -3677,16 +3685,16 @@ class BaseWorker(abc.ABC):
             coarse_time = "busy"
 
         elif isinstance(stim, ExecuteSuccessEvent):
-            activity = ("execute", key_split(stim.key))
+            activity = ("execute", span_id, key_split(stim.key))
             ts = self.state.tasks.get(stim.key)
             coarse_time = False if ts and ts.state == "memory" else "cancelled"
 
         elif isinstance(stim, ExecuteFailureEvent):
-            activity = ("execute", key_split(stim.key))
+            activity = ("execute", span_id, key_split(stim.key))
             coarse_time = "failed"
 
         elif isinstance(stim, RescheduleEvent):
-            activity = ("execute", key_split(stim.key))
+            activity = ("execute", span_id, key_split(stim.key))
             coarse_time = "cancelled"
 
         else:
@@ -3729,10 +3737,12 @@ class BaseWorker(abc.ABC):
                 )
 
             elif isinstance(inst, Execute):
+                ts = self.state.tasks[inst.key]
                 self._start_async_instruction(
                     f"execute({inst.key})",
                     self.execute,
                     inst.key,
+                    span_id=ts.span_id,
                     stimulus_id=inst.stimulus_id,
                 )
 


Test patch:
diff --git a/distributed/tests/test_spans.py b/distributed/tests/test_spans.py
index d7c467b7..efd24a6a 100644
--- a/distributed/tests/test_spans.py
+++ b/distributed/tests/test_spans.py
@@ -4,7 +4,8 @@ import pytest
 
 from dask import delayed
 
-from distributed import Client, Event, Future
+from distributed import Client, Event, Future, wait
+from distributed.compatibility import WINDOWS
 from distributed.metrics import time
 from distributed.spans import span
 from distributed.utils_test import (
@@ -427,3 +428,105 @@ async def test_scatter_creates_tg(c, s, a, b):
     x1 = c.submit(inc, 1, key="x-1")
     assert await x1 == 2
     assert s.tasks["x-0"].group.span_id is not None
+
+
+@gen_cluster(client=True)
+async def test_worker_metrics(c, s, a, b):
+    """Test that Scheduler.cumulative_worker_metrics and Span.cumulative_worker_metrics
+    are sync'ed for all 'execute' activities
+
+    See also
+    --------
+    test_worker_metrics.py::test_send_metrics_to_scheduler
+    """
+    sid_y = []
+    with span("foo") as sid_x:
+        # Make sure that cumulative sync over multiple heartbeats works as expected
+        for _ in range(2):
+            # Span 'foo' remains the same across for iterations
+            x0 = c.submit(inc, 1, key=("x", 0), workers=[a.address])
+            x1 = c.submit(inc, 1, key=("x", 1), workers=[b.address])
+            await wait([x0, x1])
+
+            # Span 'bar' is opened and closed between for iterations
+            with span("bar") as sid_yi:
+                sid_y.append(sid_yi)
+                y0 = c.submit(inc, 1, key=("y", 0), workers=[a.address])
+                y1 = c.submit(inc, 1, key=("y", 1), workers=[b.address])
+            await wait([y0, y1])
+
+            # Populate gather_dep and get_data metrics.
+            # They will be ignored by the spans extension because they have no span_id.
+            x2 = c.submit(inc, x0, key=("x", 2), workers=[b.address])
+            x3 = c.submit(inc, x1, key=("x", 3), workers=[a.address])
+            await wait([x2, x3])
+
+            # Flush metrics from workers to scheduler
+            await a.heartbeat()
+            await b.heartbeat()
+
+            # Cleanup
+            del x0, x1, x2, x3, y0, y1
+            await async_poll_for(
+                lambda: not s.tasks and not a.state.tasks and not b.state.tasks,
+                timeout=5,
+            )
+
+    metrics = {
+        k: v
+        for k, v in s.cumulative_worker_metrics.items()
+        if isinstance(k, tuple) and k[0] == "execute"
+    }
+
+    assert list(metrics) == [
+        ("execute", sid_x, "x", "deserialize", "seconds"),
+        ("execute", sid_x, "x", "thread-cpu", "seconds"),
+        ("execute", sid_x, "x", "thread-noncpu", "seconds"),
+        ("execute", sid_x, "x", "executor", "seconds"),
+        ("execute", sid_x, "x", "other", "seconds"),
+        ("execute", sid_y[0], "y", "deserialize", "seconds"),
+        ("execute", sid_y[0], "y", "thread-cpu", "seconds"),
+        ("execute", sid_y[0], "y", "thread-noncpu", "seconds"),
+        ("execute", sid_y[0], "y", "executor", "seconds"),
+        ("execute", sid_y[0], "y", "other", "seconds"),
+        ("execute", sid_x, "x", "memory-read", "count"),
+        ("execute", sid_x, "x", "memory-read", "bytes"),
+        ("execute", sid_y[1], "y", "deserialize", "seconds"),
+        ("execute", sid_y[1], "y", "thread-cpu", "seconds"),
+        ("execute", sid_y[1], "y", "thread-noncpu", "seconds"),
+        ("execute", sid_y[1], "y", "executor", "seconds"),
+        ("execute", sid_y[1], "y", "other", "seconds"),
+    ]
+
+    if not WINDOWS:
+        assert all(v > 0 for v in metrics.values())
+
+    ext = s.extensions["spans"]
+
+    # Metrics have been synchronized from scheduler to spans
+    for (context, span_id, prefix, activity, unit), v in metrics.items():
+        assert (
+            ext.spans[span_id]._cumulative_worker_metrics[
+                context, prefix, activity, unit
+            ]
+            == v
+        )
+
+    # foo metrics include self and its child bar
+    cum_metrics = ext.spans_search_by_name["foo",][0].cumulative_worker_metrics
+    assert list(cum_metrics) == [
+        ("execute", "x", "deserialize", "seconds"),
+        ("execute", "x", "thread-cpu", "seconds"),
+        ("execute", "x", "thread-noncpu", "seconds"),
+        ("execute", "x", "executor", "seconds"),
+        ("execute", "x", "other", "seconds"),
+        ("execute", "x", "memory-read", "count"),
+        ("execute", "x", "memory-read", "bytes"),
+        ("execute", "y", "deserialize", "seconds"),
+        ("execute", "y", "thread-cpu", "seconds"),
+        ("execute", "y", "thread-noncpu", "seconds"),
+        ("execute", "y", "executor", "seconds"),
+        ("execute", "y", "other", "seconds"),
+    ]
+    for k, v in cum_metrics.items():
+        assert v == sum(sp._cumulative_worker_metrics[k] for sp in ext.spans.values())
diff --git a/distributed/tests/test_worker_metrics.py b/distributed/tests/test_worker_metrics.py
index 56a7e6ed..c7e2e3a0 100644
--- a/distributed/tests/test_worker_metrics.py
+++ b/distributed/tests/test_worker_metrics.py
@@ -42,6 +42,14 @@ def get_digests(
     return digests
 
 
+def span_id(s: Scheduler) -> str | None:
+    ext = s.extensions["spans"]
+    defaults = ext.spans_search_by_name["default",]
+    # This is an arbitrary constraint enforced by the tests below
+    assert len(defaults) == 1
+    return defaults[0].id
+
+
 @gen_cluster(client=True, config={"distributed.worker.memory.target": 1e-9})
 async def test_task_lifecycle(c, s, a, b):
     x = (await c.scatter({"x": "x" * 20_000}, workers=[a.address]))["x"]
@@ -74,25 +82,25 @@ async def test_task_lifecycle(c, s, a, b):
         ("gather-dep", "other", "seconds"),
         # a.execute()
         # -> Deserialize run_spec
-        ("execute", "z", "deserialize", "seconds"),
+        ("execute", span_id(s), "z", "deserialize", "seconds"),
         # -> Unspill inputs
         # (There's also another execute-deserialize-seconds entry)
-        ("execute", "z", "disk-read", "seconds"),
-        ("execute", "z", "disk-read", "count"),
-        ("execute", "z", "disk-read", "bytes"),
-        ("execute", "z", "decompress", "seconds"),
+        ("execute", span_id(s), "z", "disk-read", "seconds"),
+        ("execute", span_id(s), "z", "disk-read", "count"),
+        ("execute", span_id(s), "z", "disk-read", "bytes"),
+        ("execute", span_id(s), "z", "decompress", "seconds"),
         # -> Run in thread
-        ("execute", "z", "thread-cpu", "seconds"),
-        ("execute", "z", "thread-noncpu", "seconds"),
-        ("execute", "z", "executor", "seconds"),
+        ("execute", span_id(s), "z", "thread-cpu", "seconds"),
+        ("execute", span_id(s), "z", "thread-noncpu", "seconds"),
+        ("execute", span_id(s), "z", "executor", "seconds"),
         # Spill output; added by _transition_to_memory
-        ("execute", "z", "serialize", "seconds"),
-        ("execute", "z", "compress", "seconds"),
-        ("execute", "z", "disk-write", "seconds"),
-        ("execute", "z", "disk-write", "count"),
-        ("execute", "z", "disk-write", "bytes"),
+        ("execute", span_id(s), "z", "serialize", "seconds"),
+        ("execute", span_id(s), "z", "compress", "seconds"),
+        ("execute", span_id(s), "z", "disk-write", "seconds"),
+        ("execute", span_id(s), "z", "disk-write", "count"),
+        ("execute", span_id(s), "z", "disk-write", "bytes"),
         # Delta to end-to-end runtime as seen from the worker state machine
-        ("execute", "z", "other", "seconds"),
+        ("execute", span_id(s), "z", "other", "seconds"),
         # a.get_data() (triggered by the client retrieving the Future for z)
         # Unspill
         ("get-data", "disk-read", "seconds"),
@@ -108,8 +116,8 @@ async def test_task_lifecycle(c, s, a, b):
     assert list(get_digests(a)) == expect
 
     assert get_digests(a, allow="count") == {
-        ("execute", "z", "disk-read", "count"): 2,
-        ("execute", "z", "disk-write", "count"): 1,
+        ("execute", span_id(s), "z", "disk-read", "count"): 2,
+        ("execute", span_id(s), "z", "disk-write", "count"): 1,
         ("gather-dep", "disk-write", "count"): 1,
         ("get-data", "disk-read", "count"): 1,
     }
@@ -121,8 +129,10 @@ async def test_task_lifecycle(c, s, a, b):
 async def test_async_task(c, s, a):
     """Test that async tasks are metered"""
     await c.submit(asyncio.sleep, 0.1, key=("x-123", 0))
-    assert a.digests_total["execute", "x" "thread-cpu", "seconds"] == 0
-    assert 0 < a.digests_total["execute", "x", "thread-noncpu", "seconds"] < 1
+    assert a.digests_total["execute", span_id(s), "x" "thread-cpu", "seconds"] == 0
+    assert (
+        0 < a.digests_total["execute", span_id(s), "x", "thread-noncpu", "seconds"] < 1
+    )
 
 
 @gen_cluster(client=True, nthreads=[("", 1)])
@@ -139,19 +149,21 @@ async def test_custom_executor(c, s, a):
             await c.submit(sleep, 0.1)
 
     assert list(get_digests(a, "execute")) == [
-        ("execute", "sleep", "deserialize", "seconds"),
-        ("execute", "sleep", "executor", "seconds"),
-        ("execute", "sleep", "other", "seconds"),
+        ("execute", span_id(s), "sleep", "deserialize", "seconds"),
+        ("execute", span_id(s), "sleep", "executor", "seconds"),
+        ("execute", span_id(s), "sleep", "other", "seconds"),
     ]
 
-    assert 0 < a.digests_total["execute", "sleep", "executor", "seconds"] < 1
+    assert (
+        0 < a.digests_total["execute", span_id(s), "sleep", "executor", "seconds"] < 1
+    )
 
 
 @gen_cluster(client=True, nthreads=[("", 1)])
 async def test_run_spec_deserialization(c, s, a):
     """Test that deserialization of run_spec is metered"""
     await c.submit(inc, 1, key="x")
-    assert 0 < a.digests_total["execute", "x", "deserialize", "seconds"] < 1
+    assert 0 < a.digests_total["execute", span_id(s), "x", "deserialize", "seconds"] < 1
 
 
 @gen_cluster(client=True)
@@ -167,8 +179,8 @@ async def test_offload(c, s, a, b, monkeypatch):
     assert list(get_digests(b, {"offload", "serialize", "deserialize"})) == [
         ("gather-dep", "offload", "seconds"),
         ("gather-dep", "deserialize", "seconds"),
-        ("execute", "y", "offload", "seconds"),
-        ("execute", "y", "deserialize", "seconds"),
+        ("execute", span_id(s), "y", "offload", "seconds"),
+        ("execute", span_id(s), "y", "deserialize", "seconds"),
         ("get-data", "offload", "seconds"),
         ("get-data", "serialize", "seconds"),
     ]
@@ -180,7 +192,7 @@ async def test_execute_failed(c, s, a):
     x = c.submit(lambda: 1 / 0, key="x")
     await wait(x)
 
-    assert list(get_digests(a)) == [("execute", "x", "failed", "seconds")]
+    assert list(get_digests(a)) == [("execute", span_id(s), "x", "failed", "seconds")]
 
 
 @gen_cluster(client=True, nthreads=[("", 1)])
@@ -194,7 +206,9 @@ async def test_cancelled_execute(c, s, a):
     await ev.set()
     await async_poll_for(lambda: not a.state.tasks, timeout=5)
 
-    assert list(get_digests(a)) == [("execute", "x", "cancelled", "seconds")]
+    assert list(get_digests(a)) == [
+        ("execute", span_id(s), "x", "cancelled", "seconds")
+    ]
 
 
 @gen_cluster(client=True, nthreads=[("", 1)])
@@ -349,18 +363,18 @@ async def test_user_metrics_sync(c, s, a):
     await wait(c.submit(f, key="x"))
 
     assert list(get_digests(a)) == [
-        ("execute", "x", "deserialize", "seconds"),
-        ("execute", "x", "I/O", "seconds"),
-        ("execute", "x", "thread-cpu", "seconds"),
-        ("execute", "x", "thread-noncpu", "seconds"),
-        ("execute", "x", "executor", "seconds"),
-        ("execute", "x", "other", "seconds"),
+        ("execute", span_id(s), "x", "deserialize", "seconds"),
+        ("execute", span_id(s), "x", "I/O", "seconds"),
+        ("execute", span_id(s), "x", "thread-cpu", "seconds"),
+        ("execute", span_id(s), "x", "thread-noncpu", "seconds"),
+        ("execute", span_id(s), "x", "executor", "seconds"),
+        ("execute", span_id(s), "x", "other", "seconds"),
     ]
-    assert get_digests(a)["execute", "x", "I/O", "seconds"] == 5
-    assert get_digests(a)["execute", "x", "thread-cpu", "seconds"] == 0
-    assert get_digests(a)["execute", "x", "thread-noncpu", "seconds"] == 0
-    assert get_digests(a)["execute", "x", "executor", "seconds"] == 0
-    assert get_digests(a)["execute", "x", "other", "seconds"] == 0
+    assert get_digests(a)["execute", span_id(s), "x", "I/O", "seconds"] == 5
+    assert get_digests(a)["execute", span_id(s), "x", "thread-cpu", "seconds"] == 0
+    assert get_digests(a)["execute", span_id(s), "x", "thread-noncpu", "seconds"] == 0
+    assert get_digests(a)["execute", span_id(s), "x", "executor", "seconds"] == 0
+    assert get_digests(a)["execute", span_id(s), "x", "other", "seconds"] == 0
 
 
 @gen_cluster(client=True, nthreads=[("", 1)])
@@ -372,14 +386,14 @@ async def test_user_metrics_async(c, s, a):
     await wait(c.submit(f, key="x"))
 
     assert list(get_digests(a)) == [
-        ("execute", "x", "deserialize", "seconds"),
-        ("execute", "x", "I/O", "seconds"),
-        ("execute", "x", "thread-noncpu", "seconds"),
-        ("execute", "x", "other", "seconds"),
+        ("execute", span_id(s), "x", "deserialize", "seconds"),
+        ("execute", span_id(s), "x", "I/O", "seconds"),
+        ("execute", span_id(s), "x", "thread-noncpu", "seconds"),
+        ("execute", span_id(s), "x", "other", "seconds"),
     ]
-    assert get_digests(a)["execute", "x", "I/O", "seconds"] == 5
-    assert get_digests(a)["execute", "x", "thread-noncpu", "seconds"] == 0
-    assert get_digests(a)["execute", "x", "other", "seconds"] == 0
+    assert get_digests(a)["execute", span_id(s), "x", "I/O", "seconds"] == 5
+    assert get_digests(a)["execute", span_id(s), "x", "thread-noncpu", "seconds"] == 0
+    assert get_digests(a)["execute", span_id(s), "x", "other", "seconds"] == 0
 
 
 @gen_cluster(client=True, nthreads=[("", 1)])
@@ -392,11 +406,11 @@ async def test_user_metrics_fail(c, s, a):
     await wait(c.submit(f, key="x"))
 
     assert list(get_digests(a)) == [
-        ("execute", "x", "I/O", "bytes"),
-        ("execute", "x", "failed", "seconds"),
+        ("execute", span_id(s), "x", "I/O", "bytes"),
+        ("execute", span_id(s), "x", "failed", "seconds"),
     ]
-    assert get_digests(a)["execute", "x", "I/O", "bytes"] == 100
-    assert get_digests(a)["execute", "x", "failed", "seconds"] < 1
+    assert get_digests(a)["execute", span_id(s), "x", "I/O", "bytes"] == 100
+    assert get_digests(a)["execute", span_id(s), "x", "failed", "seconds"] < 1
 
 
 @gen_cluster(client=True, nthreads=[("", 1)])
@@ -416,13 +430,13 @@ async def test_user_metrics_weird(c, s, a):
         list(get_digests(a))
         == list(get_digests(s))
         == [
-            ("execute", "x", "deserialize", "seconds"),
-            ("execute", "x", ("foo", 1), "seconds"),
-            ("execute", "x", None, "custom"),
-            ("execute", "x", "thread-cpu", "seconds"),
-            ("execute", "x", "thread-noncpu", "seconds"),
-            ("execute", "x", "executor", "seconds"),
-            ("execute", "x", "other", "seconds"),
+            ("execute", span_id(s), "x", "deserialize", "seconds"),
+            ("execute", span_id(s), "x", ("foo", 1), "seconds"),
+            ("execute", span_id(s), "x", None, "custom"),
+            ("execute", span_id(s), "x", "thread-cpu", "seconds"),
+            ("execute", span_id(s), "x", "thread-noncpu", "seconds"),
+            ("execute", span_id(s), "x", "executor", "seconds"),
+            ("execute", span_id(s), "x", "other", "seconds"),
         ]
     )
 
@@ -446,9 +460,9 @@ async def test_do_not_leak_metrics(c, s, a, b):
     await wait([x, y, z])
 
     assert get_digests(a, "count") == {
-        ("execute", "x", "ping-x", "count"): 10,
-        ("execute", "y", "ping-y", "count"): 10,
-        ("execute", "z", "ping-z", "count"): 10,
+        ("execute", span_id(s), "x", "ping-x", "count"): 10,
+        ("execute", span_id(s), "y", "ping-y", "count"): 10,
+        ("execute", span_id(s), "z", "ping-z", "count"): 10,
     }
 
 
@@ -477,15 +491,21 @@ async def test_reschedule(c, s, a, b):
     assert all(f.key in b.data for f in futures)
 
     evs = get_digests(a, "x")
-    k = ("execute", "x", "cancelled", "seconds")
+    k = ("execute", span_id(s), "x", "cancelled", "seconds")
     assert list(evs) == [k]
     assert evs[k] > 0
 
 
-@gen_cluster(client=True)
+@gen_cluster(
+    client=True, scheduler_kwargs={"extensions": {}}, worker_kwargs={"extensions": {}}
+)
 async def test_send_metrics_to_scheduler(c, s, a, b):
     """Test that Worker.digests_total are sync'ed by the heartbeat to
     Scheduler.cumulative_worker_metrics
+
+    See also
+    --------
+    test_spans.py::test_worker_metrics
     """
     # Race condition: metrics that are updated while idle may or may not be there already
     assert s.cumulative_worker_metrics.keys() in (set(), {"latency"})
@@ -515,11 +535,11 @@ async def test_send_metrics_to_scheduler(c, s, a, b):
         == list(b_metrics)
         == list(s_metrics)
         == [
-            ("execute", "x", "deserialize", "seconds"),
-            ("execute", "x", "thread-cpu", "seconds"),
-            ("execute", "x", "thread-noncpu", "seconds"),
-            ("execute", "x", "executor", "seconds"),
-            ("execute", "x", "other", "seconds"),
+            ("execute", None, "x", "deserialize", "seconds"),
+            ("execute", None, "x", "thread-cpu", "seconds"),
+            ("execute", None, "x", "thread-noncpu", "seconds"),
+            ("execute", None, "x", "executor", "seconds"),
+            ("execute", None, "x", "other", "seconds"),
             ("get-data", "memory-read", "count"),
             ("get-data", "memory-read", "bytes"),
             ("get-data", "serialize", "seconds"),
@@ -529,8 +549,8 @@ async def test_send_metrics_to_scheduler(c, s, a, b):
             ("gather-dep", "network", "seconds"),
             ("gather-dep", "other", "seconds"),
             ("get-data", "network", "seconds"),
-            ("execute", "x", "memory-read", "count"),
-            ("execute", "x", "memory-read", "bytes"),
+            ("execute", None, "x", "memory-read", "count"),
+            ("execute", None, "x", "memory-read", "bytes"),
         ]
     )
     for k in s_metrics:
@@ -538,3 +558,27 @@ async def test_send_metrics_to_scheduler(c, s, a, b):
             assert a_metrics[k] > 0
             assert b_metrics[k] > 0
         assert s_metrics[k] == a_metrics[k] + b_metrics[k]
+
+
+@gen_cluster(
+    client=True,
+    nthreads=[("", 1)],
+    scheduler_kwargs={"extensions": {}},
+    worker_kwargs={"extensions": {}},
+)
+async def test_no_spans_extension(c, s, a):
+    await wait(c.submit(lambda: 1 / 0, key="x"))
+    await wait(c.submit(inc, 1, key="y"))
+    await a.heartbeat()
+
+    digests = get_digests(a)
+    assert list(digests) == [
+        ("execute", None, "x", "failed", "seconds"),
+        ("execute", None, "y", "deserialize", "seconds"),
+        ("execute", None, "y", "thread-cpu", "seconds"),
+        ("execute", None, "y", "thread-noncpu", "seconds"),
+        ("execute", None, "y", "executor", "seconds"),
+        ("execute", None, "y", "other", "seconds"),
+    ]
+    for k, v in digests.items():
+        assert s.cumulative_worker_metrics[k] == v