Add fail_hard decorator for worker methods

distributed/tests/test_utils_test.py

@@ -679,3 +679,17 @@ async def test_log_invalid_worker_task_states(c, s, a):
 
     assert "released" in out + err
     assert "task-name" in out + err
+
+
+def test_worker_fail_hard(capsys):
+    @gen_cluster(client=True, nthreads=[("127.0.0.1", 1)])
+    async def test_fail_hard(c, s, a):
+        with pytest.raises(Exception):
+            await a.gather_dep(
+                worker="abcd", to_gather=["x"], total_nbytes=0, stimulus_id="foo"
+            )
+
+    with pytest.raises(Exception) as info:
+        test_fail_hard()
+
+    assert "abcd" in str(info.value)

distributed/utils_test.py

@@ -50,13 +50,21 @@
 from distributed.comm.tcp import TCP
 from distributed.compatibility import WINDOWS
 from distributed.config import initialize_logging
-from distributed.core import CommClosedError, ConnectionPool, Status, connect, rpc
+from distributed.core import (
+    CommClosedError,
+    ConnectionPool,
+    Status,
+    clean_exception,
+    connect,
+    rpc,
+)
 from distributed.deploy import SpecCluster
 from distributed.diagnostics.plugin import WorkerPlugin
 from distributed.metrics import time
 from distributed.nanny import Nanny
 from distributed.node import ServerNode
 from distributed.proctitle import enable_proctitle_on_children
+from distributed.protocol import deserialize
 from distributed.security import Security
 from distributed.utils import (
     DequeHandler,
@@ -867,6 +875,7 @@ async def start_cluster(
             await s.close(fast=True)
             check_invalid_worker_transitions(s)
             check_invalid_task_states(s)
+            check_worker_fail_hard(s)
             raise TimeoutError("Cluster creation timeout")
     return s, workers
 
@@ -898,6 +907,20 @@ def check_invalid_task_states(s: Scheduler) -> None:
     raise ValueError("Invalid worker task state")
 
 
+def check_worker_fail_hard(s: Scheduler) -> None:
+    if not s.events.get("worker-fail-hard"):
+        return
+
+    for timestamp, msg in s.events["worker-fail-hard"]:
+        msg = msg.copy()
+        worker = msg.pop("worker")
+        msg["exception"] = deserialize(msg["exception"].header, msg["exception"].frames)
+        msg["traceback"] = deserialize(msg["traceback"].header, msg["traceback"].frames)
+        print("Failed worker", worker)
+        typ, exc, tb = clean_exception(**msg)
+        raise exc.with_traceback(tb)
+
+
 async def end_cluster(s, workers):
     logger.debug("Closing out test cluster")
 
@@ -910,6 +933,7 @@ async def end_worker(w):
     s.stop()
     check_invalid_worker_transitions(s)
     check_invalid_task_states(s)
+    check_worker_fail_hard(s)
 
 
 def gen_cluster(

distributed/worker.py

@@ -157,6 +157,32 @@
 DEFAULT_STARTUP_INFORMATION: dict[str, Callable[[Worker], Any]] = {}
 
 
+def fail_hard(method):
+    """
+    Decorator to close the worker if this method encounters an exception
+    """
+    assert iscoroutinefunction(method)
+
+    @functools.wraps(method)
+    async def wrapper(self, *args, **kwargs):
+        try:
+            return await method(self, *args, **kwargs)
+        except Exception as e:
+            logger.exception(e)
+            self.log_event(
+                "worker-fail-hard",
+                {
+                    **error_message(e),
+                    "worker": self.address,
+                },
+            )
+            # TODO: send event to scheduler
+            await self.close(nanny=False, executor_wait=False)
+            raise
+
+    return wrapper
+
+
 class Worker(ServerNode):
     """Worker node in a Dask distributed cluster
 
@@ -900,14 +926,15 @@ def logs(self):
         return self._deque_handler.deque
 
     def log_event(self, topic, msg):
-        self.loop.add_callback(
-            self.batched_stream.send,
-            {
-                "op": "log-event",
-                "topic": topic,
-                "msg": msg,
-            },
-        )
+        full_msg = {
+            "op": "log-event",
+            "topic": topic,
+            "msg": msg,
+        }
+        if self.thread_id == threading.get_ident():
+            self.batched_stream.send(full_msg)
+        else:
+            self.loop.add_callback(self.batched_stream.send, full_msg)
 
     @property
     def executing_count(self) -> int:
@@ -1199,6 +1226,7 @@ async def heartbeat(self):
         finally:
             self.heartbeat_active = False
 
+    @fail_hard
     async def handle_scheduler(self, comm):
         try:
             await self.handle_stream(comm, every_cycle=[self.ensure_communicating])
@@ -3020,6 +3048,7 @@ def _update_metrics_received_data(
         self.counters["transfer-count"].add(len(data))
         self.incoming_count += 1
 
+    @fail_hard
     @log_errors
     async def gather_dep(
         self,
@@ -4104,7 +4133,6 @@ def validate_state(self):
 
         except Exception as e:
             logger.error("Validate state failed.  Closing.", exc_info=e)
-            self.loop.add_callback(self.close)
             logger.exception(e)
             if LOG_PDB:
                 import pdb