Repository: dask/distributed

Issue:
Restructure `Worker.close`
At the moment, the order of operations in `Worker.close` is somewhat random and does not appear intentionally structured. 

It reads roughly as follows:
1. Update `self.status` to `Status.closing`
    1. Inform scheduler of status change
    2. Handle `PauseEvent` to prevent *further* tasks from being executed/gathered
2. Stop `self.periodic_callbacks`
3. `self.stop()`
    1. Purge workspace
    2. Close monitor
    3. Stop sync listeners
    4. Abort handshaking comms
    5. Stop async listeners in the background
4. Run `BaseWorker.close(timeout=...)`
    1. Cancel `self.async_instructions`
    2. Wait for them to finish (with timeout)
5. Teardown preloads
6. Close extensions
7. *Maybe* tell nanny to `close_gracefully`
8. Teardown plugins
9. Close worker-initialized clients
10. Close scheduler RPC
11. Stop services
12. Wait a bit in case we’re using UCX
13. Close batched stream by first sending `close-stream`, then closing with a timeout
14. Shutdown executors (maybe wait with timeout)
15. Close RPC
16. Update status to `Status.closed`
17. `ServerNode.close()`
    1. stop periodic_callbacks (******again******)
    2. stop monitor (*again*)
    3. stop listeners (******again******)
    4. Abort handshaking comms *******(again)*******
    5. Stop async listeners (*again but blocking this time*)
    6. Stop background_tasks
    7. close RPC (******again******)
    8. Close comms
    9. Remove `local_directory` from `sys.path`
18. Unwind `exit_stack`

There are multiple issues with this:
* The workspace gets purged before plugins get torn down, which may remove data of an in-progress P2P shuffle and cause exceptions.
* It's generally hard to understand what state a worker is in at any given point in time
* Just to add to the confusion, some stuff is stopped multiple times.

As an alternative, I suggest the following rough order of operations:

1. Inform the cluster that the worker is closing
    * Inform scheduler
    * Inform nanny
    * Set `self.status = Status.closing`
2. Tear down any running "add-ins" while resources are still working
    * plugins
    * extensions
    * services
    * preloads
5. Prevent new communications from being established
6. Close existing comms
7. Close local resources and clean up internal state

In particular, add-ins such as plugins should be torn down while the worker is still functional to allow them to take arbitrary actions such as informing the scheduler or initiating communications with other workers.


Note:
Timeouts are also a problem but should be tackled in a dedicated PR (see #7318, #7320).


Discussion:
> In particular, add-ins such as plugins should be torn down while the worker is still functional to allow them to take arbitrary actions such as informing the scheduler or initiating communications with other workers.

An alternative to this would be to remove the worker from the scheduler ASAP, mimicking the case where a worker dies. Any necessary cluster-wide coordination would then be left to `SchedulerPlugin`s, and `WorkerPlugin`s would be scoped to local teardown.
Fun fact: By stopping `periodic_callbacks` in step 2, we also stop heartbeats and keep-alives. If closing the worker takes too long, its TTL will expire and the scheduler will trigger `remove_worker` regardless of how far along the worker is in its shutdown.
> Inform the cluster that the worker is closing

+1

> Inform scheduler

this is done by setting the status. I don't think we should do anything else at this stage.

---

overall the proposed ordering makes sense to me. I think the PCs have to be stopped between 2 and 3. PCs do need comms and we should stop them before closing them.

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -1535,23 +1535,24 @@ async def close(  # type: ignore
             logger.info("Closed worker has not yet started: %s", self.status)
         if not executor_wait:
             logger.info("Not waiting on executor to close")
-        self.status = Status.closing
 
-        # Stop callbacks before giving up control in any `await`.
-        # We don't want to heartbeat while closing.
-        for pc in self.periodic_callbacks.values():
-            pc.stop()
+        # This also informs the scheduler about the status update
+        self.status = Status.closing
+        setproctitle("dask worker [closing]")
 
-        self.stop()
+        if nanny and self.nanny:
+            with self.rpc(self.nanny) as r:
+                await r.close_gracefully(reason=reason)
 
         # Cancel async instructions
         await BaseWorker.close(self, timeout=timeout)
 
-        for preload in self.preloads:
-            try:
-                await preload.teardown()
-            except Exception:
-                logger.exception("Failed to tear down preload")
+        teardowns = [
+            plugin.teardown(self)
+            for plugin in self.plugins.values()
+            if hasattr(plugin, "teardown")
+        ]
+        await asyncio.gather(*(td for td in teardowns if isawaitable(td)))
 
         for extension in self.extensions.values():
             if hasattr(extension, "close"):
@@ -1559,19 +1560,18 @@ async def close(  # type: ignore
                 if isawaitable(result):
                     await result
 
-        if nanny and self.nanny:
-            with self.rpc(self.nanny) as r:
-                await r.close_gracefully(reason=reason)
+        self.stop_services()
 
-        setproctitle("dask worker [closing]")
+        for preload in self.preloads:
+            try:
+                await preload.teardown()
+            except Exception:
+                logger.exception("Failed to tear down preload")
 
-        teardowns = [
-            plugin.teardown(self)
-            for plugin in self.plugins.values()
-            if hasattr(plugin, "teardown")
-        ]
+        for pc in self.periodic_callbacks.values():
+            pc.stop()
 
-        await asyncio.gather(*(td for td in teardowns if isawaitable(td)))
+        self.stop()
 
         if self._client:
             # If this worker is the last one alive, clean up the worker
@@ -1597,8 +1597,6 @@ async def close(  # type: ignore
 
         await self.scheduler.close_rpc()
 
-        self.stop_services()
-
         # Give some time for a UCX scheduler to complete closing endpoints
         # before closing self.batched_stream, otherwise the local endpoint
         # may be closed too early and errors be raised on the scheduler when
@@ -1649,10 +1647,11 @@ def _close(executor, wait):
         await self.rpc.close()
 
         self.status = Status.closed
+        setproctitle("dask worker [closed]")
+
         await ServerNode.close(self)
 
         self.__exit_stack.__exit__(None, None, None)
-        setproctitle("dask worker [closed]")
         return "OK"
 
     async def close_gracefully(


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
@@ -3678,7 +3678,7 @@ async def test_gather_on_worker_bad_recipient(c, s, a, b):
     """The recipient is missing"""
     x = await c.scatter("x")
     await b.close()
-    assert s.workers.keys() == {a.address}
+    await async_poll_for(lambda: s.workers.keys() == {a.address}, timeout=5)
     out = await s.gather_on_worker(b.address, {x.key: [a.address]})
     assert out == {x.key}