Repository: dask/distributed

Issue:
exclude ipython code from computations
Run `ipython` and then...

```python
import dask
from distributed import Client
client = Client()
dask.config.set({"distributed.diagnostics.computations.nframes": 2})

def foo(x): print(x); return x;
def zip(x): return foo(x);
N = client.map(zip, range(2))

# wait just a sec otherwise there won't be computation yet

print(client.cluster.scheduler.computations[-1].code[0][0])
```

and you get

```
    async def run_code(self, code_obj, result=None, *, async_=False):
        """Execute a code object.

        When an exception occurs, self.showtraceback() is called to display a
        traceback.

        Parameters
        ----------
        code_obj : code object
          A compiled code object, to be executed
        result : ExecutionResult, optional
          An object to store exceptions that occur during execution.
        async_ :  Bool (Experimental)
          Attempt to run top-level asynchronous code in a default loop.

        Returns
        -------
        False : successful execution.
        True : an error occurred.
        """
        # special value to say that anything above is IPython and should be
        # hidden.
        __tracebackhide__ = "__ipython_bottom__"
        # Set our own excepthook in case the user code tries to call it
        # directly, so that the IPython crash handler doesn't get triggered
        old_excepthook, sys.excepthook = sys.excepthook, self.excepthook

        # we save the original sys.excepthook in the instance, in case config
        # code (such as magics) needs access to it.
        self.sys_excepthook = old_excepthook
        outflag = True  # happens in more places, so it's easier as default
        try:
            try:
                if async_:
                    await eval(code_obj, self.user_global_ns, self.user_ns)
                else:
                    exec(code_obj, self.user_global_ns, self.user_ns)
            finally:
                # Reset our crash handler in place
                sys.excepthook = old_excepthook
        except SystemExit as e:
            if result is not None:
                result.error_in_exec = e
            self.showtraceback(exception_only=True)
            warn("To exit: use 'exit', 'quit', or Ctrl-D.", stacklevel=1)
        except bdb.BdbQuit:
            etype, value, tb = sys.exc_info()
            if result is not None:
                result.error_in_exec = value
            # the BdbQuit stops here
        except self.custom_exceptions:
            etype, value, tb = sys.exc_info()
            if result is not None:
                result.error_in_exec = value
            self.CustomTB(etype, value, tb)
        except:
            if result is not None:
                result.error_in_exec = sys.exc_info()[1]
            self.showtraceback(running_compiled_code=True)
        else:
            outflag = False
        return outflag
```

which is from [ipython](https://github.com/ipython/ipython/blob/9c07131ad450854fd265895e0b79cc0b89ef72d4/IPython/core/interactiveshell.py#L3469) and doesn't seem useful to collect.


Discussion:
In case it's worth mentioning, run `print(client.cluster.scheduler.computations[-1].code[0][1])` and you get the more interesting

```
import dask
from distributed import Client
client = Client()
dask.config.set({"distributed.diagnostics.computations.nframes": 2})

def foo(x): print(x); return x;
def zip(x): return foo(x);
N = client.map(zip, range(2))
```

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -3013,12 +3013,10 @@ def _get_computation_code(
                 "Ignored modules must be a list. Instead got "
                 f"({type(ignore_modules)}, {ignore_modules})"
             )
+        pattern: re.Pattern | None = None
         if stacklevel is None:
-            pattern: re.Pattern | None
             if ignore_modules:
                 pattern = re.compile("|".join([f"(?:{mod})" for mod in ignore_modules]))
-            else:
-                pattern = None
         else:
             # stacklevel 0 or less - shows dask internals which likely isn't helpful
             stacklevel = stacklevel if stacklevel > 0 else 1
@@ -3027,9 +3025,8 @@ def _get_computation_code(
         for i, (fr, _) in enumerate(traceback.walk_stack(sys._getframe().f_back), 1):
             if len(code) >= nframes:
                 break
-            if stacklevel is not None:
-                if i != stacklevel:
-                    continue
+            elif stacklevel is not None and i != stacklevel:
+                continue
             elif pattern is not None and (
                 pattern.match(fr.f_globals.get("__name__", ""))
                 or fr.f_code.co_name in ("<listcomp>", "<dictcomp>")
@@ -3038,19 +3035,23 @@ def _get_computation_code(
             try:
                 code.append(inspect.getsource(fr))
             except OSError:
-                # Try to fine the source if we are in %%time or %%timeit magic.
-                if (
-                    fr.f_code.co_filename in {"<timed exec>", "<magic-timeit>"}
-                    and "IPython" in sys.modules
-                ):
+                try:
                     from IPython import get_ipython
 
                     ip = get_ipython()
                     if ip is not None:
                         # The current cell
                         code.append(ip.history_manager._i00)
+                except ImportError:
+                    pass  # No IPython
+
                 break
 
+            # Ignore IPython related wrapping functions to user code
+            if hasattr(fr.f_back, "f_globals"):
+                module_name = sys.modules[fr.f_back.f_globals["__name__"]].__name__  # type: ignore
+                if module_name.endswith("interactiveshell"):
+                    break
         return tuple(reversed(code))
 
     def _graph_to_futures(


Test patch:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -6,11 +6,13 @@
 import gc
 import inspect
 import logging
+import multiprocessing
 import operator
 import os
 import pathlib
 import pickle
 import random
+import re
 import subprocess
 import sys
 import threading
@@ -7152,6 +7154,58 @@ def nested_call():
         assert nested_call()[-1] == upper_frame_code
 
 
+def run_in_ipython(code):
+    from IPython.testing.globalipapp import start_ipython
+
+    shell = start_ipython()
+    return shell.run_cell(code)
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("nframes", (2, 3))
+@gen_cluster()
+async def test_computation_ignore_ipython_frames(s, a, b, nframes):
+    pytest.importorskip("IPython")
+
+    source_code = f"""
+        import time
+        import dask
+        from distributed import Client
+
+        dask.config.set({{"distributed.diagnostics.computations.nframes": {nframes}}})
+        with Client("{s.address}") as client:
+            def foo(x): print(x); return x;
+            def bar(x): return client.map(foo, range(x))
+
+            N = client.gather(bar(3))
+    """
+    # When not running IPython in a new process, it does not shutdown
+    # properly and leaks a thread. There should be a way to fix this.
+    # Seems to be another (deeper) issue that this shouldn't need
+    # a subprocess/thread/@gen_cluster/test at all, and ought to be able to run
+    # directly in InteractiveShell (and does) but requires `--reruns=1`
+    # due to some underlying lag in the asyncio if ran by itself, but will
+    # otherwise run fine in the suite of tests.
+    ctx = multiprocessing.get_context("spawn")
+    with concurrent.futures.ProcessPoolExecutor(1, mp_context=ctx) as executor:
+        loop = asyncio.get_running_loop()
+        result = await loop.run_in_executor(executor, run_in_ipython, source_code)
+
+    result.raise_error()
+    computations = s.computations
+
+    assert len(computations) == 1
+    assert len(computations[0].code) == 1
+    code = computations[0].code[0]
+    assert len(code) == 2  # 2 frames when ignoring IPython frames
+
+    def normalize(s):
+        return re.sub(r"\s+", " ", s).strip()
+
+    assert normalize(code[0]) == normalize(source_code)
+    assert normalize(code[1]) == "def bar(x): return client.map(foo, range(x))"
+
+
 @gen_cluster(client=True, nthreads=[("", 1)])
 async def test_computation_store_annotations(c, s, a):
     # We do not want to store layer annotations