Repository: dask/distributed

Issue:
Static Performance Report: Allow stacklevel optional argument for gathering source code
Allow the user to specific which level of the call stack is used for capturing source code in the static report generated by `dask.distributed.performance_report`. In some cases, the method might be wrapped by another function so allowing users to ability to select which level to include in the report might be helpful to them.

```python
    with dask.distributed.performance_report(filename=filename, stacklevel=2):
        ...
```

I am happy to pick this up as I'm working on related functionality. 


Discussion:
I'm slightly conflicted work-wise, but this seems fine to me.  I think that the use of the `stacklevel=` keyword name is nice because it matches existing Python API with `warnings.warn`
This seems reasonable to me, but I have a similar conflict. @quasiben would this be useful for any of the RAPIDS folks? You and your group seem to produce a lot of `performance_report`s
Conflicted meaning it is useful for other work projects or something else?

Anyways I don't see any issue with this. People frequently are interested in recommendations on how to profile in a distributed setting. Continuing to improve our tooling on this front makes sense to me
Conflicted meaning that both James and I work for Coiled, and this is necessary for internal Coiled work.  This seems pretty innocuous, but we'd feel more comfortable giving the green light if some non-Coiled person signed off on it.  I think that @jakirkham 's comment does that.  Thanks!
> I am happy to pick this up as I'm working on related functionality

Woo! @ndanielsen a PR adding `stacklevel=` would be very welcome
Sorry to be late here.  This seems like a solid improvement for everyone +1

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -4670,17 +4670,29 @@ class performance_report:
     browser.  Locally we recommend using ``python -m http.server`` or hosting
     the file live online.
 
+    Parameters
+    ----------
+    filename: str (optional)
+        The filename to save the performance report locally
+
+    stacklevel: int (optional)
+        The code execution frame utilized for populating the Calling Code section
+        of the report. Defaults to `1` which is the frame calling ``performance_report``
+
+
     Examples
     --------
-    >>> with performance_report(filename="myfile.html"):
+    >>> with performance_report(filename="myfile.html", stacklevel=1):
     ...     x.compute()
 
     $ python -m http.server
     $ open myfile.html
     """
 
-    def __init__(self, filename="dask-report.html"):
+    def __init__(self, filename="dask-report.html", stacklevel=1):
         self.filename = filename
+        # stacklevel 0 or less - shows dask internals which likely isn't helpful
+        self._stacklevel = stacklevel if stacklevel > 0 else 1
 
     async def __aenter__(self):
         self.start = time()
@@ -4689,7 +4701,7 @@ async def __aenter__(self):
     async def __aexit__(self, typ, value, traceback, code=None):
         if not code:
             try:
-                frame = inspect.currentframe().f_back
+                frame = sys._getframe(self._stacklevel)
                 code = inspect.getsource(frame)
             except Exception:
                 code = ""
@@ -4704,7 +4716,7 @@ def __enter__(self):
 
     def __exit__(self, typ, value, traceback):
         try:
-            frame = inspect.currentframe().f_back
+            frame = sys._getframe(self._stacklevel)
             code = inspect.getsource(frame)
         except Exception:
             code = ""


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
@@ -6178,7 +6178,7 @@ async def test_performance_report(c, s, a, b):
     pytest.importorskip("bokeh")
     da = pytest.importorskip("dask.array")
 
-    async def f():
+    async def f(stacklevel):
         """
         We wrap this in a function so that the assertions aren't in the
         performanace report itself
@@ -6187,14 +6187,15 @@ async def f():
         """
         x = da.random.random((1000, 1000), chunks=(100, 100))
         with tmpfile(extension="html") as fn:
-            async with performance_report(filename=fn):
+            async with performance_report(filename=fn, stacklevel=stacklevel):
                 await c.compute((x + x.T).sum())
 
             with open(fn) as f:
                 data = f.read()
         return data
 
-    data = await f()
+    # Ensure default kwarg maintains backward compatability
+    data = await f(stacklevel=1)
 
     assert "Also, we want this comment to appear" in data
     assert "bokeh" in data
@@ -6204,6 +6205,18 @@ async def f():
     assert "Threads: 4" in data
     assert dask.__version__ in data
 
+    # Stacklevel two captures code two frames back -- which in this case
+    # is the testing function
+    data = await f(stacklevel=2)
+    assert "async def test_performance_report(c, s, a, b):" in data
+    assert "Dask Performance Report" in data
+
+    # Stacklevel zero or lower is overridden to stacklevel=1 so we don't see
+    # distributed internals
+    data = await f(stacklevel=0)
+    assert "Also, we want this comment to appear" in data
+    assert "Dask Performance Report" in data
+
 
 @pytest.mark.asyncio
 async def test_client_gather_semaphore_loop(cleanup):