Repository: dask/distributed

Issue:
Paused workers should not be able to fetch any more data
The current implementation of the worker `ensure_communicating` will continue to fetch dependencies for as long as there are dependencies to fetch. This can push an already overloaded worker over the edge and cause it to fail.

https://github.com/dask/distributed/blob/8734c9db78a68f6bea5d2a1deb39412a9e7b739a/distributed/worker.py#L2684-L2687

A paused worker should not be allowed to fetch more data.

There are two possible ways to achieve this

1. Add another guard to `ensure_communicating` to stop scheduling additional `gather_dep` coroutines
2. Remove all tasks from a paused worker that aren't in memory. This would indirectly empty the `data_needed` heap and cause a worker to stabilize. This could be achieved by either aggressively stealing or by implementing a custom scheduler handler.

I think both options have a certain appeal. I'm wondering which one is the best to choose, specifically in context of the latest changes to AMM / retirement / pause. 

cc @crusaderky 

Note: right now, network traffic is only restricted for egress, i.e. incoming `get_data` requests from other workers, see https://github.com/dask/distributed/blob/8734c9db78a68f6bea5d2a1deb39412a9e7b739a/distributed/worker.py#L1712-L1716 


Discussion:
> Remove all tasks from a paused worker that aren't in memory. 

This approach would fix at the same time this issue and
- #3761

> aggressively stealing

This is my preferred choice, as it's most likely the one that adds the least complexity. However it also means that if anybody were to disable stealing, they would also face both this issue and #3761. I think #3761 is fairly intuitive to correlate to stealing; this issue less so. So the question is who, if anybody, ever disables stealing.

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -2748,6 +2748,9 @@ def stimulus_story(
         return [ev for ev in self.stimulus_log if getattr(ev, "key", None) in keys]
 
     def ensure_communicating(self) -> None:
+        if self.status != Status.running:
+            return
+
         stimulus_id = f"ensure-communicating-{time()}"
         skipped_worker_in_flight_or_busy = []
 
@@ -3489,7 +3492,7 @@ async def _maybe_deserialize_task(
             raise
 
     def _ensure_computing(self) -> RecsInstrs:
-        if self.status in (Status.paused, Status.closing_gracefully):
+        if self.status != Status.running:
             return {}, []
 
         recs: Recs = {}


Test patch:
diff --git a/distributed/tests/test_worker_memory.py b/distributed/tests/test_worker_memory.py
--- a/distributed/tests/test_worker_memory.py
+++ b/distributed/tests/test_worker_memory.py
@@ -478,6 +478,73 @@ def f(ev):
         assert "Resuming worker" in logger.getvalue()
 
 
+@gen_cluster(
+    client=True,
+    nthreads=[("", 1), ("", 1)],
+    config={
+        "distributed.worker.memory.target": False,
+        "distributed.worker.memory.spill": False,
+        "distributed.worker.memory.pause": False,
+    },
+)
+async def test_pause_prevents_deps_fetch(c, s, a, b):
+    """A worker is paused while there are dependencies ready to fetch, but all other
+    workers are in flight
+    """
+    a_addr = a.address
+
+    class X:
+        def __sizeof__(self):
+            return 2**40  # Disable clustering in select_keys_for_gather
+
+        def __reduce__(self):
+            return X.pause_on_unpickle, ()
+
+        @staticmethod
+        def pause_on_unpickle():
+            # Note: outside of task execution, distributed.get_worker()
+            # returns a random worker running in the process
+            for w in Worker._instances:
+                if w.address == a_addr:
+                    w.status = Status.paused
+                    return X()
+            assert False
+
+    x = c.submit(X, key="x", workers=[b.address])
+    y = c.submit(inc, 1, key="y", workers=[b.address])
+    await wait([x, y])
+    w = c.submit(lambda _: None, x, key="w", priority=1, workers=[a.address])
+    z = c.submit(inc, y, key="z", priority=0, workers=[a.address])
+
+    # - w and z reach worker a within the same message
+    # - w and z respectively make x and y go into fetch state.
+    #   w has a higher priority than z, therefore w's dependency x has a higher priority
+    #   than z's dependency y.
+    #   a.data_needed = ["x", "y"]
+    # - ensure_communicating decides to fetch x but not to fetch y together with it, as
+    #   it thinks x is 1TB in size
+    # - x fetch->flight; a is added to in_flight_workers
+    # - y is skipped by ensure_communicating since all workers that hold a replica are
+    #   in flight
+    # - x reaches a and sends a into paused state
+    # - x flight->memory; a is removed from in_flight_workers
+    # - ensure_communicating is triggered again
+    # - ensure_communicating refuses to fetch y because the worker is paused
+
+    while "y" not in a.tasks or a.tasks["y"].state != "fetch":
+        await asyncio.sleep(0.01)
+    await asyncio.sleep(0.1)
+    assert a.tasks["y"].state == "fetch"
+    assert "y" not in a.data
+    assert [ts.key for ts in a.data_needed] == ["y"]
+
+    # Unpausing kicks off ensure_communicating again
+    a.status = Status.running
+    assert await z == 3
+    assert a.tasks["y"].state == "memory"
+    assert "y" in a.data
+
+
 @gen_cluster(
     client=True,
     nthreads=[("", 1)],