Repository: dask/dask

Issue:
Importing dask breaks ipython's %pdb command
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**Describe the issue**:

In an ipython shell, running the `%pdb` command turns on automatic pdb calling. If any later commands cause an exception to be raised, the debugger should be entered after the traceback is printed. If dask is imported, the debugger is no longer entered.

**Minimal Complete Verifiable Example**:

In an ipython shell:

```python
In [1]: %pdb
Automatic pdb calling has been turned ON

In [2]: import dask

In [3]: raise ValueError("foo")
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[3], line 1
----> 1 raise ValueError("foo")

ValueError: foo

In [4]:
```

If dask is not imported, the debugger is entered as expected:

```
In [1]: %pdb
Automatic pdb calling has been turned ON

In [2]: raise ValueError("foo")
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[2], line 1
----> 1 raise ValueError("foo")

ValueError: foo
> <ipython-input-2-66d99343b0b5>(1)<module>()
----> 1 raise ValueError("foo")

ipdb>
```

**Anything else we need to know?**:

**Environment**:

- Dask version: This happens with versions 2023.6.1, 2023.7.0 and 2023.7.1. This does not happen with 2023.6.0
- Python version: 3.11.4
- Operating System: Debian Bookworm (in python:3.11.4 docker container running on macOS 13.5)
- Install method (conda, pip, source): pip



Discussion:
Thanks @laurence-kobold for the issue. I'm able to reproduce and it looks like https://github.com/dask/dask/pull/10354 is where behavior started to change. @crusaderky do you have bandwidth to look into this? 

Resolving patch:
diff --git a/dask/base.py b/dask/base.py
--- a/dask/base.py
+++ b/dask/base.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import atexit
 import dataclasses
 import datetime
 import hashlib
@@ -8,7 +7,6 @@
 import os
 import pathlib
 import pickle
-import sys
 import threading
 import uuid
 import warnings
@@ -18,7 +16,7 @@
 from contextlib import contextmanager
 from contextvars import ContextVar
 from enum import Enum
-from functools import partial, wraps
+from functools import partial
 from numbers import Integral, Number
 from operator import getitem
 from typing import TYPE_CHECKING, Any, Literal, Protocol
@@ -63,43 +61,6 @@
 if TYPE_CHECKING:
     from _typeshed import ReadableBuffer
 
-
-def _clean_traceback_hook(func):
-    @wraps(func)
-    def wrapper(exc_type, exc, tb):
-        tb = shorten_traceback(tb)
-        return func(exc_type, exc.with_traceback(tb), tb)
-
-    return wrapper
-
-
-def _clean_ipython_traceback(self, etype, value, tb, tb_offset=None):
-    short_tb = shorten_traceback(tb)
-    short_exc = value.with_traceback(short_tb)
-    stb = self.InteractiveTB.structured_traceback(
-        etype, short_exc, short_tb, tb_offset=tb_offset
-    )
-    self._showtraceback(etype, short_exc, stb)
-
-
-try:
-    from IPython import get_ipython
-except ImportError:
-    pass
-else:
-    # if we're running in ipython, customize exception handling
-    ip = get_ipython()
-    if ip is not None:
-        ip.set_custom_exc((Exception,), _clean_ipython_traceback)
-
-
-def _restore_excepthook():
-    sys.excepthook = original_excepthook
-
-
-original_excepthook = sys.excepthook
-sys.excepthook = _clean_traceback_hook(sys.excepthook)
-
 _annotations: ContextVar[dict[str, Any]] = ContextVar("annotations", default={})
 
 
@@ -663,7 +624,9 @@ def compute(
         keys.append(x.__dask_keys__())
         postcomputes.append(x.__dask_postcompute__())
 
-    results = schedule(dsk, keys, **kwargs)
+    with shorten_traceback():
+        results = schedule(dsk, keys, **kwargs)
+
     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
 
 
@@ -964,7 +927,9 @@ def persist(*args, traverse=True, optimize_graph=True, scheduler=None, **kwargs)
         keys.extend(a_keys)
         postpersists.append((rebuild, a_keys, state))
 
-    results = schedule(dsk, keys, **kwargs)
+    with shorten_traceback():
+        results = schedule(dsk, keys, **kwargs)
+
     d = dict(zip(keys, results))
     results2 = [r({k: d[k] for k in ks}, *s) for r, ks, s in postpersists]
     return repack(results2)
@@ -1603,6 +1568,3 @@ def clone_key(key, seed):
         prefix = key_split(key)
         return prefix + "-" + tokenize(key, seed)
     raise TypeError(f"Expected str or tuple[str, Hashable, ...]; got {key}")
-
-
-atexit.register(_restore_excepthook)
diff --git a/dask/utils.py b/dask/utils.py
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -8,7 +8,7 @@
 import shutil
 import sys
 import tempfile
-import traceback
+import types
 import uuid
 import warnings
 from collections.abc import Hashable, Iterable, Iterator, Mapping, Set
@@ -2130,52 +2130,53 @@ def get_meta_library(like):
     return import_module(typename(like).partition(".")[0])
 
 
-def shorten_traceback(exc_traceback):
-    """Remove irrelevant stack elements from traceback.
+class shorten_traceback:
+    """Context manager that removes irrelevant stack elements from traceback.
 
-    * only shortens traceback if any of the traceback lines match
-      `admin.traceback.shorten.when`
-    * omits frames from modules that match `admin.traceback.shorten.what`
+    * omits frames from modules that match `admin.traceback.shorten`
     * always keeps the first and last frame.
+    """
 
-    Parameters
-    ----------
-    exc_traceback : types.TracebackType
-        Original traceback
+    __slots__ = ()
 
-    Returns
-    -------
-    types.TracebackType
-        Shortened traceback
-    """
-    when_paths = config.get("admin.traceback.shorten.when")
-    what_paths = config.get("admin.traceback.shorten.what")
-    if not when_paths or not what_paths:
-        return exc_traceback
-
-    when_exp = re.compile(".*(" + "|".join(when_paths) + ")")
-    for f, _ in traceback.walk_tb(exc_traceback):
-        if when_exp.match(f.f_code.co_filename):
-            break
-    else:
-        return exc_traceback
-
-    what_exp = re.compile(".*(" + "|".join(what_paths) + ")")
-    curr = exc_traceback
-    prev = None
-
-    while curr:
-        if prev is None:
-            # always keep first frame
-            prev = curr
-        elif not curr.tb_next:
-            # always keep last frame
-            prev.tb_next = curr
-            prev = prev.tb_next
-        elif not what_exp.match(curr.tb_frame.f_code.co_filename):
-            # keep if module is not listed in what
-            prev.tb_next = curr
-            prev = prev.tb_next
-        curr = curr.tb_next
-
-    return exc_traceback
+    def __enter__(self) -> None:
+        pass
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None:
+        if exc_val and exc_tb:
+            exc_val.__traceback__ = self.shorten(exc_tb)
+
+    @staticmethod
+    def shorten(exc_tb: types.TracebackType) -> types.TracebackType:
+        paths = config.get("admin.traceback.shorten")
+        if not paths:
+            return exc_tb
+
+        exp = re.compile(".*(" + "|".join(paths) + ")")
+        curr: types.TracebackType | None = exc_tb
+        prev: types.TracebackType | None = None
+
+        while curr:
+            if prev is None:
+                prev = curr  # first frame
+            elif not curr.tb_next:
+                # always keep last frame
+                prev.tb_next = curr
+                prev = prev.tb_next
+            elif not exp.match(curr.tb_frame.f_code.co_filename):
+                # keep if module is not listed in config
+                prev.tb_next = curr
+                prev = curr
+            curr = curr.tb_next
+
+        # Uncomment to remove the first frame, which is something you don't want to keep
+        # if it matches the regexes. Requires Python >=3.11.
+        # if exc_tb.tb_next and exp.match(exc_tb.tb_frame.f_code.co_filename):
+        #     return exc_tb.tb_next
+
+        return exc_tb


Test patch:
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -6,7 +6,6 @@
 
 import asyncio
 import os
-import subprocess
 import sys
 from functools import partial
 from operator import add
@@ -1036,71 +1035,3 @@ async def test_bag_groupby_default(c, s, a, b):
     b = db.range(100, npartitions=10)
     b2 = b.groupby(lambda x: x % 13)
     assert not any("partd" in k[0] for k in b2.dask)
-
-
-def test_shorten_traceback_excepthook(tmp_path):
-    """
-    See Also
-    --------
-    test_distributed.py::test_shorten_traceback_ipython
-    test_utils.py::test_shorten_traceback
-    """
-    client_script = """
-from dask.distributed import Client
-if __name__ == "__main__":
-    f1 = lambda: 2 / 0
-    f2 = lambda: f1() + 5
-    f3 = lambda: f2() + 1
-    with Client() as client:
-        client.submit(f3).result()
-    """
-    with open(tmp_path / "script.py", mode="w") as f:
-        f.write(client_script)
-
-    proc_args = [sys.executable, os.path.join(tmp_path, "script.py")]
-    with popen(proc_args, capture_output=True) as proc:
-        out, err = proc.communicate(timeout=60)
-
-    lines = out.decode("utf-8").split("\n")
-    lines = [line for line in lines if line.startswith("  File ")]
-
-    assert len(lines) == 4
-    assert 'script.py", line 8, in <module>' in lines[0]
-    assert 'script.py", line 6, in <lambda>' in lines[1]
-    assert 'script.py", line 5, in <lambda>' in lines[2]
-    assert 'script.py", line 4, in <lambda>' in lines[3]
-
-
-def test_shorten_traceback_ipython(tmp_path):
-    """
-    See Also
-    --------
-    test_distributed.py::test_shorten_traceback_excepthook
-    test_utils.py::test_shorten_traceback
-    """
-    pytest.importorskip("IPython", reason="Requires IPython")
-
-    client_script = """
-from dask.distributed import Client
-f1 = lambda: 2 / 0
-f2 = lambda: f1() + 5
-f3 = lambda: f2() + 1
-with Client() as client: client.submit(f3).result()
-"""
-    with popen(["ipython"], capture_output=True, stdin=subprocess.PIPE) as proc:
-        out, err = proc.communicate(input=client_script.encode(), timeout=60)
-
-    lines = out.decode("utf-8").split("\n")
-    lines = [
-        line
-        for line in lines
-        if line.startswith("File ")
-        or line.startswith("Cell ")
-        or "<ipython-input" in line
-    ]
-
-    assert len(lines) == 4
-    assert "In[5]" in lines[0] or "<ipython-input-5-" in lines[0]
-    assert "In[4]" in lines[1] or "<ipython-input-4-" in lines[1]
-    assert "In[3]" in lines[2] or "<ipython-input-3-" in lines[2]
-    assert "In[2]" in lines[3] or "<ipython-input-2-" in lines[3]
diff --git a/dask/tests/test_traceback.py b/dask/tests/test_traceback.py
new file mode 100644
--- /dev/null
+++ b/dask/tests/test_traceback.py
@@ -0,0 +1,134 @@
+"""Tests on traceback shortening heuristics
+
+See Also
+--------
+distributed/tests/test_client.py::test_short_tracebacks
+distributed/tests/test_client.py::test_short_tracebacks_async
+"""
+from __future__ import annotations
+
+import traceback
+from contextlib import contextmanager
+
+import pytest
+
+import dask
+from dask.utils import shorten_traceback
+
+
+@contextmanager
+def assert_tb_levels(expect):
+    with pytest.raises(ZeroDivisionError) as e:
+        yield
+    frames = list(traceback.walk_tb(e.tb))
+    frame_names = [frame[0].f_code.co_name for frame in frames]
+    assert frame_names[0] == "assert_tb_levels", frame_names
+    assert frame_names[1:] == expect, frame_names
+
+
+def f1():
+    return 1 / 0
+
+
+def f2():
+    return f1() + 5
+
+
+def f3():
+    return f2() + 1
+
+
+@pytest.mark.parametrize(
+    "regexes,expect",
+    [
+        (None, ["test_shorten_traceback", "f3", "f2", "f1"]),  # Disabled
+        ([], ["test_shorten_traceback", "f3", "f2", "f1"]),  # Disabled
+        (["nomatch"], ["test_shorten_traceback", "f3", "f2", "f1"]),
+        # Remove absolutely everything, but keep the first and last frame
+        # Python <= 3.11 can't remove the first frame of a traceback from a __exit__
+        # function
+        ([".*"], ["test_shorten_traceback", "f1"]),
+        (["tests"], ["test_shorten_traceback", "f1"]),
+    ],
+)
+def test_shorten_traceback(regexes, expect):
+    """
+    See also
+    --------
+    test_distributed.py::test_shorten_traceback_excepthook
+    test_distributed.py::test_shorten_traceback_ipython
+    """
+    with dask.config.set({"admin.traceback.shorten": regexes}):
+        with assert_tb_levels(expect):
+            with shorten_traceback():
+                f3()
+
+
+try:
+    import tblib
+except ImportError:
+    tblib = None
+
+
+@pytest.mark.parametrize("scheduler", ["threads", "processes", "sync"])
+def test_compute_shorten_traceback(scheduler):
+    d = dask.delayed(f3)()
+
+    TEST_NAME = "test_compute_shorten_traceback"
+    if scheduler == "processes" and not tblib:
+        remote_stack = ["reraise"]
+    else:
+        remote_stack = ["f3", "f2", "f1"]
+
+    expect = [TEST_NAME, "compute", *remote_stack]
+    with assert_tb_levels(expect):
+        dask.compute(d(), scheduler=scheduler)
+
+    expect = [TEST_NAME, "compute", "compute", *remote_stack]
+    with assert_tb_levels(expect):
+        d.compute(scheduler=scheduler)
+
+
+@pytest.mark.parametrize("scheduler", ["threads", "processes", "sync"])
+def test_persist_shorten_traceback(scheduler):
+    d = dask.delayed(f3)()
+
+    TEST_NAME = "test_persist_shorten_traceback"
+    if scheduler == "processes" and not tblib:
+        remote_stack = ["reraise"]
+    else:
+        remote_stack = ["f3", "f2", "f1"]
+
+    expect = [TEST_NAME, "persist", *remote_stack]
+    with assert_tb_levels(expect):
+        dask.persist(d(), scheduler=scheduler)
+
+    expect = [TEST_NAME, "persist", "persist", *remote_stack]
+    with assert_tb_levels(expect):
+        d.persist(scheduler=scheduler)
+
+
+def test_distributed_shorten_traceback():
+    distributed = pytest.importorskip("distributed")
+    with distributed.Client(processes=False, dashboard_address=":0"):
+        d = dask.delayed(f3)()
+
+        (dp1,) = dask.persist(d)
+        dp2 = d.persist()
+
+        TEST_NAME = "test_distributed_shorten_traceback"
+        expect = [TEST_NAME, "compute", "f3", "f2", "f1"]
+        with assert_tb_levels(expect):
+            dask.compute(d())
+
+        expect = [TEST_NAME, "compute", "compute", "f3", "f2", "f1"]
+        with assert_tb_levels(expect):
+            d.compute()
+
+        expect = [TEST_NAME, "compute", "f3", "f2", "f1"]
+        with assert_tb_levels(expect):
+            dask.compute(dp1)
+
+        expect = [TEST_NAME, "compute", "compute", "f3", "f2", "f1"]
+        with assert_tb_levels(expect):
+            dp2.compute()
diff --git a/dask/tests/test_utils.py b/dask/tests/test_utils.py
--- a/dask/tests/test_utils.py
+++ b/dask/tests/test_utils.py
@@ -4,13 +4,12 @@
 import functools
 import operator
 import pickle
-import traceback
 from array import array
 
 import pytest
 from tlz import curry
 
-from dask import config, get
+from dask import get
 from dask.highlevelgraph import HighLevelGraph
 from dask.optimization import SubgraphCallable
 from dask.utils import (
@@ -42,7 +41,6 @@
     parse_timedelta,
     partial_by_order,
     random_state_data,
-    shorten_traceback,
     skip_doctest,
     stringify,
     stringify_collection_keys,
@@ -939,49 +937,3 @@ def test_get_meta_library_gpu():
     assert get_meta_library(cp.ndarray([])) == get_meta_library(
         da.from_array([]).to_backend("cupy")
     )
-
-
-@pytest.mark.parametrize(
-    "when,what,expect",
-    [
-        ([], [], 4),
-        ([".*"], [], 4),
-        ([], [".*"], 4),
-        ([r"nomatch"], [".*"], 4),
-        ([r".*"], ["nomatch"], 4),
-        ([".*"], [".*"], 2),
-        ([r"dask[\\\/]tests"], [], 4),
-        ([r"dask[\\\/]tests"], [r"dask[\\\/]tests"], 2),
-        ([], [r"dask[\\\/]tests"], 4),
-    ],
-)
-def test_shorten_traceback(when, what, expect):
-    """
-    See also
-    --------
-    test_distributed.py::test_shorten_traceback_excepthook
-    test_distributed.py::test_shorten_traceback_ipython
-    """
-
-    def f1():
-        return 2 / 0
-
-    def f2():
-        return f1() + 5
-
-    def f3():
-        return f2() + 1
-
-    with pytest.raises(ZeroDivisionError) as ex:
-        f3()
-
-    tb = ex.value.__traceback__
-    with config.set(
-        {
-            "admin.traceback.shorten.when": when,
-            "admin.traceback.shorten.what": what,
-        }
-    ):
-        tb = shorten_traceback(tb)
-    frame_count = len(list(traceback.walk_tb(tb)))
-    assert frame_count == expect