Support two NIXL bench output formats

Author: amaslenn

src/cloudai/workloads/nixl_bench/report_generation_strategy.py

@@ -37,20 +37,22 @@ def extract_data(stdout_file: Path) -> pd.DataFrame:
 
     header_present, data = False, []
     for line in stdout_file.read_text().splitlines():
-        if "Block Size (B)      Batch Size     Avg Lat. (us)  B/W (MiB/Sec)  B/W (GiB/Sec)  B/W (GB/Sec)" in line:
+        if not header_present and (
+            "Block Size (B)      Batch Size     " in line and "Avg Lat. (us)" in line and "B/W (GB/Sec)" in line
+        ):
             header_present = True
             continue
-        if header_present and len(line.split()) == 6:
-            data.append(line.split())
-
-    df = lazy.pd.DataFrame(
-        data, columns=["block_size", "batch_size", "avg_lat", "bw_mib_sec", "bw_gib_sec", "bw_gb_sec"]
-    )
+        parts = line.split()
+        if header_present and (len(parts) == 6 or len(parts) == 10):
+            if len(parts) == 6:
+                data.append([parts[0], parts[1], parts[2], parts[-1]])
+            else:
+                data.append([parts[0], parts[1], parts[3], parts[2]])
+
+    df = lazy.pd.DataFrame(data, columns=["block_size", "batch_size", "avg_lat", "bw_gb_sec"])
     df["block_size"] = df["block_size"].astype(int)
     df["batch_size"] = df["batch_size"].astype(int)
     df["avg_lat"] = df["avg_lat"].astype(float)
-    df["bw_mib_sec"] = df["bw_mib_sec"].astype(float)
-    df["bw_gib_sec"] = df["bw_gib_sec"].astype(float)
     df["bw_gb_sec"] = df["bw_gb_sec"].astype(float)
 
     return df

tests/report_generation_strategy/test_nixl_bench_report.py

@@ -0,0 +1,41 @@
+from pathlib import Path
+
+import pytest
+
+from cloudai.workloads.nixl_bench.report_generation_strategy import extract_data
+
+LEGACY_FORMAT = """
+Block Size (B)      Batch Size     Avg Lat. (us)  B/W (MiB/Sec)  B/W (GiB/Sec)  B/W (GB/Sec)
+--------------------------------------------------------------------------------
+4096                1              21.1726        184.495        0.180171       0.193457
+8192                1              21.7391        359.376        0.350953       0.376833
+33554432            1              24.6508        1.29813e+06    1267.71        1361.19
+67108864            1              39.746         1.61022e+06    1572.48        1688.44
+"""
+
+NEW_FORMAT = """
+Block Size (B)      Batch Size     B/W (GB/Sec)   Avg Lat. (us)  Avg Prep (us)  P99 Prep (us)  Avg Post (us)  P99 Post (us)  Avg Tx (us)    P99 Tx (us)
+----------------------------------------------------------------------------------------------------------------------------------------------------------------
+4096                1              0.958841       4.3            11.0           11.0           0.8            1.0            3.5            4.0
+8192                1              1.745779       4.7            13.0           13.0           0.8            1.0            3.8            5.0
+33554432            1              23.506385      1427.5         13.0           13.0           1.0            9.0            1426.0         1446.0
+67108864            1              23.582432      2845.7         13.0           13.0           0.9            9.0            2844.5         2851.0
+"""
+
+
+@pytest.mark.parametrize(
+    "sample,exp_latency,exp_bw",
+    [
+        (LEGACY_FORMAT, [21.1726, 21.7391, 24.6508, 39.746], [0.193457, 0.376833, 1361.19, 1688.44]),
+        (NEW_FORMAT, [4.3, 4.7, 1427.5, 2845.7], [0.958841, 1.745779, 23.506385, 23.582432]),
+    ],
+    ids=["LegacyFormat", "NewFormat"],
+)
+def test_nixl_bench_report_parsing(tmp_path: Path, sample: str, exp_latency: list[float], exp_bw: list[float]):
+    (tmp_path / "nixl_bench.log").write_text(sample)
+    df = extract_data(tmp_path / "nixl_bench.log")
+    assert df.shape == (4, 4)
+    assert df["block_size"].tolist() == [4096, 8192, 33554432, 67108864]
+    assert df["batch_size"].tolist() == [1, 1, 1, 1]
+    assert df["avg_lat"].tolist() == exp_latency
+    assert df["bw_gb_sec"].tolist() == exp_bw