Update compute_system_metrics.py

scripts/compute_system_metrics.py

@@ -1,5 +1,6 @@
 import os
 import csv
+from glob import glob
 
 import tyro
 import pandas as pd
@@ -14,13 +15,15 @@ def main(data_dir: str, out_file: str) -> None:
         os.makedirs(dirname, exist_ok=True)
     out_csv = csv.writer(open(out_file, "w", newline=""))
     metrics = ["throughput", "response_length", "latency", "energy"]
-    out_csv.writerow(["model"] + metrics)
+    out_csv.writerow(["model", "batch_size"] + metrics)
 
     for model_name in model_names:
-        df = pd.read_json(f"{data_dir}/{model_name}/benchmark.json")
-        out_csv.writerow(
-            [model_name.replace("--", "/")] + df[metrics].mean().to_list(),
-        )
+        for benchmark_file in glob(f"{data_dir}/{model_name}/benchmark_batch_*.json"):
+            batch_size = int(benchmark_file.split("_")[-1][:-5])
+            df = pd.read_json(benchmark_file)
+            out_csv.writerow(
+                [model_name.replace("--", "/"), str(batch_size)] + df[metrics].mean().to_list(),
+            )
 
 
 if __name__ == "__main__":