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:

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -25,7 +25,7 @@
 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 @@
 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 @@ def run(
     @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 @@ def _get_computation_code(
             # 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 @@ def _get_computation_code(
                 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 @@ def _get_computation_code(
                     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 @@ def _graph_to_futures(
                     "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 @@ def _graph_to_futures(
                     "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 @@ async def __aexit__(self, exc_type, exc_value, traceback, code=None):
         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 @@ def __enter__(self):
     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
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -68,6 +68,7 @@
 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
 


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
@@ -116,6 +116,7 @@
     popen,
     raises_with_cause,
     randominc,
+    relative_frame_linenumber,
     save_sys_modules,
     slowadd,
     slowdec,
@@ -7126,16 +7127,41 @@ def f(x):
 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 nested_call():
         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 @@ def bar(x): return client.map(foo, range(x))
     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 fetch_comp_code(dask_scheduler):
     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 @@ def func(x):
     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 @@ def func(x):
     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 @@ def func(x):
     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 @@ def func(x):
     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
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -2577,3 +2577,8 @@ def __sizeof__(self) -> int:
 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