Repository: dask/distributed

Issue:
Flaky `test_worker_doesnt_await_task_completion`
I have a feeling this is a subtle bug in the recently-changed `restart` logic
```
___________________ test_worker_doesnt_await_task_completion ___________________

loop = <tornado.platform.asyncio.AsyncIOMainLoop object at 0x7fd8662543d0>

    def test_worker_doesnt_await_task_completion(loop):
        with cluster(nanny=True, nworkers=1) as (s, [w]):
            with Client(s["address"], loop=loop) as c:
                future = c.submit(sleep, 100)
                sleep(0.1)
                start = time()
>               c.restart()

distributed/tests/test_failed_workers.py:170: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
distributed/client.py:3395: in restart
    return self.sync(
distributed/utils.py:338: in sync
    return sync(
distributed/utils.py:405: in sync
    raise exc.with_traceback(tb)
distributed/utils.py:378: in f
    result = yield future
../../../miniconda3/envs/dask-distributed/lib/python3.8/site-packages/tornado/gen.py:762: in run
    value = future.result()
distributed/client.py:3364: in _restart
    await self.scheduler.restart(timeout=timeout, wait_for_workers=wait_for_workers)
distributed/core.py:1153: in send_recv_from_rpc
    return await send_recv(comm=comm, op=key, **kwargs)
distributed/core.py:943: in send_recv
    raise exc.with_traceback(tb)
distributed/core.py:769: in _handle_comm
    result = await result
distributed/utils.py:799: in wrapper
    return await func(*args, **kwargs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

>   raise TimeoutError(msg) from None
E   asyncio.exceptions.TimeoutError: Waited for 1 worker(s) to reconnect after restarting, but after 20s, only 1 have returned. Consider a longer timeout, or `wait_for_workers=False`.

distributed/scheduler.py:5289: TimeoutError
----------------------------- Captured stderr call -----------------------------
2022-08-02 15:33:17,513 - distributed.http.proxy - INFO - To route to workers diagnostics web server please install jupyter-server-proxy: python -m pip install jupyter-server-proxy
2022-08-02 15:33:17,524 - distributed.scheduler - INFO - State start
2022-08-02 15:33:17,548 - distributed.scheduler - INFO - Clear task state
2022-08-02 15:33:17,549 - distributed.scheduler - INFO -   Scheduler at:     tcp://127.0.0.1:56665
2022-08-02 15:33:17,549 - distributed.scheduler - INFO -   dashboard at:            127.0.0.1:8787
2022-08-02 15:33:17,657 - distributed.nanny - INFO -         Start Nanny at: 'tcp://127.0.0.1:56666'
2022-08-02 15:33:19,841 - distributed.diskutils - INFO - Found stale lock file and directory '/var/folders/24/8k48jl6d249_n_qfxwsl6xvm0000gn/T/dask-worker-space/worker-zgduxqdu', purging
2022-08-02 15:33:19,843 - distributed.diskutils - INFO - Found stale lock file and directory '/var/folders/24/8k48jl6d249_n_qfxwsl6xvm0000gn/T/dask-worker-space/worker-qonwlf1u', purging
2022-08-02 15:33:19,886 - distributed.worker - INFO -       Start worker at:      tcp://127.0.0.1:56669
2022-08-02 15:33:19,887 - distributed.worker - INFO -          Listening to:      tcp://127.0.0.1:56669
2022-08-02 15:33:19,887 - distributed.worker - INFO -          dashboard at:            127.0.0.1:56670
2022-08-02 15:33:19,887 - distributed.worker - INFO - Waiting to connect to:      tcp://127.0.0.1:56665
2022-08-02 15:33:19,887 - distributed.worker - INFO - -------------------------------------------------
2022-08-02 15:33:19,887 - distributed.worker - INFO -               Threads:                          1
2022-08-02 15:33:19,887 - distributed.worker - INFO -                Memory:                  14.00 GiB
2022-08-02 15:33:19,887 - distributed.worker - INFO -       Local Directory: /var/folders/24/8k48jl6d249_n_qfxwsl6xvm0000gn/T/dask-worker-space/worker-b2ocd1ce
2022-08-02 15:33:19,887 - distributed.worker - INFO - -------------------------------------------------
2022-08-02 15:33:20,649 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://127.0.0.1:56669', status: init, memory: 0, processing: 0>
2022-08-02 15:33:21,426 - distributed.scheduler - INFO - Starting worker compute stream, tcp://127.0.0.1:56669
2022-08-02 15:33:21,426 - distributed.core - INFO - Starting established connection
2022-08-02 15:33:21,426 - distributed.worker - INFO -         Registered to:      tcp://127.0.0.1:56665
2022-08-02 15:33:21,427 - distributed.worker - INFO - -------------------------------------------------
2022-08-02 15:33:21,429 - distributed.core - INFO - Starting established connection
2022-08-02 15:33:21,548 - distributed.scheduler - INFO - Receive client connection: Client-70bfb78a-[1278](https://github.com/dask/distributed/runs/7634516786?check_suite_focus=true#step:11:1279)-11ed-801c-005056abdb4e
2022-08-02 15:33:21,549 - distributed.core - INFO - Starting established connection
2022-08-02 15:33:21,739 - distributed.scheduler - INFO - Releasing all requested keys
2022-08-02 15:33:21,740 - distributed.scheduler - INFO - Clear task state
2022-08-02 15:33:21,750 - distributed.nanny - INFO - Nanny asking worker to close
2022-08-02 15:33:21,752 - distributed.worker - INFO - Stopping worker at tcp://127.0.0.1:56669
2022-08-02 15:33:21,756 - distributed.worker - INFO - Connection to scheduler broken. Closing without reporting. ID: Worker-b7cbdb91-0b6f-45f3-9c78-a460b37247fc Address tcp://127.0.0.1:56669 Status: Status.closing
2022-08-02 15:33:21,756 - distributed.scheduler - INFO - Remove worker <WorkerState 'tcp://127.0.0.1:56669', status: closing, memory: 0, processing: 0>
2022-08-02 15:33:21,757 - distributed.core - INFO - Removing comms to tcp://127.0.0.1:56669
2022-08-02 15:33:21,757 - distributed.scheduler - INFO - Lost all workers
2022-08-02 15:33:37,836 - distributed.nanny - WARNING - Worker process still alive after 15.99999580383301 seconds, killing
2022-08-02 15:33:37,851 - distributed.nanny - INFO - Worker process 53153 was killed by signal 15
2022-08-02 15:33:37,855 - distributed.nanny - WARNING - Restarting worker
2022-08-02 15:33:40,458 - distributed.worker - INFO -       Start worker at:      tcp://127.0.0.1:56680
2022-08-02 15:33:40,458 - distributed.worker - INFO -          Listening to:      tcp://127.0.0.1:56680
2022-08-02 15:33:40,458 - distributed.worker - INFO -          dashboard at:            127.0.0.1:56681
2022-08-02 15:33:40,458 - distributed.worker - INFO - Waiting to connect to:      tcp://127.0.0.1:56665
2022-08-02 15:33:40,459 - distributed.worker - INFO - -------------------------------------------------
2022-08-02 15:33:40,459 - distributed.worker - INFO -               Threads:                          1
2022-08-02 15:33:40,459 - distributed.worker - INFO -                Memory:                  14.00 GiB
2022-08-02 15:33:40,459 - distributed.worker - INFO -       Local Directory: /var/folders/24/8k48jl6d249_n_qfxwsl6xvm0000gn/T/dask-worker-space/worker-i9qscgl6
2022-08-02 15:33:40,459 - distributed.worker - INFO - -------------------------------------------------
2022-08-02 15:33:41,265 - distributed.scheduler - INFO - Register worker <WorkerState 'tcp://127.0.0.1:56680', status: init, memory: 0, processing: 0>
2022-08-02 15:33:42,754 - distributed.scheduler - INFO - Starting worker compute stream, tcp://127.0.0.1:56680
2022-08-02 15:33:42,754 - distributed.core - INFO - Starting established connection
2022-08-02 15:33:42,755 - distributed.worker - INFO -         Registered to:      tcp://127.0.0.1:56665
2022-08-02 15:33:42,755 - distributed.worker - INFO - -------------------------------------------------
2022-08-02 15:33:42,757 - distributed.core - INFO - Starting established connection
2022-08-02 15:33:42,887 - distributed.core - ERROR - Waited for 1 worker(s) to reconnect after restarting, but after 20s, only 1 have returned. Consider a longer timeout, or `wait_for_workers=False`.
Traceback (most recent call last):
  File "/Users/runner/work/distributed/distributed/distributed/utils.py", line 799, in wrapper
    return await func(*args, **kwargs)
  File "/Users/runner/work/distributed/distributed/distributed/scheduler.py", line 5289, in restart
    raise TimeoutError(msg) from None
asyncio.exceptions.TimeoutError: Waited for 1 worker(s) to reconnect after restarting, but after 20s, only 1 have returned. Consider a longer timeout, or `wait_for_workers=False`.
2022-08-02 15:33:42,888 - distributed.core - ERROR - Exception while handling op restart
Traceback (most recent call last):
  File "/Users/runner/work/distributed/distributed/distributed/core.py", line 769, in _handle_comm
    result = await result
  File "/Users/runner/work/distributed/distributed/distributed/utils.py", line 799, in wrapper
    return await func(*args, **kwargs)
  File "/Users/runner/work/distributed/distributed/distributed/scheduler.py", line 5289, in restart
    raise TimeoutError(msg) from None
asyncio.exceptions.TimeoutError: Waited for 1 worker(s) to reconnect after restarting, but after 20s, only 1 have returned. Consider a longer timeout, or `wait_for_workers=False`.
2022-08-02 15:33:42,898 - distributed.scheduler - INFO - Remove client Client-70bfb78a-1278-11ed-801c-005056abdb4e
2022-08-02 15:33:42,899 - distributed.scheduler - INFO - Remove client Client-70bfb78a-1278-11ed-801c-005056abdb4e
2022-08-02 15:33:42,899 - distributed.scheduler - INFO - Close client connection: Client-70bfb78a-1278-11ed-801c-005056abdb4e
2022-08-02 15:33:42,907 - distributed.nanny - INFO - Closing Nanny at 'tcp://127.0.0.1:56666'.
2022-08-02 15:33:42,908 - distributed.nanny - INFO - Nanny asking worker to close
2022-08-02 15:33:42,909 - distributed.worker - INFO - Stopping worker at tcp://127.0.0.1:56680
2022-08-02 15:33:42,911 - distributed.scheduler - INFO - Scheduler closing...
2022-08-02 15:33:42,912 - distributed.scheduler - INFO - Scheduler closing all comms
2022-08-02 15:33:42,913 - distributed.worker - INFO - Connection to scheduler broken. Closing without reporting. ID: Worker-9dbf444e-9fff-4972-9ed4-9e7b51bb3a89 Address tcp://127.0.0.1:56680 Status: Status.closing
2022-08-02 15:33:42,913 - distributed.scheduler - INFO - Remove worker <WorkerState 'tcp://127.0.0.1:56680', status: closing, memory: 0, processing: 0>
2022-08-02 15:33:42,913 - distributed.core - INFO - Removing comms to tcp://127.0.0.1:5668
```
https://github.com/dask/distributed/runs/7634516786?check_suite_focus=true#step:11:1264


Discussion:

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -5265,28 +5265,27 @@ async def restart(self, client=None, timeout=30, wait_for_workers=True):
         self.log_event([client, "all"], {"action": "restart", "client": client})
 
         if wait_for_workers:
-            while monotonic() < start + timeout:
+            while len(self.workers) < n_workers:
                 # NOTE: if new (unrelated) workers join while we're waiting, we may return before
                 # our shut-down workers have come back up. That's fine; workers are interchangeable.
-                if len(self.workers) >= n_workers:
-                    return
-                await asyncio.sleep(0.2)
-            else:
-                msg = (
-                    f"Waited for {n_workers} worker(s) to reconnect after restarting, "
-                    f"but after {timeout}s, only {len(self.workers)} have returned. "
-                    "Consider a longer timeout, or `wait_for_workers=False`."
-                )
-
-                if (n_nanny := len(nanny_workers)) < n_workers:
-                    msg += (
-                        f" The {n_workers - n_nanny} worker(s) not using Nannies were just shut "
-                        "down instead of restarted (restart is only possible with Nannies). If "
-                        "your deployment system does not automatically re-launch terminated "
-                        "processes, then those workers will never come back, and `Client.restart` "
-                        "will always time out. Do not use `Client.restart` in that case."
+                if monotonic() < start + timeout:
+                    await asyncio.sleep(0.2)
+                else:
+                    msg = (
+                        f"Waited for {n_workers} worker(s) to reconnect after restarting, "
+                        f"but after {timeout}s, only {len(self.workers)} have returned. "
+                        "Consider a longer timeout, or `wait_for_workers=False`."
                     )
-                raise TimeoutError(msg) from None
+
+                    if (n_nanny := len(nanny_workers)) < n_workers:
+                        msg += (
+                            f" The {n_workers - n_nanny} worker(s) not using Nannies were just shut "
+                            "down instead of restarted (restart is only possible with Nannies). If "
+                            "your deployment system does not automatically re-launch terminated "
+                            "processes, then those workers will never come back, and `Client.restart` "
+                            "will always time out. Do not use `Client.restart` in that case."
+                        )
+                    raise TimeoutError(msg) from None
 
     async def broadcast(
         self,


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
@@ -706,6 +706,43 @@ async def test_restart_not_all_workers_return(c, s, a, b):
     assert b.status in (Status.closed, Status.closing)
 
 
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_restart_worker_rejoins_after_timeout_expired(c, s, a):
+    """
+    We don't want to see an error message like:
+
+    ``Waited for 1 worker(s) to reconnect after restarting, but after 0s, only 1 have returned.``
+
+    If a worker rejoins after our last poll for new workers, but before we raise the error,
+    we shouldn't raise the error.
+    """
+    # We'll use a 0s timeout on the restart, so it always expires.
+    # And we'll use a plugin to block the restart process, and spin up a new worker
+    # in the middle of it.
+
+    class Plugin(SchedulerPlugin):
+        removed = asyncio.Event()
+        proceed = asyncio.Event()
+
+        async def remove_worker(self, *args, **kwargs):
+            self.removed.set()
+            await self.proceed.wait()
+
+    s.add_plugin(Plugin())
+
+    task = asyncio.create_task(c.restart(timeout=0))
+    await Plugin.removed.wait()
+    assert not s.workers
+
+    async with Worker(s.address, nthreads=1) as w:
+        assert len(s.workers) == 1
+        Plugin.proceed.set()
+
+        # New worker has joined, but the timeout has expired (since it was 0).
+        # Still, we should not time out.
+        await task
+
+
 @gen_cluster(client=True, nthreads=[("", 1)] * 2)
 async def test_restart_no_wait_for_workers(c, s, a, b):
     await c.restart(timeout="1s", wait_for_workers=False)