Repository: dask/distributed

Issue:
Performance report Calling Code wrong when using `%%time` 
**What happened**:
When using the `%%time` magic at the beginning of  notebook cell that will generate a performance report, the resulting report in the Summary page section **Calling Code** will display the code below instead of the corresponding code cell. 

```
    def _run(self):
        try:
            self._context.run(self._callback, *self._args)
        except (SystemExit, KeyboardInterrupt):
            raise
        except BaseException as exc:
            cb = format_helpers._format_callback_source(
                self._callback, self._args)
            msg = f'Exception in callback {cb}'
            context = {
                'message': msg,
                'exception': exc,
                'handle': self,
            }
            if self._source_traceback:
                context['source_traceback'] = self._source_traceback
            self._loop.call_exception_handler(context)
        self = None  # Needed to break cycles when an exception occurs.
```
 
**What you expected to happen**:
See the code that I executed under the `performance_report` context manager

**Minimal Complete Verifiable Example**:

```python
from dask.distributed import Client, performance_report
client = Client()
import dask.array as da
```
new cell
```python

%%time
with performance_report(filename="test_magic_err.html"):
    x = da.random.random((1000, 1000), chunks=(100, 100))
    y = x + x.T
    z = y[::2, 5000:].mean(axis=1)
    z.compute()
```
Open the performance report and check the Calling Code section. 

**Anything else we need to know?**:
If I remove the `%%time` magic I don't see this problem. 

**Environment**:

- Dask/distributeds version: '2021.09.1'
- Python version: 3.8.12
- Operating System: macOS Big Sur 11.6
- Install method (conda, pip, source): conda-forge

Performance report Calling Code wrong when using `%%time` 
**What happened**:
When using the `%%time` magic at the beginning of  notebook cell that will generate a performance report, the resulting report in the Summary page section **Calling Code** will display the code below instead of the corresponding code cell. 

```
    def _run(self):
        try:
            self._context.run(self._callback, *self._args)
        except (SystemExit, KeyboardInterrupt):
            raise
        except BaseException as exc:
            cb = format_helpers._format_callback_source(
                self._callback, self._args)
            msg = f'Exception in callback {cb}'
            context = {
                'message': msg,
                'exception': exc,
                'handle': self,
            }
            if self._source_traceback:
                context['source_traceback'] = self._source_traceback
            self._loop.call_exception_handler(context)
        self = None  # Needed to break cycles when an exception occurs.
```
 
**What you expected to happen**:
See the code that I executed under the `performance_report` context manager

**Minimal Complete Verifiable Example**:

```python
from dask.distributed import Client, performance_report
client = Client()
import dask.array as da
```
new cell
```python

%%time
with performance_report(filename="test_magic_err.html"):
    x = da.random.random((1000, 1000), chunks=(100, 100))
    y = x + x.T
    z = y[::2, 5000:].mean(axis=1)
    z.compute()
```
Open the performance report and check the Calling Code section. 

**Anything else we need to know?**:
If I remove the `%%time` magic I don't see this problem. 

**Environment**:

- Dask/distributeds version: '2021.09.1'
- Python version: 3.8.12
- Operating System: macOS Big Sur 11.6
- Install method (conda, pip, source): conda-forge



Discussion:
Thanks for raising an issue around this @ncclementi. I'm able to reproduce. It looks like the `%%time` magic (and possibly other IPython magic commands?) is influencing the frame we grab for the calling source code. Perhaps @eriknw and/or @ian-r-rose has thoughts on how we could improve the current situation?  
I see what's going on.  I had to deal with this in `afar` too (https://github.com/eriknw/afar/pull/12) as @jrbourbeau well-knows :)

I can whip up a PR tonight after dinner or first thing tomorrow.
Thanks for raising an issue around this @ncclementi. I'm able to reproduce. It looks like the `%%time` magic (and possibly other IPython magic commands?) is influencing the frame we grab for the calling source code. Perhaps @eriknw and/or @ian-r-rose has thoughts on how we could improve the current situation?  
I see what's going on.  I had to deal with this in `afar` too (https://github.com/eriknw/afar/pull/12) as @jrbourbeau well-knows :)

I can whip up a PR tonight after dinner or first thing tomorrow.

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -2518,13 +2518,16 @@ def run_coroutine(self, function, *args, **kwargs):
         return self.run(function, *args, **kwargs)
 
     @staticmethod
-    def _get_computation_code() -> str:
+    def _get_computation_code(stacklevel=None) -> str:
         """Walk up the stack to the user code and extract the code surrounding
         the compute/submit/persist call. All modules encountered which are
         blacklisted by the option
         `distributed.diagnostics.computations.ignore-modules` will be ignored.
         This can be used to blacklist commonly used libraries which wrap
         dask/distributed compute calls.
+
+        ``stacklevel`` may be used to explicitly indicate from which frame on
+        the stack to get the source code.
         """
 
         ignore_modules = dask.config.get(
@@ -2535,22 +2538,39 @@ def _get_computation_code() -> str:
                 "Ignored modules must be a list. Instead got "
                 f"({type(ignore_modules)}, {ignore_modules})"
             )
-
-        pattern: re.Pattern | None
-        if ignore_modules:
-            pattern = re.compile("|".join([f"(?:{mod})" for mod in ignore_modules]))
+        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:
-            pattern = None
-
-        for fr, _ in traceback.walk_stack(None):
-            if pattern is None or (
-                not pattern.match(fr.f_globals.get("__name__", ""))
-                and fr.f_code.co_name not in ("<listcomp>", "<dictcomp>")
+            # stacklevel 0 or less - shows dask internals which likely isn't helpful
+            stacklevel = stacklevel if stacklevel > 0 else 1
+        for i, (fr, _) in enumerate(traceback.walk_stack(None), 1):
+            if stacklevel is not None:
+                if 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>")
             ):
-                try:
-                    return inspect.getsource(fr)
-                except OSError:
-                    break
+                continue
+            try:
+                return 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
+                ):
+                    from IPython import get_ipython
+
+                    ip = get_ipython()
+                    if ip is not None:
+                        # The current cell
+                        return ip.history_manager._i00
+                break
         return "<Code not available>"
 
     def _graph_to_futures(
@@ -4843,13 +4863,10 @@ async def __aenter__(self):
         await get_client().get_task_stream(start=0, stop=0)  # ensure plugin
 
     async def __aexit__(self, typ, value, traceback, code=None):
-        if not code:
-            try:
-                frame = sys._getframe(self._stacklevel)
-                code = inspect.getsource(frame)
-            except Exception:
-                code = ""
-        data = await get_client().scheduler.performance_report(
+        client = get_client()
+        if code is None:
+            code = client._get_computation_code(self._stacklevel + 1)
+        data = await client.scheduler.performance_report(
             start=self.start, last_count=self.last_count, code=code, mode=self.mode
         )
         with open(self.filename, "w") as f:
@@ -4859,12 +4876,9 @@ def __enter__(self):
         get_client().sync(self.__aenter__)
 
     def __exit__(self, typ, value, traceback):
-        try:
-            frame = sys._getframe(self._stacklevel)
-            code = inspect.getsource(frame)
-        except Exception:
-            code = ""
-        get_client().sync(self.__aexit__, type, value, traceback, code=code)
+        client = get_client()
+        code = client._get_computation_code(self._stacklevel + 1)
+        client.sync(self.__aexit__, type, value, traceback, code=code)
 
 
 class get_task_metadata:


Test patch: