Repository: dask/distributed

Issue:
Capture line number for code frames
I'd love if the line number was captured for computations. This would allow us to add an indicator to the calling line in what is potentially a sea of function calls. The line-number should be available.

https://github.com/dask/distributed/blob/515dffe40bfceb3141c7e72df7ee4de74d35315a/distributed/client.py#L3041



Discussion:
github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/distributed/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

       26 files  ±0         26 suites  ±0   13h 11m 18s [:stopwatch:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "duration of all tests") + 32m 15s
  3 551 tests ±0    3 441 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "passed tests")  - 2     105 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "skipped / disabled tests") ±0  5 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "failed tests") +2 
44 939 runs  ±0  42 806 [:heavy_check_mark:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "passed tests")  - 5  2 127 [:zzz:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "skipped / disabled tests") +2  6 [:x:](https://github.com/EnricoMi/publish-unit-test-result-action/blob/v1.20/README.md#the-symbols "failed tests") +3 

For more details on these failures, see [this check](https://github.com/dask/distributed/runs/12827061435).

Results for commit ee2c23a4. ± Comparison against base commit b83c0557.

[test-results]:data:application/gzip;base64,H4sIAORgPmQC/03MSw7DIAxF0a1UjDsAbAJ0MxVyQELNpyIwqrr3Omp+w3tsvY9IeYiLeNx0d7+JpeV6Vt9KqHmeuNGidUx8resdjFF7PpdGtBrixV75zaakOSiFPDCdEEuZC4tkKW1adxE9+C33XdROdof9d7XSdqdt93i5ztI8jrlyiBg1aQhoe5WUBaPJkZMeOpJRkYLgAVD1SXx/HLEbsBQBAAA=

milesgranger: @ntabris IIRC, you had mentioned you wanted to check out how this would affect the platform? I'm reasonably sure, as it is now, it shouldn't affect much of anything. All callers immediately drop the use of the line number and only use the src code lines as before so I think it shouldn't have any side effects. 

I was thinking this would be strictly implementing capturing the line numbers, then in a later PR actually making use of them in.
ntabris: Okay, I didn't realize that. Doing this in multiple steps makes it unclear to me when review of design should happen. Would code from this PR be in scope when reviewing the PR which makes use of this code? (What's the reason for doing this in multiple PR? It seems like a single unit from design and UX perspective.)
milesgranger: > What's the reason for doing this in multiple PR?

Because the [original issue](https://github.com/dask/distributed/issues/7779) seemed to be, in my interpretation, a feature for the platform indicating the line number; and I'm not sure where that would be. If I'm mistaken, I'd be happy to do this in one; just need a bit of guidance on where the indicator would go, in which caller. 
shughes-uk: @milesgranger i'm a bit clueless to the internals of dask but here is where where we are capturing code frames and sending them to the platform

https://github.com/coiled/platform/blob/0b6bfe34da7a2b1e1b68a77ef4103a14e500bef6/analytics/preload_scripts/telemetry.py#L402-L464


milesgranger: Thanks @shughes-uk, seems like we'd be in the first situation, this PR would break that code. As `code` would no longer have the type `Tuple[str, ...]` but `Tuple[SourceCode, ...]` where `SourceCode` is a namedtuple with the original 'code' and 'line_number':

Alternatively, I suppose [we could update the sending to scheduler](https://github.com/dask/distributed/blob/581e91dba69b8ce7a2618619a882c9cc490f73d4/distributed/client.py#L3121) to have another `line_number` key that when zipped with `code` are aligned if you really wanted to avoid breaking anything on your end. 
hendrikmakait: I am wondering whether we should capture the line number within the file - line 7174 in `test_client.py` - or the line number within the frame - line 5 in 
```python3
def test_computation_object_code_dask_compute(client): 
    ...
```

From what I understand, in this PR we capture the former, which might not be very useful without capturing the source file (name). It might also not suffice if one wanted to add indicators to the calling line. @shughes-uk: What were you thinking about when creating #7779? Similarly, @jacobtomlinson: What exactly would you be looking for to add line numbers to performance reports?
jacobtomlinson: @hendrikmakait taking [this performance report](https://gistcdn.githack.com/jacobtomlinson/7cd1d4b0ffb67034f64385628996e38f/raw/deeaa2cb93c7f98b3d6e32c32a5ec121bb0d328c/dask-report-199.html) as an example there is a code snippet at the bottom. It would be nice if it included the file name and line numbers for where that snippet came from. 
milesgranger: Thanks for the review @hendrikmakait I've updated to make line number reference the current source code https://github.com/dask/distributed/pull/7786/commits/75cfdcbc9214e7551b1a8bed3a946a107c6b43fc, give a holler if this isn't what you had in mind.
hendrikmakait: To summarize an offline discussion with @ntabris and @milesgranger:

We care about knowing what line number triggered a computation/submission within a frame. It would also be nice to know the file number the frame's code belongs to and its absolute position within the file (as per @jacobtomlinson's comment). Whatever changes we make in this PR should accommodate that. If it's easy to do it would be nice if this PR also added those attributes. 

We do not care about deduplicating equal code frames with different line numbers (but we do care about deduplicating equal code frames with _equal_ line numbers which might be caused by a for-loop). We assume that the remaining duplicates will not cause issues on the scheduler in terms of size and could easily be deduplicated by the consumer. Given that we feel a need to generally iterate on computations, we do not want to sink resources into building a nice tree-like deduplication.

We should run an A/B test before merging this PR to avoid an unexpected performance impact.

cc @fjetter in case you see any red flags with this approach.
fjetter: > cc @fjetter in case you see any red flags with this approach.

Haven't looked at the code but what you are describing sounds good.

Patch:
diff --git a/.github/workflows/conda.yml b/.github/workflows/conda.yml
index fb289ae8..c22261ef 100644
--- a/.github/workflows/conda.yml
+++ b/.github/workflows/conda.yml
@@ -35,9 +35,10 @@ jobs:
           miniforge-variant: Mambaforge
           use-mamba: true
           python-version: 3.9
+          channel-priority: strict
       - name: Install dependencies
         run: |
-          mamba install boa conda-verify
+          mamba install -c conda-forge boa conda-verify
 
           which python
           pip list
@@ -73,5 +74,6 @@ jobs:
           ANACONDA_API_TOKEN: ${{ secrets.DASK_CONDA_TOKEN }}
         run: |
           # install anaconda for upload
-          mamba install anaconda-client
+          mamba install -c conda-forge anaconda-client
+
           anaconda upload --label dev noarch/*.tar.bz2
diff --git a/distributed/client.py b/distributed/client.py
index 13eb9466..0cd4e39c 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -25,7 +25,7 @@ from functools import partial
 from importlib.metadata import PackageNotFoundError, version
 from numbers import Number
 from queue import Queue as pyQueue
-from typing import Any, ClassVar, Literal, TypedDict
+from typing import Any, ClassVar, Literal, NamedTuple, TypedDict
 
 from packaging.version import parse as parse_version
 from tlz import first, groupby, keymap, merge, partition_all, valmap
@@ -130,6 +130,13 @@ DEFAULT_EXTENSIONS = {
 TOPIC_PREFIX_FORWARDED_LOG_RECORD = "forwarded-log-record"
 
 
+class SourceCode(NamedTuple):
+    code: str
+    lineno_frame: int
+    lineno_relative: int
+    filename: str
+
+
 def _get_global_client() -> Client | None:
     c = _current_client.get()
     if c:
@@ -2992,7 +2999,7 @@ class Client(SyncMethodMixin):
     @staticmethod
     def _get_computation_code(
         stacklevel: int | None = None, nframes: int = 1
-    ) -> tuple[str, ...]:
+    ) -> tuple[SourceCode, ...]:
         """Walk up the stack to the user code and extract the code surrounding
         the compute/submit/persist call. All modules encountered which are
         ignored through the option
@@ -3021,8 +3028,10 @@ class Client(SyncMethodMixin):
             # stacklevel 0 or less - shows dask internals which likely isn't helpful
             stacklevel = stacklevel if stacklevel > 0 else 1
 
-        code: list[str] = []
-        for i, (fr, _) in enumerate(traceback.walk_stack(sys._getframe().f_back), 1):
+        code: list[SourceCode] = []
+        for i, (fr, lineno_frame) in enumerate(
+            traceback.walk_stack(sys._getframe().f_back), 1
+        ):
             if len(code) >= nframes:
                 break
             elif stacklevel is not None and i != stacklevel:
@@ -3032,8 +3041,15 @@ class Client(SyncMethodMixin):
                 or fr.f_code.co_name in ("<listcomp>", "<dictcomp>")
             ):
                 continue
+
+            kwargs = dict(
+                lineno_frame=lineno_frame,
+                lineno_relative=lineno_frame - fr.f_code.co_firstlineno,
+                filename=fr.f_code.co_filename,
+            )
+
             try:
-                code.append(inspect.getsource(fr))
+                code.append(SourceCode(code=inspect.getsource(fr), **kwargs))  # type: ignore
             except OSError:
                 try:
                     from IPython import get_ipython
@@ -3041,7 +3057,7 @@ class Client(SyncMethodMixin):
                     ip = get_ipython()
                     if ip is not None:
                         # The current cell
-                        code.append(ip.history_manager._i00)
+                        code.append(SourceCode(code=ip.history_manager._i00, **kwargs))  # type: ignore
                 except ImportError:
                     pass  # No IPython
 
@@ -3111,6 +3127,10 @@ class Client(SyncMethodMixin):
                     "This may cause some slowdown.\n"
                     "Consider scattering data ahead of time and using futures."
                 )
+
+            computations = self._get_computation_code(
+                nframes=dask.config.get("distributed.diagnostics.computations.nframes")
+            )
             self._send_to_scheduler(
                 {
                     "op": "update-graph",
@@ -3121,11 +3141,7 @@ class Client(SyncMethodMixin):
                     "submitting_task": getattr(thread_state, "key", None),
                     "fifo_timeout": fifo_timeout,
                     "actors": actors,
-                    "code": self._get_computation_code(
-                        nframes=dask.config.get(
-                            "distributed.diagnostics.computations.nframes"
-                        )
-                    ),
+                    "code": ToPickle(computations),
                     "annotations": ToPickle(annotations),
                 }
             )
@@ -5756,7 +5772,7 @@ class performance_report:
         client = get_client()
         if code is None:
             frames = client._get_computation_code(self._stacklevel + 1, nframes=1)
-            code = frames[0] if frames else "<Code not available>"
+            code = frames[0].code if frames else "<Code not available>"
         data = await client.scheduler.performance_report(
             start=self.start, last_count=self.last_count, code=code, mode=self.mode
         )
@@ -5771,7 +5787,7 @@ class performance_report:
     def __exit__(self, exc_type, exc_value, traceback):
         client = get_client()
         frames = client._get_computation_code(self._stacklevel + 1, nframes=1)
-        code = frames[0] if frames else "<Code not available>"
+        code = frames[0].code if frames else "<Code not available>"
         client.sync(self.__aexit__, exc_type, exc_value, traceback, code=code)
 
 
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 064392fe..50550325 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -68,6 +68,7 @@ from distributed import versions as version_module
 from distributed._stories import scheduler_story
 from distributed.active_memory_manager import ActiveMemoryManagerExtension, RetireWorker
 from distributed.batched import BatchedSend
+from distributed.client import SourceCode
 from distributed.collections import HeapSet
 from distributed.comm import (
     Comm,
@@ -848,7 +849,7 @@ class Computation:
 
     start: float
     groups: set[TaskGroup]
-    code: SortedSet
+    code: SortedSet[SourceCode]
     id: uuid.UUID
     annotations: dict
 
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 64702ea5..17cac8b5 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -1,6 +1,43 @@
 Changelog
 =========
 
+.. _v2023.5.1:
+
+2023.5.1
+--------
+
+Released on May 26, 2023
+
+.. note::
+
+  This release drops support for Python 3.8. As of this release
+  Dask supports Python 3.9, 3.10, and 3.11.
+  See `this community issue <https://github.com/dask/community/issues/315>`_
+  for more details.
+
+Enhancements
+^^^^^^^^^^^^
+- Exclude IPython code from computations (:pr:`7788`) `Miles`_
+- Drop Python 3.8 support (:pr:`7840`) `Thomas Grainger`_
+- Add ``storage_options`` to ``performance_report`` (:pr:`7636`) `ypogorelova`_
+- Don't warn about mismatched ``msgpack`` (:pr:`7839`) `Irina Truong`_
+- Clean up ``sys.path`` on ``Server`` shutdown (:pr:`7838`) `James Bourbeau`_
+- Dashboard: Fine Performance Metrics (:pr:`7725`) `Miles`_
+
+Bug Fixes
+^^^^^^^^^
+- Properly handle unknown chunk sizes in P2P rechunking (:pr:`7856`) `Hendrik Makait`_
+- Minimal change to work around (:issue:`7726`) / support for UCX (:pr:`7851`) `Benjamin Zaitlen`_
+- Don't end computations until cluster is truly idle (:pr:`7790`) `crusaderky`_
+
+Maintenance
+^^^^^^^^^^^
+- Explicitly install ``anaconda-client`` from conda-forge when uploading conda nightlies (:pr:`7861`) `Charles Blackmon-Luca`_
+- Fix ``is_idle`` docs build (:pr:`7854`) `James Bourbeau`_
+- Add tests for P2P barrier fusion (:pr:`7845`) `Hendrik Makait`_
+- Avoid ``DeprecationWarning`` in ``cupy`` dispatch registration (:pr:`7836`) `Lawrence Mitchell`_
+
+
 .. _v2023.5.0:
 
 2023.5.0
@@ -4950,3 +4987,4 @@ significantly without many new features.
 .. _`Eugene Druzhynin`: https://github.com/eugene-graft
 .. _`ypogorelova`: https://github.com/ypogorelova
 .. _`Patrick Hoefler`: https://github.com/phofl
+.. _`Irina Truong`: https://github.com/j-bennet
diff --git a/pyproject.toml b/pyproject.toml
index 7cbd4465..1159f609 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,7 +27,7 @@ requires-python = ">=3.9"
 dependencies = [
     "click >= 8.0",
     "cloudpickle >= 1.5.0",
-    "dask == 2023.5.0",
+    "dask == 2023.5.1",
     "jinja2 >= 2.10.3",
     "locket >= 1.0.0",
     "msgpack >= 1.0.0",


Test patch:
diff --git a/.github/workflows/test-report.yaml b/.github/workflows/test-report.yaml
index 399f3489..613d7a5a 100644
--- a/.github/workflows/test-report.yaml
+++ b/.github/workflows/test-report.yaml
@@ -56,7 +56,7 @@ jobs:
           mv test_report.html test_short_report.html deploy/
 
       - name: Deploy 🚀
-        uses: JamesIves/github-pages-deploy-action@v4.4.1
+        uses: JamesIves/github-pages-deploy-action@v4.4.2
         with:
           branch: gh-pages
           folder: deploy
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
index 57482bc7..37e522db 100644
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -116,6 +116,7 @@ from distributed.utils_test import (
     popen,
     raises_with_cause,
     randominc,
+    relative_frame_linenumber,
     save_sys_modules,
     slowadd,
     slowdec,
@@ -7126,16 +7127,41 @@ async def test_get_client_functions_spawn_clusters(c, s, a):
 def test_computation_code_walk_frames():
     test_function_code = inspect.getsource(test_computation_code_walk_frames)
     code = Client._get_computation_code()
+    lineno_relative = relative_frame_linenumber(inspect.currentframe()) - 1
+    lineno_frame = inspect.getframeinfo(inspect.currentframe()).lineno - 2
 
-    assert code == (test_function_code,)
+    # Sanity check helper function works, called 7 lines down in this function
+    assert relative_frame_linenumber(inspect.currentframe()) == 7
+
+    assert len(code) == 1
+    code = code[0]
+
+    assert code.code == test_function_code
+    assert code.lineno_frame == lineno_frame
+    assert code.lineno_relative == lineno_relative
+    assert code.filename == __file__
 
     def nested_call():
-        return Client._get_computation_code(nframes=2)
+        code = Client._get_computation_code(nframes=2)
+        nonlocal lineno_relative, lineno_frame
+        lineno_relative = 1  # called on first line in this function
+        lineno_frame = inspect.getframeinfo(inspect.currentframe()).lineno - 3
+        return code
 
     nested = nested_call()
+    nested_call_lineno_relative = relative_frame_linenumber(inspect.currentframe()) - 1
+    nested_call_lineno_frame = inspect.getframeinfo(inspect.currentframe()).lineno - 2
+
     assert len(nested) == 2
-    assert nested[-1] == inspect.getsource(nested_call)
-    assert nested[-2] == test_function_code
+    assert nested[-1].code == inspect.getsource(nested_call)
+    assert nested[-1].lineno_frame == lineno_frame
+    assert nested[-1].lineno_relative == lineno_relative
+    assert nested[-1].filename == __file__
+
+    assert nested[-2].code == test_function_code
+    assert nested[-2].lineno_frame == nested_call_lineno_frame
+    assert nested[-2].lineno_relative == nested_call_lineno_relative
+    assert nested[-2].filename == __file__
 
     with pytest.raises(TypeError, match="Ignored modules must be a list"):
         with dask.config.set(
@@ -7149,9 +7175,17 @@ def test_computation_code_walk_frames():
         import sys
 
         upper_frame_code = inspect.getsource(sys._getframe(1))
+        lineno_relative = relative_frame_linenumber(sys._getframe(1))
+        lineno_frame = inspect.getframeinfo(sys._getframe(1)).lineno
         code = Client._get_computation_code()
-        assert code == (upper_frame_code,)
-        assert nested_call()[-1] == upper_frame_code
+
+        assert len(code) == 1
+        code = code[0]
+
+        assert code.code == upper_frame_code
+        assert code.lineno_relative == lineno_relative
+        assert code.lineno_frame == lineno_frame
+        assert nested_call()[-1].code == upper_frame_code
 
 
 def run_in_ipython(code):
@@ -7202,8 +7236,8 @@ async def test_computation_ignore_ipython_frames(s, a, b, nframes):
     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))"
+    assert normalize(code[0].code) == normalize(source_code)
+    assert normalize(code[1].code) == "def bar(x): return client.map(foo, range(x))"
 
 
 @gen_cluster(client=True, nthreads=[("", 1)])
@@ -7237,8 +7271,8 @@ def test_computation_object_code_dask_compute(client):
     code = client.run_on_scheduler(fetch_comp_code)
 
     assert len(code) == 2
-    assert code[-1] == test_function_code
-    assert code[-2] == inspect.getsource(sys._getframe(1))
+    assert code[-1].code == test_function_code
+    assert code[-2].code == inspect.getsource(sys._getframe(1))
 
 
 def test_computation_object_code_dask_compute_no_frames_default(client):
@@ -7289,8 +7323,8 @@ async def test_computation_object_code_dask_persist(c, s, a, b):
     assert len(comp.code) == 1
 
     assert len(comp.code[0]) == 2
-    assert comp.code[0][-1] == test_function_code
-    assert comp.code[0][-2] == inspect.getsource(sys._getframe(1))
+    assert comp.code[0][-1].code == test_function_code
+    assert comp.code[0][-2].code == inspect.getsource(sys._getframe(1))
 
 
 @gen_cluster(client=True, config={"distributed.diagnostics.computations.nframes": 2})
@@ -7312,8 +7346,8 @@ async def test_computation_object_code_client_submit_simple(c, s, a, b):
     assert len(comp.code) == 1
 
     assert len(comp.code[0]) == 2
-    assert comp.code[0][-1] == test_function_code
-    assert comp.code[0][-2] == inspect.getsource(sys._getframe(1))
+    assert comp.code[0][-1].code == test_function_code
+    assert comp.code[0][-2].code == inspect.getsource(sys._getframe(1))
 
 
 @gen_cluster(client=True, config={"distributed.diagnostics.computations.nframes": 2})
@@ -7336,8 +7370,8 @@ async def test_computation_object_code_client_submit_list_comp(c, s, a, b):
     assert len(comp.code) == 1
 
     assert len(comp.code[0]) == 2
-    assert comp.code[0][-1] == test_function_code
-    assert comp.code[0][-2] == inspect.getsource(sys._getframe(1))
+    assert comp.code[0][-1].code == test_function_code
+    assert comp.code[0][-2].code == inspect.getsource(sys._getframe(1))
 
 
 @gen_cluster(client=True, config={"distributed.diagnostics.computations.nframes": 2})
@@ -7360,8 +7394,8 @@ async def test_computation_object_code_client_submit_dict_comp(c, s, a, b):
     assert len(comp.code) == 1
 
     assert len(comp.code[0]) == 2
-    assert comp.code[0][-1] == test_function_code
-    assert comp.code[0][-2] == inspect.getsource(sys._getframe(1))
+    assert comp.code[0][-1].code == test_function_code
+    assert comp.code[0][-2].code == inspect.getsource(sys._getframe(1))
 
 
 @gen_cluster(client=True, config={"distributed.diagnostics.computations.nframes": 2})
@@ -7381,8 +7415,8 @@ async def test_computation_object_code_client_map(c, s, a, b):
     assert len(comp.code) == 1
 
     assert len(comp.code[0]) == 2
-    assert comp.code[0][-1] == test_function_code
-    assert comp.code[0][-2] == inspect.getsource(sys._getframe(1))
+    assert comp.code[0][-1].code == test_function_code
+    assert comp.code[0][-2].code == inspect.getsource(sys._getframe(1))
 
 
 @gen_cluster(client=True, config={"distributed.diagnostics.computations.nframes": 2})
@@ -7401,8 +7435,8 @@ async def test_computation_object_code_client_compute(c, s, a, b):
     assert len(comp.code) == 1
 
     assert len(comp.code[0]) == 2
-    assert comp.code[0][-1] == test_function_code
-    assert comp.code[0][-2] == inspect.getsource(sys._getframe(1))
+    assert comp.code[0][-1].code == test_function_code
+    assert comp.code[0][-2].code == inspect.getsource(sys._getframe(1))
 
 
 @pytest.mark.slow
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
index 605a1e33..67565eb6 100644
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -2577,3 +2577,8 @@ class SizeOf:
 def gen_nbytes(nbytes: int) -> SizeOf:
     """A function that emulates exactly nbytes on the worker data structure."""
     return SizeOf(nbytes)
+
+
+def relative_frame_linenumber(frame):
+    """Line number of call relative to the frame"""
+    return inspect.getframeinfo(frame).lineno - frame.f_code.co_firstlineno