Change data backend

bar_plot.py

@@ -86,7 +86,7 @@ def create_matplotlib_bar_plot() -> None:
             fig.suptitle(
                 f"Mismatch for batch size, sequence length and number of tokens to generate between configs: {bs} "
                 f"!= {batch_size}, {seqlen} != {sequence_length}, {n_tok} != {num_tokens_to_generate}",
-                color='white', fontsize=18, pad=20
+                color='white', fontsize=18
             )
             return None
 

data.py

@@ -10,25 +10,24 @@ class ModelBenchmarkData:
 
     def __init__(self, json_path: str) -> None:
         with open(json_path, "r") as f:
-            self.data = json.load(f)
+            self.data: dict = json.load(f)
 
-    def compute_e2e_latency(self, measures: dict) -> tuple[float, Optional[float]]:
-        return measures["e2e_latency"]
+    def compute_ttft(self, measures: dict) -> list[float]:
+        return [dts[0] for dts in measures["dt_tokens"]]
+    
+    def compute_itl(self, measures: dict) -> list[float]:
+        return [
+            (dts[-1] - dts[0]) / (len(dts) - 1) if len(dts) > 2 else 0
+            for dts in measures["dt_tokens"]
+        ]
 
-    def compute_ttft(self, measures: dict) -> float:
-        return measures["t_tokens"][0] - measures["wall_time_start"]
-
-    def compute_itl(self, measures: dict) -> Optional[float]:
-        if len(measures["t_tokens"]) < 2:
-            return None
-        delta_t = measures["t_tokens"][-1] - measures["t_tokens"][0]
-        num_tokens = len(measures["t_tokens"]) - 1
-        return delta_t / num_tokens
+    def compute_e2e_latency(self, measures: dict) -> list[float]:
+        return measures["e2e_latency"][:]
 
     def ensure_coherence(self) -> tuple[int, int, int]:
         all_hyperparams = set()
-        for cfg_name, data in self.data.items():
-            config = data["metadata"]["config"]
+        for data in self.data.values():
+            config = data["config"]
             hyperparams = (config["batch_size"], config["sequence_length"], config["num_tokens_to_generate"])
             all_hyperparams.add(hyperparams)
         if len(all_hyperparams) > 1:
@@ -40,12 +39,12 @@ class ModelBenchmarkData:
     def get_bar_plot_data(self, collapse_on_cache: bool = True, collapse_on_compile_mode: bool = True) -> dict:
         # Gather data for each scenario
         per_scenario_data = {}
-        for i, (cfg_name, data) in enumerate(self.data.items()):
+        for cfg_name, data in self.data.items():
             per_scenario_data[cfg_name] = {
-                "ttft": [self.compute_ttft(d) for d in data["measures"]],
-                "itl": [self.compute_itl(d) for d in data["measures"]],
-                "e2e": [self.compute_e2e_latency(d) for d in data["measures"]],
-                "config": data["metadata"]["config"],
+                "ttft": self.compute_ttft(data["measures"]),
+                "itl": self.compute_itl(data["measures"]),
+                "e2e": self.compute_e2e_latency(data["measures"]),
+                "config": data["config"],
             }
         # Eventually collapse on cache
         if collapse_on_cache: