Fix timeout on startup

dask · May 11, 2022 · 159bc09 · 159bc09
1 parent 622feb5
commit 159bc09
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/distributed/cli/dask_worker.py b/distributed/cli/dask_worker.py
@@ -478,10 +478,12 @@ async def wait_for_signals_and_close():
             wait_for_nannies_to_finish()
         )
 
-        await asyncio.wait(
+        done, _ = await asyncio.wait(
             [wait_for_signals_and_close_task, wait_for_nannies_to_finish_task],
             return_when=asyncio.FIRST_COMPLETED,
         )
+        # Re-raise exceptions from done tasks
+        [task.result() for task in done]
 
     try:
         asyncio.run(run())

diff --git a/distributed/cli/tests/test_dask_worker.py b/distributed/cli/tests/test_dask_worker.py
@@ -690,7 +690,14 @@ def dask_setup(worker):
 @pytest.mark.parametrize("nanny", ["--nanny", "--no-nanny"])
 def test_timeout(nanny):
     worker = subprocess.run(
-        ["dask-worker", "192.168.1.100:7777", nanny, "--death-timeout=1"],
+        [
+            "python",
+            "-m",
+            "distributed.cli.dask_worker",
+            "192.168.1.100:7777",
+            nanny,
+            "--death-timeout=1",
+        ],
         text=True,
         encoding="utf8",
         capture_output=True,