Collect average "iterations" in parallel_bench.

PiperOrigin-RevId: 731762429
Change-Id: Ia1f1b795232a9ef2939c7ee4d7fde3a58e1ea3bb

Author: liyuying0000

fleetbench/parallel/parallel_bench_lib.py

@@ -91,6 +91,8 @@ class BenchmarkMetrics:
   per_iteration_wall_time: float
   # per benchmark iteration cpu time
   per_iteration_cpu_time: float
+  # per benchmark iteration
+  per_bm_run_iteration: int
 
 
 class ParallelBench:
@@ -179,6 +181,7 @@ def _PreRun(
                 total_duration=1.0,
                 per_iteration_wall_time=0.0,
                 per_iteration_cpu_time=0.0,
+                per_bm_run_iteration=0,
             )
         ]
         for benchmark in self.benchmarks.keys()
@@ -306,6 +309,7 @@ def _RunSchedulingLoop(self) -> None:
                   total_duration=r.duration,
                   per_iteration_wall_time=r.bm_wall_time,
                   per_iteration_cpu_time=r.bm_cpu_time,
+                  per_bm_run_iteration=r.iteration,
               )
           )
 
@@ -373,6 +377,7 @@ def ConvertToDataFrame(self) -> pd.DataFrame:
             "Benchmark": benchmark,
             "WallTimes": t.per_iteration_wall_time,
             "CPUTimes": t.per_iteration_cpu_time,
+            "Iterations": t.per_bm_run_iteration,
         }
         data.append(entry)
     runtimes = pd.DataFrame(data)
@@ -443,6 +448,7 @@ def Run(
                 total_duration=r.duration,
                 per_iteration_wall_time=r.bm_wall_time,
                 per_iteration_cpu_time=r.bm_cpu_time,
+                per_bm_run_iteration=r.iteration,
             )
         )
 

fleetbench/parallel/parallel_bench_lib_test.py

@@ -49,13 +49,15 @@ def tearDown(self):
   @mock.patch.object(bm, "GetSubBenchmarks", autospec=True)
   @mock.patch.object(run.Run, "Execute", autospec=True)
   @mock.patch.object(cpu, "Utilization", autospec=True)
+  @mock.patch.object(reporter, "GenerateBenchmarkReport", autospec=True)
   @mock.patch.object(reporter, "SaveBenchmarkResults", autospec=True)
   @flagsaver.flagsaver(
       benchmark_dir=absltest.get_default_test_tmpdir(),
   )
   def testRun(
       self,
       mock_save_benchmark_results,
+      mock_generate_benchmark_report,
       mock_utilization,
       mock_execute,
       mock_get_subbenchmarks,
@@ -81,6 +83,7 @@ def fake_utilization(unused_cpus):
 
     mock_utilization.side_effect = fake_utilization
 
+    mock_generate_benchmark_report.return_value = pd.DataFrame()
     mock_save_benchmark_results.return_value = None
 
     self.pb = parallel_bench_lib.ParallelBench(
@@ -126,23 +129,27 @@ def test_convert_to_dataframe(self):
             total_duration=10,
             per_iteration_wall_time=1,
             per_iteration_cpu_time=1,
+            per_bm_run_iteration=2,
         ),
         parallel_bench_lib.BenchmarkMetrics(
             total_duration=2,
             per_iteration_wall_time=3.01,
             per_iteration_cpu_time=3,
+            per_bm_run_iteration=4,
         ),
     ]
     self.pb.runtimes["BM_Test2"] = [
         parallel_bench_lib.BenchmarkMetrics(
             total_duration=10,
             per_iteration_wall_time=1,
             per_iteration_cpu_time=1,
+            per_bm_run_iteration=10,
         ),
         parallel_bench_lib.BenchmarkMetrics(
             total_duration=4,
             per_iteration_wall_time=4,
             per_iteration_cpu_time=5,
+            per_bm_run_iteration=8,
         ),
     ]
     self.pb.utilization_samples.append((pd.Timestamp.now(), 0.5))
@@ -151,8 +158,18 @@ def test_convert_to_dataframe(self):
     self.assertEqual(
         df.to_dict("records"),
         [
-            {"Benchmark": "BM_Test1", "WallTimes": 3.01, "CPUTimes": 3},
-            {"Benchmark": "BM_Test2", "WallTimes": 4, "CPUTimes": 5},
+            {
+                "Benchmark": "BM_Test1",
+                "WallTimes": 3.01,
+                "CPUTimes": 3,
+                "Iterations": 4,
+            },
+            {
+                "Benchmark": "BM_Test2",
+                "WallTimes": 4,
+                "CPUTimes": 5,
+                "Iterations": 8,
+            },
         ],
     )
 

fleetbench/parallel/reporter.py

@@ -103,9 +103,13 @@ def GenerateBenchmarkReport(
           Count=("WallTimes", "count"),
           Mean_Wall_Time=("WallTimes", "mean"),
           Mean_CPU_Time=("CPUTimes", "mean"),
+          Mean_Iterations=("Iterations", "mean"),
       )
       .round(3)
   )
+  grouped_results["Mean_Iterations"] = grouped_results[
+      "Mean_Iterations"
+  ].astype(int)
 
   # Combine perf_counter_df and benchmark run results on the same
   # "benchmark" entry.
@@ -134,6 +138,7 @@ def SaveBenchmarkResults(output_dir, df: pd.DataFrame) -> None:
       columns={
           "Mean_Wall_Time": "real_time",
           "Mean_CPU_Time": "cpu_time",
+          "Mean_Iterations": "iterations",
       }
   )
   data = df.reset_index().to_dict(orient="records")

fleetbench/parallel/reporter_test.py

@@ -115,9 +115,24 @@ def test_aggregate_constext_inconsistent_data(self):
 
   def test_generate_benchmark_report(self):
     benchmark_df = pd.DataFrame([
-        {"Benchmark": "BM_Test1", "WallTimes": 3, "CPUTimes": 3},
-        {"Benchmark": "BM_Test2", "WallTimes": 4, "CPUTimes": 5},
-        {"Benchmark": "BM_Test1", "WallTimes": 6, "CPUTimes": 7},
+        {
+            "Benchmark": "BM_Test1",
+            "WallTimes": 3,
+            "CPUTimes": 3,
+            "Iterations": 10,
+        },
+        {
+            "Benchmark": "BM_Test2",
+            "WallTimes": 4,
+            "CPUTimes": 5,
+            "Iterations": 20,
+        },
+        {
+            "Benchmark": "BM_Test1",
+            "WallTimes": 6,
+            "CPUTimes": 7,
+            "Iterations": 50,
+        },
     ])
     perf_counter_df = pd.DataFrame([
         {"Benchmark": "BM_Test1", "instructions": 130.0, "cycles": 3.0},
@@ -134,6 +149,7 @@ def test_generate_benchmark_report(self):
             "Count": 2,
             "Mean_Wall_Time": 4.5,
             "Mean_CPU_Time": 5.0,
+            "Mean_Iterations": 30,
             "instructions": 130.0,
             "cycles": 3.0,
         },
@@ -142,6 +158,7 @@ def test_generate_benchmark_report(self):
             "Count": 1,
             "Mean_Wall_Time": 4,
             "Mean_CPU_Time": 5.0,
+            "Mean_Iterations": 20,
             "instructions": 200.0,
             "cycles": 2.0,
         },

fleetbench/parallel/result.py

@@ -30,3 +30,4 @@ class Result:
   duration: Optional[float] = None
   bm_wall_time: Optional[float] = None
   bm_cpu_time: Optional[float] = None
+  iteration: Optional[int] = None

fleetbench/parallel/run.py

@@ -52,15 +52,16 @@ def Execute(self) -> result.Result:
 
     with open(self.out_file, "r") as f:
       benchmark_output = f.read()
-    benchmark_cpu_time, benchmark_wall_time = self._GetBenchmarkRuntime(
-        benchmark_output
+    benchmark_cpu_time, benchmark_wall_time, iteration = (
+        self._GetBenchmarkRuntime(benchmark_output)
     )
     return result.Result(
         benchmark=self.benchmark.Name(),
         rc=proc.returncode,
         duration=end - start,
         bm_wall_time=benchmark_wall_time,
         bm_cpu_time=benchmark_cpu_time,
+        iteration=iteration,
         result=benchmark_output,
         stdout=proc.stdout,
         stderr=proc.stderr,
@@ -72,6 +73,12 @@ def _CommandLine(self) -> list[str]:
         f"--benchmark_out={self.out_file}",
     ]
 
-  def _GetBenchmarkRuntime(self, benchmark_output: str) -> tuple[float, float]:
+  def _GetBenchmarkRuntime(
+      self, benchmark_output: str
+  ) -> tuple[float, float, int]:
     data = json.loads(benchmark_output)
-    return data["benchmarks"][0]["cpu_time"], data["benchmarks"][0]["real_time"]
+    return (
+        data["benchmarks"][0]["cpu_time"],
+        data["benchmarks"][0]["real_time"],
+        data["benchmarks"][0]["iterations"],
+    )

fleetbench/parallel/run_test.py

@@ -33,7 +33,11 @@ class RunTest(absltest.TestCase):
   def testRun(self, mock_run):
     output_file = self.create_tempfile()
 
-    data = {"benchmarks": [{"cpu_time": 12.345, "real_time": 12.3}]}
+    data = {
+        "benchmarks": [
+            {"cpu_time": 12.345, "real_time": 12.3, "iterations": 10}
+        ]
+    }
     json_object = json.dumps(data, indent=4)
     with open(output_file.full_path, "w") as f:
       f.write(json_object)
@@ -56,6 +60,7 @@ def testRun(self, mock_run):
     self.assertGreater(result.duration, 0)
     self.assertEqual(result.bm_cpu_time, 12.345)
     self.assertEqual(result.bm_wall_time, 12.3)
+    self.assertEqual(result.iteration, 10)
     self.assertEqual(result.rc, 0)
 
   @flagsaver.flagsaver(