Repository: dask/distributed

Issue:
Migrate `ensure_computing` transitions to new `WorkerState` event mechanism
The task executing transitions should be migrated to the `WorkerState` event mechanism as outlined in https://github.com/dask/distributed/issues/5736#issuecomment-1040599751

- The `Worker.execute` method is modified such that it no longer performs any transition but instead returns appropriate `StateMachineEvents` that trigger the necessary handlers. For instance
  - TaskFinished
  - Rescheduled
  - TaskErred
- The `ensure_computing` method is removed from the list of `handle_comm` `every_cycle` callbacks
- The logic of `ensure_computing` is moved to a private method (perspectively a method of the new class `WorkerState`)
- The new private method will not perform any transitions but rather return a list of recommendations
- This private method is then called as part of the transition system, e.g.
  - `unpause handler`
  - `transition_waiting_ready`
  - `transition_executing_*`
- The implicit goal is to remove all invocations of `self.loop.add_callback(self.execute, ...)` and replace this with the new callback method


### Blocked by 
- [x] https://github.com/dask/distributed/issues/5894

### Related work 
- [x] https://github.com/dask/distributed/issues/5891
- [x] https://github.com/dask/distributed/issues/5896



Discussion:

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -421,6 +421,7 @@ class Worker(ServerNode):
     execution_state: dict[str, Any]
     plugins: dict[str, WorkerPlugin]
     _pending_plugins: tuple[WorkerPlugin, ...]
+    _async_instructions: set[asyncio.Task]
 
     def __init__(
         self,
@@ -838,6 +839,8 @@ def __init__(
             self.lifetime += (random.random() * 2 - 1) * lifetime_stagger
             self.io_loop.call_later(self.lifetime, self.close_gracefully)
 
+        self._async_instructions = set()
+
         Worker._instances.add(self)
 
     ################
@@ -1406,6 +1409,19 @@ async def close(
                 logger.info("Closed worker has not yet started: %s", self.status)
             self.status = Status.closing
 
+            if self._async_instructions:
+                for task in self._async_instructions:
+                    task.cancel()
+                # async tasks can handle cancellation and could take an arbitrary amount
+                # of time to terminate
+                _, pending = await asyncio.wait(
+                    self._async_instructions, timeout=timeout
+                )
+                for task in pending:
+                    logger.error(
+                        f"Failed to cancel asyncio task after {timeout} seconds: {task}"
+                    )
+
             for preload in self.preloads:
                 await preload.teardown()
 
@@ -2603,12 +2619,16 @@ def handle_stimulus(self, stim: StateMachineEvent) -> None:
             self.transitions(recs, stimulus_id=stim.stimulus_id)
             self._handle_instructions(instructions)
 
-    def _handle_stimulus_from_future(
-        self, future: asyncio.Future[StateMachineEvent | None]
+    def _handle_stimulus_from_task(
+        self, task: asyncio.Task[StateMachineEvent | None]
     ) -> None:
+        self._async_instructions.remove(task)
         with log_errors():
-            # This *should* never raise
-            stim = future.result()
+            try:
+                # This *should* never raise any other exceptions
+                stim = task.result()
+            except asyncio.CancelledError:
+                return
         if stim:
             self.handle_stimulus(stim)
 
@@ -2619,10 +2639,12 @@ def _handle_instructions(self, instructions: Instructions) -> None:
             if isinstance(inst, SendMessageToScheduler):
                 self.batched_stream.send(inst.to_dict())
             elif isinstance(inst, Execute):
-                coro = self.execute(inst.key, stimulus_id=inst.stimulus_id)
-                task = asyncio.create_task(coro)
-                # TODO track task (at the moment it's fire-and-forget)
-                task.add_done_callback(self._handle_stimulus_from_future)
+                task = asyncio.create_task(
+                    self.execute(inst.key, stimulus_id=inst.stimulus_id),
+                    name=f"execute({inst.key})",
+                )
+                self._async_instructions.add(task)
+                task.add_done_callback(self._handle_stimulus_from_task)
             else:
                 raise TypeError(inst)  # pragma: nocover
 


Test patch:
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -26,6 +26,7 @@
 import distributed
 from distributed import (
     Client,
+    Event,
     Nanny,
     Reschedule,
     default_client,
@@ -1455,6 +1456,61 @@ async def test_close_gracefully(c, s, a, b):
         assert_amm_transfer_story(key, b, a)
 
 
+@pytest.mark.parametrize("sync", [False, pytest.param(True, marks=[pytest.mark.slow])])
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_close_while_executing(c, s, a, sync):
+    ev = Event()
+
+    if sync:
+
+        def f(ev):
+            ev.set()
+            sleep(2)
+
+    else:
+
+        async def f(ev):
+            await ev.set()
+            await asyncio.Event().wait()  # Block indefinitely
+
+    f1 = c.submit(f, ev, key="f1")
+    await ev.wait()
+    task = next(
+        task for task in asyncio.all_tasks() if "execute(f1)" in task.get_name()
+    )
+    await a.close()
+    assert task.cancelled()
+    assert s.tasks["f1"].state == "no-worker"
+
+
+@pytest.mark.slow
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_close_async_task_handles_cancellation(c, s, a):
+    ev = Event()
+
+    async def f(ev):
+        await ev.set()
+        try:
+            await asyncio.Event().wait()  # Block indefinitely
+        except asyncio.CancelledError:
+            await asyncio.Event().wait()  # Ignore the first cancel()
+
+    f1 = c.submit(f, ev, key="f1")
+    await ev.wait()
+    task = next(
+        task for task in asyncio.all_tasks() if "execute(f1)" in task.get_name()
+    )
+    start = time()
+    with captured_logger("distributed.worker", level=logging.ERROR) as logger:
+        await a.close(timeout=1)
+    assert "Failed to cancel asyncio task" in logger.getvalue()
+    assert time() - start < 5
+    assert not task.cancelled()
+    assert s.tasks["f1"].state == "no-worker"
+    task.cancel()
+    await asyncio.wait({task})
+
+
 @pytest.mark.slow
 @gen_cluster(client=True, nthreads=[("", 1)], timeout=10)
 async def test_lifetime(c, s, a):