Repository: dask/distributed

Issue:
`test_compute_per_key` flaky
Ubuntu, 3.8: https://github.com/dask/distributed/pull/4925/checks?check_run_id=2888102571#step:10:2318

```python
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

c = <Client: No scheduler connected>
s = <Scheduler: "tcp://127.0.0.1:52494" workers: 0 cores: 0, tasks: 0>
a = <Worker: 'tcp://127.0.0.1:52495', 0, Status.closed, stored: 0, running: 0/1, ready: 0, comm: 0, waiting: -4>
b = <Worker: 'tcp://127.0.0.1:52497', 1, Status.closed, stored: 0, running: 0/2, ready: 0, comm: 0, waiting: -2>

    @pytest.mark.flaky(reruns=10, reruns_delay=5, condition=MACOS)
    @gen_cluster(client=True, scheduler_kwargs={"dashboard": True})
    async def test_compute_per_key(c, s, a, b):
        mbk = ComputePerKey(s)
    
        da = pytest.importorskip("dask.array")
        x = (da.ones((20, 20), chunks=(10, 10)) + 1).persist(optimize_graph=False)
    
        await x
        y = await dask.delayed(inc)(1).persist()
        z = (x + x.T) - x.mean(axis=0)
        await c.compute(z.sum())
    
        mbk.update()
        http_client = AsyncHTTPClient()
        response = await http_client.fetch(
            "http://localhost:%d/individual-compute-time-per-key" % s.http_server.port
        )
        assert response.code == 200
>       assert ("sum-aggregate") in mbk.compute_source.data["names"]
E       AssertionError: assert 'sum-aggregate' in ['mean_chunk', 'sum', 'ones', 'add', 'mean_agg-aggregate']

distributed\dashboard\tests\test_scheduler_bokeh.py:808: AssertionError
```


Discussion:
Haven't seen a failure in the past 30days. Closing
Happened again on Windows 3.10 on 07/13:
```
____________________________ test_compute_per_key _____________________________
c = <Client: No scheduler connected>
s = <Scheduler 'tcp://127.0.0.1:54512', workers: 0, cores: 0, tasks: 0>
a = <Worker 'tcp://127.0.0.1:54513', name: 0, status: closed, stored: 0, running: 0/1, ready: 0, comm: 0, waiting: 0>
b = <Worker 'tcp://127.0.0.1:54515', name: 1, status: closed, stored: 0, running: 0/2, ready: 0, comm: 0, waiting: 0>
@gen_cluster(client=True, scheduler_kwargs={"dashboard": True})
asyncdeftest_compute_per_key(c, s, a, b):
        mbk = ComputePerKey(s)
        da = pytest.importorskip("dask.array")
        x = (da.ones((20, 20), chunks=(10, 10)) + 1).persist(optimize_graph=False)
await x
        y = await dask.delayed(inc)(1).persist()
        z = (x + x.T) - x.mean(axis=0)
await c.compute(z.sum())
        mbk.update()
        http_client = AsyncHTTPClient()
        response = await http_client.fetch(
"http://localhost:%d/individual-compute-time-per-key" % s.http_server.port
        )
assert response.code == 200
>       assert ("sum-aggregate") in mbk.compute_source.data["names"]
E       AssertionError: assert 'sum-aggregate' in ['ones_like', 'sum', 'mean_chunk', 'add', 'mean_agg-aggregate']
distributed\dashboard\tests\test_scheduler_bokeh.py:1022: AssertionError
---------------------------- Captured stdout call -----------------------------
Dumped cluster state to test_cluster_dump\test_compute_per_key.yaml
---------------------------- Captured stderr call -----------------------------
2022-07-13 10:35:41,533 - distributed.scheduler - WARNING - Received heartbeat from unregistered worker 'tcp://127.0.0.1:54515'.
2022-07-13 10:35:41,535 - distributed.worker - ERROR - Scheduler was unaware of this worker 'tcp://127.0.0.1:54515'. Shutting down.
```
https://github.com/dask/distributed/runs/7318976477?check_suite_focus=true#step:11:1840

The `Received heartbeat from unregistered worker` is suspicious and probably related to the issue.

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
@@ -51,7 +51,7 @@
 from bokeh.themes import Theme
 from bokeh.transform import cumsum, factor_cmap, linear_cmap, stack
 from jinja2 import Environment, FileSystemLoader
-from tlz import curry, pipe, valmap
+from tlz import curry, pipe, second, valmap
 from tlz.curried import concat, groupby, map
 from tornado import escape
 
@@ -1523,42 +1523,44 @@ def __init__(self, scheduler, **kwargs):
     def update(self):
         compute_times = defaultdict(float)
 
-        for key, ts in self.scheduler.task_prefixes.items():
-            name = key_split(key)
-            for action, t in ts.all_durations.items():
+        for name, tp in self.scheduler.task_prefixes.items():
+            for action, t in tp.all_durations.items():
                 if action == "compute":
                     compute_times[name] += t
 
-        # order by largest time first
-        compute_times = sorted(compute_times.items(), key=lambda x: x[1], reverse=True)
-
-        # keep only time which are 2% of max or greater
-        if compute_times:
-            max_time = compute_times[0][1] * 0.02
-            compute_times = [(n, t) for n, t in compute_times if t > max_time]
-            compute_colors = list()
-            compute_names = list()
-            compute_time = list()
-            total_time = 0
-            for name, t in compute_times:
-                compute_names.append(name)
-                compute_colors.append(ts_color_of(name))
-                compute_time.append(t)
-                total_time += t
-
-            angles = [t / total_time * 2 * math.pi for t in compute_time]
-
-            self.fig.x_range.factors = compute_names
-
-            compute_result = dict(
-                angles=angles,
-                times=compute_time,
-                color=compute_colors,
-                names=compute_names,
-                formatted_time=[format_time(t) for t in compute_time],
-            )
+        if not compute_times:
+            return
 
-            update(self.compute_source, compute_result)
+        # order by largest time first
+        compute_times = sorted(compute_times.items(), key=second, reverse=True)
+
+        # Keep only times which are 2% of max or greater
+        max_time = compute_times[0][1] * 0.02
+        compute_colors = []
+        compute_names = []
+        compute_time = []
+        total_time = 0
+        for name, t in compute_times:
+            if t < max_time:
+                break
+            compute_names.append(name)
+            compute_colors.append(ts_color_of(name))
+            compute_time.append(t)
+            total_time += t
+
+        angles = [t / total_time * 2 * math.pi for t in compute_time]
+
+        self.fig.x_range.factors = compute_names
+
+        compute_result = dict(
+            angles=angles,
+            times=compute_time,
+            color=compute_colors,
+            names=compute_names,
+            formatted_time=[format_time(t) for t in compute_time],
+        )
+
+        update(self.compute_source, compute_result)
 
 
 class AggregateAction(DashboardComponent):


Test patch:
diff --git a/distributed/dashboard/tests/test_scheduler_bokeh.py b/distributed/dashboard/tests/test_scheduler_bokeh.py
--- a/distributed/dashboard/tests/test_scheduler_bokeh.py
+++ b/distributed/dashboard/tests/test_scheduler_bokeh.py
@@ -1286,7 +1286,6 @@ async def test_compute_per_key(c, s, a, b):
 
     da = pytest.importorskip("dask.array")
     x = (da.ones((20, 20), chunks=(10, 10)) + 1).persist(optimize_graph=False)
-
     await x
     y = await dask.delayed(inc)(1).persist()
     z = (x + x.T) - x.mean(axis=0)
@@ -1294,14 +1293,18 @@ async def test_compute_per_key(c, s, a, b):
     await c.compute(zsum)
 
     mbk.update()
+
+    # Keep only times which are 2% of max or greater.
+    # This means that the list of names is not stable (but max time is always preserved)
+    assert mbk.compute_source.data["names"]
+    assert set(mbk.compute_source.data["names"]).issubset(s.task_prefixes)
+    assert "angles" in mbk.compute_source.data
+
     http_client = AsyncHTTPClient()
     response = await http_client.fetch(
         "http://localhost:%d/individual-compute-time-per-key" % s.http_server.port
     )
     assert response.code == 200
-    assert ("sum-aggregate") in mbk.compute_source.data["names"]
-    assert ("add") in mbk.compute_source.data["names"]
-    assert "angles" in mbk.compute_source.data.keys()
 
 
 @gen_cluster(scheduler_kwargs={"http_prefix": "foo-bar", "dashboard": True})