Ensure actors set erred state properly in case of worker failure

Author: fjetter

distributed/actor.py

@@ -141,7 +141,7 @@ def __dir__(self):
 
     def __getattr__(self, key):
         if self._future and self._future.status not in ("finished", "pending"):
-            raise ValueError(
+            raise RuntimeError(
                 "Worker holding Actor was lost.  Status: " + self._future.status
             )
         self._try_bind_worker_client()

distributed/scheduler.py

@@ -2599,7 +2599,7 @@ def _transition_memory_released(self, key: Key, stimulus_id: str) -> RecsMsgs:
             if ts.who_wants:
                 ts.exception_blame = ts
                 ts.exception = Serialized(
-                    *serialize(ValueError("Worker holding Actor was lost"))
+                    *serialize(RuntimeError("Worker holding Actor was lost"))
                 )
                 return {ts.key: "erred"}, {}, {}  # don't try to recreate
 
@@ -2652,7 +2652,7 @@ def _transition_released_erred(self, key: Key, stimulus_id: str) -> RecsMsgs:
 
         if self.validate:
             assert ts.exception_blame
-            assert not ts.who_has
+            assert not ts.who_has or ts.actor
             assert not ts.waiting_on
 
         failing_ts = ts.exception_blame
@@ -2772,8 +2772,8 @@ def _transition_processing_erred(
         self,
         key: Key,
         stimulus_id: str,
-        worker: str,
         *,
+        worker: str | None = None,
         cause: Key | None = None,
         exception: Serialized | None = None,
         traceback: Serialized | None = None,
@@ -2988,6 +2988,45 @@ def _remove_key(self, key: Key) -> None:
         ts.exception_blame = ts.exception = ts.traceback = None
         self.task_metadata.pop(key, None)
 
+    def _transition_memory_erred(self, key: Key, stimulus_id: str) -> RecsMsgs:
+        ts = self.tasks[key]
+        if self.validate:
+            assert ts.actor
+        recommendations: Recs = {}
+        client_msgs: Msgs = {}
+        worker_msgs: Msgs = {}
+        # XXX factor this out?
+        worker_msg = {
+            "op": "free-keys",
+            "keys": [key],
+            "stimulus_id": stimulus_id,
+        }
+        for ws in ts.who_has or ():
+            worker_msgs[ws.address] = [worker_msg]
+        self.remove_all_replicas(ts)
+
+        for dts in ts.dependents:
+            if not dts.who_has:
+                dts.exception_blame = ts
+                recommendations[dts.key] = "erred"
+        exception = Serialized(
+            *serialize(RuntimeError("Worker holding Actor was lost"))
+        )
+        report_msg = {
+            "op": "task-erred",
+            "key": key,
+            "exception": exception,
+        }
+        for cs in ts.who_wants or ():
+            client_msgs[cs.client_key] = [report_msg]
+
+        ts.state = "erred"
+        return self._propagate_erred(
+            ts,
+            cause=ts.key,
+            exception=exception,
+        )
+
     def _transition_memory_forgotten(self, key: Key, stimulus_id: str) -> RecsMsgs:
         ts = self.tasks[key]
 
@@ -3078,6 +3117,7 @@ def _transition_released_forgotten(self, key: Key, stimulus_id: str) -> RecsMsgs
         ("no-worker", "processing"): _transition_no_worker_processing,
         ("no-worker", "erred"): _transition_no_worker_erred,
         ("released", "forgotten"): _transition_released_forgotten,
+        ("memory", "erred"): _transition_memory_erred,
         ("memory", "forgotten"): _transition_memory_forgotten,
         ("erred", "released"): _transition_erred_released,
         ("memory", "released"): _transition_memory_released,
@@ -5521,7 +5561,9 @@ async def remove_worker(
 
         for ts in list(ws.has_what):
             self.remove_replica(ts, ws)
-            if not ts.who_has:
+            if ts in ws.actors:
+                recommendations[ts.key] = "erred"
+            elif not ts.who_has:
                 if ts.run_spec:
                     recompute_keys.add(ts.key)
                     recommendations[ts.key] = "released"

distributed/tests/test_actor.py

@@ -13,6 +13,7 @@
     Actor,
     BaseActorFuture,
     Client,
+    Event,
     Future,
     Nanny,
     as_completed,
@@ -23,7 +24,7 @@
 from distributed.metrics import time
 from distributed.utils import LateLoopEvent
 from distributed.utils_test import cluster, double, gen_cluster, inc
-from distributed.worker import get_worker
+from distributed.worker import Worker, get_worker
 
 
 class Counter:
@@ -290,7 +291,7 @@ async def test_failed_worker(c, s, a, b):
 
     await a.close()
 
-    with pytest.raises(ValueError, match="Worker holding Actor was lost"):
+    with pytest.raises(RuntimeError, match="Worker holding Actor was lost"):
         await counter.increment()
 
 
@@ -824,3 +825,68 @@ def demo(self):
 
     actor = await c.submit(Actor, actor=True, workers=[a.address])
     assert await actor.demo() == a.address
+
+
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_actor_worker_host_leaves_gracefully(c, s, a):
+    # see also test_actor_worker_host_dies
+    async with Worker(s.address, nthreads=1) as b:
+
+        counter = await c.submit(
+            Counter, actor=True, workers=[b.address], allow_other_workers=True
+        )
+
+        enter_ev = Event()
+        wait_ev = Event()
+
+        def foo(couner, enter_ev, wait_ev):
+            enter_ev.set()
+            wait_ev.wait()
+
+        fut = c.submit(
+            foo,
+            counter,
+            enter_ev,
+            wait_ev,
+            workers=[a.address],
+            allow_other_workers=True,
+        )
+
+        await enter_ev.wait()
+    await wait_ev.set()
+    with pytest.raises(RuntimeError, match="Worker holding Actor was lost"):
+        await fut.result()
+
+
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_actor_worker_host_dies(c, s, a):
+    # see also test_actor_worker_host_leaves_gracefully
+    async with Worker(s.address, nthreads=1) as b:
+
+        counter = await c.submit(
+            Counter, actor=True, workers=[b.address], allow_other_workers=True
+        )
+
+        enter_ev = Event()
+        wait_ev = Event()
+
+        def foo(couner, enter_ev, wait_ev):
+            enter_ev.set()
+            wait_ev.wait()
+
+        fut = c.submit(
+            foo,
+            counter,
+            enter_ev,
+            wait_ev,
+            workers=[a.address],
+            allow_other_workers=True,
+        )
+
+        await enter_ev.wait()
+        # Simulate the worker going down
+        s.stream_comms[b.address].close()
+        await b.finished()
+        await wait_ev.set()
+        with pytest.raises(RuntimeError, match="Worker holding Actor was lost"):
+            await fut.result()