Repository: dask/distributed

Issue:
Add `needs triage` label to new issues
We do this over in `dask/dask` with [this GHA workflow](https://github.com/dask/dask/blob/3edd87f6f582979264dee4f2f0cda52b39b23e78/.github/workflows/label-all.yml#L11) and it's reasonably useful 

Discussion:

Patch:
diff --git a/.github/workflows/label-all.yml b/.github/workflows/label-all.yml
new file mode 100644
index 00000000..9d09c42e
--- /dev/null
+++ b/.github/workflows/label-all.yml
@@ -0,0 +1,14 @@
+name: "Issue and PR Labeler"
+on:
+  pull_request:
+    types: [opened]
+  issues:
+    types: [opened, reopened]
+jobs:
+  label-all-on-open:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: andymckay/labeler@1.0.4
+        with:
+          add-labels: "needs triage"
+          ignore-if-labeled: false
diff --git a/distributed/dashboard/components/scheduler.py b/distributed/dashboard/components/scheduler.py
index 238f1acf..d7731437 100644
--- a/distributed/dashboard/components/scheduler.py
+++ b/distributed/dashboard/components/scheduler.py
@@ -6,8 +6,7 @@ import operator
 import os
 from collections import OrderedDict, defaultdict
 from collections.abc import Iterable
-from datetime import datetime, timedelta
-from itertools import chain
+from datetime import datetime
 from numbers import Number
 from typing import Any, TypeVar
 
@@ -47,7 +46,7 @@ from bokeh.models import (
 )
 from bokeh.models.widgets import DataTable, TableColumn
 from bokeh.models.widgets.markups import Div
-from bokeh.palettes import Viridis11, YlGnBu9, interp_palette
+from bokeh.palettes import Viridis11
 from bokeh.plotting import figure
 from bokeh.themes import Theme
 from bokeh.transform import cumsum, factor_cmap, linear_cmap, stack
@@ -91,6 +90,7 @@ from distributed.diagnostics.task_stream import color_of as ts_color_of
 from distributed.diagnostics.task_stream import colors as ts_color_lookup
 from distributed.metrics import time
 from distributed.scheduler import Scheduler
+from distributed.spans import Span, SpansSchedulerExtension
 from distributed.utils import Log, log_errors
 
 if dask.config.get("distributed.dashboard.export-tool"):
@@ -3383,167 +3383,117 @@ class FinePerformanceMetrics(DashboardComponent):
     The main overview of the Fine Performance Metrics page.
     """
 
+    BASE_TOOLS = ["pan", "wheel_zoom", "box_zoom", "reset"]
+    scheduler: Scheduler
+    task_exec_by_prefix_src: ColumnDataSource
+    task_exec_by_prefix_ymax: float
+    task_exec_by_activity_src: ColumnDataSource
+    get_data_by_activity_src: ColumnDataSource
+    substantial_change: bool
+    visible_functions: list[str]
+    visible_activities: list[str]
+    function_selector: MultiChoice
+    span_tag_selector: MultiChoice
+    unit_selector: Select
+    task_exec_by_activity_chart: figure | None
+    task_exec_by_prefix_chart: figure | None
+    get_data_by_activity_chart: figure | None
+
     @log_errors
-    def __init__(self, scheduler, **kwargs):
+    def __init__(self, scheduler: Scheduler, **kwargs: Any):
         self.scheduler = scheduler
-        self.senddata = defaultdict(list)
-        self.sendsrc = ColumnDataSource(data=dict())
-        self.task_exec_data = defaultdict(list)
-        self.task_exec_data_limited = defaultdict(list)
-        self.task_exec_by_prefix_src = ColumnDataSource(data=dict())
-        self.task_exec_by_activity_src = ColumnDataSource(data=dict())
+        self.task_exec_by_prefix_src = ColumnDataSource(data={})
+        self.task_exec_by_prefix_ymax = 0.0
+        self.task_exec_by_activity_src = ColumnDataSource(data={})
+        self.get_data_by_activity_src = ColumnDataSource(data={})
         self.substantial_change = False
-        self.task_activities = []
-        self.init_root()
-
-    def init_root(self):
-        def handle_selector_chng(attr, old, new):
-            self.unit_selected = new
-            self.substantial_change = True
-
-        self.function_selector = MultiChoice(value=[], options=[])
-        self.function_selector.placeholder = "Select specific functions"
-        self.unit_selector = Select(title="Unit selection", options=["seconds"])
-        self.unit_selector.on_change("value", handle_selector_chng)
-        self.unit_selected = "seconds"
+        self.visible_functions = []
+        self.visible_activities = []
         self.task_exec_by_activity_chart = None
         self.task_exec_by_prefix_chart = None
-        self.senddata_by_activity_chart = None
+        self.get_data_by_activity_chart = None
+
+        # Selectors
+        self.function_selector = MultiChoice(
+            title="Filter by function",
+            placeholder="Select specific functions",
+            value=[],
+            options=[],
+        )
+        self.span_tag_selector = MultiChoice(
+            title="Filter by span tag",
+            placeholder="Select specific span tags",
+            value=[],
+            options=[],
+        )
+        self.unit_selector = Select(title="Unit selection", options=["seconds"])
+        self.unit_selector.value = "seconds"
+        self.unit_selector.on_change("value", self._handle_change_unit)
+
+        selectors_row = row(
+            children=[
+                self.span_tag_selector,
+                self.function_selector,
+                self.unit_selector,
+            ],
+            sizing_mode="stretch_width",
+        )
         self.root = column(
-            self.function_selector,
-            self.unit_selector,
+            selectors_row,
             sizing_mode="scale_width",
         )
 
-    def format(self, unit: str, val: Any) -> str:
-        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
+    def _handle_change_unit(self, attr: str, old: str, new: str) -> None:
+        self.substantial_change = True
 
-        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),
-        )
+        if new == "seconds":
+            yfmt = "00:00:00"
+        elif new == "bytes":
+            yfmt = "0.00b"
+        elif new == "count":
+            yfmt = "0"
+        else:
+            yfmt = "0.000"
+        assert self.task_exec_by_prefix_chart
+        self.task_exec_by_prefix_chart.yaxis[0].formatter.format = yfmt
 
     @without_property_validation
     @log_errors
     def update(self):
-        execute_metrics, get_data_metrics = self.get_metrics()
+        self._update_selectors()
+        self._build_data_sources()
 
-        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)
-
-        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:
+        needs_figures_row = self.task_exec_by_activity_chart is None
+        if needs_figures_row:
             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 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.
-        if self.function_selector.value:
-            indexes = [data["functions"].index(f) for f in self.function_selector.value]
-            for key, values in data.items():
-                data[key] = [values[idx] for idx in indexes]
-
-        # Otherwise limit those being shown which have 'expired' to be displayed
-        else:
-            cutoff = datetime.utcnow() - timedelta(seconds=10)
-            n_show = len([d for d in data["timestamp"] if d > cutoff]) or 5
-            for key in data:
-                data[key] = data[key][-n_show:]
-        self.task_exec_data_limited = data.copy()
 
-        # Show total number of functions to choose from
-        self.function_selector.title = (
-            f"Filter by function ({len(self.function_selector.options)}):"
-        )
+        if needs_figures_row:
+            # First call to update()
+            self.task_exec_by_prefix_chart = (
+                self._build_task_execution_by_prefix_chart()
+            )
+            self.task_exec_by_activity_chart = self._build_pie_chart(
+                source=self.task_exec_by_activity_src,
+                title="Task execution, by activity",
+            )
+            self.get_data_by_activity_chart = self._build_pie_chart(
+                source=self.get_data_by_activity_src,
+                # get_data is called via RPC for data pull; the metrics for it are
+                # recorded on the worker *donating* the data
+                title="Send data, by activity",
+            )
 
-        needs_figures_row = self.task_exec_by_activity_chart is None
-        self._build_task_execution_by_activity_chart(self.task_exec_data_limited.copy())
-        self._build_task_execution_by_prefix_chart(self.task_exec_data_limited.copy())
-        self._build_senddata_chart(self.senddata.copy())
+        self.task_exec_by_prefix_chart.y_range.end = self.task_exec_by_prefix_ymax * 1.1
 
-        if self.substantial_change or needs_figures_row:
+        if self.substantial_change:
+            # Visible activities and/or functions changed
             self.substantial_change = False
-
+            self._update_task_execution_by_prefix_chart()
             figures_row = row(
                 children=[
                     self.task_exec_by_prefix_chart,
                     self.task_exec_by_activity_chart,
-                    self.senddata_by_activity_chart,
+                    self.get_data_by_activity_chart,
                 ],
                 sizing_mode="stretch_width",
             )
@@ -3555,164 +3505,284 @@ class FinePerformanceMetrics(DashboardComponent):
                 # Otherwise needs forced refresh by replacing the figures row
                 self.root.children[-1] = figures_row
 
-    def _build_task_execution_by_activity_chart(
-        self, task_exec_data: defaultdict[str, list]
-    ) -> None:
-        piechart_data = {}
-        piechart_data["value"] = [
-            sum(task_exec_data[f"{op}_{self.unit_selected}"])
-            for op in self.task_activities
-        ]
-        piechart_data["text"] = [
-            self.format(self.unit_selected, v) for v in piechart_data["value"]
-        ]
-        piechart_data["angle"] = [
-            (
-                sum(task_exec_data[f"{activity}_{self.unit_selected}"])
-                / sum(piechart_data["value"])
-                if sum(piechart_data["value"])
-                else 0  # may not have any bytes movement reported, avoid divide by zero
-            )
-            * 2
-            * math.pi
-            for activity in self.task_activities
-        ]
-        piechart_data["color"] = self._get_palette(len(self.task_activities))
-        piechart_data["activity"] = self.task_activities
-        self.task_exec_by_activity_src.data = piechart_data
-
-        if self.task_exec_by_activity_chart is None:
-            piechart = figure(
-                height=500,
-                sizing_mode="scale_both",
-                title="Task execution, by activity",
-                tools="hover",
-                tooltips="@{activity}: @text",
-                x_range=(-0.5, 1.0),
-            )
-            piechart.axis.axis_label = None
-            piechart.axis.visible = False
-            piechart.grid.grid_line_color = None
-
-            piechart.wedge(
-                x=0,
-                y=1,
-                radius=0.4,
-                start_angle=cumsum("angle", include_zero=True),
-                end_angle=cumsum("angle"),
-                line_color="white",
-                fill_color="color",
-                legend_field="activity",
-                source=self.task_exec_by_activity_src,
-            )
-            self.task_exec_by_activity_chart = piechart
-
-    def _build_task_execution_by_prefix_chart(
-        self, task_exec_data: defaultdict[str, list]
-    ) -> None:
-        functions = task_exec_data["functions"]
-        base_tools = "pan,wheel_zoom,box_zoom,reset"
-
-        if self.task_exec_by_prefix_chart is None:
-            barchart = figure(
-                x_range=functions,
-                height=500,
-                sizing_mode="scale_both",
-                title="Task execution, by function",
-                tools=base_tools,
-            )
-            barchart.yaxis.visible = False
-            barchart.xaxis.major_label_orientation = 0.2
-            barchart.grid.grid_line_color = None
-            self.task_exec_by_prefix_chart = barchart
+    def _update_selectors(self) -> None:
+        """Update choices available in
 
-        if self.task_exec_by_prefix_chart.x_range.factors != functions:
-            self.substantial_change = True
-            self.task_exec_by_prefix_chart.x_range = FactorRange(*functions)
+        - self.unit_selector
+        - self.function_selector
+        - self.span_tag_selector
+        """
+        units = set()
+        functions = set()
 
-        stackers = [
-            name for name in task_exec_data if name.endswith(self.unit_selected)
-        ]
+        for k in self.scheduler.cumulative_worker_metrics:
+            if not isinstance(k, tuple):
+                continue
+            context, *other, activity, unit = k
 
-        if stackers and stackers != getattr(self, "_prev_stackers", []):
-            self._prev_stackers = stackers
-            renderers = self.task_exec_by_prefix_chart.vbar_stack(
-                stackers,
-                x="functions",
-                width=0.9,
-                source=self.task_exec_by_prefix_src,
-                color=self._get_palette(len(self.task_activities)),
-                legend_label=self.task_activities,
+            assert isinstance(unit, str)
+            units.add(unit)
+
+            if context == "execute":
+                (function,) = other
+                assert isinstance(function, str)
+                functions.add(function)
+
+        units.difference_update(self.unit_selector.options)
+        if units:
+            self.unit_selector.options.extend(units)
+
+        functions.difference_update(self.function_selector.options)
+        if functions:
+            self.function_selector.options.extend(functions)
+            self.function_selector.title = (
+                f"Filter by function ({len(self.function_selector.options)}):"
             )
 
-            # Create hovertools ontop of base tools
-            tools = self.task_exec_by_prefix_chart.tools
-            self.task_exec_by_prefix_chart.tools = tools[: len(base_tools.split(","))]
-
-            for vbar in renderers:
-                tooltips = [
-                    (
-                        vbar.name,
-                        f"@{{{vbar.name}_text}}",
-                    ),
-                    ("function", "@functions"),
-                ]
-                self.task_exec_by_prefix_chart.add_tools(
-                    HoverTool(tooltips=tooltips, renderers=[vbar])
+        spans_ext: SpansSchedulerExtension | None = self.scheduler.extensions.get(
+            "spans"
+        )
+        if spans_ext:
+            tags = set(spans_ext.spans_search_by_tag)
+            tags.difference_update(self.span_tag_selector.options)
+            if tags:
+                self.span_tag_selector.options.extend(tags)
+                self.span_tag_selector.title = (
+                    f"Filter by span tag ({len(self.span_tag_selector.options)}):"
                 )
 
+    def _format(self, val: float) -> str:
+        unit = self.unit_selector.value
+        assert isinstance(unit, str)
+        if unit == "seconds":
+            return format_time(val)
+        elif unit == "bytes":
+            return format_bytes(int(val))
+        # count or custom user-defined metric
+        elif (ival := int(val)) == val:
+            return str(ival)
+        else:
+            return str(val)
+
+    def _get_palette(self) -> list[str]:
+        n = len(self.visible_activities)
+        try:
+            from bokeh.palettes import interp_palette
+
+            return list(interp_palette(Viridis11, n))
+        except ImportError:
+            # Bokeh 2.4
+            return [Viridis11[i % len(Viridis11)] for i in range(n)]
+
+    def _build_data_sources(self) -> None:
+        """Pre-process and filter fine performance metrics; build data tables in
+        Bokeh format
+
+        Updates:
+
+        - self.substantial_change
+        - self.visible_activities
+        - self.visible_functions
+        - self.task_exec_by_prefix_src.data
+        - self.task_exec_by_activity_src.data
+        - self.get_data_by_activity_src.data
+        """
+        visible_functions = set()
+        visible_activities = set()
+        execute_by_func: defaultdict[tuple[str, str], float] = defaultdict(float)
+        execute: defaultdict[str, float] = defaultdict(float)
+        get_data: defaultdict[str, float] = defaultdict(float)
+
+        function_sel = set(self.function_selector.value)
+
+        span: Span | Scheduler
+        if self.span_tag_selector.value:
+            spans_ext: SpansSchedulerExtension = self.scheduler.extensions["spans"]
+            span = spans_ext.merge_by_tags(*self.span_tag_selector.value)
+        else:
+            span = self.scheduler
+
+        for k, v in span.cumulative_worker_metrics.items():
+            if not isinstance(k, tuple):
+                continue  # Only happens in global metrics
+            context, *other, activity, unit = k
+            assert isinstance(unit, str)
+            assert self.unit_selector.value
+            if unit != self.unit_selector.value:
+                continue
+
+            if context == "execute":
+                (function,) = other
+                assert isinstance(function, str)
+                if not function_sel or function in function_sel:
+                    # Custom metrics can provide any hashable as the label
+                    activity = str(activity)
+                    execute_by_func[function, activity] += v
+                    execute[activity] += v
+                    visible_functions.add(function)
+                    visible_activities.add(activity)
+
+            elif context == "get-data" and not function_sel:
+                # Note: this will always be empty when a span is selected
+                assert isinstance(activity, str)
+                visible_activities.add(activity)
+                get_data[activity] += v
+
+            # Ignore memory-monitor and gather-dep metrics
+
+        if visible_functions != set(self.visible_functions):
             self.substantial_change = True
-            self.task_exec_by_prefix_chart.renderers = renderers
-        self.task_exec_by_prefix_src.data = dict(task_exec_data)
-
-    def _build_senddata_chart(self, senddata: defaultdict[str, list]) -> None:
-        piedata = {}
-        piedata["activity"] = senddata["activity"]
-        piedata["value"] = [
-            (sum(senddata[f"{op}_{self.unit_selected}"])) for op in senddata["activity"]
-        ]
-        piedata["text"] = [self.format(self.unit_selected, v) for v in piedata["value"]]
-        piedata["angle"] = [
-            (
-                (sum(senddata[f"{op}_{self.unit_selected}"]) / sum(piedata["value"]))
-                if sum(piedata["value"])
-                else 0.0
-            )
-            * 2
-            * math.pi
-            for op in piedata["activity"]
+            self.visible_functions = sorted(visible_functions)
+
+        if visible_activities != set(self.visible_activities):
+            self.substantial_change = True
+            self.visible_activities = sorted(visible_activities)
+
+        (
+            self.task_exec_by_prefix_src.data,
+            self.task_exec_by_prefix_ymax,
+        ) = self._build_task_execution_by_prefix_data(execute_by_func)
+        self.task_exec_by_activity_src.data = self._build_pie_data(execute)
+        self.get_data_by_activity_src.data = self._build_pie_data(get_data)
+
+    def _build_pie_data(self, data: defaultdict[str, float]) -> dict[str, list]:
+        """Build the data source for a pie chart by activity
+
+        See also
+        --------
+        _build_pie_chart
+        """
+        total_value = sum(data.values())
+        percent_k = 100.0 / total_value if total_value else 0.0
+        angle_k = 2.0 * math.pi / total_value if total_value else 0.0
+        activities = self.visible_activities
+        values = [data[activity] for activity in activities]
+        total_text = self._format(sum(values))
+        return {
+            "activity": activities,
+            "value": values,
+            "text": [self._format(v) + f" ({v * percent_k:.0f}%)" for v in values],
+            "angle": [v * angle_k for v in values],
+            "color": self._get_palette(),
+            "total_text": [total_text] * len(values),
+        }
+
+    def _build_pie_chart(self, source: ColumnDataSource, title: str) -> figure:
+        """Create pie chart by activity
+
+        See also
+        --------
+        _build_pie_data
+        """
+        piechart = figure(
+            height=500,
+            sizing_mode="scale_both",
+            title=title,
+            tools="hover",
+            tooltips="@{activity}: @text<br>total: @{total_text}",
+            x_range=(-0.5, 1.0),
+        )
+        piechart.axis.axis_label = None
+        piechart.axis.visible = False
+        piechart.grid.grid_line_color = None
+
+        piechart.wedge(
+            x=0,
+            y=1,
+            radius=0.4,
+            start_angle=cumsum("angle", include_zero=True),
+            end_angle=cumsum("angle"),
+            line_color="white",
+            fill_color="color",
+            legend_field="activity",
+            source=source,
+        )
+        return piechart
+
+    def _build_task_execution_by_prefix_data(
+        self, data: defaultdict[tuple[str, str], float]
+    ) -> tuple[dict[str, list], float]:
+        """Build the data source for the execute by function stacked chart
+
+        See also
+        --------
+        _build_task_execution_by_prefix_chart
+        _update_task_execution_by_prefix_chart
+        """
+        func_totals = [
+            sum(data[function, activity] for activity in self.visible_activities)
+            for function in self.visible_functions
         ]
-        piedata["color"] = self._get_palette(len(piedata["activity"]))
+        perc_k = [100.0 / v if v else 0.0 for v in func_totals]
+        out: dict[str, list] = {
+            "__functions": self.visible_functions,
+            "____total_text": [self._format(v) for v in func_totals],
+        }
 
-        self.sendsrc.data = piedata
+        for activity in self.visible_activities:
+            values = [data[function, activity] for function in self.visible_functions]
+            out[activity] = values
+            out[f"__{activity}_text"] = [
+                self._format(v) + f" ({v * perc_ki:.0f}%)"
+                for v, perc_ki in zip(values, perc_k)
+            ]
+        return out, max(func_totals, default=0.0)
 
-        if self.senddata_by_activity_chart is None:
-            senddata_piechart = figure(
-                height=500,
-                sizing_mode="scale_both",
-                title="Send data, by activity",
-                tools="hover",
-                tooltips="@{activity}: @text",
-                x_range=(-0.5, 1.0),
-            )
-            senddata_piechart.wedge(
-                x=0,
-                y=1,
-                radius=0.4,
-                start_angle=cumsum("angle", include_zero=True),
-                end_angle=cumsum("angle"),
-                line_color="white",
-                fill_color="color",
-                legend_field="activity",
-                source=self.sendsrc,
-            )
-            senddata_piechart.axis.axis_label = None
-            senddata_piechart.axis.visible = False
-            senddata_piechart.grid.grid_line_color = None
-            self.senddata_by_activity_chart = senddata_piechart
+    def _build_task_execution_by_prefix_chart(self) -> figure:
+        """Create empty stacked bar chart for execute by function
+
+        See also
+        --------
+        _build_task_execution_by_prefix_data
+        _update_task_execution_by_prefix_chart
+        """
+        barchart = figure(
+            x_range=[],
+            height=500,
+            sizing_mode="scale_both",
+            title="Task execution, by function",
+            tools=",".join(self.BASE_TOOLS),
+        )
+        barchart.yaxis.visible = True
+        # As of Bokeh 3.1, DataRange1D (the default) does not work when switching back
+        # from bytes (GiBs) to seconds (hundreds). So we need to manually update it.
+        barchart.y_range = Range1d(0, 1)
+        barchart.yaxis[0].formatter = NumeralTickFormatter(format="00:00:00")
+        barchart.xaxis.major_label_orientation = 0.2
+        barchart.grid.grid_line_color = None
+        return barchart
+
+    def _update_task_execution_by_prefix_chart(self) -> None:
+        """Rebuild X axis and tooltips of execution by prefix stacked chart
+
+        See also
+        --------
+        _build_task_execution_by_prefix_data
+        _build_task_execution_by_prefix_chart
+        """
+        barchart = self.task_exec_by_prefix_chart
+        assert barchart is not None
+        barchart.x_range = FactorRange(*self.visible_functions)
+        renderers = barchart.vbar_stack(
+            self.visible_activities,
+            x="__functions",
+            width=0.9,
+            source=self.task_exec_by_prefix_src,
+            color=self._get_palette(),
+            legend_label=self.visible_activities,
+        )
 
-    def _get_palette(self, n: int) -> list[str]:
-        return interp_palette(YlGnBu9, n)
+        # Create or refresh hovertools on top of base tools
+        barchart.tools = barchart.tools[: len(self.BASE_TOOLS)]
+
+        for vbar in renderers:
+            tooltips = [
+                ("function", "@__functions"),
+                (vbar.name, f"@{{__{vbar.name}_text}}"),
+                ("total", "@____total_text"),
+            ]
+            barchart.add_tools(HoverTool(tooltips=tooltips, renderers=[vbar]))
+        barchart.renderers = renderers
 
 
 class Contention(DashboardComponent):
diff --git a/distributed/spans.py b/distributed/spans.py
index a6b8d8cc..56757b01 100644
--- a/distributed/spans.py
+++ b/distributed/spans.py
@@ -282,6 +282,16 @@ class Span:
                 out[k] += v
         return out
 
+    @staticmethod
+    def merge(*items: Span) -> Span:
+        """Merge multiple spans into a synthetic one.
+        The input spans must not be related with each other.
+        """
+        out = Span(name=("(merged)",), id_="(merged)", parent=None)
+        out.children.extend(items)
+        out.enqueued = min(child.enqueued for child in items)
+        return out
+
 
 class SpansSchedulerExtension:
     """Scheduler extension for spans support"""
@@ -300,6 +310,11 @@ class SpansSchedulerExtension:
     #: All spans, keyed by the individual tags that make up their name and sorted by
     #: creation time.
     #: This is a convenience helper structure to speed up searches.
+    #:
+    #: See Also
+    #: --------
+    #: find_by_tags
+    #: merge_by_tags
     spans_search_by_tag: defaultdict[str, list[Span]]
 
     def __init__(self, scheduler: Scheduler):
@@ -379,6 +394,29 @@ class SpansSchedulerExtension:
 
         return span
 
+    def find_by_tags(self, *tags: str) -> Iterator[Span]:
+        """Yield all spans that contain any of the given tags.
+        When a tag is shared both by a span and its (grand)children, only return the
+        parent.
+        """
+        by_level = defaultdict(list)
+        for tag in tags:
+            for sp in self.spans_search_by_tag[tag]:
+                by_level[len(sp.name)].append(sp)
+
+        seen = set()
+        for _, level in sorted(by_level.items()):
+            seen.update(level)
+            for sp in level:
+                if sp.parent not in seen:
+                    yield sp
+
+    def merge_by_tags(self, *tags: str) -> Span:
+        """Return a synthetic Span which is the sum of all spans containing the given
+        tags
+        """
+        return Span.merge(*self.find_by_tags(*tags))
+
     def heartbeat(
         self, ws: WorkerState, data: dict[str, dict[tuple[str, ...], float]]
     ) -> None:
diff --git a/distributed/worker.py b/distributed/worker.py
index 6adeb083..9e7bce3b 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1034,6 +1034,14 @@ class Worker(BaseWorker, ServerNode):
             # spilling is disabled
             spilled_memory, spilled_disk = 0, 0
 
+        # Squash span_id in metrics.
+        # SpansWorkerExtension, if loaded, will send them out disaggregated.
+        digests: defaultdict[Hashable, float] = defaultdict(float)
+        for k, v in self.digests_total_since_heartbeat.items():
+            if isinstance(k, tuple) and k[0] == "execute":
+                k = k[:1] + k[2:]
+            digests[k] += v
+
         out = dict(
             task_counts=self.state.task_counter.current_count(by_prefix=False),
             bandwidth={
@@ -1041,7 +1049,7 @@ class Worker(BaseWorker, ServerNode):
                 "workers": dict(self.bandwidth_workers),
                 "types": keymap(typename, self.bandwidth_types),
             },
-            digests_total_since_heartbeat=dict(self.digests_total_since_heartbeat),
+            digests_total_since_heartbeat=dict(digests),
             managed_bytes=self.state.nbytes,
             spilled_bytes={
                 "memory": spilled_memory,


Test patch:
diff --git a/distributed/dashboard/tests/test_scheduler_bokeh.py b/distributed/dashboard/tests/test_scheduler_bokeh.py
index b8c393c5..57497f04 100644
--- a/distributed/dashboard/tests/test_scheduler_bokeh.py
+++ b/distributed/dashboard/tests/test_scheduler_bokeh.py
@@ -342,6 +342,10 @@ async def test_WorkersMemory(c, s, a, b):
 async def test_FinePerformanceMetrics(c, s, a, b):
     cl = FinePerformanceMetrics(s)
 
+    # Test with no metrics
+    cl.update()
+    assert not cl.visible_functions
+
     # execute on default span; multiple tasks in same TaskGroup
     x0 = c.submit(inc, 0, key="x-0", workers=[a.address])
     x1 = c.submit(inc, 1, key="x-1", workers=[a.address])
@@ -368,8 +372,9 @@ async def test_FinePerformanceMetrics(c, s, a, b):
 
     # Custom metric with non-string label and custom unit
     def f():
-        context_meter.digest_metric(None, 1, "custom")
-        context_meter.digest_metric(("foo", 1), 2, "custom")
+        context_meter.digest_metric(None, 1, "seconds")
+        context_meter.digest_metric(("foo", 1), 2, "seconds")
+        context_meter.digest_metric("hideme", 1.1, "custom")
 
     v = c.submit(f, key="v")
     await wait([y0, y1, z, v])
@@ -377,14 +382,35 @@ async def test_FinePerformanceMetrics(c, s, a, b):
     await a.heartbeat()
     await b.heartbeat()
 
-    assert not cl.task_exec_data
     cl.update()
-    assert cl.task_exec_data
-    assert set(cl.task_exec_data["functions"]) == {"v", "w", "x", "y", "z"}
-    assert set(cl.unit_selector.options) == {"seconds", "count", "bytes", "custom"}
-    assert "thread-cpu" in cl.task_activities
-    assert "('foo', 1)" in cl.task_activities
-    assert "None" in cl.task_activities
+    assert sorted(cl.visible_functions) == ["v", "w", "x", "y", "z"]
+    assert sorted(cl.unit_selector.options) == ["bytes", "count", "custom", "seconds"]
+    assert "thread-cpu" in cl.visible_activities
+    assert "('foo', 1)" in cl.visible_activities
+    assert "None" in cl.visible_activities
+    assert "hideme" not in cl.visible_activities
+
+    orig_activities = cl.visible_activities[:]
+
+    cl.unit_selector.value = "bytes"
+    cl.update()
+    assert sorted(cl.visible_activities) == ["disk-read", "disk-write", "memory-read"]
+
+    cl.unit_selector.value = "count"
+    cl.update()
+    assert sorted(cl.visible_activities) == ["disk-read", "disk-write", "memory-read"]
+
+    cl.unit_selector.value = "custom"
+    cl.update()
+    assert sorted(cl.visible_activities) == ["hideme"]
+
+    cl.unit_selector.value = "seconds"
+    cl.update()
+    assert cl.visible_activities == orig_activities
+
+    cl.span_tag_selector.value = ["foo"]
+    cl.update()
+    assert sorted(cl.visible_functions) == ["y", "z"]
 
 
 @gen_cluster(client=True)
diff --git a/distributed/tests/test_spans.py b/distributed/tests/test_spans.py
index efd24a6a..fa1e0bc2 100644
--- a/distributed/tests/test_spans.py
+++ b/distributed/tests/test_spans.py
@@ -13,6 +13,7 @@ from distributed.utils_test import (
     async_poll_for,
     gen_cluster,
     inc,
+    slowinc,
     wait_for_state,
 )
 
@@ -439,8 +440,8 @@ async def test_worker_metrics(c, s, a, b):
     --------
     test_worker_metrics.py::test_send_metrics_to_scheduler
     """
-    sid_y = []
-    with span("foo") as sid_x:
+    bar_sid = []
+    with span("foo") as foo_sid:
         # Make sure that cumulative sync over multiple heartbeats works as expected
         for _ in range(2):
             # Span 'foo' remains the same across for iterations
@@ -449,11 +450,9 @@ async def test_worker_metrics(c, s, a, b):
             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])
+            with span("bar") as bar_sid_i:
+                bar_sid.append(bar_sid_i)
+                await c.submit(inc, 1, key=("y", 0))
 
             # Populate gather_dep and get_data metrics.
             # They will be ignored by the spans extension because they have no span_id.
@@ -466,55 +465,20 @@ async def test_worker_metrics(c, s, a, b):
             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"),
-    ]
+            del x0, x1, x2, x3
+            await async_poll_for(lambda: not s.tasks, timeout=5)
 
-    if not WINDOWS:
-        assert all(v > 0 for v in metrics.values())
+            # Have metrics with 'y' task prefix in foo too
+            await c.submit(inc, 1, key=("y", 1))
 
     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 = ext.spans[foo_sid].cumulative_worker_metrics
+    bar0_metrics = ext.spans[bar_sid[0]].cumulative_worker_metrics
+    bar1_metrics = ext.spans[bar_sid[1]].cumulative_worker_metrics
 
-    # foo metrics include self and its child bar
-    cum_metrics = ext.spans_search_by_name["foo",][0].cumulative_worker_metrics
-    assert list(cum_metrics) == [
+    # metrics for foo include self and its child bar
+    assert list(foo_metrics) == [
         ("execute", "x", "deserialize", "seconds"),
         ("execute", "x", "thread-cpu", "seconds"),
         ("execute", "x", "thread-noncpu", "seconds"),
@@ -528,5 +492,94 @@ async def test_worker_metrics(c, s, a, b):
         ("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())
+    assert (
+        list(bar0_metrics)
+        == list(bar1_metrics)
+        == [
+            ("execute", "y", "deserialize", "seconds"),
+            ("execute", "y", "thread-cpu", "seconds"),
+            ("execute", "y", "thread-noncpu", "seconds"),
+            ("execute", "y", "executor", "seconds"),
+            ("execute", "y", "other", "seconds"),
+        ]
+    )
+
+    if not WINDOWS:
+        for metrics in (foo_metrics, bar0_metrics, bar1_metrics):
+            assert all(v > 0 for v in metrics.values()), metrics
+
+    # Metrics have been synchronized from scheduler to spans
+    for k, v in foo_metrics.items():
+        assert s.cumulative_worker_metrics[k] == v
+
+    # Metrics for foo contain the sum of metrics from itself and for bar
+    for k in bar0_metrics:
+        assert (
+            foo_metrics[k]
+            == bar0_metrics[k]
+            + bar1_metrics[k]
+            + ext.spans[foo_sid]._cumulative_worker_metrics[k]
+        )
+
+
+@gen_cluster(client=True)
+async def test_merge_by_tags(c, s, a, b):
+    with span("foo") as foo1:
+        await c.submit(inc, 1, key="x1")
+        with span("bar") as bar1:  # foo, bar
+            await c.submit(inc, 2, key="x2")
+            with span("foo") as foo2:  # foo, bar, foo
+                await c.submit(inc, 3, key="x3")
+        with span("foo") as foo3:  # foo, foo
+            await c.submit(inc, 4, key="x4")
+    with span("bar") as bar2:  # bar
+        await c.submit(inc, 5, key="x5")
+
+    ext = s.extensions["spans"]
+    assert {s.id for s in ext.find_by_tags("foo")} == {foo1}
+    assert {s.id for s in ext.find_by_tags("foo", "bar")} == {foo1, bar2}
+    assert {s.id for s in ext.find_by_tags("bar", "foo")} == {foo1, bar2}
+    assert {s.id for s in ext.find_by_tags("bar")} == {bar1, bar2}
+
+    def tgnames(*tags):
+        return [tg.name for tg in ext.merge_by_tags(*tags).traverse_groups()]
+
+    assert tgnames("foo") == ["x1", "x2", "x3", "x4"]
+    assert tgnames("foo", "bar") == ["x1", "x2", "x3", "x4", "x5"]
+    assert tgnames("bar", "foo") == ["x5", "x1", "x2", "x3", "x4"]
+    assert tgnames("bar") == ["x5", "x2", "x3"]
+
+
+@gen_cluster(client=True)
+async def test_merge_by_tags_metrics(c, s, a, b):
+    with span("foo") as foo1:
+        await c.submit(slowinc, 1, delay=0.05, key="x-1")
+    await async_poll_for(lambda: not s.task_groups, timeout=5)
+
+    with span("foo") as foo2:
+        await c.submit(slowinc, 2, delay=0.06, key="x-2")
+    await async_poll_for(lambda: not s.task_groups, timeout=5)
+
+    with span("bar") as bar1:
+        await c.submit(slowinc, 3, delay=0.07, key="x-3")
+    await async_poll_for(lambda: not s.task_groups, timeout=5)
+
+    await a.heartbeat()
+    await b.heartbeat()
+
+    ext = s.extensions["spans"]
+    k = ("execute", "x", "thread-noncpu", "seconds")
+    t_foo = ext.merge_by_tags("foo").cumulative_worker_metrics[k]
+    t_bar = ext.merge_by_tags("bar").cumulative_worker_metrics[k]
+    t_foo1 = ext.spans[foo1]._cumulative_worker_metrics[k]
+    t_foo2 = ext.spans[foo2]._cumulative_worker_metrics[k]
+    t_bar1 = ext.spans[bar1]._cumulative_worker_metrics[k]
+    assert t_foo1 > 0
+    assert t_foo2 > 0
+    assert t_bar1 > 0
+    assert t_foo == t_foo1 + t_foo2
+    assert t_bar == t_bar1
+
+    assert ext.merge_by_tags("foo").enqueued == min(
+        ext.spans[foo1].enqueued, ext.spans[foo2].enqueued
+    )
diff --git a/distributed/tests/test_worker_metrics.py b/distributed/tests/test_worker_metrics.py
index c7e2e3a0..cfc0bbdb 100644
--- a/distributed/tests/test_worker_metrics.py
+++ b/distributed/tests/test_worker_metrics.py
@@ -426,19 +426,20 @@ async def test_user_metrics_weird(c, s, a):
     await wait(c.submit(f, key="x"))
     await a.heartbeat()
 
-    assert (
-        list(get_digests(a))
-        == list(get_digests(s))
-        == [
-            ("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"),
-        ]
-    )
+    s_metrics = get_digests(s)
+    a_metrics = get_digests(a)
+
+    assert list(s_metrics) == [
+        ("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"),
+    ]
+    for (context, prefix, activity, unit), v in s_metrics.items():
+        assert a_metrics[context, span_id(s), prefix, activity, unit] == v
 
 
 @gen_cluster(client=True, nthreads=[("", 3)])
@@ -530,34 +531,36 @@ async def test_send_metrics_to_scheduler(c, s, a, b):
     a_metrics = get_digests(a)
     b_metrics = get_digests(b)
     s_metrics = get_digests(s)
-    assert (
-        list(a_metrics)
-        == list(b_metrics)
-        == list(s_metrics)
-        == [
-            ("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"),
-            ("get-data", "compress", "seconds"),
-            ("gather-dep", "decompress", "seconds"),
-            ("gather-dep", "deserialize", "seconds"),
-            ("gather-dep", "network", "seconds"),
-            ("gather-dep", "other", "seconds"),
-            ("get-data", "network", "seconds"),
-            ("execute", None, "x", "memory-read", "count"),
-            ("execute", None, "x", "memory-read", "bytes"),
-        ]
-    )
-    for k in s_metrics:
+
+    expect_worker = [
+        ("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"),
+        ("get-data", "compress", "seconds"),
+        ("gather-dep", "decompress", "seconds"),
+        ("gather-dep", "deserialize", "seconds"),
+        ("gather-dep", "network", "seconds"),
+        ("gather-dep", "other", "seconds"),
+        ("get-data", "network", "seconds"),
+        ("execute", None, "x", "memory-read", "count"),
+        ("execute", None, "x", "memory-read", "bytes"),
+    ]
+    assert list(a_metrics) == list(b_metrics) == expect_worker
+    expect_scheduler = [
+        k[:1] + k[2:] if k[0] == "execute" else k for k in expect_worker
+    ]
+    assert list(s_metrics) == expect_scheduler
+
+    for wk, sk in zip(expect_worker, expect_scheduler):
         if not WINDOWS:
-            assert a_metrics[k] > 0
-            assert b_metrics[k] > 0
-        assert s_metrics[k] == a_metrics[k] + b_metrics[k]
+            assert a_metrics[wk] > 0
+            assert b_metrics[wk] > 0
+        assert s_metrics[sk] == a_metrics[wk] + b_metrics[wk]
 
 
 @gen_cluster(
@@ -571,8 +574,9 @@ async def test_no_spans_extension(c, s, a):
     await wait(c.submit(inc, 1, key="y"))
     await a.heartbeat()
 
-    digests = get_digests(a)
-    assert list(digests) == [
+    w_metrics = get_digests(a)
+    s_metrics = get_digests(s)
+    expect_worker = [
         ("execute", None, "x", "failed", "seconds"),
         ("execute", None, "y", "deserialize", "seconds"),
         ("execute", None, "y", "thread-cpu", "seconds"),
@@ -580,5 +584,13 @@ async def test_no_spans_extension(c, s, a):
         ("execute", None, "y", "executor", "seconds"),
         ("execute", None, "y", "other", "seconds"),
     ]
-    for k, v in digests.items():
-        assert s.cumulative_worker_metrics[k] == v
+    expect_scheduler = [
+        k[:1] + k[2:] if k[0] == "execute" else k for k in expect_worker
+    ]
+    assert list(w_metrics) == expect_worker
+    assert list(s_metrics) == expect_scheduler
+
+    for wk, sk in zip(expect_worker, expect_scheduler):
+        if not WINDOWS:
+            assert w_metrics[wk] > 0
+        assert s_metrics[sk] == w_metrics[wk]