Repository: dask/distributed

Issue:
`client.restart()` may cause workers to shut down instead of restarting
`Client.restart()` immediately removes WorkerStates before the workers have fully shut down (https://github.com/dask/distributed/issues/6390) [by calling `remove_worker`](https://github.com/dask/distributed/blob/a3414329458e0bb1eb7efa9d948dfe1c91d53bf9/distributed/scheduler.py#L5096-L5106). However, this doesn't flush the `BatchedSend` or wait for confirmation that the worker has received the message. So if the worker heartbeats in this interval after the `WorkerState` has been removed, but before the `op: close` has reached the worker, the scheduler will [get mad at it](https://github.com/dask/distributed/blob/a3414329458e0bb1eb7efa9d948dfe1c91d53bf9/distributed/scheduler.py#L3463-L3466) and the worker will [shut itself down](https://github.com/dask/distributed/blob/a3414329458e0bb1eb7efa9d948dfe1c91d53bf9/distributed/worker.py#L1239-L1245) instead of restarting. Only after it starts to shut down will it receive the `op: close, nanny=True` message from the scheduler, which it will effectively ignore.

There are a few ways to address this:
- short-term: the worker should restart when it gets the `missing` message instead of shutting down. This is reasonable to do anyway. https://github.com/dask/distributed/issues/6387
- medium-term: https://github.com/dask/distributed/issues/6390

@jrbourbeau discovered this, and we initially thought it was an issue with the fact that `Worker.close` `await`s a bunch of things before [turning off its heartbeats](https://github.com/dask/distributed/blob/a3414329458e0bb1eb7efa9d948dfe1c91d53bf9/distributed/worker.py#L1537-L1538) to the scheduler. We thought that the worker was partway through the closing process, but still heartbeating. It's true that this can happen, and it probably shouldn't. However, if an extraneous heartbeat does occur while closing, and the scheduler replies with `missing`, then the `Worker.close()` call in response to that will just [jump on the bandwagon](https://github.com/dask/distributed/blob/a3414329458e0bb1eb7efa9d948dfe1c91d53bf9/distributed/worker.py#L1473-L1475) of the first close call that's already running, so it won't actually cause a shutdown if the first call was doing a restart.

Therefore, I think this is purely about the race condition on the scheduler.

cc @fjetter
`client.restart()` may cause workers to shut down instead of restarting
`Client.restart()` immediately removes WorkerStates before the workers have fully shut down (https://github.com/dask/distributed/issues/6390) [by calling `remove_worker`](https://github.com/dask/distributed/blob/a3414329458e0bb1eb7efa9d948dfe1c91d53bf9/distributed/scheduler.py#L5096-L5106). However, this doesn't flush the `BatchedSend` or wait for confirmation that the worker has received the message. So if the worker heartbeats in this interval after the `WorkerState` has been removed, but before the `op: close` has reached the worker, the scheduler will [get mad at it](https://github.com/dask/distributed/blob/a3414329458e0bb1eb7efa9d948dfe1c91d53bf9/distributed/scheduler.py#L3463-L3466) and the worker will [shut itself down](https://github.com/dask/distributed/blob/a3414329458e0bb1eb7efa9d948dfe1c91d53bf9/distributed/worker.py#L1239-L1245) instead of restarting. Only after it starts to shut down will it receive the `op: close, nanny=True` message from the scheduler, which it will effectively ignore.

There are a few ways to address this:
- short-term: the worker should restart when it gets the `missing` message instead of shutting down. This is reasonable to do anyway. https://github.com/dask/distributed/issues/6387
- medium-term: https://github.com/dask/distributed/issues/6390

@jrbourbeau discovered this, and we initially thought it was an issue with the fact that `Worker.close` `await`s a bunch of things before [turning off its heartbeats](https://github.com/dask/distributed/blob/a3414329458e0bb1eb7efa9d948dfe1c91d53bf9/distributed/worker.py#L1537-L1538) to the scheduler. We thought that the worker was partway through the closing process, but still heartbeating. It's true that this can happen, and it probably shouldn't. However, if an extraneous heartbeat does occur while closing, and the scheduler replies with `missing`, then the `Worker.close()` call in response to that will just [jump on the bandwagon](https://github.com/dask/distributed/blob/a3414329458e0bb1eb7efa9d948dfe1c91d53bf9/distributed/worker.py#L1473-L1475) of the first close call that's already running, so it won't actually cause a shutdown if the first call was doing a restart.

Therefore, I think this is purely about the race condition on the scheduler.

cc @fjetter


Discussion:
Wouldn't something like below be a very-short-term release?
I think we should get this immediate problem fixed asap since it's messing with out CI **and** is failing the coiled-runtime integration tests


(untested)
```diff
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index d4a84c184..149fa0819 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -5095,10 +5095,10 @@ class Scheduler(SchedulerState, ServerNode):

         for addr in list(self.workers):
             try:
-                # Ask the worker to close if it doesn't have a nanny,
-                # otherwise the nanny will kill it anyway
-                await self.remove_worker(
-                    address=addr, close=addr not in nannies, stimulus_id=stimulus_id
+                await self.close_worker(
+                    worker=addr,
+                    stimulus_id=stimulus_id,
+                    safe=True,
                 )
             except Exception:
                 logger.info(
@@ -5141,6 +5141,13 @@ class Scheduler(SchedulerState, ServerNode):
                         exc_info=True,
                     )

+        for addr in list(self.workers):
+            await self.remove_worker(
+                address=addr,
+                stimulus_id=stimulus_id,
+                safe=True,
+                close=False,
+            )
         self.clear_task_state()

         with suppress(AttributeError):
```

If https://github.com/dask/distributed/issues/6387 is a quick fix that also works for me
No, because `close_worker` just calls `remove_worker` internally. It's a redundant method to `remove_worker(close=True)` (and `close=False` should never be used):

https://github.com/dask/distributed/blob/6d85a8536ee02fd2a9ae52f50de030d1cc7afe7e/distributed/scheduler.py#L3432-L3442

We should either remove `close_worker` and have just 1 method, or make `close_worker` do what it sounds like it should do.

But the approach you're suggesting here is correct. Having `restart` just tell workers to close, then flushing and closing their stream comms, and only then removing the state would be fix. That's basically what's proposed to do for https://github.com/dask/distributed/issues/6390. That will be pretty quick, but it's not the quickest possible path.

I think the quickest fix would be 
```diff
diff --git a/distributed/worker.py b/distributed/worker.py
index a8f4b6ba..1277a78b 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1245,7 +1245,8 @@ class Worker(ServerNode):
                 logger.error(
                     f"Scheduler was unaware of this worker {self.address!r}. Shutting down."
                 )
-                await self.close()
+                # Something is out of sync; have the nanny restart us if possible.
+                await self.close(nanny=False)
                 return
 
             self.scheduler_delay = response["time"] - middle

```
especially since we should do it anyway https://github.com/dask/distributed/issues/6387.
Wouldn't something like below be a very-short-term release?
I think we should get this immediate problem fixed asap since it's messing with out CI **and** is failing the coiled-runtime integration tests


(untested)
```diff
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index d4a84c184..149fa0819 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -5095,10 +5095,10 @@ class Scheduler(SchedulerState, ServerNode):

         for addr in list(self.workers):
             try:
-                # Ask the worker to close if it doesn't have a nanny,
-                # otherwise the nanny will kill it anyway
-                await self.remove_worker(
-                    address=addr, close=addr not in nannies, stimulus_id=stimulus_id
+                await self.close_worker(
+                    worker=addr,
+                    stimulus_id=stimulus_id,
+                    safe=True,
                 )
             except Exception:
                 logger.info(
@@ -5141,6 +5141,13 @@ class Scheduler(SchedulerState, ServerNode):
                         exc_info=True,
                     )

+        for addr in list(self.workers):
+            await self.remove_worker(
+                address=addr,
+                stimulus_id=stimulus_id,
+                safe=True,
+                close=False,
+            )
         self.clear_task_state()

         with suppress(AttributeError):
```

If https://github.com/dask/distributed/issues/6387 is a quick fix that also works for me
No, because `close_worker` just calls `remove_worker` internally. It's a redundant method to `remove_worker(close=True)` (and `close=False` should never be used):

https://github.com/dask/distributed/blob/6d85a8536ee02fd2a9ae52f50de030d1cc7afe7e/distributed/scheduler.py#L3432-L3442

We should either remove `close_worker` and have just 1 method, or make `close_worker` do what it sounds like it should do.

But the approach you're suggesting here is correct. Having `restart` just tell workers to close, then flushing and closing their stream comms, and only then removing the state would be fix. That's basically what's proposed to do for https://github.com/dask/distributed/issues/6390. That will be pretty quick, but it's not the quickest possible path.

I think the quickest fix would be 
```diff
diff --git a/distributed/worker.py b/distributed/worker.py
index a8f4b6ba..1277a78b 100644
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1245,7 +1245,8 @@ class Worker(ServerNode):
                 logger.error(
                     f"Scheduler was unaware of this worker {self.address!r}. Shutting down."
                 )
-                await self.close()
+                # Something is out of sync; have the nanny restart us if possible.
+                await self.close(nanny=False)
                 return
 
             self.scheduler_delay = response["time"] - middle

```
especially since we should do it anyway https://github.com/dask/distributed/issues/6387.

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1245,7 +1245,8 @@ async def heartbeat(self):
                 logger.error(
                     f"Scheduler was unaware of this worker {self.address!r}. Shutting down."
                 )
-                await self.close()
+                # Something is out of sync; have the nanny restart us if possible.
+                await self.close(nanny=False)
                 return
 
             self.scheduler_delay = response["time"] - middle


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
@@ -1783,6 +1783,28 @@ async def close(self):
         assert not s.workers
 
 
+@gen_cluster(
+    client=True,
+    nthreads=[("", 1)],
+    Worker=Nanny,
+    worker_kwargs={"heartbeat_interval": "1ms"},
+)
+async def test_heartbeat_missing_restarts(c, s: Scheduler, n: Nanny):
+    old_heartbeat_handler = s.handlers["heartbeat_worker"]
+    s.handlers["heartbeat_worker"] = lambda *args, **kwargs: {"status": "missing"}
+
+    assert n.process
+    await n.process.stopped.wait()
+
+    assert not s.workers
+    s.handlers["heartbeat_worker"] = old_heartbeat_handler
+
+    await n.process.running.wait()
+    assert n.status == Status.running
+
+    await c.wait_for_workers(1)
+
+
 @gen_cluster(nthreads=[])
 async def test_bad_local_directory(s):
     try: