shield kill

distributed/nanny.py

@@ -516,7 +516,7 @@ async def _():
                 await self.instantiate()
 
         try:
-            await wait_for(_(), timeout)
+            await wait_for(asyncio.shield(_()), timeout)
         except asyncio.TimeoutError:
             logger.error(
                 f"Restart timed out after {timeout}s; returning before finished"
@@ -769,7 +769,6 @@ async def start(self) -> Status:
         self.worker_dir = msg["dir"]
         assert self.worker_address
         self.status = Status.running
-        self.running.set()
 
         return self.status
 
@@ -804,6 +803,7 @@ def mark_stopped(self):
                 msg = self._death_message(self.process.pid, r)
                 logger.info(msg)
             self.status = Status.stopped
+            self.running.clear()
             self.stopped.set()
             # Release resources
             self.process.close()

distributed/tests/test_nanny.py

@@ -208,25 +208,47 @@ async def test_scheduler_file():
         s.stop()
 
 
-@pytest.mark.xfail(
-    os.environ.get("MINDEPS") == "true",
-    reason="Timeout errors with mindeps environment",
-)
-@gen_cluster(client=True, Worker=Nanny, nthreads=[("127.0.0.1", 2)])
-async def test_nanny_timeout(c, s, a):
+@gen_cluster(client=True, Worker=Nanny, nthreads=[("", 1)])
+async def test_nanny_restart(c, s, a):
+    x = await c.scatter(123)
+    assert await c.submit(lambda: 1) == 1
+
+    await a.restart()
+
+    while x.status != "cancelled":
+        await asyncio.sleep(0.1)
+
+    assert await c.submit(lambda: 1) == 1
+
+
+@gen_cluster(client=True, Worker=Nanny, nthreads=[("", 1)])
+async def test_nanny_restart_timeout(c, s, a):
     x = await c.scatter(123)
     with captured_logger(
         logging.getLogger("distributed.nanny"), level=logging.ERROR
     ) as logger:
-        await a.restart(timeout=0.1)
+        await a.restart(timeout=0)
 
     out = logger.getvalue()
     assert "timed out" in out.lower()
 
-    start = time()
     while x.status != "cancelled":
         await asyncio.sleep(0.1)
-        assert time() < start + 7
+
+    assert await c.submit(lambda: 1) == 1
+
+
+@gen_cluster(client=True, Worker=Nanny, nthreads=[("", 1)])
+async def test_nanny_restart_timeout_stress(c, s, a):
+    x = await c.scatter(123)
+    restarts = [a.restart(timeout=random.random() * 0.1) for _ in range(100)]
+    await asyncio.gather(*restarts)
+
+    while x.status != "cancelled":
+        await asyncio.sleep(0.1)
+
+    assert await c.submit(lambda: 1) == 1
+    assert len(s.workers) == 1
 
 
 @gen_cluster(