Repository: dask/distributed

Issue:
Nanny restarts worker once if it fails to start / doesn't clean up processes on startup failure
If the worker fails during startup, the Nanny will try to restart it once. This is not too important, but it's a bug and it makes testing difficult.

If we slightly modify `test_worker_start_exception` to verify that `"Restarting worker"` isn't in the logs...
https://github.com/dask/distributed/blob/4211e5cddbc59e31e338ddcd4e497a68893e6f14/distributed/tests/test_nanny.py#L485-L490

```python
@gen_cluster(nthreads=[])
async def test_worker_start_exception(s):
    with captured_logger(logger="distributed.nanny", level=logging.DEBUG) as logs:
        # make sure this raises the right Exception:
        with raises_with_cause(RuntimeError, None, StartException, None):
            async with Nanny(s.address, worker_class=BrokenWorker) as n:
                pass

    print("HERE")
    print(logs.getvalue())
    assert "Restarting worker" not in logs.getvalue()
    assert False # fail test so we can see stdout
```

...it passes.

Yet if we run it, the message we didn't want to see actually _is_ in the logs (but not in the captured logs)!
(😬 emojis added for visibility to the relevant spots)
```
(dask-distributed) gabe dev/distributed ‹fix-6363*› » pytest distributed/tests/test_nanny.py::test_worker_start_exception -s                                               1 ↵
============================================================================= test session starts =============================================================================
platform darwin -- Python 3.9.5, pytest-6.2.4, py-1.10.0, pluggy-0.13.1 -- /Users/gabe/miniconda3/envs/dask-distributed/bin/python
cachedir: .pytest_cache
rootdir: /Users/gabe/dev/distributed, configfile: setup.cfg
plugins: repeat-0.8.0, xdist-2.3.0, asyncio-0.12.0, timeout-1.4.2, forked-1.3.0, cov-3.0.0, rerunfailures-10.0
timeout: 300.0s
timeout method: thread
timeout func_only: False
collected 1 item                                                                                                                                                              

distributed/tests/test_nanny.py::test_worker_start_exception 2022-05-23 16:11:06,515 - distributed.http.proxy - INFO - To route to workers diagnostics web server please install jupyter-server-proxy: python -m pip install jupyter-server-proxy
2022-05-23 16:11:06,519 - distributed.scheduler - INFO - State start
2022-05-23 16:11:06,522 - distributed.scheduler - INFO - Clear task state
2022-05-23 16:11:06,523 - distributed.scheduler - INFO -   Scheduler at:     tcp://127.0.0.1:61231
2022-05-23 16:11:06,523 - distributed.scheduler - INFO -   dashboard at:           127.0.0.1:61230
2022-05-23 16:11:06,531 - distributed.nanny - INFO -         Start Nanny at: 'tcp://127.0.0.1:61232'
2022-05-23 16:11:07,260 - distributed.nanny - ERROR - Failed to start worker
Traceback (most recent call last):
  File "/Users/gabe/dev/distributed/distributed/nanny.py", line 872, in run
    await worker
  File "/Users/gabe/dev/distributed/distributed/tests/test_nanny.py", line 482, in start
    raise StartException("broken")
test_nanny.StartException: broken
2022-05-23 16:11:07,283 - distributed.nanny - ERROR - Failed while trying to start worker process: broken
2022-05-23 16:11:07,284 - distributed.nanny - ERROR - Failed to connect to process
Traceback (most recent call last):
  File "/Users/gabe/dev/distributed/distributed/nanny.py", line 681, in start
    msg = await self._wait_until_connected(uid)
  File "/Users/gabe/dev/distributed/distributed/nanny.py", line 801, in _wait_until_connected
    raise msg["exception"]
test_nanny.StartException: broken
2022-05-23 16:11:07,285 - distributed.nanny - ERROR - Failed to start process
Traceback (most recent call last):
  File "/Users/gabe/dev/distributed/distributed/nanny.py", line 425, in instantiate
    result = await self.process.start()
  File "/Users/gabe/dev/distributed/distributed/nanny.py", line 681, in start
    msg = await self._wait_until_connected(uid)
  File "/Users/gabe/dev/distributed/distributed/nanny.py", line 801, in _wait_until_connected
    raise msg["exception"]
test_nanny.StartException: broken
2022-05-23 16:11:07,285 - distributed.nanny - INFO - Closing Nanny at 'tcp://127.0.0.1:61232'.
HERE  # <--- 😬😬😬 OUR LOGGING CAPTURE ENDS. `async with` block has exited
        Start Nanny at: 'tcp://127.0.0.1:61232'
Failed while trying to start worker process: broken
Failed to connect to process
Traceback (most recent call last):
  File "/Users/gabe/dev/distributed/distributed/nanny.py", line 681, in start
    msg = await self._wait_until_connected(uid)
  File "/Users/gabe/dev/distributed/distributed/nanny.py", line 801, in _wait_until_connected
    raise msg["exception"]
test_nanny.StartException: broken
Failed to start process
Traceback (most recent call last):
  File "/Users/gabe/dev/distributed/distributed/nanny.py", line 425, in instantiate
    result = await self.process.start()
  File "/Users/gabe/dev/distributed/distributed/nanny.py", line 681, in start
    msg = await self._wait_until_connected(uid)
  File "/Users/gabe/dev/distributed/distributed/nanny.py", line 801, in _wait_until_connected
    raise msg["exception"]
test_nanny.StartException: broken
Closing Nanny at 'tcp://127.0.0.1:61232'.

Dumped cluster state to test_cluster_dump/test_worker_start_exception.yaml
2022-05-23 16:11:07,599 - distributed.nanny - INFO - Worker process 37919 was killed by signal 15
2022-05-23 16:11:07,601 - distributed.nanny - WARNING - Restarting worker  # <--- 😬😬😬 WORKER IS RESTARTING
2022-05-23 16:11:07,602 - distributed.scheduler - INFO - Scheduler closing...
2022-05-23 16:11:07,612 - distributed.scheduler - INFO - Scheduler closing all comms
FAILED

================================================================================== FAILURES ===================================================================================
_________________________________________________________________________ test_worker_start_exception _________________________________________________________________________

s = <Scheduler 'tcp://127.0.0.1:61231', workers: 0, cores: 0, tasks: 0>

    @gen_cluster(nthreads=[])
    async def test_worker_start_exception(s):
        with captured_logger(logger="distributed.nanny", level=logging.DEBUG) as logs:
            # make sure this raises the right Exception:
            with raises_with_cause(RuntimeError, None, StartException, None):
                async with Nanny(s.address, worker_class=BrokenWorker) as n:
                    pass
    
        print("HERE")
        print(logs.getvalue())
        assert "Restarting worker" not in logs.getvalue()
>       assert False # fail test so we can see stdout
E       assert False

distributed/tests/test_nanny.py:511: AssertionError
============================================================================ slowest 20 durations =============================================================================
1.11s call     distributed/tests/test_nanny.py::test_worker_start_exception

(2 durations < 0.005s hidden.  Use -vv to show these durations.)
=========================================================================== short test summary info ===========================================================================
FAILED distributed/tests/test_nanny.py::test_worker_start_exception - assert False
============================================================================== 1 failed in 1.37s ==============================================================================
```


Discussion:

Resolving patch:
diff --git a/distributed/core.py b/distributed/core.py
--- a/distributed/core.py
+++ b/distributed/core.py
@@ -15,7 +15,7 @@
 from contextlib import suppress
 from enum import Enum
 from functools import partial
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, TypedDict, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, TypedDict, TypeVar, final
 
 import tblib
 from tlz import merge
@@ -462,6 +462,7 @@ async def start_unsafe(self):
         await self.rpc.start()
         return self
 
+    @final
     async def start(self):
         async with self._startup_lock:
             if self.status == Status.failed:
diff --git a/distributed/nanny.py b/distributed/nanny.py
--- a/distributed/nanny.py
+++ b/distributed/nanny.py
@@ -398,7 +398,7 @@ async def instantiate(self) -> Status:
             self.process = WorkerProcess(
                 worker_kwargs=worker_kwargs,
                 silence_logs=self.silence_logs,
-                on_exit=self._on_exit_sync,
+                on_exit=self._on_worker_exit_sync,
                 worker=self.Worker,
                 env=self.env,
                 config=self.config,
@@ -490,19 +490,20 @@ def is_alive(self):
     def run(self, comm, *args, **kwargs):
         return run(self, comm, *args, **kwargs)
 
-    def _on_exit_sync(self, exitcode):
+    def _on_worker_exit_sync(self, exitcode):
         try:
-            self._ongoing_background_tasks.call_soon(self._on_exit, exitcode)
+            self._ongoing_background_tasks.call_soon(self._on_worker_exit, exitcode)
         except AsyncTaskGroupClosedError:  # Async task group has already been closed, so the nanny is already clos(ed|ing).
             pass
 
     @log_errors
-    async def _on_exit(self, exitcode):
+    async def _on_worker_exit(self, exitcode):
         if self.status not in (
             Status.init,
             Status.closing,
             Status.closed,
             Status.closing_gracefully,
+            Status.failed,
         ):
             try:
                 await self._unregister()
@@ -517,6 +518,7 @@ async def _on_exit(self, exitcode):
                 Status.closing,
                 Status.closed,
                 Status.closing_gracefully,
+                Status.failed,
             ):
                 logger.warning("Restarting worker")
                 await self.instantiate()
@@ -577,7 +579,7 @@ async def close(self, timeout=5):
             if self.process is not None:
                 await self.kill(timeout=timeout)
         except Exception:
-            pass
+            logger.exception("Error in Nanny killing Worker subprocess")
         self.process = None
         await self.rpc.close()
         self.status = Status.closed
@@ -662,15 +664,15 @@ async def start(self) -> Status:
             await self.process.start()
         except OSError:
             logger.exception("Nanny failed to start process", exc_info=True)
-            self.process.terminate()
+            # NOTE: doesn't wait for process to terminate, just for terminate signal to be sent
+            await self.process.terminate()
             self.status = Status.failed
-            return self.status
         try:
             msg = await self._wait_until_connected(uid)
         except Exception:
-            logger.exception("Failed to connect to process")
+            # NOTE: doesn't wait for process to terminate, just for terminate signal to be sent
+            await self.process.terminate()
             self.status = Status.failed
-            self.process.terminate()
             raise
         if not msg:
             return self.status
@@ -731,7 +733,12 @@ def mark_stopped(self):
     async def kill(self, timeout: float = 2, executor_wait: bool = True) -> None:
         """
         Ensure the worker process is stopped, waiting at most
-        *timeout* seconds before terminating it abruptly.
+        ``timeout * 0.8`` seconds before killing it abruptly.
+
+        When `kill` returns, the worker process has been joined.
+
+        If the worker process does not terminate within ``timeout`` seconds,
+        even after being killed, `asyncio.TimeoutError` is raised.
         """
         deadline = time() + timeout
 
@@ -740,32 +747,38 @@ async def kill(self, timeout: float = 2, executor_wait: bool = True) -> None:
         if self.status == Status.stopping:
             await self.stopped.wait()
             return
-        assert self.status in (Status.starting, Status.running)
+        assert self.status in (
+            Status.starting,
+            Status.running,
+            Status.failed,  # process failed to start, but hasn't been joined yet
+        ), self.status
         self.status = Status.stopping
         logger.info("Nanny asking worker to close")
 
         process = self.process
+        assert self.process
+        wait_timeout = timeout * 0.8
         self.child_stop_q.put(
             {
                 "op": "stop",
-                "timeout": max(0, deadline - time()) * 0.8,
+                "timeout": wait_timeout,
                 "executor_wait": executor_wait,
             }
         )
         await asyncio.sleep(0)  # otherwise we get broken pipe errors
         self.child_stop_q.close()
 
-        while process.is_alive() and time() < deadline:
-            await asyncio.sleep(0.05)
+        try:
+            await process.join(wait_timeout)
+            return
+        except asyncio.TimeoutError:
+            pass
 
-        if process.is_alive():
-            logger.warning(
-                f"Worker process still alive after {timeout} seconds, killing"
-            )
-            try:
-                await process.terminate()
-            except Exception as e:
-                logger.error("Failed to kill worker process: %s", e)
+        logger.warning(
+            f"Worker process still alive after {wait_timeout} seconds, killing"
+        )
+        await process.kill()
+        await process.join(max(0, deadline - time()))
 
     async def _wait_until_connected(self, uid):
         while True:
@@ -783,9 +796,6 @@ async def _wait_until_connected(self, uid):
                 continue
 
             if "exception" in msg:
-                logger.error(
-                    "Failed while trying to start worker process: %s", msg["exception"]
-                )
                 raise msg["exception"]
             else:
                 return msg
diff --git a/distributed/process.py b/distributed/process.py
--- a/distributed/process.py
+++ b/distributed/process.py
@@ -17,7 +17,7 @@
 
 import dask
 
-from distributed.utils import TimeoutError, get_mp_context
+from distributed.utils import get_mp_context
 
 logger = logging.getLogger(__name__)
 
@@ -299,15 +299,9 @@ async def join(self, timeout=None):
         assert self._state.pid is not None, "can only join a started process"
         if self._state.exitcode is not None:
             return
-        if timeout is None:
-            await self._exit_future
-        else:
-            try:
-                # Shield otherwise the timeout cancels the future and our
-                # on_exit callback will try to set a result on a canceled future
-                await asyncio.wait_for(asyncio.shield(self._exit_future), timeout)
-            except TimeoutError:
-                pass
+        # Shield otherwise the timeout cancels the future and our
+        # on_exit callback will try to set a result on a canceled future
+        await asyncio.wait_for(asyncio.shield(self._exit_future), timeout)
 
     def close(self):
         """


Test patch:
diff --git a/distributed/tests/test_asyncprocess.py b/distributed/tests/test_asyncprocess.py
--- a/distributed/tests/test_asyncprocess.py
+++ b/distributed/tests/test_asyncprocess.py
@@ -80,7 +80,8 @@ async def test_simple():
     assert proc.exitcode is None
 
     t1 = time()
-    await proc.join(timeout=0.02)
+    with pytest.raises(asyncio.TimeoutError):
+        await proc.join(timeout=0.02)
     dt = time() - t1
     assert 0.2 >= dt >= 0.001
     assert proc.is_alive()
diff --git a/distributed/tests/test_nanny.py b/distributed/tests/test_nanny.py
--- a/distributed/tests/test_nanny.py
+++ b/distributed/tests/test_nanny.py
@@ -13,6 +13,8 @@
 import psutil
 import pytest
 
+from distributed.diagnostics.plugin import WorkerPlugin
+
 pytestmark = pytest.mark.gpu
 
 from tlz import first, valmap
@@ -27,7 +29,7 @@
 from distributed.diagnostics import SchedulerPlugin
 from distributed.metrics import time
 from distributed.protocol.pickle import dumps
-from distributed.utils import TimeoutError, parse_ports
+from distributed.utils import TimeoutError, get_mp_context, parse_ports
 from distributed.utils_test import (
     captured_logger,
     gen_cluster,
@@ -474,21 +476,29 @@ async def test_nanny_closed_by_keyboard_interrupt(ucx_loop, protocol):
             assert "remove-worker" in str(s.events)
 
 
-class StartException(Exception):
-    pass
-
-
 class BrokenWorker(worker.Worker):
-    async def start(self):
-        raise StartException("broken")
+    async def start_unsafe(self):
+        raise ValueError("broken")
 
 
 @gen_cluster(nthreads=[])
 async def test_worker_start_exception(s):
-    # make sure this raises the right Exception:
-    with raises_with_cause(RuntimeError, None, StartException, None):
-        async with Nanny(s.address, worker_class=BrokenWorker) as n:
-            pass
+    nanny = Nanny(s.address, worker_class=BrokenWorker)
+    with captured_logger(logger="distributed.nanny", level=logging.WARNING) as logs:
+        with raises_with_cause(
+            RuntimeError,
+            "Nanny failed to start",
+            RuntimeError,
+            "BrokenWorker failed to start",
+        ):
+            async with nanny:
+                pass
+    assert nanny.status == Status.failed
+    # ^ NOTE: `Nanny.close` sets it to `closed`, then `Server.start._close_on_failure` sets it to `failed`
+    assert nanny.process is None
+    assert "Restarting worker" not in logs.getvalue()
+    # Avoid excessive spewing. (It's also printed once extra within the subprocess, which is okay.)
+    assert logs.getvalue().count("ValueError: broken") == 1, logs.getvalue()
 
 
 @gen_cluster(nthreads=[])
@@ -571,6 +581,44 @@ async def test_restart_memory(c, s, n):
         await asyncio.sleep(0.1)
 
 
+class BlockClose(WorkerPlugin):
+    def __init__(self, close_happened):
+        self.close_happened = close_happened
+
+    async def teardown(self, worker):
+        # Never let the worker cleanly shut down, so it has to be killed
+        self.close_happened.set()
+        while True:
+            await asyncio.sleep(10)
+
+
+@pytest.mark.slow
+@gen_cluster(nthreads=[])
+async def test_close_joins(s):
+    close_happened = get_mp_context().Event()
+
+    nanny = Nanny(s.address, plugins=[BlockClose(close_happened)])
+    async with nanny:
+        p = nanny.process
+        assert p
+        close_t = asyncio.create_task(nanny.close())
+
+        while not close_happened.wait(0):
+            await asyncio.sleep(0.01)
+
+        assert not close_t.done()
+        assert nanny.status == Status.closing
+        assert nanny.process and nanny.process.status == Status.stopping
+
+        await close_t
+
+        assert nanny.status == Status.closed
+        assert not nanny.process
+
+        assert p.status == Status.stopped
+        assert not p.process
+
+
 @gen_cluster(Worker=Nanny, nthreads=[("", 1)])
 async def test_scheduler_crash_doesnt_restart(s, a):
     # Simulate a scheduler crash by disconnecting it first