Log worker info on error

distributed/client.py

@@ -333,6 +333,9 @@
         await self._state.wait()
         if self.status == "error":
             exc = clean_exception(self._state.exception, self._state.traceback)
+            logger.error(
+                f"in task: {self._state.key}\non worker:{self._state.erred_on}"
+            )
             if raiseit:
                 typ, exc, tb = exc
                 raise exc.with_traceback(tb)
@@ -351,6 +354,9 @@
     async def _exception(self):
         await self._state.wait()
         if self.status == "error":
+            logger.error(
+                f"in task: {self._state.key}\non worker:{self._state.erred_on}"
+            )
             return self._state.exception
         else:
             return None
@@ -545,7 +551,15 @@
     This is shared between all Futures with the same key and client.
     """
 
-    __slots__ = ("_event", "status", "type", "exception", "traceback")
+    __slots__ = (
+        "_event",
+        "status",
+        "type",
+        "exception",
+        "traceback",
+        "key",
+        "erred_on",
+    )
 
     def __init__(self):
         self._event = None
@@ -590,7 +604,7 @@
         self.status = "pending"
         self._get_event().clear()
 
-    def set_error(self, exception, traceback):
+    def set_error(self, exception, traceback, erred_on, key):
         """Sets the error data
 
         Sets the status to 'error'. Sets the exception, the traceback,
@@ -602,9 +616,15 @@
             The exception
         traceback: Exception
             The traceback
+        erred_on: set
+            Workers where it errored
+        key: Hashable
+            Task key
         """
         _, exception, traceback = clean_exception(exception, traceback)
 
+        self.key = key
+        self.erred_on = erred_on
         self.status = "error"
         self.exception = exception
         self.traceback = traceback
@@ -1618,10 +1638,12 @@
         if state is not None:
             state.retry()
 
-    def _handle_task_erred(self, key=None, exception=None, traceback=None):
+    def _handle_task_erred(
+        self, key=None, exception=None, traceback=None, erred_on=None
+    ):
         state = self.futures.get(key)
         if state is not None:
-            state.set_error(exception, traceback)
+            state.set_error(exception, traceback, erred_on, key)
 
     def _handle_restart(self):
         logger.info("Receive restart signal from scheduler")
@@ -2228,6 +2250,7 @@
                         except (KeyError, AttributeError):
                             exc = CancelledError(key)
                         else:
+                            logger.error(f"in task: {st.key}\non worker:{st.erred_on}")
                             raise exception.with_traceback(traceback)
                         raise exc
                     if errors == "skip":
@@ -3241,6 +3264,9 @@
                     should_rejoin = False
             try:
                 results = self.gather(packed, asynchronous=asynchronous, direct=direct)
+            except Exception:
+                exc_type, exc_value, exc_traceback = sys.exc_info()
+                raise exc_value.with_traceback(exc_traceback)
             finally:
                 for f in futures.values():
                     f.release()

distributed/queues.py

@@ -267,7 +267,7 @@ def process(d):
             if d["type"] == "Future":
                 value = Future(d["value"], self.client, inform=True, state=d["state"])
                 if d["state"] == "erred":
-                    value._state.set_error(d["exception"], d["traceback"])
+                    value._state.set_error(d["exception"], d["traceback"], None, None)
                 self.client._send_to_scheduler(
                     {"op": "queue-future-release", "name": self.name, "key": d["value"]}
                 )

distributed/scheduler.py

@@ -2483,6 +2483,7 @@ def transition_released_erred(self, key: str, stimulus_id: str) -> RecsMsgs:
             "key": key,
             "exception": failing_ts.exception,
             "traceback": failing_ts.traceback,
+            "erred_on": failing_ts.erred_on,
         }
         for cs in ts.who_wants:
             client_msgs[cs.client_key] = [report_msg]
@@ -2677,6 +2678,7 @@ def transition_processing_erred(
             "key": key,
             "exception": failing_ts.exception,
             "traceback": failing_ts.traceback,
+            "erred_on": failing_ts.erred_on,
         }
         for cs in ts.who_wants:
             client_msgs[cs.client_key] = [report_msg]

distributed/tests/test_client.py

@@ -400,18 +400,34 @@ async def test_future_tuple_repr(c, s, a, b):
 
 @gen_cluster(client=True)
 async def test_Future_exception(c, s, a, b):
-    x = c.submit(div, 1, 0)
-    result = await x.exception()
-    assert isinstance(result, ZeroDivisionError)
+    with captured_logger("distributed.client") as logger:
+        x = c.submit(div, 1, 0)
+        exc = await x.exception()
+        assert isinstance(exc, ZeroDivisionError)
+    assert logger.getvalue().startswith("in task: ")
+
+    with captured_logger("distributed.client") as logger:
+        x = c.submit(div, 1, 0)
+        with pytest.raises(ZeroDivisionError):
+            _ = await x.result()
+    assert logger.getvalue().startswith("in task: ")
 
     x = c.submit(div, 1, 1)
     result = await x.exception()
     assert result is None
 
 
 def test_Future_exception_sync(c):
-    x = c.submit(div, 1, 0)
-    assert isinstance(x.exception(), ZeroDivisionError)
+    with captured_logger("distributed.client") as logger:
+        x = c.submit(div, 1, 0)
+        assert isinstance(x.exception(), ZeroDivisionError)
+    assert logger.getvalue().startswith("in task: ")
+
+    with captured_logger("distributed.client") as logger:
+        x = c.submit(div, 1, 0)
+        with pytest.raises(ZeroDivisionError):
+            _ = x.result()
+    assert logger.getvalue().startswith("in task: ")
 
     x = c.submit(div, 1, 1)
     assert x.exception() is None

distributed/variable.py

@@ -208,7 +208,7 @@ async def _get(self, timeout=None):
         if d["type"] == "Future":
             value = Future(d["value"], self.client, inform=True, state=d["state"])
             if d["state"] == "erred":
-                value._state.set_error(d["exception"], d["traceback"])
+                value._state.set_error(d["exception"], d["traceback"], None, None)
             self.client._send_to_scheduler(
                 {
                     "op": "variable-future-release",