tweaks

Author: crusaderky

distributed/tests/test_worker_state_machine.py

@@ -475,9 +475,9 @@ async def test_forget_data_needed(c, s, a, b):
        popped from data_needed
     3. Task is forgotten
     4. Task is recreated from scratch and transitioned to fetch again
-    5. BUG: at the moment of writing this test, data_needed.push silently did nothing,
-       because it still contained the forgotten task, which is a different TaskState
-       instance which will be no longer updated.
+    5. BUG: at the moment of writing this test, adding to data_needed silently did
+       nothing, because it still contained the forgotten task, which is a different
+       TaskState instance which will be no longer updated.
     6. _ensure_communicating runs. It pops the forgotten task and discards it.
     7. We now have a task stuck in fetch state.
     """

distributed/worker.py

@@ -2998,13 +2998,14 @@ def _ensure_communicating(self, *, stimulus_id: str) -> RecsInstrs:
 
             ts = self.data_needed.pop()
 
-            if not ts.who_has:
-                recommendations[ts] = "missing"
-                continue
             if self.validate:
                 assert ts.state == "fetch"
                 assert self.address not in ts.who_has
 
+            if not ts.who_has:
+                recommendations[ts] = "missing"
+                continue
+
             workers = [
                 w
                 for w in ts.who_has
@@ -3315,7 +3316,7 @@ def done_event():
             )
             for d in has_what:
                 ts = self.tasks[d]
-                ts.who_has.discard(worker)
+                ts.who_has.remove(worker)
                 if not ts.who_has and ts.state in (
                     "fetch",
                     "flight",
@@ -4639,7 +4640,7 @@ async def _get_data():
 cache_loads = LRU(maxsize=100)
 
 
-def loads_function(bytes_object: bytes) -> Callable:
+def loads_function(bytes_object):
     """Load a function from bytes, cache bytes"""
     if len(bytes_object) < 100000:
         try:
@@ -4691,7 +4692,7 @@ def execute_task(task):
 _cache_lock = threading.Lock()
 
 
-def dumps_function(func: Callable) -> bytes:
+def dumps_function(func) -> bytes:
     """Dump a function to bytes, cache functions"""
     try:
         with _cache_lock: