display_graph_for_all moved to the extension

Author: OutlyingWest

src/jumper/context.py

@@ -1,9 +1,12 @@
 from dataclasses import dataclass
 
+from jumper.perfdatahandler import PerformanceDataHandler
+
 
 @dataclass
 class KernelContext:
     nodelist: list = None
+    perfdata_handler = PerformanceDataHandler()
 
 
 kernel_context = KernelContext()

src/jumper/kernel.py

@@ -95,8 +95,10 @@ def __init__(self, **kwargs):
 
         # will be set to True as soon as GPU data is received
         self.gpu_avail = False
-        self.perfdata_handler = PerformanceDataHandler()
-        kernel_context.nodelist = self.perfdata_handler.get_nodelist()
+        # TODO: Temporary share perfdata_handler instance with an ipython extension
+        #  as it contains data that should be shared with the extension.
+        kernel_context.perfdata_handler = PerformanceDataHandler()
+        kernel_context.nodelist = kernel_context.perfdata_handler.get_nodelist()
 
         self.scorep_available_ = shutil.which("scorep")
         self.scorep_python_available_ = True
@@ -717,7 +719,7 @@ async def scorep_execute(
             cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=proc_env
         )
 
-        self.perfdata_handler.start_perfmonitor(proc.pid)
+        kernel_context.perfdata_handler.start_perfmonitor(proc.pid)
         # For memory mode jupyter_dump and jupyter_update must be awaited
         # concurrently to the running subprocess
         if self.pershelper.mode == "memory":
@@ -779,7 +781,7 @@ async def scorep_execute(
         time_indices = None
         if len(multicellmode_timestamps):
             # retrieve the index this cell will have in the global history
-            sub_idx = len(self.perfdata_handler.get_code_history())
+            sub_idx = len(kernel_context.perfdata_handler.get_code_history())
             # append to have end of last code fragment
             multicellmode_timestamps.append("MCM_TS" + str(time.time()))
             time_indices = [[]]
@@ -819,7 +821,7 @@ async def scorep_execute(
                 time_indices[0][idx] = (sub_idx, ms)
 
         performance_data_nodes, duration = (
-            self.perfdata_handler.end_perfmonitor()
+            kernel_context.perfdata_handler.end_perfmonitor()
         )
 
         # In disk mode, subprocess already terminated
@@ -909,7 +911,7 @@ async def scorep_execute(
         self.pershelper.postprocess()
         if performance_data_nodes:
             self.report_perfdata(performance_data_nodes, duration)
-            self.perfdata_handler.append_code(
+            kernel_context.perfdata_handler.append_code(
                 datetime.datetime.now(), code, time_indices
             )
         return self.standard_reply()
@@ -957,9 +959,9 @@ async def do_execute(
             return self.standard_reply()
         """
         if code.startswith("%%display_graph_for_last"):
-            if not len(self.perfdata_handler.get_perfdata_history()):
+            if not len(kernel_context.perfdata_handler.get_perfdata_history()):
                 self.cell_output("No performance data available.")
-            time_indices = self.perfdata_handler.get_time_indices()[-1]
+            time_indices = kernel_context.perfdata_handler.get_time_indices()[-1]
             if time_indices:
                 sub_idxs = [x[0] for x in time_indices[0]]
                 self.cell_output(
@@ -969,7 +971,7 @@ async def do_execute(
                 )
             perfvis.draw_performance_graph(
                 kernel_context.nodelist,
-                self.perfdata_handler.get_perfdata_history()[-1],
+                kernel_context.perfdata_handler.get_perfdata_history()[-1],
                 self.gpu_avail,
                 time_indices,
             )
@@ -981,14 +983,14 @@ async def do_execute(
                     "stdout",
                 )
             index = int(code.split(" ")[1])
-            if index >= len(self.perfdata_handler.get_perfdata_history()):
+            if index >= len(kernel_context.perfdata_handler.get_perfdata_history()):
                 self.cell_output(
                     "Tracked only "
-                    + str(len(self.perfdata_handler.get_perfdata_history()))
+                    + str(len(kernel_context.perfdata_handler.get_perfdata_history()))
                     + " cells. This index is not available."
                 )
             else:
-                time_indices = self.perfdata_handler.get_time_indices()[index]
+                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(
@@ -998,22 +1000,11 @@ async def do_execute(
                     )
                 perfvis.draw_performance_graph(
                     kernel_context.nodelist,
-                    self.perfdata_handler.get_perfdata_history()[index],
+                    kernel_context.perfdata_handler.get_perfdata_history()[index],
                     self.gpu_avail,
                     time_indices,
                 )
             return self.standard_reply()
-        elif code.startswith("%%display_graph_for_all"):
-            data, time_indices = (
-                self.perfdata_handler.get_perfdata_aggregated()
-            )
-            perfvis.draw_performance_graph(
-                kernel_context.nodelist,
-                data,
-                self.gpu_avail,
-                time_indices,
-            )
-            return self.standard_reply()
 
         elif code.startswith("%%display_code_for_index"):
             if len(code.split(" ")) == 1:
@@ -1022,28 +1013,28 @@ async def do_execute(
                     "stdout",
                 )
             index = int(code.split(" ")[1])
-            if index >= len(self.perfdata_handler.get_perfdata_history()):
+            if index >= len(kernel_context.perfdata_handler.get_perfdata_history()):
                 self.cell_output(
                     "Tracked only "
-                    + str(len(self.perfdata_handler.get_perfdata_history()))
+                    + str(len(kernel_context.perfdata_handler.get_perfdata_history()))
                     + " cells. This index is not available."
                 )
             else:
                 self.cell_output(
                     "Cell timestamp: "
-                    + str(self.perfdata_handler.get_code_history()[index][0])
+                    + str(kernel_context.perfdata_handler.get_code_history()[index][0])
                     + "\n--\n",
                     "stdout",
                 )
                 self.cell_output(
-                    self.perfdata_handler.get_code_history()[index][1],
+                    kernel_context.perfdata_handler.get_code_history()[index][1],
                     "stdout",
                 )
             return self.standard_reply()
         elif code.startswith("%%display_code_history"):
             show(
                 pd.DataFrame(
-                    self.perfdata_handler.get_code_history(),
+                    kernel_context.perfdata_handler.get_code_history(),
                     columns=["timestamp", "code"],
                 ).reset_index(),
                 layout={"topStart": "search", "topEnd": None},
@@ -1065,14 +1056,14 @@ async def do_execute(
                 # In addition, it has a counter for the number of measurements
                 # each sub cell corresponds to in the list of performance data
                 # measurements, e.g. (2_0, 5), (2_1, 3), (2_2, 7)
-                mcm_time_indices = self.perfdata_handler.get_time_indices()
+                mcm_time_indices = kernel_context.perfdata_handler.get_time_indices()
                 mcm_time_indices = list(
                     filter(lambda item: item is not None, mcm_time_indices)
                 )
 
                 code = (
                     f"{varname}="
-                    f"{self.perfdata_handler.get_perfdata_history()}"
+                    f"{kernel_context.perfdata_handler.get_perfdata_history()}"
                 )
 
                 if mcm_time_indices:
@@ -1107,13 +1098,13 @@ async def do_execute(
                 filename = code.split(" ")[1]
                 with open(f"{filename}_perfdata.json", "w") as f:
                     json.dump(
-                        self.perfdata_handler.get_perfdata_history(),
+                        kernel_context.perfdata_handler.get_perfdata_history(),
                         default=str,
                         fp=f,
                     )
                 with open(f"{filename}_code.json", "w") as f:
                     json.dump(
-                        self.perfdata_handler.get_code_history(),
+                        kernel_context.perfdata_handler.get_code_history(),
                         default=str,
                         fp=f,
                     )
@@ -1212,7 +1203,7 @@ async def do_execute(
         else:
             if self.mode == KernelMode.DEFAULT:
                 self.pershelper.parse(magics_cleanup(code)[1], "jupyter")
-                self.perfdata_handler.start_perfmonitor(os.getpid())
+                kernel_context.perfdata_handler.start_perfmonitor(os.getpid())
                 parent_ret = await super().do_execute(
                     code,
                     silent,
@@ -1222,11 +1213,11 @@ async def do_execute(
                     cell_id=cell_id,
                 )
                 performance_data_nodes, duration = (
-                    self.perfdata_handler.end_perfmonitor()
+                    kernel_context.perfdata_handler.end_perfmonitor()
                 )
                 if performance_data_nodes:
                     self.report_perfdata(performance_data_nodes, duration)
-                    self.perfdata_handler.append_code(
+                    kernel_context.perfdata_handler.append_code(
                         datetime.datetime.now(), code
                     )
                 return parent_ret

src/magic_extension/magic.py

@@ -7,19 +7,35 @@
 from jumper.kernel import KernelMode
 
 from jumper.perfdatahandler import PerformanceDataHandler
+import jumper.visualization as perfvis
 
 
 @magics_class
 class KernelMagics(Magics):
     def __init__(self, shell):
         super(KernelMagics, self).__init__(shell)
-        self.perfdata_handler = PerformanceDataHandler()
         self.mode = KernelMode.DEFAULT
 
+        # will be set to True as soon as GPU data is received
+        self.gpu_avail = False
+
     @cell_magic
     def abra(self, line, cell):
         print('HELLO!\n', cell)
 
+    @line_magic
+    def display_graph_for_all(self, line):
+        data, time_indices = (
+            kernel_context.perfdata_handler.get_perfdata_aggregated()
+        )
+        perfvis.draw_performance_graph(
+            kernel_context.nodelist,
+            data,
+            self.gpu_avail,
+            time_indices,
+        )
+
+
     @cell_magic
     def set_perfmonitor(self, line, code):
         """
@@ -33,8 +49,8 @@ def set_perfmonitor(self, line, code):
                 )
             else:
                 try:
-                    self.perfdata_handler.set_monitor(monitor)
-                    kernel_context.nodelist = self.perfdata_handler.get_nodelist()
+                    kernel_context.perfdata_handler.set_monitor(monitor)
+                    kernel_context.nodelist = kernel_context.perfdata_handler.get_nodelist()
                     if len(kernel_context.nodelist) <= 1:
                         kernel_context.nodelist = None
                         self.cell_output(