Send worker validation errors to scheduler and err on test completion (…

…#6192)

distributed/tests/test_utils_test.py

@@ -640,3 +640,24 @@ async def test_log_invalid_transitions(c, s, a):
 
     assert "foo" in out + err
     assert "task-name" in out + err
+
+
+def test_invalid_worker_states(capsys):
+    @gen_cluster(client=True, nthreads=[("127.0.0.1", 1)])
+    async def test_log_invalid_worker_task_states(c, s, a):
+        x = c.submit(inc, 1, key="task-name")
+        await x
+        a.tasks[x.key].state = "released"
+        with pytest.raises(Exception):
+            a.validate_task(a.tasks[x.key])
+
+        while not s.events["invalid-worker-task-states"]:
+            await asyncio.sleep(0.01)
+
+    with pytest.raises(Exception) as info:
+        test_log_invalid_worker_task_states()
+
+    out, err = capsys.readouterr()
+
+    assert "released" in out + err
+    assert "task-name" in out + err

distributed/utils_test.py

@@ -866,11 +866,12 @@ async def start_cluster(
             await asyncio.gather(*(w.close(timeout=1) for w in workers))
             await s.close(fast=True)
             check_invalid_worker_transitions(s)
+            check_invalid_task_states(s)
             raise TimeoutError("Cluster creation timeout")
     return s, workers
 
 
-def check_invalid_worker_transitions(s):
+def check_invalid_worker_transitions(s: Scheduler) -> None:
     if not s.events.get("invalid-worker-transition"):
         return
 
@@ -884,6 +885,19 @@ def check_invalid_worker_transitions(s):
     )
 
 
+def check_invalid_task_states(s: Scheduler) -> None:
+    if not s.events.get("invalid-worker-task-states"):
+        return
+
+    for timestamp, msg in s.events["invalid-worker-task-states"]:
+        print("Worker:", msg["worker"])
+        print("State:", msg["state"])
+        for line in msg["story"]:
+            print(line)
+
+    raise ValueError("Invalid worker task state")
+
+
 async def end_cluster(s, workers):
     logger.debug("Closing out test cluster")
 
@@ -895,6 +909,7 @@ async def end_worker(w):
     await s.close()  # wait until scheduler stops completely
     s.stop()
     check_invalid_worker_transitions(s)
+    check_invalid_task_states(s)
 
 
 def gen_cluster(

distributed/worker.py

@@ -4009,6 +4009,16 @@ def validate_task(self, ts):
 
                 pdb.set_trace()
 
+            self.log_event(
+                "invalid-worker-task-states",
+                {
+                    "key": ts.key,
+                    "state": ts.state,
+                    "story": self.story(ts),
+                    "worker": self.address,
+                },
+            )
+
             raise AssertionError(
                 f"Invalid TaskState encountered for {ts!r}.\nStory:\n{self.story(ts)}\n"
             ) from e