Repository: dask/distributed

Issue:
`@log_errors` can't decorate asyncio tasks
This logs the exception:
```python
@log_errors
async def f():
    assert False

await f()
```
but this doesn't:
```python
await asyncio.create_task(f())
```

From debugging, I see that here:
https://github.com/dask/distributed/blob/095a09d4e71385f9ceed45de0780480dd26b9015/distributed/utils.py#L831-L836

`modname` is `asyncio.events` instead of the module that raised the error.


Discussion:

Resolving patch:
diff --git a/distributed/dashboard/components/worker.py b/distributed/dashboard/components/worker.py
--- a/distributed/dashboard/components/worker.py
+++ b/distributed/dashboard/components/worker.py
@@ -51,7 +51,7 @@
 
 def standard_doc(title, active_page, *, template="simple.html"):
     def decorator(f):
-        @log_errors(unroll_stack=2)
+        @log_errors(unroll_stack=1)
         def wrapper(arg, extra, doc):
             doc.title = title
             doc.template = env.get_template(template)
diff --git a/distributed/utils.py b/distributed/utils.py
--- a/distributed/utils.py
+++ b/distributed/utils.py
@@ -16,6 +16,7 @@
 import sys
 import tempfile
 import threading
+import traceback
 import warnings
 import weakref
 import xml.etree.ElementTree
@@ -755,7 +756,7 @@ def log_errors(*, pdb: bool = False, unroll_stack: int = 1) -> _LogErrors:
     ...
 
 
-def log_errors(func=None, /, *, pdb=False, unroll_stack=1):
+def log_errors(func=None, /, *, pdb=False, unroll_stack=0):
     """Log any errors and then reraise them.
 
     This can be used:
@@ -783,13 +784,34 @@ def func(...):
         Set to True to break into the debugger in case of exception
     unroll_stack: int, optional
         Number of levels of stack to unroll when determining the module's name for the
-        purpose of logging. Normally you should omit this. Set to 2 if you are writing a
+        purpose of logging. Normally you should omit this. Set to 1 if you are writing a
         helper function, context manager, or decorator.
     """
     le = _LogErrors(pdb=pdb, unroll_stack=unroll_stack)
     return le(func) if func else le
 
 
+_getmodulename_with_path_map: dict[str, str] = {}
+
+
+def _getmodulename_with_path(fname: str) -> str:
+    """Variant of inspect.getmodulename that returns the full module path"""
+    try:
+        return _getmodulename_with_path_map[fname]
+    except KeyError:
+        pass
+
+    for modname, mod in sys.modules.items():
+        fname2 = getattr(mod, "__file__", None)
+        if fname2:
+            _getmodulename_with_path_map[fname2] = modname
+
+    try:
+        return _getmodulename_with_path_map[fname]
+    except KeyError:  # pragma: nocover
+        return os.path.splitext(os.path.basename(fname))[0]
+
+
 class _LogErrors:
     __slots__ = ("pdb", "unroll_stack")
 
@@ -820,16 +842,15 @@ def wrapper(*args, **kwargs):
     def __enter__(self):
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(self, exc_type, exc_value, tb):
         from distributed.comm import CommClosedError
 
         if not exc_type or issubclass(exc_type, (CommClosedError, gen.Return)):
             return
 
-        stack = inspect.stack()
+        stack = traceback.extract_tb(tb)
         frame = stack[self.unroll_stack]
-        mod = inspect.getmodule(frame[0])
-        modname = mod.__name__
+        modname = _getmodulename_with_path(frame.filename)
 
         try:
             logger = logging.getLogger(modname)


Test patch:
diff --git a/distributed/tests/test_utils.py b/distributed/tests/test_utils.py
--- a/distributed/tests/test_utils.py
+++ b/distributed/tests/test_utils.py
@@ -9,6 +9,7 @@
 import os
 import queue
 import socket
+import sys
 import traceback
 import warnings
 import xml
@@ -934,6 +935,9 @@ async def _():
 
         assert await _() == 123
 
+        task = asyncio.create_task(_())
+        assert await task == 123
+
         @log_errors
         async def _():
             raise CustomError("err6")
@@ -941,6 +945,29 @@ async def _():
         with pytest.raises(CustomError):
             await _()
 
+        task = asyncio.create_task(_())
+        with pytest.raises(CustomError):
+            await task
+
+        @log_errors()
+        async def _():
+            raise CustomError("err7")
+
+        with pytest.raises(CustomError):
+            await _()
+
+        task = asyncio.create_task(_())
+        with pytest.raises(CustomError):
+            await task
+
+        async def _():
+            with log_errors():
+                raise CustomError("err8")
+
+        task = asyncio.create_task(_())
+        with pytest.raises(CustomError):
+            await task
+
     assert [row for row in caplog.getvalue().splitlines() if row.startswith("err")] == [
         "err1",
         "err2",
@@ -951,15 +978,35 @@ async def _():
         "err5",
         "err5",
         "err6",
+        "err6",
+        "err7",
+        "err7",
+        "err8",
     ]
 
-    # Test unroll_stack
-    with captured_logger("distributed.utils") as caplog:
-        with pytest.raises(CustomError):
-            with log_errors(unroll_stack=0):
-                raise CustomError("err7")
 
-    assert caplog.getvalue().startswith("err7\n")
+@pytest.mark.parametrize("unroll_stack,logger_name", [(0, "test_utils"), (1, "a.b.c")])
+def test_log_errors_unroll_stack(unroll_stack, logger_name, tmp_path):
+    (tmp_path / "a").mkdir()
+    (tmp_path / "a").touch("__init__.py")
+    (tmp_path / "a" / "b").mkdir()
+    (tmp_path / "a" / "b").touch("__init__.py")
+    with (tmp_path / "a" / "b" / "c.py").open("w") as fh:
+        fh.write("def f():\n    raise ValueError('myerr')\n")
+
+    sys.path.insert(0, str(tmp_path))
+    import a.b.c
+
+    del sys.path[0]
+
+    with (
+        captured_logger(logger_name) as caplog,
+        pytest.raises(ValueError),
+        log_errors(unroll_stack=unroll_stack),
+    ):
+        a.b.c.f()
+
+    assert "myerr" in caplog.getvalue()
 
 
 def test_load_json_robust_timeout(tmp_path):