Repository: dask/distributed

Issue:
Graceful shutdown should not increase suspicious counter
When a worker dies unexpectedly, the scheduler increases a `suspicious` counter on all tasks that were computing on it at the moment, then reschedules them somewhere else. When the counter reaches `distributed.scheduler.allowed-failures`, the task is marked as failed.

Graceful worker shutdown is when the user or an agent (e.g. dynamic scaling, or Coiled reacting to the imminent demise of a EC2 spot instance) _asks_ the scheduler to terminate a worker. The worker will be shut down as soon as there are no unique keys in memory on it; if there are running tasks, the scheduler won't wait for them to finish (as they could take arbitrarily long) and just reschedules somewhere else, wasting any accrued progress.

When this happens, the scheduler is currently increasing the `suspicious` counter. This is wrong, as the reason the worker died while it was running the task is that the scheduler itself asked it to.

This is particularly glaring if the user disables retries by setting `distributed.scheduler.allowed-failures: 0`; for example this happens in 
https://github.com/coiled/benchmarks/blob/main/tests/benchmarks/test_spill.py, which just failed with

> FAILED tests/benchmarks/test_spill.py::test_dot_product_spill[compressible] - distributed.scheduler.KilledWorker: Attempted to run task ('chunk_sum-6f8083eb44862f3ee7669497a35c8a25', 8, 7, 7) on 0 different workers, but all those workers died while running it. The last worker that attempt to run the task was tls://10.0.34.36:41345. Inspecting worker logs is often a good next step to diagnose what went wrong. For more information see https://distributed.dask.org/en/stable/killed.html.

(note the off-by-one in the log - it should read `on 1 different workers`).

On the worker log I read:
```
distributed.worker - INFO - asking worker to shutdown gracefully, spot interruption at 2023-09-21T23:11:01Z
distributed.worker - INFO - Closing worker gracefully: tls://10.0.34.36:41345. Reason: instance-preemption
distributed.worker - INFO - Stopping worker at tls://10.0.34.36:41345. Reason: instance-preemption
distributed.nanny - INFO - Closing Nanny gracefully at 'tls://10.0.34.36:38745'. Reason: instance-preemption
distributed.core - INFO - Connection to tls://10.0.34.214:8786 has been closed.
distributed.worker - INFO - done asking worker to close gracefully
distributed.nanny - INFO - Worker closed
distributed.nanny - INFO - Closing Nanny at 'tls://10.0.34.36:38745'. Reason: nanny-close-gracefully
```


Discussion:
FYI @ntabris this should _rarely_ be an issue on Coiled with spot instances turned on, as the chances of a task hitting 3 times in a row a worker just before it's shut down are very low. However users that disabled retries for whatever reason (typically, to increase dev velocity when tasks have a tendency to go OOM) will definitely notice.
Well, once upon a time this worked as expected. The graceful shutdown should close the worker using `safe=True` (which it is. Is suspect this is an ordering thing and the worker removal triggered by the closed connection happens first. I'm a bit surprised this isn't tested.

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1686,7 +1686,7 @@ async def close_gracefully(
         await self.scheduler.retire_workers(
             workers=[self.address],
             close_workers=False,
-            remove=False,
+            remove=True,
             stimulus_id=f"worker-close-gracefully-{time()}",
         )
         if restart is None:


Test patch:
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
@@ -1576,6 +1576,37 @@ async def test_close_gracefully(c, s, a, b):
         assert_amm_transfer_story(key, b, a)
 
 
+@pytest.mark.slow
+@gen_cluster(client=True, Worker=Nanny)
+async def test_close_gracefully_no_suspicious_tasks(c, s, a, b):
+    workers = {a.worker_address: a, b.worker_address: b}
+    started = Event()
+    block_task = Event()
+
+    def blocking_task():
+        started.set()
+        block_task.wait()
+
+    fut = c.submit(blocking_task)
+    await started.wait()
+
+    to_close = s.tasks[fut.key].processing_on.address
+
+    async def close_gracefully(dask_worker):
+        await asyncio.shield(dask_worker.close_gracefully())
+
+    try:
+        await c.run(close_gracefully, workers=[to_close])
+    except CommClosedError:
+        pass
+    await async_poll_for(lambda: to_close not in s.workers, 5)
+
+    assert b.address not in s.workers
+    assert s.tasks[fut.key].suspicious == 0
+    await block_task.set()
+    await fut
+
+
 @pytest.mark.parametrize("sync", [False, pytest.param(True, marks=[pytest.mark.slow])])
 @gen_cluster(client=True, nthreads=[("", 1)])
 async def test_close_while_executing(c, s, a, sync):