display_graph_for_all moved to extension; gpu_avail moved to context

Author: OutlyingWest

src/jumper/context.py

@@ -5,7 +5,8 @@
 
 @dataclass
 class KernelContext:
-    nodelist: list = None
+    # will be set to True as soon as GPU data is received
+    gpu_avail = False
     perfdata_handler = PerformanceDataHandler()
 
 

src/jumper/kernel.py

@@ -93,10 +93,6 @@ def __init__(self, **kwargs):
         self.writefile_scorep_binding_args = []
         self.writefile_multicell = False
 
-        # will be set to True as soon as GPU data is received
-        self.gpu_avail = False
-        # TODO: Temporary share perfdata_handler instance with an ipython extension
-        #  as it contains data that should be shared with the extension.
         self.nodelist = kernel_context.perfdata_handler.get_nodelist()
 
         self.scorep_available_ = shutil.which("scorep")
@@ -558,7 +554,7 @@ def report_perfdata(self, performance_data_nodes, duration):
                     )
 
                 if gpu_util[0] and gpu_mem[0]:
-                    self.gpu_avail = True
+                    kernel_context.gpu_avail = True
                     self.cell_output(
                         "--GPU Util and Mem per GPU--\n", "stdout"
                     )
@@ -971,39 +967,10 @@ async def do_execute(
             perfvis.draw_performance_graph(
                 self.nodelist,
                 kernel_context.perfdata_handler.get_perfdata_history()[-1],
-                self.gpu_avail,
+                kernel_context.gpu_avail,
                 time_indices,
             )
             return self.standard_reply()
-        elif code.startswith("%%display_graph_for_index"):
-            if len(code.split(" ")) == 1:
-                self.cell_output(
-                    "No index specified. Use: %%display_graph_for_index index",
-                    "stdout",
-                )
-            index = int(code.split(" ")[1])
-            if index >= len(kernel_context.perfdata_handler.get_perfdata_history()):
-                self.cell_output(
-                    "Tracked only "
-                    + str(len(kernel_context.perfdata_handler.get_perfdata_history()))
-                    + " cells. This index is not available."
-                )
-            else:
-                time_indices = kernel_context.perfdata_handler.get_time_indices()[index]
-                if time_indices:
-                    sub_idxs = [x[0] for x in time_indices[0]]
-                    self.cell_output(
-                        f"Cell seemed to be tracked in multi cell"
-                        " mode. Got performance data for the"
-                        f" following sub cells: {sub_idxs}"
-                    )
-                perfvis.draw_performance_graph(
-                    self.nodelist,
-                    kernel_context.perfdata_handler.get_perfdata_history()[index],
-                    self.gpu_avail,
-                    time_indices,
-                )
-            return self.standard_reply()
 
         elif code.startswith("%%display_code_for_index"):
             if len(code.split(" ")) == 1:

src/magic_extension/magic.py

@@ -16,8 +16,6 @@ def __init__(self, shell):
         super(KernelMagics, self).__init__(shell)
         self.mode = KernelMode.DEFAULT
 
-        # will be set to True as soon as GPU data is received
-        self.gpu_avail = False
         self.nodelist = kernel_context.perfdata_handler.get_nodelist()
 
     @cell_magic
@@ -32,13 +30,57 @@ def display_graph_for_all(self, line):
         perfvis.draw_performance_graph(
             self.nodelist,
             data,
-            self.gpu_avail,
+            kernel_context.gpu_avail,
             time_indices,
         )
 
     @line_magic
     def display_graph_for_index(self, line):
-        pass
+        """
+        Display performance graph for a given index.
+        Usage:
+            %display_graph_for_index 2
+        """
+
+        if not line.strip():
+            self.cell_output(
+                "No index specified. Use: %display_graph_for_index index",
+                "stderr"
+            )
+            return
+
+        try:
+            index = int(line.strip())
+        except ValueError:
+            self.cell_output(
+                "Invalid index. Please provide an integer.",
+                "stderr"
+            )
+            return
+
+        history = kernel_context.perfdata_handler.get_perfdata_history()
+        if index >= len(history):
+            self.cell_output(
+                f"Tracked only {len(history)} cells. This index is not available.",
+                "stderr"
+            )
+            return
+
+        time_indices = kernel_context.perfdata_handler.get_time_indices()[index]
+
+        if time_indices:
+            sub_idxs = [x[0] for x in time_indices[0]]
+            self.cell_output(
+                f"📈 Cell seemed to be tracked in multi-cell mode. "
+                f"Got performance data for the following sub-cells: {sub_idxs}"
+            )
+
+        perfvis.draw_performance_graph(
+            self.nodelist,
+            history[index],
+            kernel_context.gpu_avail,
+            time_indices,
+        )
 
 
     @cell_magic
@@ -71,15 +113,18 @@ def set_perfmonitor(self, line, code):
                             + str(self.nodelist)
                         )
                 except Exception as e:
-                    self.cell_output(f"Error setting monitor\n{e}", "stderr")
+                    self.cell_output(
+                        f"Error setting monitor\n{e}",
+                        "stderr"
+                    )
         else:
             self.cell_output(
                 f"KernelWarning: Currently in {self.mode}, command ignored.",
                 "stderr",
             )
 
     @staticmethod
-    def cell_output(string, stream="stdout"):
+    def cell_output(string: str, stream="stdout"):
         if stream == "stderr":
             print(string, file=sys.stderr)
         else: