Repository: dask/distributed

Issue:
Worker stuck in closing_gracefully state
(Note: @mrocklin heavily edited this issue, the words are his (sorry Brett to impersonate you) but the logs are coming from Brett)

We've found that workers can be stuck in a `closing_gracefully` state indefinitely.  Here is some context:

### Cluster
We're running an adaptively scaled cluster of roughly 4000 workers.  

```
cluster = KubeHelmCluster(release_name=release_name)  # custom wrapper around dask_kubernetes.KubeCluster
cluster.adapt(minimum=0, maximum=4000, interval="10s", target_duration="60s")
```

### Computation

Our workload is pretty simple.  We're just reading a bunch of data, and then calling a map_partitions call and that's it.  We see that the cluster scales up pretty fast, and then scales down pretty fast.

```
dd.read_parquet("gs://.../*.parquet").map_partitions(len).compute()  # ~11k parquet files
```

### Some findings

Both the scheduler and the worker agree that it is in a closing gracefully state (this is after a long while)

```
In [30]: c.run_on_scheduler(lambda dask_scheduler: dask_scheduler.workers[w].status)
Out[30]: <Status.closing_gracefully: 'closing_gracefully'>

In [31]: c.run(lambda dask_worker: str(dask_worker), workers=["tcp://10.36.228.19:34273"])
Out[31]: {'tcp://10.36.228.19:34273': "<Worker 'tcp://10.36.228.19:34273', name: 128, status: closing_gracefully, stored: 1, running: 0/4, ready: 3, comm: 0, waiting: 0>"}
```

Interestingly, they both also agree that there are three tasks that are ready to go.  An artificial call to `Worker._ensure_computing` at this point doesn't do anything because the first check in that method is `if self.status != Status.Running: return`.  

Here are the logs on the scheduler that refer to this worker.  

```
(.venv) ➜  model git:(ta_backfill) ✗ kl  brett-f4d63eb0-daskscheduler-667b4fc95-f897l | grep '10.36.228.19:34273'
2022-04-27 13:22:17,287 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://10.36.228.19:34273', name: 128, status: undefined, memory: 0, processing: 0>
2022-04-27 13:22:17,291 - distributed.scheduler - INFO - Starting worker compute stream, tcp://10.36.228.19:34273
2022-04-27 13:23:01,001 - distributed.scheduler - INFO - Retiring worker tcp://10.36.228.19:34273
2022-04-27 13:23:01,448 - distributed.active_memory_manager - INFO - Retiring worker tcp://10.36.228.19:34273; 3 keys are being moved away.
2022-04-27 13:23:02,341 - distributed.active_memory_manager - INFO - Retiring worker tcp://10.36.228.19:34273; 3 keys are being moved away.
2022-04-27 13:23:03,716 - distributed.active_memory_manager - INFO - Retiring worker tcp://10.36.228.19:34273; 2 keys are being moved away.
2022-04-27 13:23:05,544 - distributed.active_memory_manager - INFO - Retiring worker tcp://10.36.228.19:34273; no unique keys need to be moved away.
```

And the events from the scheduler's perspective

```
In [32]: c.run_on_scheduler(lambda dask_scheduler: dask_scheduler.events[w])
Out[32]:
deque([(1651065737.2876427, {'action': 'add-worker'}),
       (1651065737.629527,
        {'action': 'worker-status-change',
         'prev-status': 'undefined',
         'status': 'running'}),
       (1651065793.225713,
        {'action': 'missing-data',
         'key': "('_split-659a14d513522415a58d9ec3d470e072', 8087)"}),
       (1651065793.2970057,
        {'action': 'missing-data',
         'key': "('_split-659a14d513522415a58d9ec3d470e072', 8087)"})])
```

And the logs on the worker side

```
(.venv) ➜  model git:(ta_backfill) ✗ kl -f brett-f4d63eb0-daskworkers-584d467f76-9zctq
/usr/src/python/distributed/distributed/cli/dask_worker.py:319: FutureWarning: The --nprocs flag will be removed in a future release. It has been renamed to --nworkers.
  warnings.warn(
2022-04-27 13:22:12,055 - distributed.nanny - INFO -         Start Nanny at: 'tcp://10.36.228.19:46595'
2022-04-27 13:22:13,553 - distributed.worker - INFO -       Start worker at:   tcp://10.36.228.19:34273
2022-04-27 13:22:13,553 - distributed.worker - INFO -          Listening to:   tcp://10.36.228.19:34273
2022-04-27 13:22:13,553 - distributed.worker - INFO -          dashboard at:         10.36.228.19:33865
2022-04-27 13:22:13,553 - distributed.worker - INFO - Waiting to connect to: tcp://brett-f4d63eb0-daskscheduler:8786
2022-04-27 13:22:13,553 - distributed.worker - INFO - -------------------------------------------------
2022-04-27 13:22:13,553 - distributed.worker - INFO -               Threads:                          4
2022-04-27 13:22:13,553 - distributed.worker - INFO -                Memory:                   7.82 GiB
2022-04-27 13:22:13,553 - distributed.worker - INFO -       Local Directory: /src/dask-worker-space/worker-ykvkqgn_
2022-04-27 13:22:13,554 - distributed.worker - INFO - -------------------------------------------------
2022-04-27 13:22:17,292 - distributed.worker - INFO -         Registered to: tcp://brett-f4d63eb0-daskscheduler:8786
2022-04-27 13:22:17,292 - distributed.worker - INFO - -------------------------------------------------
2022-04-27 13:22:17,293 - distributed.core - INFO - Starting established connection
2022-04-27 13:22:56,885 - distributed.utils_perf - INFO - full garbage collection released 105.93 MiB from 0 reference cycles (threshold: 9.54 MiB)
```

### Thoughts

My guess is that when entering a `closing_gracefully` state on the worker side we don't immediately clear out ready work, when maybe we should.  

Also, we should have some sanity checks so that `close_gracefully` can't go on for forever.  

cc @crusaderky @fjetter 


Discussion:
I hope that this is easy to resolve (we should be able to manage state in the `close_gracefully` method I hope).

In the meantime, a plugin like this might be helpful.

```python
class CloseUngracefully(WorkerPlugin):
    def setup(self, worker):
        self.worker = worker
        self.count = 0

        pc = PeriodicCallback(self.check, 1000)
        pc.start()
    
    def check(self):
        if self.worker.status == Status.closing_gracefully:
            self.count += 1

        if self.count >= 10:
            sys.exit(0)
```
To be clear, the code above is not actually tested.  It is likely wrong (but hopefully right directionally)
I don't think this is the worker's fault. `Status.closing_gracefully` is sort of just a sentinel value on the worker, where the scheduler is telling the worker, "don't do any more work, but don't shut down yet until I tell you it's okay". The reason being that [`retire_workers`](https://github.com/dask/distributed/blob/198522bf767d7d8c1927df07951cd34a79109bd0/distributed/scheduler.py#L5945-L6064) (what adaptive scaling uses to scale down) doesn't tell the worker to shut down until its data has been copied to other workers in the cluster. The worker in question obviously can't know when all those keys have been replicated elsewhere; the scheduler is the one that has to orchestrate this.

Having a sanity timeout so that `close_gracefully` can't go on forever probably makes sense. But I'm guessing it should be on the scheduler side, not the workers, since the scheduler is the one coordinating anyway—fewer race conditions to worry about if we do it there. In fact, the scheduler [may tell a worker](https://github.com/dask/distributed/blob/198522bf767d7d8c1927df07951cd34a79109bd0/distributed/scheduler.py#L6075-L6082) that's `closing_gracefully` "actually wait no, don't go! we still need you!". This gets tricky to reason about if the worker might also be shutting itself too due to a deadline.

----

Anyway, my guess is that there's a discrepancy between `RetireWorker.done` https://github.com/dask/distributed/blob/198522bf767d7d8c1927df07951cd34a79109bd0/distributed/active_memory_manager.py#L662-L667
and the cases in which the `RetireWorker` AMM policy [removes](https://github.com/dask/distributed/blob/198522bf767d7d8c1927df07951cd34a79109bd0/distributed/active_memory_manager.py#L651) [itself](https://github.com/dask/distributed/blob/198522bf767d7d8c1927df07951cd34a79109bd0/distributed/active_memory_manager.py#L660) from the AMM policies.

Because `Scheduler._track_retire_worker` is [waiting for `policy.done()`](https://github.com/dask/distributed/blob/198522bf767d7d8c1927df07951cd34a79109bd0/distributed/scheduler.py#L6074) to flip before it tells the worker to actually shut down, if the policy gets removed before it's `done()`, it may never become done. Or, because there's a poll interval, maybe `done()` actually was true at the point when the policy removed itself, but somehow things changed in between (a different worker holding the other replica suddenly died, for example), and by the time `done()` gets called, it no longer matches the `done()` criteria (and without a running policy, never will)?

@crusaderky what do you think? It still feels odd to me that the criteria for `self.done()` isn't just "the policy has been removed". (Some previous discussion in https://github.com/dask/distributed/pull/5381/files#r794919819.) What about:
```diff
diff --git a/distributed/active_memory_manager.py b/distributed/active_memory_manager.py
index 448a7581..b549d300 100644
--- a/distributed/active_memory_manager.py
+++ b/distributed/active_memory_manager.py
@@ -661,6 +661,8 @@ class RetireWorker(ActiveMemoryManagerPolicy):
 
     def done(self) -> bool:
         """Return True if it is safe to close the worker down; False otherwise"""
+        if self not in self.manager.policies:
+            return True
         ws = self.manager.scheduler.workers.get(self.address)
         if ws is None:
             return True
```

Also, the obvious place to add a timeout would be in this `while not policy.done()` loop. But if the timeout expires, do we want to close the worker (even if it still has essential data), or switch it back to `running`?

> My guess is that when entering a closing_gracefully state on the worker side we don't immediately clear out ready work, when maybe we should.

This might be a nice optimization (steal tasks immediately from workers that are closing gracefully, instead of waiting for them to close, and letting the normal remove_worker handling reassign the tasks). But if the tasks on the worker depend on data that's on the worker, having work-stealing replicating data around at the same time as a `RetireWorkers` policy might be counterproductive. Data replication has to happen before those tasks can run elsewhere. Once `RetireWorkers` is done, the data is replicated. So I don't think we gain much from this. We just have to be patient and wait for data transfer.
> I don't think this is the worker's fault. Status.closing_gracefully is sort of just a sentinel value on the worker, where the scheduler is telling the worker, "don't do any more work, but don't shut down yet until I tell you it's okay". 

Should the scheduler immediately reschedule processing tasks on other workers?
@gjoseph92 do you think it is doable to get a fix that would resolve this issue in by the release on Friday?

Resolving patch:
diff --git a/distributed/active_memory_manager.py b/distributed/active_memory_manager.py
--- a/distributed/active_memory_manager.py
+++ b/distributed/active_memory_manager.py
@@ -661,6 +661,12 @@ def run(self) -> SuggestionGenerator:
 
     def done(self) -> bool:
         """Return True if it is safe to close the worker down; False otherwise"""
+        if self not in self.manager.policies:
+            # Either the no_recipients flag has been raised, or there were no unique replicas
+            # as of the latest AMM run. Note that due to tasks transitioning from running to
+            # memory there may be some now; it's OK to lose them and just recompute them
+            # somewhere else.
+            return True
         ws = self.manager.scheduler.workers.get(self.address)
         if ws is None:
             return True
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -6184,24 +6184,24 @@ async def _track_retire_worker(
         stimulus_id: str,
     ) -> tuple:  # tuple[str | None, dict]
         while not policy.done():
-            if policy.no_recipients:
-                # Abort retirement. This time we don't need to worry about race
-                # conditions and we can wait for a scheduler->worker->scheduler
-                # round-trip.
-                self.stream_comms[ws.address].send(
-                    {
-                        "op": "worker-status-change",
-                        "status": prev_status.name,
-                        "stimulus_id": stimulus_id,
-                    }
-                )
-                return None, {}
-
             # Sleep 0.01s when there are 4 tasks or less
             # Sleep 0.5s when there are 200 or more
             poll_interval = max(0.01, min(0.5, len(ws.has_what) / 400))
             await asyncio.sleep(poll_interval)
 
+        if policy.no_recipients:
+            # Abort retirement. This time we don't need to worry about race
+            # conditions and we can wait for a scheduler->worker->scheduler
+            # round-trip.
+            self.stream_comms[ws.address].send(
+                {
+                    "op": "worker-status-change",
+                    "status": prev_status.name,
+                    "stimulus_id": stimulus_id,
+                }
+            )
+            return None, {}
+
         logger.debug(
             "All unique keys on worker %s have been replicated elsewhere", ws.address
         )


Test patch: