From 89a84e5d64db0df5183c0ecc75a80c97bd3a9f02 Mon Sep 17 00:00:00 2001
From: John Kirkham <jakirkham@gmail.com>
Date: Sat, 23 Jan 2021 18:02:07 -0800
Subject: [PATCH] Intersect `TaskGroup` states and `ALL_TASK_STATES`

As `dict.values` supports intersections and `ALL_TASK_STATES` is a
`set`, we can just compute their intersection directly. This winds up
being significantly faster and simpler.
---
 distributed/scheduler.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 0d9080c79df..b87715c49b0 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -5908,7 +5908,7 @@ def transition(self, key, finish: str, *args, **kwargs):
             tg: TaskGroup = ts._group
             if ts._state == "forgotten" and tg._name in parent._task_groups:
                 # Remove TaskGroup if all tasks are in the forgotten state
-                if not any([tg._states.get(s) for s in ALL_TASK_STATES]):
+                if not (tg._states.values() & ALL_TASK_STATES):
                     ts._prefix._groups.remove(tg)
                     del parent._task_groups[tg._name]