Repository: dask/distributed

Issue:
Incorrect task counts in scheduler status page
Surprised this hasn't already been posted as I started noticing it at least a month ago:
```
c = distributed.Client()
f = c.submit(sum, [1, 2])
g = c.submit(len, [1, 2])
h = c.submit(sum, [3, 4])
```
Previously the plot would show 2 `sum`s and 1 `len`; now (on 2.9.1 and master, `distributed==2.9.1+10.g44d24098`) it shows 3 of each:
![image](https://user-images.githubusercontent.com/903655/72282921-91a3a480-35f2-11ea-97d9-6b2a79bcd146.png)
This behavior continues indefinitely, so with e.g. 1 `len` and 10000 `sums` we'd see 10001 of each for a total of 20002. Observed the same behavior w/ bokeh 1.0.4 and 1.4.0.


Discussion:

Resolving patch:
diff --git a/distributed/dashboard/scheduler.py b/distributed/dashboard/scheduler.py
--- a/distributed/dashboard/scheduler.py
+++ b/distributed/dashboard/scheduler.py
@@ -280,7 +280,11 @@ def collect(self):
 
         yield CounterMetricFamily(
             "dask_scheduler_tasks_forgotten",
-            "Total number of processed tasks no longer in memory and already removed from the scheduler job queue.",
+            (
+                "Total number of processed tasks no longer in memory and already "
+                "removed from the scheduler job queue. Note task groups on the "
+                "scheduler which have all tasks in the forgotten state are not included."
+            ),
             value=task_counter.get("forgotten", 0.0),
         )
 
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -833,7 +833,7 @@ def duration(self):
 
     @property
     def types(self):
-        return set.union(*[tg.types for tg in self.groups])
+        return set().union(*[tg.types for tg in self.groups])
 
 
 class _StateLegacyMapping(Mapping):
@@ -1956,20 +1956,21 @@ def new_task(self, key, spec, state):
         """ Create a new task, and associated states """
         ts = TaskState(key, spec)
         ts._state = state
-        try:
-            tg = self.task_groups[ts.group_key]
-        except KeyError:
-            tg = self.task_groups[ts.group_key] = TaskGroup(ts.group_key)
-        tg.add(ts)
         prefix_key = key_split(key)
         try:
             tp = self.task_prefixes[prefix_key]
         except KeyError:
-            tp = TaskPrefix(prefix_key)
-            tp.groups.append(tg)
-            self.task_prefixes[prefix_key] = tp
+            tp = self.task_prefixes[prefix_key] = TaskPrefix(prefix_key)
         ts.prefix = tp
-        tg.prefix = tp
+
+        group_key = ts.group_key
+        try:
+            tg = self.task_groups[group_key]
+        except KeyError:
+            tg = self.task_groups[group_key] = TaskGroup(group_key)
+            tg.prefix = tp
+            tp.groups.append(tg)
+        tg.add(ts)
         self.tasks[key] = ts
         return ts
 
@@ -4603,6 +4604,13 @@ def transition(self, key, finish, *args, **kwargs):
                 if ts.state == "forgotten":
                     del self.tasks[ts.key]
 
+            if ts.state == "forgotten":
+                # Remove TaskGroup if all tasks are in the forgotten state
+                tg = ts.group
+                if not any(tg.states.get(s) for s in ALL_TASK_STATES):
+                    ts.prefix.groups.remove(tg)
+                    del self.task_groups[tg.name]
+
             return recommendations
         except Exception as e:
             logger.exception("Error transitioning %r from %r to %r", key, start, finish)
diff --git a/distributed/utils.py b/distributed/utils.py
--- a/distributed/utils.py
+++ b/distributed/utils.py
@@ -637,11 +637,11 @@ def key_split_group(x):
         elif x[0] == "<":
             return x.strip("<>").split()[0].split(".")[-1]
         else:
-            return ""
+            return key_split(x)
     elif typ is bytes:
         return key_split_group(x.decode())
     else:
-        return ""
+        return key_split(x)
 
 
 @contextmanager


Test patch:
diff --git a/distributed/dashboard/tests/test_scheduler_bokeh_html.py b/distributed/dashboard/tests/test_scheduler_bokeh_html.py
--- a/distributed/dashboard/tests/test_scheduler_bokeh_html.py
+++ b/distributed/dashboard/tests/test_scheduler_bokeh_html.py
@@ -171,7 +171,7 @@ async def fetch_metrics():
     active_metrics, forgotten_tasks = await fetch_metrics()
     assert active_metrics.keys() == expected
     assert sum(active_metrics.values()) == 0.0
-    assert sum(forgotten_tasks) == 1.0
+    assert sum(forgotten_tasks) == 0.0
 
 
 @gen_cluster(
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -1740,6 +1740,8 @@ async def test_task_groups(c, s, a, b):
 
     await c.replicate(y)
     assert tg.nbytes_in_memory == y.nbytes
+    assert "array" in str(tg.types)
+    assert "array" in str(tp.types)
 
     del y
 
@@ -1748,8 +1750,9 @@ async def test_task_groups(c, s, a, b):
 
     assert tg.nbytes_in_memory == 0
     assert tg.states["forgotten"] == 5
-    assert "array" in str(tg.types)
-    assert "array" in str(tp.types)
+    # Ensure TaskGroup is removed once all tasks are in forgotten state
+    assert tg.name not in s.task_groups
+    assert sys.getrefcount(tg) == 2
 
 
 @gen_cluster(client=True)
@@ -1761,6 +1764,44 @@ async def test_task_prefix(c, s, a, b):
 
     assert s.task_prefixes["sum-aggregate"].states["memory"] == 1
 
+    a = da.arange(101, chunks=(20,))
+    b = (a + 1).sum().persist()
+    b = await b
+
+    assert s.task_prefixes["sum-aggregate"].states["memory"] == 2
+
+
+@gen_cluster(client=True)
+async def test_task_group_non_tuple_key(c, s, a, b):
+    da = pytest.importorskip("dask.array")
+    np = pytest.importorskip("numpy")
+    x = da.arange(100, chunks=(20,))
+    y = (x + 1).sum().persist()
+    y = await y
+
+    assert s.task_prefixes["sum"].states["released"] == 4
+    assert "sum" not in s.task_groups
+
+    f = c.submit(np.sum, [1, 2, 3])
+    await f
+
+    assert s.task_prefixes["sum"].states["released"] == 4
+    assert s.task_prefixes["sum"].states["memory"] == 1
+    assert "sum" in s.task_groups
+
+
+@gen_cluster(client=True)
+async def test_task_unique_groups(c, s, a, b):
+    """ This test ensure that task groups remain unique when using submit
+    """
+    x = c.submit(sum, [1, 2])
+    y = c.submit(len, [1, 2])
+    z = c.submit(sum, [3, 4])
+    await asyncio.wait([x, y, z])
+
+    assert s.task_prefixes["len"].states["memory"] == 1
+    assert s.task_prefixes["sum"].states["memory"] == 2
+
 
 class BrokenComm(Comm):
     peer_address = None