fix bali pid and bali result dict keys finding

Author: Lena-Jurkschat

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:
@@ -247,6 +249,7 @@ def _draw_bali_segments(
         if not show_bali:
             return
         segments = self._load_bali_segments()
+
         if not segments:
             return
 

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,11 +298,16 @@ 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.info(f"cached segments: {self._cached_bali_segments}")
+            logger.info(f"cached segments: {self._cached_bali_segments}")
+        
         return self._cached_bali_segments
 
     def _invalidate_bali_cache(self):

jumper_extension/bali_hook.py

@@ -16,7 +16,7 @@ def __init__(self, base_search_path: str = "."):
         )
         self.colormap_energy = mpl.colors.LinearSegmentedColormap.from_list(
             "yellow_to_red",
-            ['#EADFB4','#F6995C','#874C62']
+            ['#EADFB4','#874C62']
         )
 
     def _find_bali_directories(self, pid: int) -> List[str]:
@@ -52,8 +52,8 @@ def extract_segment(
             framework_data = benchmark_data[framework]
 
             for iteration_key, iteration_data in framework_data.items():
-                start_time = iteration_data.get("start_time")
-                end_time = iteration_data.get("end_time")
+                start_time = iteration_data.get("start_timestamp")
+                end_time = iteration_data.get("end_timestamp")
                 # ``generation_time`` is the duration of the text-generation
                 # phase; ``tokenize_time`` and ``setup_time`` are also
                 # durations (not absolute timestamps).
@@ -114,8 +114,8 @@ def extract_error_segments(self, error_data: Dict, config_data: Dict) -> List[Di
         segments = []
         if error_data:
             for framework, error_info in error_data.items():
-                start_time = error_info.get("start_time")
-                end_time = error_info.get("end_time")
+                start_time = error_info.get("start_timestamp")
+                end_time = error_info.get("end_timestamp")
 
                 segments.append({
                     "start_time": start_time,
@@ -202,4 +202,4 @@ def get_color_for_energy_efficiency(
         normalized = max(
             0.0, min(1.0, (tokens_per_sec - vmin) / (vmax - vmin))
         )
-        return self.colormap_energy(normalized)
+        return self.colormap_energy(normalized)