Repository: dask/distributed

Issue:
test_stress_scatter_death
``test_stress_scatter_death`` has suddenly started hanging very frequently on Windows and MacOSX.
According to https://dask.org/distributed/test_report.html, the potential culprit is either #6217 or #6248.
CC @fjetter 

Note: in the test report above, the failures on Windows are marked as a white box instead of red due to #6304.


Discussion:
Update: it hangs on Ubuntu too: https://github.com/dask/distributed/runs/6349119823?check_suite_focus=true

It looks like an infinite transition loop (but I may be wrong)

```
ERROR    asyncio:base_events.py:1744 Exception in callback IOLoop.add_future.<locals>.<lambda>(<Task finishe...meout >300.0s>) at /usr/share/miniconda3/envs/dask-distributed/lib/python3.10/site-packages/tornado/ioloop.py:688
handle: <Handle IOLoop.add_future.<locals>.<lambda>(<Task finishe...meout >300.0s>) at /usr/share/miniconda3/envs/dask-distributed/lib/python3.10/site-packages/tornado/ioloop.py:688>
Traceback (most recent call last):
  File "/usr/share/miniconda3/envs/dask-distributed/lib/python3.10/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
  File "/usr/share/miniconda3/envs/dask-distributed/lib/python3.10/site-packages/tornado/ioloop.py", line 688, in <lambda>
    lambda f: self._run_callback(functools.partial(callback, future))
  File "/usr/share/miniconda3/envs/dask-distributed/lib/python3.10/site-packages/tornado/ioloop.py", line 741, in _run_callback
    ret = callback()
  File "/usr/share/miniconda3/envs/dask-distributed/lib/python3.10/site-packages/tornado/ioloop.py", line 765, in _discard_future_result
    future.result()
  File "/home/runner/work/distributed/distributed/distributed/worker.py", line 175, in wrapper
    return await method(self, *args, **kwargs)
  File "/home/runner/work/distributed/distributed/distributed/worker.py", line 1278, in handle_scheduler
    await self.handle_stream(comm, every_cycle=[self.ensure_communicating])
  File "/home/runner/work/distributed/distributed/distributed/core.py", line 651, in handle_stream
    handler(**merge(extra, msg))
  File "/home/runner/work/distributed/distributed/distributed/worker.py", line 2013, in handle_compute_task
    self.transitions(recommendations, stimulus_id=stimulus_id)
  File "/home/runner/work/distributed/distributed/distributed/worker.py", line 2764, in transitions
    a_recs, a_instructions = self._transition(
  File "/home/runner/work/distributed/distributed/distributed/worker.py", line 2651, in _transition
    func = self._transitions_table.get((start, cast(str, finish)))
  File "/usr/share/miniconda3/envs/dask-distributed/lib/python3.10/typing.py", line 1736, in cast
    def cast(typ, val):
  File "/usr/share/miniconda3/envs/dask-distributed/lib/python3.10/site-packages/pytest_timeout.py", line 241, in handler
    timeout_sigalrm(item, settings.timeout)
  File "/usr/share/miniconda3/envs/dask-distributed/lib/python3.10/site-packages/pytest_timeout.py", line 409, in timeout_sigalrm
    pytest.fail("Timeout >%ss" % timeout)
  File "/usr/share/miniconda3/envs/dask-distributed/lib/python3.10/site-packages/_pytest/outcomes.py", line 196, in fail
    raise Failed(msg=reason, pytrace=pytrace)
Failed: Timeout >300.0s
```
Also seen in [this CI run]( https://github.com/dask/distributed/runs/6358362174?check_suite_focus=true ).
Replicated on high-powered desktop after 51 runs

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -1262,6 +1262,7 @@ class SchedulerState:
         "validate",
         "workers",
         "transition_counter",
+        "transition_counter_max",
         "plugins",
         "UNKNOWN_TASK_DURATION",
         "MEMORY_RECENT_TO_OLD_TIME",
@@ -1354,6 +1355,9 @@ def __init__(
             / 2.0
         )
         self.transition_counter = 0
+        self.transition_counter_max = dask.config.get(
+            "distributed.admin.transition-counter-max"
+        )
 
     @property
     def memory(self) -> MemoryState:
@@ -1430,16 +1434,24 @@ def _transition(self, key, finish: str, stimulus_id: str, *args, **kwargs):
         Scheduler.transitions : transitive version of this function
         """
         try:
-            recommendations = {}  # type: ignore
-            worker_msgs = {}  # type: ignore
-            client_msgs = {}  # type: ignore
-
             ts: TaskState = self.tasks.get(key)  # type: ignore
             if ts is None:
-                return recommendations, client_msgs, worker_msgs
+                return {}, {}, {}
             start = ts._state
             if start == finish:
-                return recommendations, client_msgs, worker_msgs
+                return {}, {}, {}
+
+            # Notes:
+            # - in case of transition through released, this counter is incremented by 2
+            # - this increase happens before the actual transitions, so that it can
+            #   catch potential infinite recursions
+            self.transition_counter += 1
+            if self.transition_counter_max:
+                assert self.transition_counter < self.transition_counter_max
+
+            recommendations = {}  # type: ignore
+            worker_msgs = {}  # type: ignore
+            client_msgs = {}  # type: ignore
 
             if self.plugins:
                 dependents = set(ts.dependents)
@@ -1451,7 +1463,7 @@ def _transition(self, key, finish: str, stimulus_id: str, *args, **kwargs):
                 recommendations, client_msgs, worker_msgs = func(
                     key, stimulus_id, *args, **kwargs
                 )  # type: ignore
-                self.transition_counter += 1
+
             elif "released" not in start_finish:
                 assert not args and not kwargs, (args, kwargs, start_finish)
                 a_recs: dict
@@ -3173,6 +3185,7 @@ def _to_dict(self, *, exclude: Container[str] = ()) -> dict:
         info = super()._to_dict(exclude=exclude)
         extra = {
             "transition_log": self.transition_log,
+            "transition_counter": self.transition_counter,
             "log": self.log,
             "tasks": self.tasks,
             "task_groups": self.task_groups,
@@ -4496,6 +4509,8 @@ def validate_state(self, allow_overlap: bool = False) -> None:
             actual_total_occupancy,
             self.total_occupancy,
         )
+        if self.transition_counter_max:
+            assert self.transition_counter < self.transition_counter_max
 
     ###################
     # Manage Messages #
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -132,6 +132,7 @@
     TaskFinishedMsg,
     TaskState,
     TaskStateState,
+    TransitionCounterMaxExceeded,
     UniqueTaskHeap,
     UnpauseEvent,
     merge_recs_instructions,
@@ -452,7 +453,8 @@ class Worker(ServerNode):
     target_message_size: int
     validate: bool
     _transitions_table: dict[tuple[str, str], Callable]
-    _transition_counter: int
+    transition_counter: int
+    transition_counter_max: int | Literal[False]
     incoming_count: int
     outgoing_count: int
     outgoing_current_count: int
@@ -659,7 +661,10 @@ def __init__(
             ("waiting", "released"): self.transition_generic_released,
         }
 
-        self._transition_counter = 0
+        self.transition_counter = 0
+        self.transition_counter_max = dask.config.get(
+            "distributed.admin.transition-counter-max"
+        )
         self.incoming_transfer_log = deque(maxlen=100000)
         self.incoming_count = 0
         self.outgoing_transfer_log = deque(maxlen=100000)
@@ -1116,6 +1121,7 @@ def _to_dict(self, *, exclude: Container[str] = ()) -> dict:
             "busy_workers": self.busy_workers,
             "log": self.log,
             "stimulus_log": self.stimulus_log,
+            "transition_counter": self.transition_counter,
             "tasks": self.tasks,
             "logs": self.get_logs(),
             "config": dask.config.config,
@@ -2654,8 +2660,28 @@ def _transition(
         start = ts.state
         func = self._transitions_table.get((start, cast(str, finish)))
 
+        # Notes:
+        # - in case of transition through released, this counter is incremented by 2
+        # - this increase happens before the actual transitions, so that it can
+        #   catch potential infinite recursions
+        self.transition_counter += 1
+        if (
+            self.transition_counter_max
+            and self.transition_counter >= self.transition_counter_max
+        ):
+            self.log_event(
+                "transition-counter-max-exceeded",
+                {
+                    "key": ts.key,
+                    "start": start,
+                    "finish": finish,
+                    "story": self.story(ts),
+                    "worker": self.address,
+                },
+            )
+            raise TransitionCounterMaxExceeded(ts.key, start, finish, self.story(ts))
+
         if func is not None:
-            self._transition_counter += 1
             recs, instructions = func(ts, *args, stimulus_id=stimulus_id, **kwargs)
             self._notify_plugins("transition", ts.key, start, finish, **kwargs)
 
@@ -4294,7 +4320,8 @@ def validate_state(self):
                         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))
-            assert self.waiting_for_data_count == waiting_for_data_count
+            # FIXME https://github.com/dask/distributed/issues/6319
+            # assert self.waiting_for_data_count == waiting_for_data_count
             for worker, keys in self.has_what.items():
                 for k in keys:
                     assert worker in self.tasks[k].who_has
@@ -4302,6 +4329,9 @@ def validate_state(self):
             for ts in self.tasks.values():
                 self.validate_task(ts)
 
+            if self.transition_counter_max:
+                assert self.transition_counter < self.transition_counter_max
+
         except Exception as e:
             logger.error("Validate state failed.  Closing.", exc_info=e)
             logger.exception(e)
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
@@ -75,7 +75,7 @@ def __init__(self, key, start, finish, story):
 
     def __repr__(self):
         return (
-            f"InvalidTransition: {self.key} :: {self.start}->{self.finish}"
+            f"{self.__class__.__name__}: {self.key} :: {self.start}->{self.finish}"
             + "\n"
             + "  Story:\n    "
             + "\n    ".join(map(str, self.story))
@@ -84,6 +84,10 @@ def __repr__(self):
     __str__ = __repr__
 
 
+class TransitionCounterMaxExceeded(InvalidTransition):
+    pass
+
+
 @lru_cache
 def _default_data_size() -> int:
     return parse_bytes(dask.config.get("distributed.scheduler.default-data-size"))


Test patch:
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -21,6 +21,7 @@
 from dask.utils import apply, parse_timedelta, stringify, tmpfile, typename
 
 from distributed import (
+    CancelledError,
     Client,
     Event,
     Lock,
@@ -3215,11 +3216,67 @@ async def test_computations_futures(c, s, a, b):
     assert "inc" in str(computation.groups)
 
 
-@gen_cluster(client=True)
-async def test_transition_counter(c, s, a, b):
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_transition_counter(c, s, a):
     assert s.transition_counter == 0
+    assert a.transition_counter == 0
     await c.submit(inc, 1)
     assert s.transition_counter > 1
+    assert a.transition_counter > 1
+
+
+@pytest.mark.slow
+@gen_cluster(client=True)
+async def test_transition_counter_max_scheduler(c, s, a, b):
+    # This is set by @gen_cluster; it's False in production
+    assert s.transition_counter_max > 0
+    s.transition_counter_max = 1
+    with captured_logger("distributed.scheduler") as logger:
+        with pytest.raises(CancelledError):
+            await c.submit(inc, 2)
+    assert s.transition_counter > 1
+    with pytest.raises(AssertionError):
+        s.validate_state()
+    assert "transition_counter_max" in logger.getvalue()
+    # Scheduler state is corrupted. Avoid test failure on gen_cluster teardown.
+    s.validate = False
+
+
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_transition_counter_max_worker(c, s, a):
+    # This is set by @gen_cluster; it's False in production
+    assert s.transition_counter_max > 0
+    a.transition_counter_max = 1
+    with captured_logger("distributed.core") as logger:
+        fut = c.submit(inc, 2)
+        while True:
+            try:
+                a.validate_state()
+            except AssertionError:
+                break
+            await asyncio.sleep(0.01)
+
+    assert "TransitionCounterMaxExceeded" in logger.getvalue()
+    # Worker state is corrupted. Avoid test failure on gen_cluster teardown.
+    a.validate = False
+
+
+@gen_cluster(
+    client=True,
+    nthreads=[("", 1)],
+    config={"distributed.admin.transition-counter-max": False},
+)
+async def test_disable_transition_counter_max(c, s, a, b):
+    """Test that the cluster can run indefinitely if transition_counter_max is disabled.
+    This is the default outside of @gen_cluster.
+    """
+    assert s.transition_counter_max is False
+    assert a.transition_counter_max is False
+    assert await c.submit(inc, 1) == 2
+    assert s.transition_counter > 1
+    assert a.transition_counter > 1
+    s.validate_state()
+    a.validate_state()
 
 
 @gen_cluster(
@@ -3339,6 +3396,7 @@ async def test_Scheduler__to_dict(c, s, a):
         "status",
         "thread_id",
         "transition_log",
+        "transition_counter",
         "log",
         "memory",
         "tasks",
diff --git a/distributed/tests/test_stress.py b/distributed/tests/test_stress.py
--- a/distributed/tests/test_stress.py
+++ b/distributed/tests/test_stress.py
@@ -228,7 +228,12 @@ async def test_stress_steal(c, s, *workers):
 
 
 @pytest.mark.slow
-@gen_cluster(nthreads=[("127.0.0.1", 1)] * 10, client=True, timeout=180)
+@gen_cluster(
+    nthreads=[("", 1)] * 10,
+    client=True,
+    timeout=180,
+    config={"distributed.admin.transition-counter-max": 500_000},
+)
 async def test_close_connections(c, s, *workers):
     da = pytest.importorskip("dask.array")
     x = da.random.random(size=(1000, 1000), chunks=(1000, 1))
@@ -291,7 +296,12 @@ async def test_no_delay_during_large_transfer(c, s, w):
 
 
 @pytest.mark.slow
-@gen_cluster(client=True, Worker=Nanny, nthreads=[("127.0.0.1", 2)] * 6)
+@gen_cluster(
+    client=True,
+    Worker=Nanny,
+    nthreads=[("", 2)] * 6,
+    config={"distributed.admin.transition-counter-max": 500_000},
+)
 async def test_chaos_rechunk(c, s, *workers):
     s.allowed_failures = 10000
 
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
@@ -3434,6 +3434,7 @@ async def test_Worker__to_dict(c, s, a):
         "busy_workers",
         "log",
         "stimulus_log",
+        "transition_counter",
         "tasks",
         "logs",
         "config",
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -27,14 +27,6 @@
 from time import sleep
 from typing import Any, Generator, Literal
 
-from distributed.compatibility import MACOS
-from distributed.scheduler import Scheduler
-
-try:
-    import ssl
-except ImportError:
-    ssl = None  # type: ignore
-
 import pytest
 import yaml
 from tlz import assoc, memoize, merge
@@ -43,12 +35,12 @@
 
 import dask
 
-from distributed import system
+from distributed import Scheduler, system
 from distributed import versions as version_module
 from distributed.client import Client, _global_clients, default_client
 from distributed.comm import Comm
 from distributed.comm.tcp import TCP
-from distributed.compatibility import WINDOWS
+from distributed.compatibility import MACOS, WINDOWS
 from distributed.config import initialize_logging
 from distributed.core import (
     CommClosedError,
@@ -79,6 +71,11 @@
 )
 from distributed.worker import WORKER_ANY_RUNNING, InvalidTransition, Worker
 
+try:
+    import ssl
+except ImportError:
+    ssl = None  # type: ignore
+
 try:
     import dask.array  # register config
 except ImportError:
@@ -447,8 +444,6 @@ async def background_read():
 
 def run_scheduler(q, nputs, config, port=0, **kwargs):
     with dask.config.set(config):
-        from distributed import Scheduler
-
         # On Python 2.7 and Unix, fork() is used to spawn child processes,
         # so avoid inheriting the parent's IO loop.
         with pristine_loop() as loop:
@@ -999,6 +994,7 @@ async def test_foo(scheduler, worker1, worker2, pytest_fixture_a, pytest_fixture
     worker_kwargs = merge(
         {"memory_limit": system.MEMORY_LIMIT, "death_timeout": 15}, worker_kwargs
     )
+    config = merge({"distributed.admin.transition-counter-max": 50_000}, config)
 
     def _(func):
         if not iscoroutinefunction(func):
@@ -1052,8 +1048,7 @@ async def coro():
                             task = asyncio.create_task(coro)
                             coro2 = asyncio.wait_for(asyncio.shield(task), timeout)
                             result = await coro2
-                            if s.validate:
-                                s.validate_state()
+                            validate_state(s, *workers)
 
                         except asyncio.TimeoutError:
                             assert task
@@ -1073,6 +1068,10 @@ async def coro():
                             while not task.cancelled():
                                 await asyncio.sleep(0.01)
 
+                            # Hopefully, the hang has been caused by inconsistent state,
+                            # which should be much more meaningful than the timeout
+                            validate_state(s, *workers)
+
                             # Remove as much of the traceback as possible; it's
                             # uninteresting boilerplate from utils_test and asyncio and
                             # not from the code being tested.
@@ -1205,6 +1204,15 @@ async def dump_cluster_state(
     print(f"Dumped cluster state to {fname}")
 
 
+def validate_state(*servers: Scheduler | Worker | Nanny) -> None:
+    """Run validate_state() on the Scheduler and all the Workers of the cluster.
+    Excludes workers wrapped by Nannies and workers manually started by the test.
+    """
+    for s in servers:
+        if s.validate and hasattr(s, "validate_state"):
+            s.validate_state()  # type: ignore
+
+
 def raises(func, exc=Exception):
     try:
         func()