diff --git a/distributed/nanny.py b/distributed/nanny.py index 7d51ba105a..395e4725f7 100644 --- a/distributed/nanny.py +++ b/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 diff --git a/distributed/tests/test_nanny.py b/distributed/tests/test_nanny.py index cd578534e4..4dc6dd3cce 100644 --- a/distributed/tests/test_nanny.py +++ b/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