Repository: dask/distributed

Issue:
Deadlock when a flaky task is stolen
1. Task x is running on worker A
2. x is stolen to worker B
3. x completes successfully on worker B
4. y, which depends on x, is scheduled on worker A
5. x eventually completes on worker A, raising an exception. To reiterate: x is a task that does not behave deterministically, and/or changes behaviour depending on which worker it's executed on.

#### Expected behaviour
Either of these is acceptable:
- the entire computation falls over, re-raising the exception on the client
-  worker A logs the incident, without informing the scheduler, and fetches the task from worker B

#### Actual behaviour
Cluster deadlock; corrupted worker state

#### Reproducer
```python
def test_worker_state_executing_failure_to_fetch(ws_with_running_task):
    ws = ws_with_running_task
    ws2 = "127.0.0.1:2"
    instructions = ws.handle_stimulus(
        FreeKeysEvent(keys=["x"], stimulus_id="s1"),
        ComputeTaskEvent.dummy(key="y", who_has={"x": [ws2]}, stimulus_id="s2"),
        ExecuteFailureEvent(
            key="x",
            start=0.0,
            stop=1.0,
            exception=Serialize(Exception()),
            traceback=None,
            exception_text="",
            traceback_text="",
            stimulus_id="s3",
        ),
    )
    assert instructions == [
        GatherDep(worker=ws2, to_gather={"x"}, total_nbytes=1, stimulus_id="s3")
    ]
    assert ws.tasks["x"].state == "flight"
```
Output:
```
Expected :[GatherDep(stimulus_id='s3', worker='127.0.0.1:2', to_gather={'x'}, total_nbytes=1)]
Actual   :[TaskErredMsg(stimulus_id='s3', key='x', ...)]

            for ts_wait in ts.waiting_for_data:
                assert ts_wait.key in self.tasks
>               assert (
                    ts_wait.state in READY | {"executing", "flight", "fetch", "missing"}
                    or ts_wait in self.missing_dep_flight
                    or ts_wait.who_has.issubset(self.in_flight_workers)
                ), (ts, ts_wait, self.story(ts), self.story(ts_wait))
E               AssertionError: (<TaskState 'y' waiting>, <TaskState 'x' error>, [('y', 'compute-task', 'released', 's2', 1657213114.0749152), ('y', 'released', 'waiting', 'waiting', {'x': 'fetch'}, 's2', 1657213114.0749364)], [('x', 'compute-task', 'released', 'compute', 1657213114.0746293), ('x', 'released', 'waiting', 'waiting', {'x': 'constrained'}, 'compute', 1657213114.0746458), ('x', 'waiting', 'constrained', 'constrained', {'x': 'executing'}, 'compute', 1657213114.0746567), ('x', 'constrained', 'executing', 'executing', {}, 'compute', 1657213114.074662), ('free-keys', ['x'], 's1', 1657213114.0748878), ('x', 'executing', 'released', 'cancelled', {}, 's1', 1657213114.0748944), ('x', 'ensure-task-exists', 'cancelled', 's2', 1657213114.0749264), ('x', 'cancelled', 'fetch', 'resumed', {}, 's2', 1657213114.0749395), ('x', 'resumed', 'error', 'error', {}, 's3', 1657213114.0749602)])

../worker_state_machine.py:3081: AssertionError
```



Discussion:
CC @hendrikmakait 

Resolving patch:
diff --git a/distributed/worker_state_machine.py b/distributed/worker_state_machine.py
--- a/distributed/worker_state_machine.py
+++ b/distributed/worker_state_machine.py
@@ -81,6 +81,18 @@
     "resumed",
 }
 READY: set[TaskStateState] = {"ready", "constrained"}
+# Valid states for a task that is found in TaskState.waiting_for_data
+WAITING_FOR_DATA: set[TaskStateState] = {
+    "constrained",
+    "executing",
+    "fetch",
+    "flight",
+    "long-running",
+    "missing",
+    "ready",
+    "resumed",
+    "waiting",
+}
 
 NO_VALUE = "--no-value-sentinel--"
 
@@ -3081,7 +3093,7 @@ def _validate_task_released(self, ts: TaskState) -> None:
     def validate_task(self, ts: TaskState) -> None:
         try:
             if ts.key in self.tasks:
-                assert self.tasks[ts.key] == ts
+                assert self.tasks[ts.key] is ts
             if ts.state == "memory":
                 self._validate_task_memory(ts)
             elif ts.state == "waiting":
@@ -3109,10 +3121,7 @@ def validate_task(self, ts: TaskState) -> None:
             ) from e
 
     def validate_state(self) -> None:
-        assert len(self.executing) >= 0
-        waiting_for_data_count = 0
         for ts in self.tasks.values():
-            assert ts.state is not None
             # check that worker has task
             for worker in ts.who_has:
                 assert worker != self.address
@@ -3125,22 +3134,18 @@ def validate_state(self) -> None:
                 # so we may have popped the key out of `self.tasks` but the
                 # dependency can still be in `memory` before GC grabs it...?
                 # Might need better bookkeeping
-                assert dep.state is not None
+                assert self.tasks[dep.key] is dep
                 assert ts in dep.dependents, ts
-            if ts.waiting_for_data:
-                waiting_for_data_count += 1
+
             for ts_wait in ts.waiting_for_data:
-                assert ts_wait.key in self.tasks
-                assert ts_wait.state in READY | {
-                    "executing",
-                    "long-running",
-                    "resumed",
-                    "flight",
-                    "fetch",
-                    "missing",
-                }, (ts, ts_wait, self.story(ts), self.story(ts_wait))
+                assert self.tasks[ts_wait.key] is ts_wait
+                assert ts_wait.state in WAITING_FOR_DATA, ts_wait
+
         # FIXME https://github.com/dask/distributed/issues/6319
-        # assert self.waiting_for_data_count == waiting_for_data_count
+        # assert self.waiting_for_data_count == sum(
+        #     bool(ts.waiting_for_data) for ts in self.tasks.values()
+        # )
+
         for worker, keys in self.has_what.items():
             assert worker != self.address
             for k in keys:


Test patch:
diff --git a/distributed/tests/test_cancelled_state.py b/distributed/tests/test_cancelled_state.py
--- a/distributed/tests/test_cancelled_state.py
+++ b/distributed/tests/test_cancelled_state.py
@@ -20,9 +20,13 @@
 from distributed.worker_state_machine import (
     ComputeTaskEvent,
     Execute,
+    ExecuteFailureEvent,
+    ExecuteSuccessEvent,
     FreeKeysEvent,
     GatherDep,
     GatherDepNetworkFailureEvent,
+    GatherDepSuccessEvent,
+    TaskFinishedMsg,
 )
 
 
@@ -544,11 +548,13 @@ async def release_all_futures():
         )
 
 
-def test_resume_executing_worker_state(ws_with_running_task):
+def test_workerstate_executing_to_executing(ws_with_running_task):
     """Test state loops:
 
-    - executing -> cancelled -> resumed -> executing
-    - executing -> long-running -> cancelled -> resumed -> long-running
+    - executing -> cancelled -> executing
+    - executing -> long-running -> cancelled -> long-running
+
+    Test that the task immediately reverts to its original state.
     """
     ws = ws_with_running_task
     ts = ws.tasks["x"]
@@ -561,3 +567,129 @@ def test_resume_executing_worker_state(ws_with_running_task):
     assert not instructions
     assert ws.tasks["x"] is ts
     assert ts.state == prev_state
+
+
+def test_workerstate_flight_to_flight(ws):
+    """Test state loop:
+
+    flight -> cancelled -> flight
+
+    Test that the task immediately reverts to its original state.
+    """
+    ws2 = "127.0.0.1:2"
+    instructions = ws.handle_stimulus(
+        ComputeTaskEvent.dummy("y", who_has={"x": [ws2]}, stimulus_id="s1"),
+        FreeKeysEvent(keys=["y", "x"], stimulus_id="s2"),
+        ComputeTaskEvent.dummy("z", who_has={"x": [ws2]}, stimulus_id="s3"),
+    )
+    assert instructions == [
+        GatherDep(worker=ws2, to_gather={"x"}, total_nbytes=1, stimulus_id="s1")
+    ]
+    assert ws.tasks["x"].state == "flight"
+
+
+def test_workerstate_executing_skips_fetch_on_success(ws_with_running_task):
+    """Test state loops:
+
+    - executing -> cancelled -> resumed (fetch) -> memory
+    - executing -> long-running -> cancelled -> resumed (fetch) -> memory
+
+    The task execution later terminates successfully.
+    Test that the task is never fetched and that dependents are unblocked.
+
+    See also: test_workerstate_executing_failure_to_fetch
+    """
+    ws = ws_with_running_task
+    ws2 = "127.0.0.1:2"
+    instructions = ws.handle_stimulus(
+        FreeKeysEvent(keys=["x"], stimulus_id="s1"),
+        ComputeTaskEvent.dummy("y", who_has={"x": [ws2]}, stimulus_id="s2"),
+        ExecuteSuccessEvent.dummy("x", 123, stimulus_id="s3"),
+    )
+    assert len(instructions) == 2
+    assert isinstance(instructions[0], TaskFinishedMsg)
+    assert instructions[0].key == "x"
+    assert instructions[0].stimulus_id == "s3"
+    assert instructions[1] == Execute(key="y", stimulus_id="s3")
+    assert ws.tasks["x"].state == "memory"
+    assert ws.data["x"] == 123
+
+
+@pytest.mark.xfail(reason="distributed#6565, distributed#6689")
+def test_workerstate_executing_failure_to_fetch(ws_with_running_task):
+    """Test state loops:
+
+    - executing -> cancelled -> resumed (fetch)
+    - executing -> long-running -> cancelled -> resumed (fetch)
+
+    The task execution later terminates with a failure.
+    This is an edge case interaction between work stealing and a task that does not
+    deterministically succeed or fail when run multiple times or on different workers.
+
+    Test that the task is fetched from the other worker. This is to avoid having to deal
+    with cancelling the dependent, which would require interaction with the scheduler
+    and increase the complexity of the use case.
+
+    See also: test_workerstate_executing_success_to_fetch
+    """
+    ws = ws_with_running_task
+    ws2 = "127.0.0.1:2"
+    instructions = ws.handle_stimulus(
+        FreeKeysEvent(keys=["x"], stimulus_id="s1"),
+        ComputeTaskEvent.dummy("y", who_has={"x": [ws2]}, stimulus_id="s2"),
+        ExecuteFailureEvent.dummy("x", stimulus_id="s3"),
+    )
+    assert instructions == [
+        GatherDep(worker=ws2, to_gather={"x"}, total_nbytes=1, stimulus_id="s3")
+    ]
+    assert ws.tasks["x"].state == "flight"
+
+
+def test_workerstate_flight_skips_executing_on_success(ws):
+    """Test state loop
+
+    flight -> cancelled -> resumed (waiting) -> memory
+
+    gather_dep later terminates successfully.
+    Test that the task is not executed and is in memory.
+    """
+    ws2 = "127.0.0.1:2"
+    instructions = ws.handle_stimulus(
+        ComputeTaskEvent.dummy("y", who_has={"x": [ws2]}, stimulus_id="s1"),
+        FreeKeysEvent(keys=["y", "x"], stimulus_id="s2"),
+        ComputeTaskEvent.dummy("x", stimulus_id="s3"),
+        GatherDepSuccessEvent(
+            worker=ws2, total_nbytes=1, data={"x": 123}, stimulus_id="s4"
+        ),
+    )
+    assert len(instructions) == 2
+    assert instructions[0] == GatherDep(
+        worker=ws2, to_gather={"x"}, total_nbytes=1, stimulus_id="s1"
+    )
+    assert isinstance(instructions[1], TaskFinishedMsg)
+    assert instructions[1].stimulus_id == "s4"
+    assert ws.tasks["x"].state == "memory"
+    assert ws.data["x"] == 123
+
+
+def test_workerstate_flight_failure_to_executing(ws):
+    """Test state loop
+
+    flight -> cancelled -> resumed (waiting)
+
+    gather_dep later terminates with a failure.
+    Test that the task is executed.
+    """
+    ws2 = "127.0.0.1:2"
+    instructions = ws.handle_stimulus(
+        ComputeTaskEvent.dummy("y", who_has={"x": [ws2]}, stimulus_id="s1"),
+        FreeKeysEvent(keys=["y", "x"], stimulus_id="s2"),
+        ComputeTaskEvent.dummy("x", stimulus_id="s3"),
+        # Peer worker does not have the data
+        GatherDepSuccessEvent(worker=ws2, total_nbytes=1, data={}, stimulus_id="s4"),
+    )
+    assert instructions == [
+        GatherDep(stimulus_id="s1", worker=ws2, to_gather={"x"}, total_nbytes=1),
+        Execute(stimulus_id="s4", key="x"),
+    ]
+    assert ws.tasks["x"].state == "executing"