Don't lose worker startup exception when its process exit is observed first (flaky test_local_cluster_redundant_kwarg)

distributed/nanny.py

@@ -721,6 +721,11 @@ async def start(self) -> Status:
         # See note in Nanny docstring.
         os.environ.update(self.pre_spawn_env)
 
+        # If the process dies, mark_stopped() (fired by the process exit callback)
+        # releases self.init_result_q. Hold a reference so that we can read the
+        # exception message that the process may have sent just before dying.
+        init_result_q = self.init_result_q
+
         try:
             try:
                 await self.process.start()
@@ -734,7 +739,7 @@ async def start(self) -> Status:
                 self.status = Status.failed
 
             try:
-                msg = await self._wait_until_connected(uid)
+                msg = await self._wait_until_connected(uid, init_result_q)
             except Exception:
                 logger.error("Worker failed to connect")
                 # The process may have already exited and been released by
@@ -758,7 +763,7 @@ async def start(self) -> Status:
         finally:
             self.running.set()
 
-        if msg:
+        if msg and self.status == Status.starting:
             self.worker_address = msg["address"]
             self.worker_dir = msg["dir"]
             assert self.worker_address
@@ -879,15 +884,23 @@ async def kill(
                 return
             raise
 
-    async def _wait_until_connected(self, uid):
+    async def _wait_until_connected(self, uid, init_result_q):
         while True:
-            if self.status != Status.starting:
-                return
+            # Read the status *before* polling the queue: if the process failed to
+            # start a worker, it sends the exception through init_result_q, flushes
+            # the queue, and terminates. If mark_stopped() (fired by the process
+            # exit callback) flipped the status away from Status.starting, the
+            # message the process may have sent while dying is guaranteed to be
+            # readable now, so one last get_nowait() will not lose it.
+            stopped = self.status != Status.starting
+
             # This is a multiprocessing queue and we'd block the event loop if
             # we simply called get
             try:
-                msg = self.init_result_q.get_nowait()
+                msg = init_result_q.get_nowait()
             except Empty:
+                if stopped:
+                    return None
                 await asyncio.sleep(self._init_msg_interval)
                 continue
 

distributed/tests/test_nanny.py

@@ -26,6 +26,7 @@
 from distributed.diagnostics import SchedulerPlugin
 from distributed.diagnostics.plugin import NannyPlugin, WorkerPlugin
 from distributed.metrics import time
+from distributed.nanny import WorkerProcess
 from distributed.protocol.pickle import dumps
 from distributed.utils import TimeoutError, get_mp_context, parse_ports
 from distributed.utils_test import (
@@ -638,6 +639,33 @@ async def test_failure_during_worker_initialization(s):
     assert "Restarting worker" not in logs.getvalue()
 
 
+@gen_cluster(nthreads=[])
+async def test_worker_start_exception_after_process_exit(s, monkeypatch):
+    """If the worker fails to start, the process sends the exception to the nanny
+    through init_result_q and terminates. The exception must not be lost if the
+    nanny observes the process exit before it reads the message from the queue
+    (flaky test_local_cluster_redundant_kwarg).
+    """
+    orig_wait_until_connected = WorkerProcess._wait_until_connected
+
+    async def wait_until_connected(self, *args, **kwargs):
+        # Force mark_stopped(), fired by the process exit callback, to win the
+        # race against the polling of init_result_q
+        await self.stopped.wait()
+        return await orig_wait_until_connected(self, *args, **kwargs)
+
+    monkeypatch.setattr(WorkerProcess, "_wait_until_connected", wait_until_connected)
+
+    with raises_with_cause(
+        RuntimeError,
+        "Nanny failed to start",
+        TypeError,
+        "unexpected keyword argument",
+    ):
+        async with Nanny(s.address, foo="bar"):
+            pass
+
+
 @gen_cluster(client=True, Worker=Nanny)
 async def test_environ_plugin(c, s, a, b):
     from dask.distributed import Environ