Move perf counters from Run() to a ParallelBench class variable.

This change is intended to be a refactoring, and shouldn't change the behavior of the program.

PiperOrigin-RevId: 746145404
Change-Id: I870110b0af524769ec5b76ac6386b7cf3442ed4b

Author: liyuying0000

fleetbench/parallel/parallel_bench.py

@@ -199,6 +199,7 @@ def main(argv: Sequence[str]) -> None:
       repetitions=_REPETITIONS.value,
       temp_parent_root=_TEMP_ROOT.value,
       keep_raw_data=_KEEP_RAW_DATA.value,
+      benchmark_perf_counters=_BENCHMARK_PERF_COUNTERS.value,
   )
 
   bench.SetWeights(
@@ -210,7 +211,6 @@ def main(argv: Sequence[str]) -> None:
   )
 
   bench.Run(
-      benchmark_perf_counters=_BENCHMARK_PERF_COUNTERS.value,
       benchmark_repetitions=_BENCHMARK_REPETITIONS.value,
       benchmark_min_time=_BENCHMARK_MIN_TIME.value,
   )

fleetbench/parallel/parallel_bench_lib.py

@@ -35,16 +35,15 @@
 
 
 def _SetExtraBenchmarkFlags(
-    benchmark_perf_counters: str,
+    benchmark_perf_counters: list[str],
     benchmark_repetitions: int,
     benchmark_min_time: str,
 ) -> list[str]:
   """Set extra benchmark flags."""
   benchmark_flags = []
   if benchmark_perf_counters:
-    benchmark_flags.append(
-        f"--benchmark_perf_counters={benchmark_perf_counters}"
-    )
+    perf_counters_str = ",".join(benchmark_perf_counters)
+    benchmark_flags.append(f"--benchmark_perf_counters={perf_counters_str}")
   if benchmark_min_time:
     benchmark_flags.append(f"--benchmark_min_time={benchmark_min_time}")
   if benchmark_repetitions:
@@ -82,6 +81,8 @@ class ParallelBench:
     temp_root: Child directory to store results in. It is in the format of
       "temp_parent_root/run_{repetition_id}".
     keep_raw_data: Whether to keep the raw results from each run.
+    perf_counters: A list of performance counters to collect during benchmark
+      run.
     runtimes: Dictionary of benchmark name -> history of benchmark runtimes.
     workers: Dictionary of CPU ID -> Worker thread.
     results: List of results from all runs.
@@ -90,7 +91,6 @@ class ParallelBench:
     target_ratios: List of target ratios for each benchmark. This is used to
       calculate the probability of each benchmark being selected, and determined
       by the benchmark weights.
-    keep_raw_data: Whether to keep the raw results from each run.
     first_run: Boolean indicating if this is the first run. We use this to
       determine if we can randomly select benchmarks or if we need to run all
       benchmarks at least once.
@@ -105,6 +105,7 @@ def __init__(
       repetitions: int,
       temp_parent_root: str,
       keep_raw_data: bool,
+      benchmark_perf_counters: str,
   ):
     """Initialize the parallel benchmark runner."""
 
@@ -123,6 +124,9 @@ def __init__(
     self.temp_parent_root = temp_parent_root
     self.temp_root = ""
     self.keep_raw_data = keep_raw_data
+    self.perf_counters: list[str] = (
+        benchmark_perf_counters.split(",") if benchmark_perf_counters else []
+    )
     self.runtimes: dict[str, list[BenchmarkMetrics]] = {}
     self.workers: dict[int, worker.Worker] = {}
     self.utilization_samples: list[tuple[pd.Timestamp, float]] = []
@@ -161,7 +165,6 @@ def SetWeights(
 
   def _PreRun(
       self,
-      benchmark_perf_counters: str,
       benchmark_repetitions: int,
       benchmark_min_time: str,
       repetition: int,
@@ -171,7 +174,7 @@ def _PreRun(
     logging.info("Initializing benchmarks and worker threads...")
 
     benchmark_flags = _SetExtraBenchmarkFlags(
-        benchmark_perf_counters, benchmark_repetitions, benchmark_min_time
+        self.perf_counters, benchmark_repetitions, benchmark_min_time
     )
 
     if benchmark_flags:
@@ -395,24 +398,16 @@ def _RunSchedulingLoop(self) -> None:
               )
           )
 
-  def GeneratePerfCounterDataFrame(
-      self, benchmark_perf_counters: str
-  ) -> pd.DataFrame | None:
+  def GeneratePerfCounterDataFrame(self) -> pd.DataFrame | None:
     """Generates a DataFrame of performance counter results for each benchmark.
 
-    Args:
-      benchmark_perf_counters: A comma-separated list of performance counters to
-        collect.
-
     Returns:
       A DataFrame of performance counter results for each benchmark, or None if
       no performance counters were specified.
     """
-    if not benchmark_perf_counters:
+    if not self.perf_counters:
       return None
 
-    counters = benchmark_perf_counters.split(",")
-
     performance_data = []
     for filename in os.listdir(self.temp_root):
 
@@ -426,7 +421,7 @@ def GeneratePerfCounterDataFrame(
       entry = {
           "Benchmark": benchmark_result["name"],
       }
-      for counter in counters:
+      for counter in self.perf_counters:
         if counter in benchmark_result:
           entry[counter] = benchmark_result[counter]
       performance_data.append(entry)
@@ -438,7 +433,7 @@ def GeneratePerfCounterDataFrame(
     # Group the results by benchmark and counter, and calculate the mean of each
     # counter for each benchmark.
     aggregations = {}
-    for counter in counters:
+    for counter in self.perf_counters:
       aggregations[counter] = pd.NamedAgg(column=counter, aggfunc=np.mean)
 
     perf_counters_results = (
@@ -481,25 +476,20 @@ def _RemoveRawData(self) -> None:
           logging.exception("Failed to remove %s: %s", file_path, e)
 
   def PostProcessBenchmarkResults(
-      self, benchmark_perf_counters: str
+      self,
   ) -> tuple[reporter.ContextInfo, reporter.BenchmarkRuntimeInfo]:
     """Generate benchmark reports and save results to a JSON file.
 
     If benchmark_perf_counters is specified, the report will include perf
     counters for each benchmark. We will also check to see if we want to remove
     the raw data.
-
-    Args:
-      benchmark_perf_counters: A comma-separated list of performance counters to
-        collect.
-
     Returns:
       A tuple containing the context and benchmark data dictionaries.
     """
 
     df = self.ConvertToDataFrame()
 
-    perf_counter_df = self.GeneratePerfCounterDataFrame(benchmark_perf_counters)
+    perf_counter_df = self.GeneratePerfCounterDataFrame()
     df = reporter.GenerateBenchmarkReport(df, perf_counter_df)
     context, data = reporter.SaveBenchmarkResults(self.temp_root, df)
 
@@ -509,7 +499,6 @@ def PostProcessBenchmarkResults(
 
   def Run(
       self,
-      benchmark_perf_counters: str = "",
       benchmark_repetitions: int = 0,
       benchmark_min_time: str = "",
   ):
@@ -526,7 +515,6 @@ def Run(
       print(f"Running trial {i}.......")
 
       self._PreRun(
-          benchmark_perf_counters,
           benchmark_repetitions,
           benchmark_min_time,
           i,
@@ -561,7 +549,7 @@ def Run(
         w.join()
 
       # Post-process benchmark results
-      context, data = self.PostProcessBenchmarkResults(benchmark_perf_counters)
+      context, data = self.PostProcessBenchmarkResults()
       context_list.append(context)
       data_list.append(data)
 

fleetbench/parallel/parallel_bench_lib_test.py

@@ -43,6 +43,7 @@ def setUp(self):
         repetitions=1,
         temp_parent_root=absltest.get_default_test_tmpdir(),
         keep_raw_data=True,
+        benchmark_perf_counters="",
     )
 
   def tearDown(self):
@@ -102,6 +103,7 @@ def fake_utilization(unused_cpus):
         repetitions=1,
         temp_parent_root=absltest.get_default_test_tmpdir(),
         keep_raw_data=True,
+        benchmark_perf_counters="",
     )
     self.pb.SetWeights(
         benchmark_target="fake_bench",
@@ -164,7 +166,7 @@ def test_run_multiple_repetitions(
   def test_set_extra_benchmark_flags(self):
     self.assertEqual(
         parallel_bench_lib._SetExtraBenchmarkFlags(
-            benchmark_perf_counters="instructions",
+            benchmark_perf_counters=["instructions"],
             benchmark_repetitions=10,
             benchmark_min_time="10s",
         ),
@@ -177,12 +179,12 @@ def test_set_extra_benchmark_flags(self):
 
     self.assertEqual(
         parallel_bench_lib._SetExtraBenchmarkFlags(
-            benchmark_perf_counters="instructions",
+            benchmark_perf_counters=["instructions", "cycles"],
             benchmark_repetitions=0,
             benchmark_min_time="",
         ),
         [
-            "--benchmark_perf_counters=instructions",
+            "--benchmark_perf_counters=instructions,cycles",
         ],
     )
 
@@ -274,8 +276,8 @@ def test_generate_perf_counter_dataframe(self):
     with open(os.path.join(self.pb.temp_root, "run_3"), "w") as f:
       json.dump(mock_data3, f)
 
-    counters = "instructions,cycles"
-    df = self.pb.GeneratePerfCounterDataFrame(counters)
+    self.pb.perf_counters = ["instructions", "cycles"]
+    df = self.pb.GeneratePerfCounterDataFrame()
 
     self.assertIsInstance(df, pd.DataFrame)
     self.assertLen(df, 2)  # There are two benchmarks
@@ -311,7 +313,7 @@ def test_post_processing_benchmark_results(
     ])
 
     mock_save_benchmark_results.return_value = (None, None)
-    self.pb.PostProcessBenchmarkResults("instructions,cycles")
+    self.pb.PostProcessBenchmarkResults()
     mock_generate_perf_counter_dataframe.assert_called_once()
     mock_generate_benchmark_report.assert_called_once()
     mock_save_benchmark_results.assert_called_once()