Repository: dask/distributed

Issue:
Migrate `ensure_communicating` transitions to new WorkerState event mechanism
The `ensure_communicating` transitions should be migrated to the `WorkerState` event mechanism as outlined in https://github.com/dask/distributed/issues/5736#issuecomment-1040599751

- The `Worker._gather_dep` method is broken up into
  - A method `Worker._gather_remote_data` which accepts an instruction indicating what keys to fetch, the remote worker address, etc. It performs the actual communication and returns a list of events about the outcome of the remote fetch (e.g. Success, DataMissing, RemoteBusy, Error, ...)
  - Handlers, e.g. `Worker(State)._handle_gather_remote_busy` that deal with the result of the fetch attempt
- The `ensure_communicating` method is removed from the list of `handle_comm` `every_cycle` callbacks
- The logic of `ensure_communicating` 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 and/or instructions as necessary
- This private method is then called as part of the transition system, e.g.
  - `transition_flight_fetch`
  - `unpause`
- The implicit goal is to remove all invocations of `self.loop.add_callback(self.gather_dep, ...)` and replace this with the new callback method

###  Subtasks

- [x] https://github.com/dask/distributed/pull/6342
- [x] https://github.com/dask/distributed/pull/6348
- [x] https://github.com/dask/distributed/pull/6379
- [ ] https://github.com/dask/distributed/pull/6442
- [ ] https://github.com/dask/distributed/pull/6388


### 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/5895




Discussion:
Is cancellation of `gather_dep`s in scope for this? How does that fit in? Just curious.
> Is cancellation of `gather_dep`s in scope for this? How does that fit in? Just curious.

It's out of scope.
gather_dep will become a task that is explicitly tracked by the worker. It should be straightforward to cancel the task when that happens. Of course many other things need to happen (sender-side cancellation, ignore spurious incoming data, etc. etc.)

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -4726,7 +4726,9 @@ def release_worker_data(self, key: str, worker: str, stimulus_id: str) -> None:
         if not ts.who_has:
             self.transitions({key: "released"}, stimulus_id)
 
-    def handle_long_running(self, key=None, worker=None, compute_duration=None):
+    def handle_long_running(
+        self, key: str, worker: str, compute_duration: float, stimulus_id: str
+    ) -> None:
         """A task has seceded from the thread pool
 
         We stop the task from being stolen in the future, and change task
@@ -4735,27 +4737,23 @@ def handle_long_running(self, key=None, worker=None, compute_duration=None):
         if key not in self.tasks:
             logger.debug("Skipping long_running since key %s was already released", key)
             return
-        ts: TaskState = self.tasks[key]
+        ts = self.tasks[key]
         steal = self.extensions.get("stealing")
         if steal is not None:
             steal.remove_key_from_stealable(ts)
 
-        ws: WorkerState = ts.processing_on
+        ws = ts.processing_on
         if ws is None:
             logger.debug("Received long-running signal from duplicate task. Ignoring.")
             return
 
-        if compute_duration:
-            old_duration: float = ts.prefix.duration_average
-            new_duration: float = compute_duration
-            if old_duration < 0:
-                avg_duration = new_duration
-            else:
-                avg_duration = 0.5 * old_duration + 0.5 * new_duration
-
-            ts.prefix.duration_average = avg_duration
+        old_duration = ts.prefix.duration_average
+        if old_duration < 0:
+            ts.prefix.duration_average = compute_duration
+        else:
+            ts.prefix.duration_average = (old_duration + compute_duration) / 2
 
-        occ: float = ws.processing[ts]
+        occ = ws.processing[ts]
         ws.occupancy -= occ
         self.total_occupancy -= occ
         # Cannot remove from processing since we're using this for things like
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -2555,11 +2555,11 @@ def transition_executing_long_running(
         self._executing.discard(ts)
         self.long_running.add(ts.key)
 
+        smsg = LongRunningMsg(
+            key=ts.key, compute_duration=compute_duration, stimulus_id=stimulus_id
+        )
         return merge_recs_instructions(
-            (
-                {},
-                [LongRunningMsg(key=ts.key, compute_duration=compute_duration)],
-            ),
+            ({}, [smsg]),
             self._ensure_computing(),
         )
 
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
@@ -257,46 +257,47 @@ def __repr__(self) -> str:
         return f"<{type(self).__name__}: {len(self)} items>"
 
 
+@dataclass
 class Instruction:
     """Command from the worker state machine to the Worker, in response to an event"""
 
-    __slots__ = ()
+    __slots__ = ("stimulus_id",)
+    stimulus_id: str
 
 
 @dataclass
 class GatherDep(Instruction):
+    __slots__ = ("worker", "to_gather", "total_nbytes")
     worker: str
     to_gather: set[str]
     total_nbytes: int
-    stimulus_id: str
-    __slots__ = tuple(__annotations__)  # type: ignore
 
 
 @dataclass
 class Execute(Instruction):
-    __slots__ = ("key", "stimulus_id")
+    __slots__ = ("key",)
     key: str
-    stimulus_id: str
 
 
-class SendMessageToScheduler(Instruction):
+@dataclass
+class EnsureCommunicatingAfterTransitions(Instruction):
     __slots__ = ()
+
+
+@dataclass
+class SendMessageToScheduler(Instruction):
     #: Matches a key in Scheduler.stream_handlers
     op: ClassVar[str]
+    __slots__ = ()
 
     def to_dict(self) -> dict[str, Any]:
         """Convert object to dict so that it can be serialized with msgpack"""
         d = {k: getattr(self, k) for k in self.__annotations__}
         d["op"] = self.op
+        d["stimulus_id"] = self.stimulus_id
         return d
 
 
-@dataclass
-class EnsureCommunicatingAfterTransitions(Instruction):
-    __slots__ = ("stimulus_id",)
-    stimulus_id: str
-
-
 @dataclass
 class TaskFinishedMsg(SendMessageToScheduler):
     op = "task-finished"
@@ -308,7 +309,6 @@ class TaskFinishedMsg(SendMessageToScheduler):
     metadata: dict
     thread: int | None
     startstops: list[StartStop]
-    stimulus_id: str
     __slots__ = tuple(__annotations__)  # type: ignore
 
     def to_dict(self) -> dict[str, Any]:
@@ -328,7 +328,6 @@ class TaskErredMsg(SendMessageToScheduler):
     traceback_text: str
     thread: int | None
     startstops: list[StartStop]
-    stimulus_id: str
     __slots__ = tuple(__annotations__)  # type: ignore
 
     def to_dict(self) -> dict[str, Any]:
@@ -341,19 +340,17 @@ def to_dict(self) -> dict[str, Any]:
 class ReleaseWorkerDataMsg(SendMessageToScheduler):
     op = "release-worker-data"
 
-    __slots__ = ("key", "stimulus_id")
+    __slots__ = ("key",)
     key: str
-    stimulus_id: str
 
 
 @dataclass
 class MissingDataMsg(SendMessageToScheduler):
     op = "missing-data"
 
-    __slots__ = ("key", "errant_worker", "stimulus_id")
+    __slots__ = ("key", "errant_worker")
     key: str
     errant_worker: str
-    stimulus_id: str
 
 
 # Not to be confused with RescheduleEvent below or the distributed.Reschedule Exception
@@ -361,9 +358,8 @@ class MissingDataMsg(SendMessageToScheduler):
 class RescheduleMsg(SendMessageToScheduler):
     op = "reschedule"
 
-    __slots__ = ("key", "stimulus_id")
+    __slots__ = ("key",)
     key: str
-    stimulus_id: str
 
 
 @dataclass
@@ -379,9 +375,8 @@ class LongRunningMsg(SendMessageToScheduler):
 class AddKeysMsg(SendMessageToScheduler):
     op = "add-keys"
 
-    __slots__ = ("keys", "stimulus_id")
+    __slots__ = ("keys",)
     keys: list[str]
-    stimulus_id: str
 
 
 @dataclass


Test patch: