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
- [ ] https://github.com/dask/distributed/issues/5896


Discussion:
github-actions[bot]: ## Unit Test Results
       12 files  +       1         12 suites  +1   6h 45m 2s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "duration of all tests") + 58m 38s
  2 675 tests +       2    2 590 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "passed tests")  -        1    82 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "skipped / disabled tests") ±  0  3 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "failed tests") +3 
15 954 runs  +1 495  15 085 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "passed tests") +1 415  861 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "skipped / disabled tests") +72  8 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "failed tests") +8 

For more details on these failures, see [this check](https://github.com/dask/distributed/runs/5723736961).

Results for commit 5dcee7e7. ± Comparison against base commit 06170d56.

crusaderky: @fjetter I did the work on _ensure_computing (see https://github.com/dask/distributed/pull/6003/commits/9d28bdeca23ab47776fffbcbb0a093e9b77801a2) but there are still some pretty conceptual failures caused by resource counting, so I'm moderately inclined to merge this PR as it is. I reverted it for now.
All tests should be green.

### In scope for this PR
- Remove all ``self.loop.add_callback(self.execute, ...)``
### Out of scope for the PR but in scope for the issue (to be implemented immediately after this PR)
- Refactor ``Worker.ensure_computing() -> None`` into ``Worker._ensure_computing() -> RecsInstr``
- Implement ``Worker.stimulus_history``
- Track tasks spawned by ``Worker._handle_instructions`` and cancel them in ``Worker.close``
### Out of scope for the issue
- Send Event Python objects from the Scheduler
- Other events (replicate etc.)

Patch:
diff --git a/distributed/cli/dask_worker.py b/distributed/cli/dask_worker.py
index f13672a4..376f2a1c 100755
--- a/distributed/cli/dask_worker.py
+++ b/distributed/cli/dask_worker.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import atexit
 import gc
@@ -6,7 +8,9 @@ import os
 import signal
 import sys
 import warnings
+from collections.abc import Iterator
 from contextlib import suppress
+from typing import Any
 
 import click
 from tlz import valmap
@@ -24,7 +28,7 @@ from distributed.proctitle import (
     enable_proctitle_on_children,
     enable_proctitle_on_current,
 )
-from distributed.utils import import_term
+from distributed.utils import import_term, parse_ports
 
 logger = logging.getLogger("distributed.dask_worker")
 
@@ -254,10 +258,10 @@ pem_file_option_type = click.Path(exists=True, resolve_path=True)
 def main(
     scheduler,
     host,
-    worker_port,
+    worker_port: str | None,
     listen_address,
     contact_address,
-    nanny_port,
+    nanny_port: str | None,
     nthreads,
     nprocs,
     n_workers,
@@ -364,7 +368,8 @@ def main(
 
     try:
         if listen_address:
-            (host, worker_port) = get_address_host_port(listen_address, strict=True)
+            host, _ = get_address_host_port(listen_address, strict=True)
+            worker_port = str(_)
             if ":" in host:
                 # IPv6 -- bracket to pass as user args
                 host = f"[{host}]"
@@ -379,11 +384,6 @@ def main(
         logger.error("Failed to launch worker. " + str(e))
         sys.exit(1)
 
-    if nanny:
-        port = nanny_port
-    else:
-        port = worker_port
-
     if not nthreads:
         nthreads = CPU_COUNT // n_workers
 
@@ -407,16 +407,16 @@ def main(
     loop = IOLoop.current()
 
     worker_class = import_term(worker_class)
+
+    port_kwargs = _apportion_ports(worker_port, nanny_port, n_workers, nanny)
+    assert len(port_kwargs) == n_workers
+
     if nanny:
         kwargs["worker_class"] = worker_class
         kwargs["preload_nanny"] = preload_nanny
-
-    if nanny:
-        kwargs.update({"worker_port": worker_port, "listen_address": listen_address})
+        kwargs["listen_address"] = listen_address
         t = Nanny
     else:
-        if nanny_port:
-            kwargs["service_ports"] = {"nanny": nanny_port}
         t = worker_class
 
     if (
@@ -442,15 +442,15 @@ def main(
             security=sec,
             contact_address=contact_address,
             host=host,
-            port=port,
             dashboard=dashboard,
             dashboard_address=dashboard_address,
             name=name
             if n_workers == 1 or name is None or name == ""
             else str(name) + "-" + str(i),
             **kwargs,
+            **port_kwargs_i,
         )
-        for i in range(n_workers)
+        for i, port_kwargs_i in enumerate(port_kwargs)
     ]
 
     async def close_all():
@@ -486,5 +486,88 @@ def main(
         logger.info("End worker")
 
 
+def _apportion_ports(
+    worker_port: str | None, nanny_port: str | None, n_workers: int, nanny: bool
+) -> list[dict[str, Any]]:
+    """Spread out evenly --worker-port and/or --nanny-port ranges to the workers and
+    nannies, avoiding overlap.
+
+    Returns
+    =======
+    List of kwargs to pass to the Worker or Nanny construtors
+    """
+    seen = set()
+
+    def parse_unique(s: str | None) -> Iterator[int | None]:
+        ports = parse_ports(s)
+        if ports in ([0], [None]):
+            for _ in range(n_workers):
+                yield ports[0]
+        else:
+            for port in ports:
+                if port not in seen:
+                    seen.add(port)
+                    yield port
+
+    worker_ports_iter = parse_unique(worker_port)
+    nanny_ports_iter = parse_unique(nanny_port)
+
+    # [(worker ports, nanny ports), ...]
+    ports: list[tuple[set[int | None], set[int | None]]] = [
+        (set(), set()) for _ in range(n_workers)
+    ]
+
+    ports_iter = iter(ports)
+    more_wps = True
+    more_nps = True
+    while more_wps or more_nps:
+        try:
+            worker_ports_i, nanny_ports_i = next(ports_iter)
+        except StopIteration:
+            # Start again in round-robin
+            ports_iter = iter(ports)
+            continue
+
+        try:
+            worker_ports_i.add(next(worker_ports_iter))
+        except StopIteration:
+            more_wps = False
+        try:
+            nanny_ports_i.add(next(nanny_ports_iter))
+        except StopIteration:
+            more_nps = False
+
+    kwargs = []
+    for worker_ports_i, nanny_ports_i in ports:
+        if not worker_ports_i or not nanny_ports_i:
+            if nanny:
+                raise ValueError(
+                    f"Not enough ports in range --worker_port {worker_port} "
+                    f"--nanny_port {nanny_port} for {n_workers} workers"
+                )
+            else:
+                raise ValueError(
+                    f"Not enough ports in range --worker_port {worker_port} "
+                    f"for {n_workers} workers"
+                )
+
+        # None and int can't be sorted together,
+        # but None and 0 are guaranteed to be alone
+        wp: Any = sorted(worker_ports_i)
+        if len(wp) == 1:
+            wp = wp[0]
+        if nanny:
+            np: Any = sorted(nanny_ports_i)
+            if len(np) == 1:
+                np = np[0]
+            kwargs_i = {"port": np, "worker_port": wp}
+        else:
+            kwargs_i = {"port": wp}
+
+        kwargs.append(kwargs_i)
+
+    return kwargs
+
+
 if __name__ == "__main__":
     main()  # pragma: no cover
diff --git a/distributed/client.py b/distributed/client.py
index 3bfd613f..86ca5f50 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -2002,8 +2002,12 @@ class Client(SyncMethodMixin):
 
         async def wait(k):
             """Want to stop the All(...) early if we find an error"""
-            st = self.futures[k]
-            await st.wait()
+            try:
+                st = self.futures[k]
+            except KeyError:
+                raise AllExit()
+            else:
+                await st.wait()
             if st.status != "finished" and errors == "raise":
                 raise AllExit()
 
diff --git a/distributed/core.py b/distributed/core.py
index 2e00c970..5d968cb4 100644
--- a/distributed/core.py
+++ b/distributed/core.py
@@ -597,7 +597,7 @@ class Server:
                         "Failed while closing connection to %r: %s", address, e
                     )
 
-    async def handle_stream(self, comm, extra=None, every_cycle=[]):
+    async def handle_stream(self, comm, extra=None, every_cycle=()):
         extra = extra or {}
         logger.info("Starting established connection")
 
diff --git a/distributed/deploy/cluster.py b/distributed/deploy/cluster.py
index 047357cc..1f7feb91 100644
--- a/distributed/deploy/cluster.py
+++ b/distributed/deploy/cluster.py
@@ -98,6 +98,7 @@ class Cluster(SyncMethodMixin):
 
     async def _start(self):
         comm = await self.scheduler_comm.live_comm()
+        comm.name = "Cluster worker status"
         await comm.write({"op": "subscribe_worker_status"})
         self.scheduler_info = SchedulerInfo(await comm.read())
         self._watch_worker_status_comm = comm
diff --git a/distributed/nanny.py b/distributed/nanny.py
index e946655f..db237152 100644
--- a/distributed/nanny.py
+++ b/distributed/nanny.py
@@ -9,6 +9,7 @@ import threading
 import uuid
 import warnings
 import weakref
+from collections.abc import Collection
 from contextlib import suppress
 from inspect import isawaitable
 from queue import Empty
@@ -24,7 +25,7 @@ from dask.system import CPU_COUNT
 from dask.utils import parse_timedelta
 
 from distributed import preloading
-from distributed.comm import get_address_host, unparse_host_port
+from distributed.comm import get_address_host
 from distributed.comm.addressing import address_from_user_args
 from distributed.core import (
     CommClosedError,
@@ -96,12 +97,16 @@ class Nanny(ServerNode):
     status = Status.undefined
     memory_manager: NannyMemoryManager
 
+    # Inputs to parse_ports()
+    _given_worker_port: int | str | Collection[int] | None
+    _start_port: int | str | Collection[int] | None
+
     def __init__(
         self,
         scheduler_ip=None,
         scheduler_port=None,
         scheduler_file=None,
-        worker_port=0,
+        worker_port: int | str | Collection[int] | None = 0,
         nthreads=None,
         loop=None,
         local_dir=None,
@@ -126,7 +131,7 @@ class Nanny(ServerNode):
         env=None,
         interface=None,
         host=None,
-        port=None,
+        port: int | str | Collection[int] | None = None,
         protocol=None,
         config=None,
         **worker_kwargs,
@@ -373,14 +378,6 @@ class Nanny(ServerNode):
 
         Blocks until the process is up and the scheduler is properly informed
         """
-        if self._listen_address:
-            start_arg = self._listen_address
-        else:
-            host = self.listener.bound_address[0]
-            start_arg = self.listener.prefix + unparse_host_port(
-                host, self._given_worker_port
-            )
-
         if self.process is None:
             worker_kwargs = dict(
                 scheduler_ip=self.scheduler_addr,
@@ -403,7 +400,6 @@ class Nanny(ServerNode):
             worker_kwargs.update(self.worker_kwargs)
             self.process = WorkerProcess(
                 worker_kwargs=worker_kwargs,
-                worker_start_args=(start_arg,),
                 silence_logs=self.silence_logs,
                 on_exit=self._on_exit_sync,
                 worker=self.Worker,
@@ -617,7 +613,6 @@ class WorkerProcess:
     def __init__(
         self,
         worker_kwargs,
-        worker_start_args,
         silence_logs,
         on_exit,
         worker,
@@ -627,7 +622,6 @@ class WorkerProcess:
         self.status = Status.init
         self.silence_logs = silence_logs
         self.worker_kwargs = worker_kwargs
-        self.worker_start_args = worker_start_args
         self.on_exit = on_exit
         self.process = None
         self.Worker = worker
@@ -658,7 +652,6 @@ class WorkerProcess:
             name="Dask Worker process (from Nanny)",
             kwargs=dict(
                 worker_kwargs=self.worker_kwargs,
-                worker_start_args=self.worker_start_args,
                 silence_logs=self.silence_logs,
                 init_result_q=self.init_result_q,
                 child_stop_q=self.child_stop_q,
@@ -808,7 +801,6 @@ class WorkerProcess:
     def _run(
         cls,
         worker_kwargs,
-        worker_start_args,
         silence_logs,
         init_result_q,
         child_stop_q,
diff --git a/distributed/utils.py b/distributed/utils.py
index 64203488..cd0d93d0 100644
--- a/distributed/utils.py
+++ b/distributed/utils.py
@@ -20,7 +20,7 @@ import weakref
 import xml.etree.ElementTree
 from asyncio import TimeoutError
 from collections import OrderedDict, UserDict, deque
-from collections.abc import Container, KeysView, ValuesView
+from collections.abc import Collection, Container, KeysView, ValuesView
 from concurrent.futures import CancelledError, ThreadPoolExecutor  # noqa: F401
 from contextlib import contextmanager, suppress
 from contextvars import ContextVar
@@ -1148,15 +1148,15 @@ def format_dashboard_link(host, port):
     )
 
 
-def parse_ports(port):
+def parse_ports(port: int | str | Collection[int] | None) -> list[int] | list[None]:
     """Parse input port information into list of ports
 
     Parameters
     ----------
-    port : int, str, None
+    port : int, str, list[int], None
         Input port or ports. Can be an integer like 8787, a string for a
-        single port like "8787", a string for a sequential range of ports like
-        "8000:8200", or None.
+        single port like "8787", string for a sequential range of ports like
+        "8000:8200", a collection of ints, or None.
 
     Returns
     -------
@@ -1188,12 +1188,7 @@ def parse_ports(port):
     [None]
 
     """
-    if isinstance(port, str) and ":" not in port:
-        port = int(port)
-
-    if isinstance(port, (int, type(None))):
-        ports = [port]
-    else:
+    if isinstance(port, str) and ":" in port:
         port_start, port_stop = map(int, port.split(":"))
         if port_stop <= port_start:
             raise ValueError(
@@ -1201,9 +1196,20 @@ def parse_ports(port):
                 "port_stop must be greater than port_start, but got "
                 f"{port_start=} and {port_stop=}"
             )
-        ports = list(range(port_start, port_stop + 1))
+        return list(range(port_start, port_stop + 1))
+
+    if isinstance(port, str):
+        return [int(port)]
+
+    if isinstance(port, int) or port is None:
+        return [port]  # type: ignore
+
+    if isinstance(port, Collection):
+        if not all(isinstance(p, int) for p in port):
+            raise TypeError(port)
+        return list(port)  # type: ignore
 
-    return ports
+    raise TypeError(port)
 
 
 is_coroutine_function = iscoroutinefunction
diff --git a/distributed/worker.py b/distributed/worker.py
index fa107e58..7a062876 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -4,6 +4,7 @@ import asyncio
 import bisect
 import builtins
 import errno
+import functools
 import heapq
 import logging
 import os
@@ -103,36 +104,38 @@ from distributed.worker_memory import (
     DeprecatedMemoryMonitor,
     WorkerMemoryManager,
 )
-from distributed.worker_state_machine import Instruction  # noqa: F401
 from distributed.worker_state_machine import (
     PROCESSING,
     READY,
     AddKeysMsg,
+    AlreadyCancelledEvent,
+    CancelComputeEvent,
+    Execute,
+    ExecuteFailureEvent,
+    ExecuteSuccessEvent,
+    Instructions,
     InvalidTransition,
     LongRunningMsg,
+    Recs,
+    RecsInstrs,
     ReleaseWorkerDataMsg,
+    RescheduleEvent,
     RescheduleMsg,
     SendMessageToScheduler,
     SerializedTask,
+    StateMachineEvent,
     TaskErredMsg,
     TaskFinishedMsg,
     TaskState,
+    TaskStateState,
     UniqueTaskHeap,
 )
 
 if TYPE_CHECKING:
-    # TODO move to typing (requires Python >=3.10)
-    from typing_extensions import TypeAlias
-
     from distributed.actor import Actor
     from distributed.client import Client
     from distributed.diagnostics.plugin import WorkerPlugin
     from distributed.nanny import Nanny
-    from distributed.worker_state_machine import TaskStateState
-
-    # {TaskState -> finish: TaskStateState | (finish: TaskStateState, transition *args)}
-    Recs: TypeAlias = "dict[TaskState, TaskStateState | tuple]"
-    Instructions: TypeAlias = "list[Instruction]"
 
 
 logger = logging.getLogger(__name__)
@@ -380,7 +383,7 @@ class Worker(ServerNode):
     bandwidth_types: defaultdict[type, tuple[float, int]]
     preloads: list[preloading.Preload]
     contact_address: str | None
-    _start_port: int | None
+    _start_port: int | str | Collection[int] | None = None
     _start_host: str | None
     _interface: str | None
     _protocol: str
@@ -446,7 +449,7 @@ class Worker(ServerNode):
         ] = DEFAULT_STARTUP_INFORMATION,
         interface: str | None = None,
         host: str | None = None,
-        port: int | None = None,
+        port: int | str | Collection[int] | None = None,
         protocol: str | None = None,
         dashboard_address: str | None = None,
         dashboard: bool = False,
@@ -1770,14 +1773,7 @@ class Worker(ServerNode):
         is in state `waiting` or `ready`.
         Nothing will happen otherwise.
         """
-        ts = self.tasks.get(key)
-        if ts and ts.state in READY | {"waiting"}:
-            self.log.append((key, "cancel-compute", stimulus_id, time()))
-            # All possible dependents of TS should not be in state Processing on
-            # scheduler side and therefore should not be assigned to a worker,
-            # yet.
-            assert not ts.dependents
-            self.transition(ts, "released", stimulus_id=stimulus_id)
+        self.handle_stimulus(CancelComputeEvent(key=key, stimulus_id=stimulus_id))
 
     def handle_acquire_replicas(
         self,
@@ -1904,7 +1900,7 @@ class Worker(ServerNode):
 
     def transition_missing_fetch(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if self.validate:
             assert ts.state == "missing"
             assert ts.priority is not None
@@ -1917,7 +1913,7 @@ class Worker(ServerNode):
 
     def transition_missing_released(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         self._missing_dep_flight.discard(ts)
         recs, instructions = self.transition_generic_released(
             ts, stimulus_id=stimulus_id
@@ -1927,7 +1923,7 @@ class Worker(ServerNode):
 
     def transition_flight_missing(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         assert ts.done
         ts.state = "missing"
         self._missing_dep_flight.add(ts)
@@ -1936,7 +1932,7 @@ class Worker(ServerNode):
 
     def transition_released_fetch(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if self.validate:
             assert ts.state == "released"
             assert ts.priority is not None
@@ -1949,7 +1945,7 @@ class Worker(ServerNode):
 
     def transition_generic_released(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         self.release_key(ts.key, stimulus_id=stimulus_id)
         recs: Recs = {}
         for dependency in ts.dependencies:
@@ -1966,7 +1962,7 @@ class Worker(ServerNode):
 
     def transition_released_waiting(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if self.validate:
             assert ts.state == "released"
             assert all(d.key in self.tasks for d in ts.dependencies)
@@ -1992,7 +1988,7 @@ class Worker(ServerNode):
 
     def transition_fetch_flight(
         self, ts: TaskState, worker, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if self.validate:
             assert ts.state == "fetch"
             assert ts.who_has
@@ -2005,7 +2001,7 @@ class Worker(ServerNode):
 
     def transition_memory_released(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         recs, instructions = self.transition_generic_released(
             ts, stimulus_id=stimulus_id
         )
@@ -2014,7 +2010,7 @@ class Worker(ServerNode):
 
     def transition_waiting_constrained(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if self.validate:
             assert ts.state == "waiting"
             assert not ts.waiting_for_data
@@ -2030,14 +2026,14 @@ class Worker(ServerNode):
 
     def transition_long_running_rescheduled(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         recs: Recs = {ts: "released"}
         smsg = RescheduleMsg(key=ts.key, worker=self.address)
         return recs, [smsg]
 
     def transition_executing_rescheduled(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         for resource, quantity in ts.resource_restrictions.items():
             self.available_resources[resource] += quantity
         self._executing.discard(ts)
@@ -2048,7 +2044,7 @@ class Worker(ServerNode):
 
     def transition_waiting_ready(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if self.validate:
             assert ts.state == "waiting"
             assert ts.key not in self.ready
@@ -2072,7 +2068,7 @@ class Worker(ServerNode):
         traceback_text: str,
         *,
         stimulus_id: str,
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         recs: Recs = {}
         instructions: Instructions = []
         if ts._previous == "executing":
@@ -2106,7 +2102,7 @@ class Worker(ServerNode):
         traceback_text: str,
         *,
         stimulus_id: str,
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         ts.exception = exception
         ts.traceback = traceback
         ts.exception_text = exception_text
@@ -2133,7 +2129,7 @@ class Worker(ServerNode):
         traceback_text: str,
         *,
         stimulus_id: str,
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         for resource, quantity in ts.resource_restrictions.items():
             self.available_resources[resource] += quantity
         self._executing.discard(ts)
@@ -2148,7 +2144,7 @@ class Worker(ServerNode):
 
     def _transition_from_resumed(
         self, ts: TaskState, finish: TaskStateState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         """`resumed` is an intermediate degenerate state which splits further up
         into two states depending on what the last signal / next state is
         intended to be. There are only two viable choices depending on whether
@@ -2186,7 +2182,7 @@ class Worker(ServerNode):
 
     def transition_resumed_fetch(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         """
         See Worker._transition_from_resumed
         """
@@ -2194,7 +2190,7 @@ class Worker(ServerNode):
 
     def transition_resumed_missing(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         """
         See Worker._transition_from_resumed
         """
@@ -2208,7 +2204,7 @@ class Worker(ServerNode):
 
     def transition_cancelled_fetch(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if ts.done:
             return {ts: "released"}, []
         elif ts._previous == "flight":
@@ -2220,14 +2216,14 @@ class Worker(ServerNode):
 
     def transition_cancelled_resumed(
         self, ts: TaskState, next: TaskStateState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         ts._next = next
         ts.state = "resumed"
         return {}, []
 
     def transition_cancelled_waiting(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if ts.done:
             return {ts: "released"}, []
         elif ts._previous == "executing":
@@ -2239,7 +2235,7 @@ class Worker(ServerNode):
 
     def transition_cancelled_forgotten(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         ts._next = "forgotten"
         if not ts.done:
             return {}, []
@@ -2247,7 +2243,7 @@ class Worker(ServerNode):
 
     def transition_cancelled_released(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if not ts.done:
             ts._next = "released"
             return {}, []
@@ -2267,7 +2263,7 @@ class Worker(ServerNode):
 
     def transition_executing_released(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         ts._previous = ts.state
         ts._next = "released"
         # See https://github.com/dask/distributed/pull/5046#discussion_r685093940
@@ -2277,13 +2273,13 @@ class Worker(ServerNode):
 
     def transition_long_running_memory(
         self, ts: TaskState, value=no_value, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         self.executed_count += 1
         return self.transition_generic_memory(ts, value=value, stimulus_id=stimulus_id)
 
     def transition_generic_memory(
         self, ts: TaskState, value=no_value, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if value is no_value and ts.key not in self.data:
             raise RuntimeError(
                 f"Tried to transition task {ts} to `memory` without data available"
@@ -2309,7 +2305,7 @@ class Worker(ServerNode):
 
     def transition_executing_memory(
         self, ts: TaskState, value=no_value, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if self.validate:
             assert ts.state == "executing" or ts.key in self.long_running
             assert not ts.waiting_for_data
@@ -2321,7 +2317,7 @@ class Worker(ServerNode):
 
     def transition_constrained_executing(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if self.validate:
             assert not ts.waiting_for_data
             assert ts.key not in self.data
@@ -2334,12 +2330,12 @@ class Worker(ServerNode):
             self.available_resources[resource] -= quantity
         ts.state = "executing"
         self._executing.add(ts)
-        self.loop.add_callback(self.execute, ts.key, stimulus_id=stimulus_id)
-        return {}, []
+        instr = Execute(key=ts.key, stimulus_id=stimulus_id)
+        return {}, [instr]
 
     def transition_ready_executing(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if self.validate:
             assert not ts.waiting_for_data
             assert ts.key not in self.data
@@ -2352,12 +2348,10 @@ class Worker(ServerNode):
 
         ts.state = "executing"
         self._executing.add(ts)
-        self.loop.add_callback(self.execute, ts.key, stimulus_id=stimulus_id)
-        return {}, []
+        instr = Execute(key=ts.key, stimulus_id=stimulus_id)
+        return {}, [instr]
 
-    def transition_flight_fetch(
-        self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    def transition_flight_fetch(self, ts: TaskState, *, stimulus_id: str) -> RecsInstrs:
         # If this transition is called after the flight coroutine has finished,
         # we can reset the task and transition to fetch again. If it is not yet
         # finished, this should be a no-op
@@ -2385,7 +2379,7 @@ class Worker(ServerNode):
         traceback_text: str,
         *,
         stimulus_id: str,
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         self._in_flight_tasks.discard(ts)
         ts.coming_from = None
         return self.transition_generic_error(
@@ -2399,7 +2393,7 @@ class Worker(ServerNode):
 
     def transition_flight_released(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if ts.done:
             # FIXME: Is this even possible? Would an assert instead be more
             # sensible?
@@ -2413,13 +2407,13 @@ class Worker(ServerNode):
 
     def transition_cancelled_memory(
         self, ts: TaskState, value, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         assert ts._next
         return {ts: ts._next}, []
 
     def transition_executing_long_running(
         self, ts: TaskState, compute_duration: float, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         ts.state = "long-running"
         self._executing.discard(ts)
         self.long_running.add(ts.key)
@@ -2429,7 +2423,7 @@ class Worker(ServerNode):
 
     def transition_released_memory(
         self, ts: TaskState, value, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         try:
             recs = self._put_key_in_memory(ts, value, stimulus_id=stimulus_id)
         except Exception as e:
@@ -2441,7 +2435,7 @@ class Worker(ServerNode):
 
     def transition_flight_memory(
         self, ts: TaskState, value, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         self._in_flight_tasks.discard(ts)
         ts.coming_from = None
         try:
@@ -2455,7 +2449,7 @@ class Worker(ServerNode):
 
     def transition_released_forgotten(
         self, ts: TaskState, *, stimulus_id: str
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         recommendations: Recs = {}
         # Dependents _should_ be released by the scheduler before this
         if self.validate:
@@ -2472,7 +2466,7 @@ class Worker(ServerNode):
 
     def _transition(
         self, ts: TaskState, finish: str | tuple, *args, stimulus_id: str, **kwargs
-    ) -> tuple[Recs, Instructions]:
+    ) -> RecsInstrs:
         if isinstance(finish, tuple):
             # the concatenated transition path might need to access the tuple
             assert not args
@@ -2592,12 +2586,35 @@ class Worker(ServerNode):
         else:
             self._handle_instructions(instructions)
 
-    def _handle_instructions(self, instructions: list[Instruction]) -> None:
+    def handle_stimulus(self, stim: StateMachineEvent) -> None:
+        with log_errors():
+            # self.stimulus_history.append(stim)  # TODO
+            recs, instructions = self.handle_event(stim)
+            self.transitions(recs, stimulus_id=stim.stimulus_id)
+            self._handle_instructions(instructions)
+            self.ensure_computing()
+            self.ensure_communicating()
+
+    def _handle_stimulus_from_future(
+        self, future: asyncio.Future[StateMachineEvent | None]
+    ) -> None:
+        with log_errors():
+            # This *should* never raise
+            stim = future.result()
+        if stim:
+            self.handle_stimulus(stim)
+
+    def _handle_instructions(self, instructions: Instructions) -> None:
         # TODO this method is temporary.
         #      See final design: https://github.com/dask/distributed/issues/5894
         for inst in instructions:
             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)
             else:
                 raise TypeError(inst)  # pragma: nocover
 
@@ -3389,24 +3406,20 @@ class Worker(ServerNode):
                 pdb.set_trace()
             raise
 
-    async def execute(self, key: str, *, stimulus_id: str) -> None:
+    async def execute(self, key: str, *, stimulus_id: str) -> StateMachineEvent | None:
         if self.status in {Status.closing, Status.closed, Status.closing_gracefully}:
-            return
-        if key not in self.tasks:
-            return
-        ts = self.tasks[key]
+            return None
+        ts = self.tasks.get(key)
+        if not ts:
+            return None
+        if ts.state == "cancelled":
+            logger.debug(
+                "Trying to execute task %s which is not in executing state anymore",
+                ts,
+            )
+            return AlreadyCancelledEvent(key=ts.key, stimulus_id=stimulus_id)
 
         try:
-            if ts.state == "cancelled":
-                # This might happen if keys are canceled
-                logger.debug(
-                    "Trying to execute task %s which is not in executing state anymore",
-                    ts,
-                )
-                ts.done = True
-                self.transition(ts, "released", stimulus_id=stimulus_id)
-                return
-
             if self.validate:
                 assert not ts.waiting_for_data
                 assert ts.state == "executing"
@@ -3430,9 +3443,7 @@ class Worker(ServerNode):
                     f"expected one of: {sorted(self.executors)}"
                 )
 
-            self.active_keys.add(ts.key)
-
-            result: dict
+            self.active_keys.add(key)
             try:
                 ts.start_time = time()
                 if iscoroutinefunction(function):
@@ -3450,7 +3461,7 @@ class Worker(ServerNode):
                         args2,
                         kwargs2,
                         self.execution_state,
-                        ts.key,
+                        key,
                         self.active_threads,
                         self.active_threads_lock,
                         self.scheduler_delay,
@@ -3465,75 +3476,137 @@ class Worker(ServerNode):
                         self.scheduler_delay,
                     )
             finally:
-                self.active_keys.discard(ts.key)
-
-            key = ts.key
-            # key *must* be still in tasks. Releasing it directly is forbidden
-            # without going through cancelled
-            ts = self.tasks.get(key)  # type: ignore
-            assert ts, self.story(key)
-            ts.done = True
-            result["key"] = ts.key
-            value = result.pop("result", None)
-            ts.startstops.append(
-                {"action": "compute", "start": result["start"], "stop": result["stop"]}
-            )
-            self.threads[ts.key] = result["thread"]
-            recommendations: Recs = {}
+                self.active_keys.discard(key)
+
+            self.threads[key] = result["thread"]
+
             if result["op"] == "task-finished":
-                ts.nbytes = result["nbytes"]
-                ts.type = result["type"]
-                recommendations[ts] = ("memory", value)
                 if self.digests is not None:
                     self.digests["task-duration"].add(result["stop"] - result["start"])
-            elif isinstance(result.pop("actual-exception"), Reschedule):
-                recommendations[ts] = "rescheduled"
-            else:
-                logger.warning(
-                    "Compute Failed\n"
-                    "Key:       %s\n"
-                    "Function:  %s\n"
-                    "args:      %s\n"
-                    "kwargs:    %s\n"
-                    "Exception: %r\n",
-                    ts.key,
-                    str(funcname(function))[:1000],
-                    convert_args_to_str(args2, max_len=1000),
-                    convert_kwargs_to_str(kwargs2, max_len=1000),
-                    result["exception_text"],
+                return ExecuteSuccessEvent(
+                    key=key,
+                    value=result["result"],
+                    start=result["start"],
+                    stop=result["stop"],
+                    nbytes=result["nbytes"],
+                    type=result["type"],
+                    stimulus_id=stimulus_id,
                 )
-                recommendations[ts] = (
-                    "error",
-                    result["exception"],
-                    result["traceback"],
-                    result["exception_text"],
-                    result["traceback_text"],
-                )
-
-            self.transitions(recommendations, stimulus_id=stimulus_id)
 
-            logger.debug("Send compute response to scheduler: %s, %s", ts.key, result)
+            if isinstance(result["actual-exception"], Reschedule):
+                return RescheduleEvent(key=ts.key, stimulus_id=stimulus_id)
 
-            if self.validate:
-                assert ts.state != "executing"
-                assert not ts.waiting_for_data
+            logger.warning(
+                "Compute Failed\n"
+                "Key:       %s\n"
+                "Function:  %s\n"
+                "args:      %s\n"
+                "kwargs:    %s\n"
+                "Exception: %r\n",
+                key,
+                str(funcname(function))[:1000],
+                convert_args_to_str(args2, max_len=1000),
+                convert_kwargs_to_str(kwargs2, max_len=1000),
+                result["exception_text"],
+            )
+            return ExecuteFailureEvent(
+                key=key,
+                start=result["start"],
+                stop=result["stop"],
+                exception=result["exception"],
+                traceback=result["traceback"],
+                exception_text=result["exception_text"],
+                traceback_text=result["traceback_text"],
+                stimulus_id=stimulus_id,
+            )
 
         except Exception as exc:
-            assert ts
-            logger.error(
-                "Exception during execution of task %s.", ts.key, exc_info=True
+            logger.error("Exception during execution of task %s.", key, exc_info=True)
+            msg = error_message(exc)
+            return ExecuteFailureEvent(
+                key=key,
+                start=None,
+                stop=None,
+                exception=msg["exception"],
+                traceback=msg["traceback"],
+                exception_text=msg["exception_text"],
+                traceback_text=msg["traceback_text"],
+                stimulus_id=stimulus_id,
             )
-            emsg = error_message(exc)
-            del emsg["status"]  # type: ignore
-            self.transition(
-                ts,
+
+    @functools.singledispatchmethod
+    def handle_event(self, ev: StateMachineEvent) -> RecsInstrs:
+        raise TypeError(ev)  # pragma: nocover
+
+    @handle_event.register
+    def _(self, ev: CancelComputeEvent) -> RecsInstrs:
+        """Scheduler requested to cancel a task"""
+        ts = self.tasks.get(ev.key)
+        if not ts or ts.state not in READY | {"waiting"}:
+            return {}, []
+
+        self.log.append((ev.key, "cancel-compute", ev.stimulus_id, time()))
+        # All possible dependents of ts should not be in state Processing on
+        # scheduler side and therefore should not be assigned to a worker, yet.
+        assert not ts.dependents
+        return {ts: "released"}, []
+
+    @handle_event.register
+    def _(self, ev: AlreadyCancelledEvent) -> RecsInstrs:
+        """Task is already cancelled by the time execute() runs"""
+        # key *must* be still in tasks. Releasing it directly is forbidden
+        # without going through cancelled
+        ts = self.tasks.get(ev.key)  # type: ignore
+        assert ts, self.story(ev.key)
+        ts.done = True
+        return {ts: "released"}, []
+
+    @handle_event.register
+    def _(self, ev: ExecuteSuccessEvent) -> RecsInstrs:
+        """Task completed successfully"""
+        # key *must* be still in tasks. Releasing it directly is forbidden
+        # without going through cancelled
+        ts = self.tasks.get(ev.key)  # type: ignore
+        assert ts, self.story(ev.key)
+
+        ts.done = True
+        ts.startstops.append({"action": "compute", "start": ev.start, "stop": ev.stop})
+        ts.nbytes = ev.nbytes
+        ts.type = ev.type
+        return {ts: ("memory", ev.value)}, []
+
+    @handle_event.register
+    def _(self, ev: ExecuteFailureEvent) -> RecsInstrs:
+        """Task execution failed"""
+        # key *must* be still in tasks. Releasing it directly is forbidden
+        # without going through cancelled
+        ts = self.tasks.get(ev.key)  # type: ignore
+        assert ts, self.story(ev.key)
+
+        ts.done = True
+        if ev.start is not None and ev.stop is not None:
+            ts.startstops.append(
+                {"action": "compute", "start": ev.start, "stop": ev.stop}
+            )
+
+        return {
+            ts: (
                 "error",
-                **emsg,
-                stimulus_id=stimulus_id,
+                ev.exception,
+                ev.traceback,
+                ev.exception_text,
+                ev.traceback_text,
             )
-        finally:
-            self.ensure_computing()
-            self.ensure_communicating()
+        }, []
+
+    @handle_event.register
+    def _(self, ev: RescheduleEvent) -> RecsInstrs:
+        """Task raised Reschedule exception while it was running"""
+        # key *must* be still in tasks. Releasing it directly is forbidden
+        # without going through cancelled
+        ts = self.tasks.get(ev.key)  # type: ignore
+        assert ts, self.story(ev.key)
+        return {ts: "rescheduled"}, []
 
     def _prepare_args_for_execution(
         self, ts: TaskState, args: tuple, kwargs: dict[str, Any]
diff --git a/distributed/worker_state_machine.py b/distributed/worker_state_machine.py
index 80016682..8ae45441 100644
--- a/distributed/worker_state_machine.py
+++ b/distributed/worker_state_machine.py
@@ -15,7 +15,7 @@ from distributed.protocol.serialize import Serialize
 from distributed.utils import recursive_to_dict
 
 if TYPE_CHECKING:
-    # TODO move to typing (requires Python >=3.10)
+    # TODO move to typing and get out of TYPE_CHECKING (requires Python >=3.10)
     from typing_extensions import TypeAlias
 
     TaskStateState: TypeAlias = Literal[
@@ -35,7 +35,8 @@ if TYPE_CHECKING:
         "resumed",
         "waiting",
     ]
-
+else:
+    TaskStateState = str
 
 # TaskState.state subsets
 PROCESSING: set[TaskStateState] = {
@@ -257,11 +258,11 @@ class Instruction:
 #    __slots__ = ()
 
 
-# @dataclass
-# class Execute(Instruction):
-#    __slots__ = ("key", "stimulus_id")
-#    key: str
-#    stimulus_id: str
+@dataclass
+class Execute(Instruction):
+    __slots__ = ("key", "stimulus_id")
+    key: str
+    stimulus_id: str
 
 
 class SendMessageToScheduler(Instruction):
@@ -322,11 +323,11 @@ class ReleaseWorkerDataMsg(SendMessageToScheduler):
     key: str
 
 
+# Not to be confused with RescheduleEvent below or the distributed.Reschedule Exception
 @dataclass
 class RescheduleMsg(SendMessageToScheduler):
     op = "reschedule"
 
-    # Not to be confused with the distributed.Reschedule Exception
     __slots__ = ("key", "worker")
     key: str
     worker: str
@@ -348,3 +349,64 @@ class AddKeysMsg(SendMessageToScheduler):
     __slots__ = ("keys", "stimulus_id")
     keys: list[str]
     stimulus_id: str
+
+
+@dataclass
+class StateMachineEvent:
+    __slots__ = ("stimulus_id",)
+    stimulus_id: str
+
+
+@dataclass
+class ExecuteSuccessEvent(StateMachineEvent):
+    key: str
+    value: object
+    start: float
+    stop: float
+    nbytes: int
+    type: type | None
+    __slots__ = tuple(__annotations__)  # type: ignore
+
+
+@dataclass
+class ExecuteFailureEvent(StateMachineEvent):
+    key: str
+    start: float | None
+    stop: float | None
+    exception: Serialize
+    traceback: Serialize | None
+    exception_text: str
+    traceback_text: str
+    __slots__ = tuple(__annotations__)  # type: ignore
+
+
+@dataclass
+class CancelComputeEvent(StateMachineEvent):
+    __slots__ = ("key",)
+    key: str
+
+
+@dataclass
+class AlreadyCancelledEvent(StateMachineEvent):
+    __slots__ = ("key",)
+    key: str
+
+
+# Not to be confused with RescheduleMsg above or the distributed.Reschedule Exception
+@dataclass
+class RescheduleEvent(StateMachineEvent):
+    __slots__ = ("key",)
+    key: str
+
+
+if TYPE_CHECKING:
+    # TODO remove quotes (requires Python >=3.9)
+    # TODO get out of TYPE_CHECKING (requires Python >=3.10)
+    # {TaskState -> finish: TaskStateState | (finish: TaskStateState, transition *args)}
+    Recs: TypeAlias = "dict[TaskState, TaskStateState | tuple]"
+    Instructions: TypeAlias = "list[Instruction]"
+    RecsInstrs: TypeAlias = "tuple[Recs, Instructions]"
+else:
+    Recs = dict
+    Instructions = list
+    RecsInstrs = tuple


Test patch:
diff --git a/distributed/cli/tests/test_dask_spec.py b/distributed/cli/tests/test_dask_spec.py
index 06a9de70..eb1a3716 100644
--- a/distributed/cli/tests/test_dask_spec.py
+++ b/distributed/cli/tests/test_dask_spec.py
@@ -5,7 +5,7 @@ import yaml
 
 from distributed import Client
 from distributed.scheduler import COMPILED
-from distributed.utils_test import gen_test, popen
+from distributed.utils_test import gen_cluster, gen_test, popen
 
 
 @pytest.mark.skipif(COMPILED, reason="Fails with cythonized scheduler")
@@ -40,7 +40,8 @@ async def test_text():
 
 @pytest.mark.skipif(COMPILED, reason="Fails with cythonized scheduler")
 @pytest.mark.asyncio
-async def test_file(cleanup, tmp_path):
+@gen_cluster(client=True, nthreads=[])
+async def test_file(c, s, tmp_path):
     fn = str(tmp_path / "foo.yaml")
     with open(fn, "w") as f:
         yaml.dump(
@@ -50,24 +51,21 @@ async def test_file(cleanup, tmp_path):
             },
             f,
         )
-
-    with popen(["dask-scheduler", "--port", "9373", "--no-dashboard"]):
-        with popen(
-            [
-                sys.executable,
-                "-m",
-                "distributed.cli.dask_spec",
-                "tcp://localhost:9373",
-                "--spec-file",
-                fn,
-            ]
-        ):
-            async with Client("tcp://localhost:9373", asynchronous=True) as client:
-                await client.wait_for_workers(1)
-                info = await client.scheduler.identity()
-                [w] = info["workers"].values()
-                assert w["name"] == "foo"
-                assert w["nthreads"] == 3
+    with popen(
+        [
+            sys.executable,
+            "-m",
+            "distributed.cli.dask_spec",
+            s.address,
+            "--spec-file",
+            fn,
+        ]
+    ):
+        await c.wait_for_workers(1)
+        info = await c.scheduler.identity()
+        [w] = info["workers"].values()
+        assert w["name"] == "foo"
+        assert w["nthreads"] == 3
 
 
 def test_errors():
diff --git a/distributed/cli/tests/test_dask_worker.py b/distributed/cli/tests/test_dask_worker.py
index 5be7d4b4..e9355c82 100644
--- a/distributed/cli/tests/test_dask_worker.py
+++ b/distributed/cli/tests/test_dask_worker.py
@@ -1,5 +1,8 @@
+from __future__ import annotations
+
 import asyncio
 import os
+import sys
 from multiprocessing import cpu_count
 from time import sleep
 
@@ -8,14 +11,152 @@ from click.testing import CliRunner
 
 from dask.utils import tmpfile
 
-import distributed.cli.dask_worker
 from distributed import Client
+from distributed.cli.dask_worker import _apportion_ports, main
 from distributed.compatibility import LINUX, to_thread
 from distributed.deploy.utils import nprocesses_nthreads
 from distributed.metrics import time
 from distributed.utils_test import gen_cluster, popen, requires_ipv6
 
 
+@pytest.mark.parametrize(
+    # args: (worker_port, nanny_port, n_workers, nanny)
+    # Passing *args tuple instead of single args is to improve readability with black
+    "args,expect",
+    [
+        # Single worker
+        (
+            (None, None, 1, False),
+            [{"port": None}],
+        ),
+        (
+            (None, None, 1, True),
+            [{"port": None, "worker_port": None}],
+        ),
+        (("123", None, 1, False), [{"port": 123}]),
+        (
+            ("123", None, 1, True),
+            [{"port": None, "worker_port": 123}],
+        ),
+        (
+            (None, "456", 1, True),
+            [{"port": 456, "worker_port": None}],
+        ),
+        (
+            ("123", "456", 1, True),
+            [{"port": 456, "worker_port": 123}],
+        ),
+        # port=None or 0 and multiple workers
+        (
+            (None, None, 2, False),
+            [
+                {"port": None},
+                {"port": None},
+            ],
+        ),
+        (
+            (None, None, 2, True),
+            [
+                {"port": None, "worker_port": None},
+                {"port": None, "worker_port": None},
+            ],
+        ),
+        (
+            (0, "0", 2, True),
+            [
+                {"port": 0, "worker_port": 0},
+                {"port": 0, "worker_port": 0},
+            ],
+        ),
+        (
+            ("0", None, 2, True),
+            [
+                {"port": None, "worker_port": 0},
+                {"port": None, "worker_port": 0},
+            ],
+        ),
+        # port ranges
+        (
+            ("100:103", None, 1, False),
+            [{"port": [100, 101, 102, 103]}],
+        ),
+        (
+            ("100:103", None, 2, False),
+            [
+                {"port": [100, 102]},  # Round robin apportion
+                {"port": [101, 103]},
+            ],
+        ),
+        # port range is not an exact multiple of n_workers
+        (
+            ("100:107", None, 3, False),
+            [
+                {"port": [100, 103, 106]},
+                {"port": [101, 104, 107]},
+                {"port": [102, 105]},
+            ],
+        ),
+        (
+            ("100:103", None, 2, True),
+            [
+                {"port": None, "worker_port": [100, 102]},
+                {"port": None, "worker_port": [101, 103]},
+            ],
+        ),
+        (
+            (None, "110:113", 2, True),
+            [
+                {"port": [110, 112], "worker_port": None},
+                {"port": [111, 113], "worker_port": None},
+            ],
+        ),
+        # port ranges have different length between nannies and workers
+        (
+            ("100:103", "110:114", 2, True),
+            [
+                {"port": [110, 112, 114], "worker_port": [100, 102]},
+                {"port": [111, 113], "worker_port": [101, 103]},
+            ],
+        ),
+        # identical port ranges
+        (
+            ("100:103", "100:103", 2, True),
+            [
+                {"port": 101, "worker_port": 100},
+                {"port": 103, "worker_port": 102},
+            ],
+        ),
+        # overlapping port ranges
+        (
+            ("100:105", "104:106", 2, True),
+            [
+                {"port": [104, 106], "worker_port": [100, 102]},
+                {"port": 105, "worker_port": [101, 103]},
+            ],
+        ),
+    ],
+)
+def test_apportion_ports(args, expect):
+    assert _apportion_ports(*args) == expect
+
+
+def test_apportion_ports_bad():
+    with pytest.raises(ValueError, match="Not enough ports in range"):
+        _apportion_ports("100:102", None, 4, False)
+    with pytest.raises(ValueError, match="Not enough ports in range"):
+        _apportion_ports(None, "100:102", 4, False)
+    with pytest.raises(ValueError, match="Not enough ports in range"):
+        _apportion_ports("100:102", "100:102", 3, True)
+    with pytest.raises(ValueError, match="Not enough ports in range"):
+        _apportion_ports("100:102", "102:104", 3, True)
+    with pytest.raises(ValueError, match="port_stop must be greater than port_start"):
+        _apportion_ports("102:100", None, 4, False)
+    with pytest.raises(ValueError, match="invalid literal for int"):
+        _apportion_ports("foo", None, 1, False)
+    with pytest.raises(ValueError, match="too many values to unpack"):
+        _apportion_ports("100:101:102", None, 1, False)
+
+
 @pytest.mark.slow
 @gen_cluster(client=True, nthreads=[])
 async def test_nanny_worker_ports(c, s):
@@ -38,8 +179,32 @@ async def test_nanny_worker_ports(c, s):
 
 
 @pytest.mark.slow
+@pytest.mark.flaky(
+    LINUX and sys.version_info == (3, 9),
+    reason="Race condition in Nanny.process.start? "
+    "See https://github.com/dask/distributed/issues/6045",
+)
 @gen_cluster(client=True, nthreads=[])
 async def test_nanny_worker_port_range(c, s):
+    async def assert_ports(min_: int, max_: int, nanny: bool) -> None:
+        port_ranges = await c.run(
+            lambda dask_worker: dask_worker._start_port, nanny=nanny
+        )
+
+        for a in port_ranges.values():
+            assert isinstance(a, list)
+            assert len(a) in (333, 334)
+            assert all(min_ <= i <= max_ for i in a)
+
+            # Test no overlap
+            for b in port_ranges.values():
+                assert a is b or not set(a) & set(b)
+
+        ports = await c.run(lambda dask_worker: dask_worker.port, nanny=nanny)
+        assert all(min_ <= p <= max_ for p in ports.values())
+        for addr, range in port_ranges.items():
+            assert ports[addr] in range
+
     with popen(
         [
             "dask-worker",
@@ -56,10 +221,8 @@ async def test_nanny_worker_port_range(c, s):
         ]
     ):
         await c.wait_for_workers(3)
-        worker_ports = await c.run(lambda dask_worker: dask_worker.port)
-        assert all(10000 <= p <= 11000 for p in worker_ports.values())
-        nanny_ports = await c.run(lambda dask_worker: dask_worker.port, nanny=True)
-        assert all(11000 <= p <= 12000 for p in nanny_ports.values())
+        await assert_ports(10000, 10999, nanny=False)
+        await assert_ports(11000, 12000, nanny=True)
 
 
 @gen_cluster(nthreads=[])
@@ -80,7 +243,9 @@ async def test_nanny_worker_port_range_too_many_workers_raises(s):
         ],
         flush_output=False,
     ) as worker:
-        assert any(b"Could not start" in worker.stdout.readline() for _ in range(100))
+        assert any(
+            b"Not enough ports in range" in worker.stdout.readline() for _ in range(100)
+        )
 
 
 @pytest.mark.slow
@@ -416,7 +581,7 @@ async def test_dashboard_non_standard_ports(c, s):
 
 def test_version_option():
     runner = CliRunner()
-    result = runner.invoke(distributed.cli.dask_worker.main, ["--version"])
+    result = runner.invoke(main, ["--version"])
     assert result.exit_code == 0
 
 
@@ -427,7 +592,7 @@ def test_worker_timeout(no_nanny):
     args = ["192.168.1.100:7777", "--death-timeout=1"]
     if no_nanny:
         args.append("--no-nanny")
-    result = runner.invoke(distributed.cli.dask_worker.main, args)
+    result = runner.invoke(main, args)
     assert result.exit_code != 0
 
 
@@ -437,14 +602,14 @@ def test_bokeh_deprecation():
     runner = CliRunner()
     with pytest.warns(UserWarning, match="dashboard"):
         try:
-            runner.invoke(distributed.cli.dask_worker.main, ["--bokeh"])
+            runner.invoke(main, ["--bokeh"])
         except ValueError:
             # didn't pass scheduler
             pass
 
     with pytest.warns(UserWarning, match="dashboard"):
         try:
-            runner.invoke(distributed.cli.dask_worker.main, ["--no-bokeh"])
+            runner.invoke(main, ["--no-bokeh"])
         except ValueError:
             # didn't pass scheduler
             pass
diff --git a/distributed/tests/test_utils.py b/distributed/tests/test_utils.py
index 931db6e0..9163458a 100644
--- a/distributed/tests/test_utils.py
+++ b/distributed/tests/test_utils.py
@@ -528,9 +528,22 @@ def test_parse_ports():
     assert parse_ports(23) == [23]
     assert parse_ports("45") == [45]
     assert parse_ports("100:103") == [100, 101, 102, 103]
+    assert parse_ports([100, 101, 102, 103]) == [100, 101, 102, 103]
+
+    out = parse_ports((100, 101, 102, 103))
+    assert out == [100, 101, 102, 103]
+    assert isinstance(out, list)
 
     with pytest.raises(ValueError, match="port_stop must be greater than port_start"):
         parse_ports("103:100")
+    with pytest.raises(TypeError):
+        parse_ports(100.5)
+    with pytest.raises(TypeError):
+        parse_ports([100, 100.5])
+    with pytest.raises(ValueError):
+        parse_ports("foo")
+    with pytest.raises(ValueError):
+        parse_ports("100.5")
 
 
 def test_lru():
diff --git a/distributed/tests/test_worker_memory.py b/distributed/tests/test_worker_memory.py
index 030daa91..0c9e88c1 100644
--- a/distributed/tests/test_worker_memory.py
+++ b/distributed/tests/test_worker_memory.py
@@ -411,7 +411,7 @@ async def test_pause_executor_manual(c, s, a):
 @gen_cluster(
     nthreads=[("", 1)],
     client=True,
-    worker_kwargs={"memory_limit": "1000 MB"},
+    worker_kwargs={"memory_limit": "10 GB"},
     config={
         "distributed.worker.memory.target": False,
         "distributed.worker.memory.spill": False,
@@ -442,8 +442,8 @@ async def test_pause_executor_with_memory_monitor(c, s, a):
         while "y" not in a.tasks:
             await asyncio.sleep(0.01)
 
-        # Hog the worker with 900MB unmanaged memory
-        mocked_rss = 900_000_000
+        # Hog the worker with 900GB unmanaged memory
+        mocked_rss = 900 * 1000**3
         while s.workers[a.address].status != Status.paused:
             await asyncio.sleep(0.01)
 
@@ -539,17 +539,17 @@ async def test_override_data_nanny(c, s, n):
 @gen_cluster(
     client=True,
     nthreads=[("", 1)],
-    worker_kwargs={"memory_limit": "1 GB", "data": UserDict},
+    worker_kwargs={"memory_limit": "10 GB", "data": UserDict},
     config={"distributed.worker.memory.monitor-interval": "10ms"},
 )
 async def test_override_data_vs_memory_monitor(c, s, a):
-    a.monitor.get_process_memory = lambda: 801_000_000 if a.data else 0
+    a.monitor.get_process_memory = lambda: 8_100_000_000 if a.data else 0
     assert memory_monitor_running(a)
 
     # Push a key that would normally trip both the target and the spill thresholds
     class C:
         def __sizeof__(self):
-            return 801_000_000
+            return 8_100_000_000
 
     # Capture output of log_errors()
     with captured_logger(logging.getLogger("distributed.utils")) as logger:
diff --git a/distributed/tests/test_worker_state_machine.py b/distributed/tests/test_worker_state_machine.py
index 2c97b217..78597a37 100644
--- a/distributed/tests/test_worker_state_machine.py
+++ b/distributed/tests/test_worker_state_machine.py
@@ -1,9 +1,13 @@
+from itertools import chain
+
 import pytest
 
 from distributed.utils import recursive_to_dict
 from distributed.worker_state_machine import (
+    Instruction,
     ReleaseWorkerDataMsg,
     SendMessageToScheduler,
+    StateMachineEvent,
     TaskState,
     UniqueTaskHeap,
 )
@@ -82,10 +86,23 @@ def test_unique_task_heap():
     assert repr(heap) == "<UniqueTaskHeap: 0 items>"
 
 
-@pytest.mark.parametrize("cls", SendMessageToScheduler.__subclasses__())
-def test_sendmsg_slots(cls):
-    smsg = cls(**dict.fromkeys(cls.__annotations__))
-    assert not hasattr(smsg, "__dict__")
+@pytest.mark.parametrize(
+    "cls",
+    chain(
+        [UniqueTaskHeap],
+        Instruction.__subclasses__(),
+        SendMessageToScheduler.__subclasses__(),
+        StateMachineEvent.__subclasses__(),
+    ),
+)
+def test_slots(cls):
+    params = [
+        k
+        for k in dir(cls)
+        if not k.startswith("_") and k != "op" and not callable(getattr(cls, k))
+    ]
+    inst = cls(**dict.fromkeys(params))
+    assert not hasattr(inst, "__dict__")
 
 
 def test_sendmsg_to_dict():