Repository: dask/distributed

Issue:
`resumed->rescheduled` is an invalid transition
```python3
def test_resumed_executing_task_releases_resources_on_reschedule(ws_with_running_task):
    ws = ws_with_running_task

    ws.handle_stimulus(FreeKeysEvent("cancel", ["x"]))
    assert ws.tasks["x"].state == "cancelled"
    assert ws.available_resources == {"R": 0}

    instructions = ws.handle_stimulus(
        ComputeTaskEvent.dummy(
            key="y",
            who_has={"x": ["127.0.0.1:1235"]},
            nbytes={"x": 8},
            stimulus_id="compute",
        )
    )
    assert ws.tasks["x"].state == "resumed"
    assert not instructions

    ws.handle_stimulus(
        RescheduleEvent(key="x", stimulus_id="reschedule")
    )
    assert ws.tasks["x"].state == "rescheduled"
    assert ws.available_resources == {"R": 1}
```

fails with `distributed.worker_state_machine.InvalidTransition: InvalidTransition: x :: resumed->rescheduled`. However, this might be a legitimate series of events to my understanding.


Discussion:
1. task x is started on w1
2. task x is cancelled
3. task x is started on w2 and terminates successfully
4. task y, which depends on x, is started on w1 -> task x is transitioned to resumed(fetch)
5. task x calls `raise Reschedule()`

There's no such thing as a `("resumed" "rescheduled")` transition.
The worker will kill itself off with `@fail_hard`, losing all data stored on itself.

This is tightly related to #6709.

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
@@ -924,6 +924,13 @@ def dummy(
 class RescheduleEvent(ExecuteDoneEvent):
     __slots__ = ()
 
+    @staticmethod
+    def dummy(key: str, *, stimulus_id: str) -> RescheduleEvent:
+        """Build an event. This method exists for compatibility with the other
+        ExecuteDoneEvent subclasses.
+        """
+        return RescheduleEvent(key=key, stimulus_id=stimulus_id)
+
 
 @dataclass
 class CancelComputeEvent(StateMachineEvent):
@@ -1962,6 +1969,24 @@ def _transition_resumed_error(
         ts.next = None
         return recs, []
 
+    def _transition_resumed_rescheduled(
+        self, ts: TaskState, *, stimulus_id: str
+    ) -> RecsInstrs:
+        """If the task raises the Reschedule() exception, but the scheduler already told
+        the worker to fetch it somewhere else, silently transition to fetch.
+
+        Note that this transition effectively duplicates the logic of
+        _transition_resumed_error.
+        """
+        assert ts.done
+        assert ts.previous in ("executing", "long-running")
+        assert ts.next == "fetch"
+        ts.state = "released"
+        ts.done = False
+        ts.previous = None
+        ts.next = None
+        return {ts: "fetch"}, []
+
     def _transition_resumed_fetch(
         self, ts: TaskState, *, stimulus_id: str
     ) -> RecsInstrs:
@@ -2297,19 +2322,20 @@ def _transition_released_forgotten(
     _TRANSITIONS_TABLE: ClassVar[
         Mapping[tuple[TaskStateState, TaskStateState], Callable[..., RecsInstrs]]
     ] = {
-        ("cancelled", "fetch"): _transition_cancelled_fetch,
         ("cancelled", "error"): _transition_cancelled_released,
+        ("cancelled", "fetch"): _transition_cancelled_fetch,
         ("cancelled", "memory"): _transition_cancelled_released,
         ("cancelled", "missing"): _transition_cancelled_released,
         ("cancelled", "released"): _transition_cancelled_released,
         ("cancelled", "rescheduled"): _transition_cancelled_released,
         ("cancelled", "waiting"): _transition_cancelled_waiting,
-        ("resumed", "memory"): _transition_resumed_memory,
         ("resumed", "error"): _transition_resumed_error,
-        ("resumed", "released"): _transition_resumed_released,
-        ("resumed", "waiting"): _transition_resumed_waiting,
         ("resumed", "fetch"): _transition_resumed_fetch,
+        ("resumed", "memory"): _transition_resumed_memory,
         ("resumed", "missing"): _transition_resumed_missing,
+        ("resumed", "released"): _transition_resumed_released,
+        ("resumed", "rescheduled"): _transition_resumed_rescheduled,
+        ("resumed", "waiting"): _transition_resumed_waiting,
         ("constrained", "executing"): _transition_constrained_executing,
         ("constrained", "released"): _transition_generic_released,
         ("error", "released"): _transition_generic_released,
@@ -2331,17 +2357,17 @@ def _transition_released_forgotten(
         ("long-running", "rescheduled"): _transition_executing_rescheduled,
         ("long-running", "released"): _transition_executing_released,
         ("memory", "released"): _transition_memory_released,
+        ("missing", "error"): _transition_generic_error,
         ("missing", "fetch"): _transition_missing_fetch,
         ("missing", "released"): _transition_missing_released,
-        ("missing", "error"): _transition_generic_error,
         ("missing", "waiting"): _transition_missing_waiting,
         ("ready", "executing"): _transition_ready_executing,
         ("ready", "released"): _transition_generic_released,
         ("released", "error"): _transition_generic_error,
         ("released", "fetch"): _transition_released_fetch,
-        ("released", "missing"): _transition_generic_missing,
         ("released", "forgotten"): _transition_released_forgotten,
         ("released", "memory"): _transition_released_memory,
+        ("released", "missing"): _transition_generic_missing,
         ("released", "waiting"): _transition_released_waiting,
         ("waiting", "constrained"): _transition_waiting_constrained,
         ("waiting", "ready"): _transition_waiting_ready,


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
@@ -30,6 +30,7 @@
     GatherDepFailureEvent,
     GatherDepNetworkFailureEvent,
     GatherDepSuccessEvent,
+    RescheduleEvent,
     TaskFinishedMsg,
     UpdateDataEvent,
 )
@@ -910,7 +911,9 @@ async def resume():
 
 
 @pytest.mark.parametrize("release_dep", [False, True])
-@pytest.mark.parametrize("done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent])
+@pytest.mark.parametrize(
+    "done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent, RescheduleEvent]
+)
 def test_cancel_with_dependencies_in_memory(ws, release_dep, done_ev_cls):
     """Cancel an executing task y with an in-memory dependency x; then simulate that x
     did not have any further dependents, so cancel x as well.
diff --git a/distributed/tests/test_reschedule.py b/distributed/tests/test_reschedule.py
--- a/distributed/tests/test_reschedule.py
+++ b/distributed/tests/test_reschedule.py
@@ -13,7 +13,9 @@
 from distributed import Event, Reschedule, get_worker, secede, wait
 from distributed.utils_test import captured_logger, gen_cluster, slowinc
 from distributed.worker_state_machine import (
+    ComputeTaskEvent,
     FreeKeysEvent,
+    GatherDep,
     RescheduleEvent,
     RescheduleMsg,
 )
@@ -129,3 +131,38 @@ def test_reschedule_releases(ws_with_running_task):
     assert instructions == [RescheduleMsg(stimulus_id="s1", key="x")]
     assert ws.available_resources == {"R": 1}
     assert "x" not in ws.tasks
+
+
+def test_reschedule_cancelled(ws_with_running_task):
+    """Test state loop:
+
+    executing -> cancelled -> rescheduled
+    executing -> long-running -> cancelled -> rescheduled
+    """
+    ws = ws_with_running_task
+    instructions = ws.handle_stimulus(
+        FreeKeysEvent(keys=["x"], stimulus_id="s1"),
+        RescheduleEvent(key="x", stimulus_id="s2"),
+    )
+    assert not instructions
+    assert "x" not in ws.tasks
+
+
+def test_reschedule_resumed(ws_with_running_task):
+    """Test state loop:
+
+    executing -> cancelled -> resumed(fetch) -> rescheduled
+    executing -> long-running -> cancelled -> resumed(fetch) -> rescheduled
+    """
+    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"),
+        RescheduleEvent(key="x", stimulus_id="s3"),
+    )
+    assert instructions == [
+        GatherDep(worker=ws2, to_gather={"x"}, total_nbytes=1, stimulus_id="s3")
+    ]
+    assert ws.tasks["x"].state == "flight"
diff --git a/distributed/tests/test_resources.py b/distributed/tests/test_resources.py
--- a/distributed/tests/test_resources.py
+++ b/distributed/tests/test_resources.py
@@ -17,6 +17,7 @@
     ExecuteFailureEvent,
     ExecuteSuccessEvent,
     FreeKeysEvent,
+    RescheduleEvent,
     TaskFinishedMsg,
 )
 
@@ -511,7 +512,9 @@ async def test_resources_from_python_override_config(c, s, a, b):
         assert info["workers"][worker.address]["resources"] == {"my_resources": 10}
 
 
-@pytest.mark.parametrize("done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent])
+@pytest.mark.parametrize(
+    "done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent, RescheduleEvent]
+)
 def test_cancelled_with_resources(ws_with_running_task, done_ev_cls):
     """Test transition loop of a task with resources:
 
@@ -528,7 +531,9 @@ def test_cancelled_with_resources(ws_with_running_task, done_ev_cls):
     assert "x" not in ws.tasks
 
 
-@pytest.mark.parametrize("done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent])
+@pytest.mark.parametrize(
+    "done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent, RescheduleEvent]
+)
 def test_resumed_with_resources(ws_with_running_task, done_ev_cls):
     """Test transition loop of a task with resources:
 
@@ -550,7 +555,9 @@ def test_resumed_with_resources(ws_with_running_task, done_ev_cls):
     assert ws.available_resources == {"R": 1}
 
 
-@pytest.mark.parametrize("done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent])
+@pytest.mark.parametrize(
+    "done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent, RescheduleEvent]
+)
 def test_resumed_with_different_resources(ws_with_running_task, done_ev_cls):
     """A task with resources is cancelled and then resumed to the same state, but with
     different resources. This is actually possible in case of manual cancellation from
diff --git a/distributed/tests/test_worker_state_machine.py b/distributed/tests/test_worker_state_machine.py
--- a/distributed/tests/test_worker_state_machine.py
+++ b/distributed/tests/test_worker_state_machine.py
@@ -1131,15 +1131,13 @@ def test_task_acquires_resources(ws, state):
 
 
 @pytest.mark.parametrize(
-    "done_ev_cls,done_status",
-    [(ExecuteSuccessEvent, "memory"), (ExecuteFailureEvent, "error")],
+    "done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent, RescheduleEvent]
 )
-def test_task_releases_resources(ws_with_running_task, done_ev_cls, done_status):
+def test_task_releases_resources(ws_with_running_task, done_ev_cls):
     ws = ws_with_running_task
     assert ws.available_resources == {"R": 0}
 
     ws.handle_stimulus(done_ev_cls.dummy("x", stimulus_id="success"))
-    assert ws.tasks["x"].state == done_status
     assert ws.available_resources == {"R": 1}
 
 
@@ -1170,12 +1168,9 @@ def test_task_with_dependencies_acquires_resources(ws):
 
 
 @pytest.mark.parametrize(
-    "done_ev_cls,done_status",
-    [(ExecuteSuccessEvent, "memory"), (ExecuteFailureEvent, "flight")],
+    "done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent, RescheduleEvent]
 )
-def test_resumed_task_releases_resources(
-    ws_with_running_task, done_ev_cls, done_status
-):
+def test_resumed_task_releases_resources(ws_with_running_task, done_ev_cls):
     ws = ws_with_running_task
     assert ws.available_resources == {"R": 0}
     ws2 = "127.0.0.1:2"
@@ -1192,7 +1187,6 @@ def test_resumed_task_releases_resources(
     assert ws.available_resources == {"R": 0}
 
     ws.handle_stimulus(done_ev_cls.dummy("x", stimulus_id="s2"))
-    assert ws.tasks["x"].state == done_status
     assert ws.available_resources == {"R": 1}
 
 
@@ -1221,28 +1215,21 @@ def test_running_task_in_all_running_tasks(ws_with_running_task):
 
 
 @pytest.mark.parametrize(
-    "done_ev_cls,done_status",
-    [(ExecuteSuccessEvent, "memory"), (ExecuteFailureEvent, "error")],
+    "done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent, RescheduleEvent]
 )
-def test_done_task_not_in_all_running_tasks(
-    ws_with_running_task, done_ev_cls, done_status
-):
+def test_done_task_not_in_all_running_tasks(ws_with_running_task, done_ev_cls):
     ws = ws_with_running_task
     ts = ws.tasks["x"]
     assert ts in ws.all_running_tasks
 
     ws.handle_stimulus(done_ev_cls.dummy("x", stimulus_id="s1"))
-    assert ts.state == done_status
     assert ts not in ws.all_running_tasks
 
 
 @pytest.mark.parametrize(
-    "done_ev_cls,done_status",
-    [(ExecuteSuccessEvent, "memory"), (ExecuteFailureEvent, "flight")],
+    "done_ev_cls", [ExecuteSuccessEvent, ExecuteFailureEvent, RescheduleEvent]
 )
-def test_done_resumed_task_not_in_all_running_tasks(
-    ws_with_running_task, done_ev_cls, done_status
-):
+def test_done_resumed_task_not_in_all_running_tasks(ws_with_running_task, done_ev_cls):
     ws = ws_with_running_task
     ws2 = "127.0.0.1:2"
 
@@ -1252,7 +1239,6 @@ def test_done_resumed_task_not_in_all_running_tasks(
         done_ev_cls.dummy("x", stimulus_id="s3"),
     )
     ts = ws.tasks["x"]
-    assert ts.state == done_status
     assert ts not in ws.all_running_tasks