declare ALL_TASK_STATES a set

dask · Jan 26, 2021 · 389bd83 · 389bd83
1 parent 496c361
commit 389bd83
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
@@ -182,7 +182,11 @@ def nogil(func):
     EventExtension,
 ]
 
-ALL_TASK_STATES = {"released", "waiting", "no-worker", "processing", "erred", "memory"}
+_ALL_TASK_STATES = declare(
+    set,
+    {"released", "waiting", "no-worker", "processing", "erred", "memory"},
+)
+ALL_TASK_STATES = _ALL_TASK_STATES
 
 
 @final
@@ -802,7 +806,7 @@ class TaskGroup:
     def __init__(self, name: str):
         self._name = name
         self._prefix = None
-        self._states = {state: 0 for state in ALL_TASK_STATES}
+        self._states = {state: 0 for state in _ALL_TASK_STATES}
         self._states["forgotten"] = 0
         self._dependencies = set()
         self._nbytes_total = 0
@@ -5967,7 +5971,7 @@ def _transition(self, key, finish: str, *args, **kwargs):
             if ts._state == "forgotten" and tg._name in parent._task_groups:
                 # Remove TaskGroup if all tasks are in the forgotten state
                 all_forgotten: bint = True
-                for s in ALL_TASK_STATES:
+                for s in _ALL_TASK_STATES:
                     if tg._states.get(s):
                         all_forgotten = False
                         break
@@ -6880,7 +6884,7 @@ def validate_task_state(ts: TaskState):
     ws: WorkerState
     dts: TaskState
 
-    assert ts._state in ALL_TASK_STATES or ts._state == "forgotten", ts
+    assert ts._state in _ALL_TASK_STATES or ts._state == "forgotten", ts
 
     if ts._waiting_on:
         assert ts._waiting_on.issubset(ts._dependencies), (