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]