Repository: dask/distributed

Issue:
Stacked area chart for TaskGroups
As part of the effort to replace the Task Stream plot I'd like to propose a stacked area chart, with each line the amount of time spent in a specific task group.  Here are some images of stacked area charts to give a sense for what I mean

![image](https://user-images.githubusercontent.com/306380/124198308-004e1980-da96-11eb-9278-e6d6ecc58ae4.png)
![image](https://user-images.githubusercontent.com/306380/124198326-07752780-da96-11eb-9a2c-d7259c93c1e0.png)

This would give us a sense for which task groups have been active, and how concentrated or diffuse they are.  We might also try to find a way to include transfers or disk access.

I think that we can probably get the information that we need for this chart by capturing the timestamp of when we reach 5% done, 10% done, 15% done, ... 95% done, and 100% done where "done" is measured by the number of tasks.  

Here is some very silly pseudocode to try to give a sense for what I mean.

```python
class TaskGroup:
    ...
    time_done: list

class Scheduler:
    def transition_processing_done(...):
        ts._group.some_counter += 1

        if ts._group.some_counter % len(ts._group) / 20:
            ts._group.time_done.append(time())
```

I think that if we then take the total duration of the task group, assume that it is uniform throughout all tasks, then we should be able to get a sense for the average number of CPUs actively working on this group during every given interval.  If the time between two successive timestamps is short, then our cluster was really concentrating on this task group during that period.  If the time between two successive timestamps is long, then our cluster wasn't concentrating on this task group.

This should give us some points in time and height values for each of our different task groups.

To merge these all into a stacked area chart my sense is that we want to select a time range, get a number of samples along that range (maybe 100) and then use some of the numpy interpolation routines on each of the task groups to sample appropriately.

I *think* that this is all somewhat normal for interpolation processes.  I'm guessing that this is moderately tricky, but shouldn't require any deep novelty on our part.


Discussion:
cc @bryevdv I think that this might an interesting non-trivial project if you're interested.
cc @randerzander (in case you have thoughts on this :)

Resolving patch:
diff --git a/distributed/dashboard/components/scheduler.py b/distributed/dashboard/components/scheduler.py
--- a/distributed/dashboard/components/scheduler.py
+++ b/distributed/dashboard/components/scheduler.py
@@ -2,7 +2,7 @@
 import math
 import operator
 import os
-from collections import defaultdict
+from collections import OrderedDict, defaultdict
 from numbers import Number
 
 import numpy as np
@@ -18,6 +18,7 @@
     CDSView,
     ColorBar,
     ColumnDataSource,
+    CustomJSHover,
     DataRange1d,
     GroupFilter,
     HoverTool,
@@ -40,8 +41,8 @@
 from bokeh.palettes import Viridis11
 from bokeh.plotting import figure
 from bokeh.themes import Theme
-from bokeh.transform import cumsum, factor_cmap, linear_cmap
-from tlz import curry, pipe
+from bokeh.transform import cumsum, factor_cmap, linear_cmap, stack
+from tlz import curry, pipe, valmap
 from tlz.curried import concat, groupby, map
 from tornado import escape
 
@@ -58,6 +59,7 @@
 )
 from distributed.dashboard.utils import BOKEH_VERSION, PROFILING, transpose, update
 from distributed.diagnostics.graph_layout import GraphLayout
+from distributed.diagnostics.progress import GroupTiming
 from distributed.diagnostics.progress_stream import color_of, progress_quads
 from distributed.diagnostics.task_stream import TaskStreamPlugin
 from distributed.diagnostics.task_stream import color_of as ts_color_of
@@ -2560,6 +2562,284 @@ def update(self):
         self.arrows_source.data.update(arrows_data)
 
 
+class TaskGroupProgress(DashboardComponent):
+    """Stacked area chart showing task groups through time"""
+
+    def __init__(self, scheduler, **kwargs):
+        self.scheduler = scheduler
+        self.source = ColumnDataSource()
+        # The length of timeseries to chart (in units of plugin.dt)
+        self.npts = 180
+
+        if GroupTiming.name not in scheduler.plugins:
+            scheduler.add_plugin(plugin=GroupTiming(scheduler))
+
+        self.plugin = scheduler.plugins[GroupTiming.name]
+
+        self.source.add(np.array(self.plugin.time) * 1000.0, "time")
+
+        x_range = DataRange1d(range_padding=0)
+        y_range = Range1d(0, max(self.plugin.nthreads))
+
+        self.root = figure(
+            id="bk-task-group-progress-plot",
+            title="Task Group Progress",
+            name="task_group_progress",
+            toolbar_location="above",
+            min_border_bottom=50,
+            x_range=x_range,
+            y_range=y_range,
+            tools="",
+            x_axis_type="datetime",
+            y_axis_location=None,
+            **kwargs,
+        )
+        self.root.yaxis.major_label_text_alpha = 0
+        self.root.yaxis.minor_tick_line_alpha = 0
+        self.root.yaxis.major_tick_line_alpha = 0
+        self.root.xgrid.visible = False
+
+        self.root.add_tools(
+            BoxZoomTool(),
+            ResetTool(),
+            PanTool(dimensions="width"),
+            WheelZoomTool(dimensions="width"),
+        )
+        self._hover = None
+        self._last_drawn = None
+        self._offset = time()
+        self._last_transition_count = scheduler.transition_counter
+        # OrderedDict so we can make a reverse iterator later and get the
+        # most-recently-added glyphs.
+        self._renderers = OrderedDict()
+        self._line_renderers = OrderedDict()
+
+    def _should_add_new_renderers(self) -> bool:
+        """
+        Whether to add new renderers to the chart.
+
+        When a new set of task groups enters the scheduler we'd like to start rendering
+        them. But it can be expensive to add new glyps, so we do it deliberately,
+        checking whether we have to do it and whether the scheduler seems busy.
+        """
+        # Always draw if we have not before
+        if not self._last_drawn:
+            return True
+        # Don't draw if there have been no new tasks completed since the last update,
+        # or if the scheduler CPU is occupied.
+        if (
+            self._last_transition_count == self.scheduler.transition_counter
+            or self.scheduler.proc.cpu_percent() > 50
+        ):
+            return False
+
+        # Only return true if there are new task groups that we have not yet added
+        # to the ColumnDataSource.
+        return not set(self.plugin.compute.keys()) <= set(self.source.data.keys())
+
+    def _should_update(self) -> bool:
+        """
+        Whether to update the ColumnDataSource. This is cheaper than redrawing,
+        but still not free, so we check whether we need it and whether the scheudler
+        is busy.
+        """
+        return (
+            self._last_transition_count != self.scheduler.transition_counter
+            and self.scheduler.proc.cpu_percent() < 50
+        )
+
+    def _get_timeseries(self, restrict_to_existing=False):
+        """
+        Update the ColumnDataSource with our time series data.
+
+        restrict_to_existing determines whether to add new task groups
+        which might have been added since the last time we rendered.
+        This is important as we want to add new stackers very deliberately.
+        """
+        # Get the front/back indices for most recent npts bins out of the timeseries
+        front = max(len(self.plugin.time) - self.npts, 0)
+        back = None
+        # Remove any periods of zero compute at the front or back of the timeseries
+        if len(self.plugin.compute):
+            agg = sum([np.array(v[front:]) for v in self.plugin.compute.values()])
+            front2 = len(agg) - len(np.trim_zeros(agg, trim="f"))
+            front += front2
+            back = len(np.trim_zeros(agg, trim="b")) - len(agg) or None
+
+        prepend = (
+            self.plugin.time[front - 1]
+            if front >= 1
+            else self.plugin.time[front] - self.plugin.dt
+        )
+        timestamps = np.array(self.plugin.time[front:back])
+        dt = np.diff(timestamps, prepend=prepend)
+
+        if restrict_to_existing:
+            new_data = {
+                k: np.array(v[front:back]) / dt
+                for k, v in self.plugin.compute.items()
+                if k in self.source.data
+            }
+        else:
+            new_data = valmap(
+                lambda x: np.array(x[front:back]) / dt,
+                self.plugin.compute,
+            )
+
+        new_data["time"] = (
+            timestamps - self._offset
+        ) * 1000.0  # bokeh likes milliseconds
+        new_data["nthreads"] = np.array(self.plugin.nthreads[front:back])
+
+        return new_data
+
+    @without_property_validation
+    def update(self):
+        """
+        Maybe update the chart. This is somewhat expensive to draw, so we update
+        it pretty defensively.
+        """
+        with log_errors():
+            if self._should_add_new_renderers():
+                # Update the chart, allowing for new task groups to be added.
+                new_data = self._get_timeseries(restrict_to_existing=False)
+                self.source.data = new_data
+
+                # Possibly update the y range if the number of threads has increased.
+                max_nthreads = max(self.plugin.nthreads)
+                if self.root.y_range.end != max_nthreads:
+                    self.root.y_range.end = max_nthreads
+
+                stackers = list(self.plugin.compute.keys())
+                colors = [color_of(key_split(k)) for k in stackers]
+
+                for i, (group, color) in enumerate(zip(stackers, colors)):
+                    # If we have already drawn the group, but it is all zero,
+                    # set it to be invisible.
+                    if group in self._renderers:
+                        if not np.count_nonzero(new_data[group]) > 0:
+                            self._renderers[group].visible = False
+                            self._line_renderers[group].visible = False
+                        else:
+                            self._renderers[group].visible = True
+                            self._line_renderers[group].visible = True
+
+                        continue
+
+                    # Draw the new area and line glyphs.
+                    renderer = self.root.varea(
+                        x="time",
+                        y1=stack(*stackers[:i]),
+                        y2=stack(*stackers[: i + 1]),
+                        color=color,
+                        alpha=0.5,
+                        source=self.source,
+                    )
+                    self._renderers[group] = renderer
+
+                    line_renderer = self.root.line(
+                        x="time",
+                        y=stack(*stackers[: i + 1]),
+                        color=color,
+                        alpha=1.0,
+                        source=self.source,
+                    )
+                    self._line_renderers[group] = line_renderer
+
+                # Don't add hover until there is something to show, as bokehjs seems to
+                # have trouble with custom hovers when there are no renderers.
+                if self.plugin.compute and self._hover is None:
+                    # Add a hover that will show occupancy for all currently active
+                    # task groups. This is a little tricky, bokeh doesn't (yet) support
+                    # hit tests for stacked area charts: https://github.com/bokeh/bokeh/issues/9182
+                    # Instead, show a single vline hover which lists the currently active task
+                    # groups. A custom formatter in JS-land pulls the relevant data index and
+                    # assembles the tooltip.
+                    formatter = CustomJSHover(code="return '';")
+                    self._hover = HoverTool(
+                        tooltips="""
+                        <div>
+                          <div style="font-size: 1.2em; font-weight: bold">
+                            <b>Worker thread occupancy</b>
+                          </div>
+                          <div>
+                            $index{custom}
+                          </div>
+                        </div>
+                        """,
+                        mode="vline",
+                        line_policy="nearest",
+                        attachment="horizontal",
+                        formatters={"$index": formatter},
+                    )
+                    self.root.add_tools(self._hover)
+
+                if self._hover:
+                    # Create a custom tooltip that:
+                    #   1. Includes nthreads
+                    #   2. Filters out inactive task groups
+                    #      (ones without any compute during the relevant dt)
+                    #   3. Colors the labels appropriately.
+                    formatter = CustomJSHover(
+                        code="""
+                        const colormap = %s;
+                        const divs = [];
+                        for (let k of Object.keys(source.data)) {
+                          const val = source.data[k][value];
+                          const color = colormap[k];
+                          if (k === "time" || k === "nthreads" || val < 1.e-3) {
+                            continue;
+                          }
+                          const label = k.length >= 20 ? k.slice(0, 20) + '…' : k;
+
+                          // Unshift so that the ordering of the labels is the same as
+                          // the ordering of the stackers.
+                          divs.unshift(
+                            '<div>'
+                              + '<span style="font-weight: bold; color:' + color + ';">'
+                                + label
+                              + '</span>'
+                              + ': '
+                              +  val.toFixed(1)
+                              + '</div>'
+                          )
+
+                        }
+                        divs.unshift(
+                          '<div>'
+                            + '<span style="font-weight: bold; color: darkgrey;">nthreads: </span>'
+                            + source.data.nthreads[value]
+                            + '</div>'
+                        );
+                        return divs.join('\\n')
+                        """
+                        % dict(
+                            zip(stackers, colors)
+                        ),  # sneak the color mapping into the callback
+                        args={"source": self.source},
+                    )
+                    # Add the HoverTool to the top line renderer.
+                    top_line = None
+                    for line in reversed(self._line_renderers.values()):
+                        if line.visible:
+                            top_line = line
+                            break
+                    self._hover.renderers = [top_line]
+                    self._hover.formatters = {"$index": formatter}
+
+                self._last_drawn = time()
+                self._last_transition_count = self.scheduler.transition_counter
+            elif self._should_update():
+                # Possibly update the y range if new threads have been added
+                max_nthreads = max(self.plugin.nthreads)
+                if self.root.y_range.end != max_nthreads:
+                    self.root.y_range.end = max_nthreads
+                # Update the data, only including existing columns, rather than redrawing
+                # the whole chart.
+                self.source.data = self._get_timeseries(restrict_to_existing=True)
+                self._last_transition_count = self.scheduler.transition_counter
+
+
 class TaskProgress(DashboardComponent):
     """Progress bars per task type"""
 
diff --git a/distributed/dashboard/scheduler.py b/distributed/dashboard/scheduler.py
--- a/distributed/dashboard/scheduler.py
+++ b/distributed/dashboard/scheduler.py
@@ -18,6 +18,7 @@
     SystemTimeseries,
     TaskGraph,
     TaskGroupGraph,
+    TaskGroupProgress,
     TaskProgress,
     TaskStream,
     WorkerNetworkBandwidth,
@@ -59,6 +60,7 @@
     "/individual-progress": individual_doc(TaskProgress, 100, height=160),
     "/individual-graph": individual_doc(TaskGraph, 200),
     "/individual-groups": individual_doc(TaskGroupGraph, 200),
+    "/individual-group-progress": individual_doc(TaskGroupProgress, 200),
     "/individual-workers-memory": individual_doc(WorkersMemory, 100),
     "/individual-cluster-memory": individual_doc(ClusterMemory, 100),
     "/individual-cpu": individual_doc(CurrentLoad, 100, fig_attr="cpu_figure"),
diff --git a/distributed/diagnostics/progress.py b/distributed/diagnostics/progress.py
--- a/distributed/diagnostics/progress.py
+++ b/distributed/diagnostics/progress.py
@@ -1,5 +1,8 @@
+from __future__ import annotations
+
 import asyncio
 import logging
+import time
 from collections import defaultdict
 from timeit import default_timer
 
@@ -291,3 +294,82 @@ def transition(self, key, start, finish, *args, **kwargs):
     def restart(self, scheduler):
         self.all.clear()
         self.state.clear()
+
+
+class GroupTiming(SchedulerPlugin):
+    """Keep track of high-level timing information for task group progress"""
+
+    name = "group-timing"
+
+    def __init__(self, scheduler):
+        self.scheduler = scheduler
+
+        # Time bin size (in seconds). TODO: make this configurable?
+        self.dt = 1.0
+
+        # Initialize our data structures.
+        self._init()
+
+    def _init(self):
+        """Shared initializatoin code between __init__ and restart"""
+        now = time.time()
+
+        # Timestamps for tracking compute durations by task group.
+        # Start with length 2 so that we always can compute a valid dt later.
+        self.time: list[float] = [now] * 2
+        # The amount of compute since the last timestamp
+        self.compute: dict[str, list[float]] = {}
+        # The number of threads at the time
+        self.nthreads: list[float] = [self.scheduler.total_nthreads] * 2
+
+    def transition(self, key, start, finish, *args, **kwargs):
+        # We are mostly interested in when tasks complete for now, so just look
+        # for when processing transitions to memory. Later we could also extend
+        # this if we can come up with useful visual channels to show it in.
+        if start == "processing" and finish == "memory":
+            startstops = kwargs.get("startstops")
+            if not startstops:
+                logger.warn(
+                    f"Task {key} finished processing, but timing information seems to be missing"
+                )
+                return
+
+            # Possibly extend the timeseries if another dt has passed
+            now = time.time()
+            self.time[-1] = now
+            while self.time[-1] - self.time[-2] > self.dt:
+                self.time[-1] = self.time[-2] + self.dt
+                self.time.append(now)
+                self.nthreads.append(self.scheduler.total_nthreads)
+                for g in self.compute.values():
+                    g.append(0.0)
+
+            # Get the task
+            task = self.scheduler.tasks[key]
+            group = task.group
+
+            # If the group is new, add it to the timeseries as if it has been
+            # here the whole time
+            if group.name not in self.compute:
+                self.compute[group.name] = [0.0] * len(self.time)
+
+            for startstop in startstops:
+                if startstop["action"] != "compute":
+                    continue
+                stop = startstop["stop"]
+                start = startstop["start"]
+                idx = len(self.time) - 1
+                # If the stop time is after the most recent bin,
+                # roll back the current index. Not clear how often this happens.
+                while idx > 0 and self.time[idx - 1] > stop:
+                    idx -= 1
+                # Allocate the timing information of the task to the time bins.
+                while idx > 0 and stop > start:
+                    delta = stop - max(self.time[idx - 1], start)
+                    self.compute[group.name][idx] += delta
+
+                    stop -= delta
+                    idx -= 1
+
+    def restart(self, scheduler):
+        self._init()


Test patch:
diff --git a/distributed/diagnostics/tests/test_progress.py b/distributed/diagnostics/tests/test_progress.py
--- a/distributed/diagnostics/tests/test_progress.py
+++ b/distributed/diagnostics/tests/test_progress.py
@@ -7,12 +7,13 @@
 from distributed.compatibility import LINUX
 from distributed.diagnostics.progress import (
     AllProgress,
+    GroupTiming,
     MultiProgress,
     Progress,
     SchedulerPlugin,
 )
 from distributed.scheduler import COMPILED
-from distributed.utils_test import dec, div, gen_cluster, inc, nodebug
+from distributed.utils_test import dec, div, gen_cluster, inc, nodebug, slowdec, slowinc
 
 
 def f(*args):
@@ -189,3 +190,35 @@ async def test_AllProgress_lost_key(c, s, a, b):
 
     while len(p.state["memory"]["inc"]) > 0:
         await asyncio.sleep(0.01)
+
+
+@gen_cluster(client=True, Worker=Nanny)
+async def test_group_timing(c, s, a, b):
+    p = GroupTiming(s)
+    s.add_plugin(p)
+
+    assert len(p.time) == 2
+    assert len(p.nthreads) == 2
+
+    futures1 = c.map(slowinc, range(10), delay=0.3)
+    futures2 = c.map(slowdec, range(10), delay=0.3)
+    await wait(futures1 + futures2)
+
+    assert len(p.time) > 2
+    assert len(p.nthreads) == len(p.time)
+    assert all([nt == s.total_nthreads for nt in p.nthreads])
+    assert "slowinc" in p.compute
+    assert "slowdec" in p.compute
+    assert all([len(v) == len(p.time) for v in p.compute.values()])
+    assert s.task_groups.keys() == p.compute.keys()
+    assert all(
+        [
+            abs(s.task_groups[k].all_durations["compute"] - sum(v)) < 1.0e-12
+            for k, v in p.compute.items()
+        ]
+    )
+
+    await s.restart()
+    assert len(p.time) == 2
+    assert len(p.nthreads) == 2
+    assert len(p.compute) == 0