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/core.py b/distributed/core.py
--- a/distributed/core.py
+++ b/distributed/core.py
@@ -662,11 +662,6 @@ def stop(self):
         if self.__stopped:
             return
 
-        if self._workdir is not None:
-            self._workdir.release()
-
-        self.monitor.close()
-
         self.__stopped = True
         _stops = set()
         for listener in self.listeners:
@@ -687,6 +682,11 @@ async def background_stops():
 
             self._ongoing_background_tasks.call_soon(background_stops)
 
+        self.monitor.close()
+
+        if self._workdir is not None:
+            self._workdir.release()
+
     @property
     def listener(self):
         if self.listeners:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -6,6 +6,7 @@
 import contextlib
 import contextvars
 import errno
+import inspect
 import logging
 import math
 import os
@@ -1571,8 +1572,6 @@ async def close(  # type: ignore
         for pc in self.periodic_callbacks.values():
             pc.stop()
 
-        self.stop()
-
         if self._client:
             # If this worker is the last one alive, clean up the worker
             # initialized clients
@@ -1595,7 +1594,27 @@ async def close(  # type: ignore
                         # otherwise
                         c.close()
 
-        await self.scheduler.close_rpc()
+        # FIXME: Copy-paste from `Server.stop`. See dask/distributed#8077
+        _stops = set()
+        for listener in self.listeners:
+            future = listener.stop()
+            if inspect.isawaitable(future):
+                _stops.add(future)
+            try:
+                abort_handshaking_comms = listener.abort_handshaking_comms
+            except AttributeError:
+                pass
+            else:
+                abort_handshaking_comms()
+
+        if _stops:
+
+            async def background_stops():
+                await asyncio.gather(*_stops)
+
+        # end copy-paste
+
+        await self.rpc.close()
 
         # Give some time for a UCX scheduler to complete closing endpoints
         # before closing self.batched_stream, otherwise the local endpoint
@@ -1644,8 +1663,7 @@ def _close(executor, wait):
                         executor=executor, wait=executor_wait
                     )  # Just run it directly
 
-        await self.rpc.close()
-
+        self.stop()
         self.status = Status.closed
         setproctitle("dask worker [closed]")
 


Test patch: