Add assertions to validate_state

dask · Jun 8, 2022 · 748dc48 · 748dc48
1 parent 8e73145
commit 748dc48
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/distributed/worker.py b/distributed/worker.py
@@ -4288,6 +4288,11 @@ def validate_state(self):
             if self.transition_counter_max:
                 assert self.transition_counter < self.transition_counter_max
 
+            # Test that there aren't multiple TaskState objects with the same key in data_needed
+            assert len({ts.key for ts in self.data_needed}) == len(self.data_needed)
+            for tss in self.data_needed_per_worker.values():
+                assert len({ts.key for ts in tss}) == len(tss)
+
         except Exception as e:
             logger.error("Validate state failed", exc_info=e)
             logger.exception(e)