Repository: dask/distributed

Issue:
Redundant remove_worker in retire_worker
`Scheduler. _track_retire_worker` has keyword arguments that allow us to close and/or remove workers.

During an unrelated debugging session we noticed that this may be redundant and we might want to clean this up. Particularly if `close_workers=True ` we will always trigger a remove_worker as well, regardless of the kwarg. This code section is _also_ marked with a TODO about redundant close calls. Is this intentional?

See https://github.com/dask/distributed/blob/a8eb3b23b8fe91f52758db155e7151e3d516cbdc/distributed/scheduler.py#L3452-L3453

https://github.com/dask/distributed/blob/a8eb3b23b8fe91f52758db155e7151e3d516cbdc/distributed/scheduler.py#L6080-L6087

cc @crusaderky 


Discussion:
@gjoseph92 added that comment.

The call definitely looks redundant in production code. A handful of tests call close_worker() without remove_worker() - they should be amended.

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -3439,19 +3439,6 @@ async def log_errors(func):
         setproctitle("dask-scheduler [closed]")
         disable_gc_diagnosis()
 
-    @log_errors
-    async def close_worker(self, worker: str, stimulus_id: str, safe: bool = False):
-        """Remove a worker from the cluster
-
-        This both removes the worker from our local state and also sends a
-        signal to the worker to shut down.  This works regardless of whether or
-        not the worker has a nanny process restarting it
-        """
-        logger.info("Closing worker %s", worker)
-        self.log_event(worker, {"action": "close-worker"})
-        self.worker_send(worker, {"op": "close"})  # TODO redundant with `remove_worker`
-        await self.remove_worker(address=worker, safe=safe, stimulus_id=stimulus_id)
-
     ###########
     # Stimuli #
     ###########
@@ -4170,16 +4157,38 @@ def stimulus_retry(self, keys, client=None):
 
         return tuple(seen)
 
+    def close_worker(self, worker: str) -> None:
+        """Ask a worker to shut itself down. Do not wait for it to take effect.
+        Note that there is no guarantee that the worker will actually accept the
+        command.
+
+        Note that :meth:`remove_worker` sends the same command internally if close=True.
+
+        See also
+        --------
+        retire_workers
+        remove_worker
+        """
+        if worker not in self.workers:
+            return
+
+        logger.info("Closing worker %s", worker)
+        self.log_event(worker, {"action": "close-worker"})
+        self.worker_send(worker, {"op": "close"})
+
     @log_errors
     async def remove_worker(
         self, address: str, *, stimulus_id: str, safe: bool = False, close: bool = True
     ) -> Literal["OK", "already-removed"]:
-        """
-        Remove worker from cluster
+        """Remove worker from cluster.
+
+        We do this when a worker reports that it plans to leave or when it appears to be
+        unresponsive. This may send its tasks back to a released state.
 
-        We do this when a worker reports that it plans to leave or when it
-        appears to be unresponsive.  This may send its tasks back to a released
-        state.
+        See also
+        --------
+        retire_workers
+        close_worker
         """
         if self.status == Status.closed:
             return "already-removed"
@@ -5935,16 +5944,16 @@ def _key(group):
     @log_errors
     async def retire_workers(
         self,
-        comm=None,
+        workers: list[str] | None = None,
         *,
-        workers: "list[str] | None" = None,
-        names: "list | None" = None,
+        names: list | None = None,
         close_workers: bool = False,
         remove: bool = True,
         stimulus_id: str | None = None,
-        **kwargs,
-    ) -> dict:
-        """Gracefully retire workers from cluster
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Gracefully retire workers from cluster. Any key that is in memory exclusively
+        on the retired workers is replicated somewhere else.
 
         Parameters
         ----------
@@ -5961,7 +5970,16 @@ async def retire_workers(
             worker.
         remove: bool (defaults to True)
             Whether or not to remove the worker metadata immediately or else
-            wait for the worker to contact us
+            wait for the worker to contact us.
+
+            If close_workers=False and remove=False, this method just flushes the tasks
+            in memory out of the workers and then returns.
+            If close_workers=True and remove=False, this method will return while the
+            workers are still in the cluster, although they won't accept new tasks.
+            If close_workers=False or for whatever reason a worker doesn't accept the
+            close command, it will be left permanently unable to accept new tasks and
+            it is expected to be closed in some other way.
+
         **kwargs: dict
             Extra options to pass to workers_to_close to determine which
             workers we should drop
@@ -5971,6 +5989,11 @@ async def retire_workers(
         Dictionary mapping worker ID/address to dictionary of information about
         that worker for each retired worker.
 
+        If there are keys that exist in memory only on the workers being retired and it
+        was impossible to replicate them somewhere else (e.g. because there aren't
+        any other running workers), the workers holding such keys won't be retired and
+        won't appear in the returned dict.
+
         See Also
         --------
         Scheduler.workers_to_close
@@ -6040,7 +6063,7 @@ async def retire_workers(
                             ws,
                             policy,
                             prev_status=prev_status,
-                            close_workers=close_workers,
+                            close=close_workers,
                             remove=remove,
                             stimulus_id=stimulus_id,
                         )
@@ -6067,7 +6090,7 @@ async def _track_retire_worker(
         ws: WorkerState,
         policy: RetireWorker,
         prev_status: Status,
-        close_workers: bool,
+        close: bool,
         remove: bool,
         stimulus_id: str,
     ) -> tuple:  # tuple[str | None, dict]
@@ -6094,14 +6117,12 @@ async def _track_retire_worker(
             "All unique keys on worker %s have been replicated elsewhere", ws.address
         )
 
-        if close_workers and ws.address in self.workers:
-            await self.close_worker(
-                worker=ws.address, safe=True, stimulus_id=stimulus_id
-            )
         if remove:
             await self.remove_worker(
-                address=ws.address, safe=True, stimulus_id=stimulus_id
+                ws.address, safe=True, close=close, stimulus_id=stimulus_id
             )
+        elif close:
+            self.close_worker(ws.address)
 
         logger.info("Retired worker %s", ws.address)
         return ws.address, ws.identity()


Test patch:
diff --git a/distributed/tests/test_cancelled_state.py b/distributed/tests/test_cancelled_state.py
--- a/distributed/tests/test_cancelled_state.py
+++ b/distributed/tests/test_cancelled_state.py
@@ -156,7 +156,7 @@ def blockable_compute(x, lock):
         await block_compute.acquire()
 
         # Close in scheduler to ensure we transition and reschedule task properly
-        await s.close_worker(worker=a.address, stimulus_id="test")
+        await s.remove_worker(a.address, stimulus_id="test")
         await wait_for_state(fut1.key, "resumed", b)
 
         block_get_data.release()
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
@@ -1243,29 +1243,28 @@ def f(dask_scheduler=None):
         await c._run_on_scheduler(f)
 
 
-@gen_cluster(client=True)
-async def test_close_worker(c, s, a, b):
+@gen_cluster()
+async def test_close_worker(s, a, b):
     assert len(s.workers) == 2
 
-    await s.close_worker(worker=a.address, stimulus_id="test")
-
-    assert len(s.workers) == 1
+    s.close_worker(a.address)
+    while len(s.workers) != 1:
+        await asyncio.sleep(0.01)
     assert a.address not in s.workers
 
-    await asyncio.sleep(0.5)
-
+    await asyncio.sleep(0.2)
     assert len(s.workers) == 1
 
 
-@pytest.mark.slow
-@gen_cluster(client=True, Worker=Nanny)
-async def test_close_nanny(c, s, a, b):
+# @pytest.mark.slow
+@gen_cluster(Worker=Nanny)
+async def test_close_nanny(s, a, b):
     assert len(s.workers) == 2
 
     assert a.process.is_alive()
     a_worker_address = a.worker_address
-    start = time()
-    await s.close_worker(worker=a_worker_address, stimulus_id="test")
+
+    await s.remove_worker(a_worker_address, stimulus_id="test")
 
     assert len(s.workers) == 1
     assert a_worker_address not in s.workers