Merge branch 'bali-hook' of https://github.com/ScaDS/jumper_jupyter_performance into bali-hook

# Conflicts:
#	jumper_extension/bali_adapter.py
#	jumper_extension/bali_hook.py

Author: Elias Werner

jumper_extension/adapters/visualizer/backends/matplotlib.py

@@ -12,7 +12,9 @@
 from jumper_extension.adapters.visualizer.visualizer import PerformanceVisualizer
 from jumper_extension.utilities import get_available_levels
 from jumper_extension.logo import jumper_colors
+import logging
 
+logger = logging.getLogger("extension")
 
 def is_ipympl_backend():
     try:
@@ -263,6 +265,7 @@ def _draw_bali_segments(
         if not show_bali:
             return
         segments = self._load_bali_segments()
+
         if not segments:
             return
 
@@ -358,7 +361,7 @@ def _draw_bali_segments(
                 "Input Length": s.get("input_len", "n/a"),
                 "Output Length": s.get("output_len", "n/a"),
                 "Output Tokens per Second": (
-                    f"{tps:.2f}" if tps else "NaN"
+                    f"{s.get('tokens_per_sec'):.2f}" if tps else "NaN"
                 ),
                 "Segment Throughput (Tok/s)": s.get(
                     "segment_throughput", "n/a"

jumper_extension/bali_adapter.py

@@ -3,6 +3,7 @@
 import numpy as np
 from itables import show
 import logging
+import os
 from jumper_extension.bali_hook import BaliResultsParser
 from jumper_extension.core.messages import EXTENSION_INFO_MESSAGES, ExtensionInfoCode
 
@@ -45,7 +46,8 @@ def refresh_segments_from_disk(self, pid: int) -> int:
         Refresh BALI segments from disk for the given process ID.
         """
         segments = self.parser.collect_all_bali_segments(pid)
-
+        with open("segments.txt", "w") as file:
+            file.write(str(segments))
         # Build DataFrame directly from segments and align to canonical column order
         df = pd.DataFrame(segments)
         self._segments_df = df.reindex(columns=self._segments_df.columns)
@@ -117,7 +119,7 @@ def _trapz(values):
             powers = np.asarray(values["gpu_power_avg"], dtype=float)
             if len(times) < 2:
                 return 0.0
-            return float(np.trapz(powers, times))
+            return float(np.trapezoid(powers, times))
 
         def _safe_div(a, b):
             return a / b if b else None
@@ -284,7 +286,7 @@ class BaliVisualizationMixin:
     def __init__(self, *args, bali_adapter=None, **kwargs):
         super().__init__(*args, **kwargs)
         self.bali_adapter = bali_adapter or BaliAdapter()
-        self._compressed_bali_segments = []
+        self._compressed_bali_segments = None
         self._cached_bali_segments = None
 
     def _load_bali_segments(self) -> List[Dict]:
@@ -296,8 +298,12 @@ def _load_bali_segments(self) -> List[Dict]:
             bali_pid = getattr(self.monitor, "bali_pid_directory", None) \
                 or getattr(self.monitor, "pid", None)
             if bali_pid is None:
-                self._cached_bali_segments = []
-                return self._cached_bali_segments
+                bali_pid = os.getpid()
+                #self._cached_bali_segments = None
+                #return self._cached_bali_segments
+
+            logger.info(f"BALI PID used: {bali_pid}")
+
             self._cached_bali_segments = self.bali_adapter.get_segments_for_visualization(
                 bali_pid)
             logging.debug("cached segments: %s", self._cached_bali_segments)