Move benchmark_filter/workload_filter functions to benchmark.py

liyuying0000 · copybara-github · commit 7a7de6a8af15 · 2025-02-18T14:49:09.000-08:00
PiperOrigin-RevId: 728366933
Change-Id: I12e37ce3c2540bc6c1a4ea01b371dbd5250ad2f5
diff --git a/fleetbench/parallel/benchmark.py b/fleetbench/parallel/benchmark.py
@@ -60,6 +60,14 @@ def GetSubBenchmarks(benchmark_path: str, workload: str = "") -> list[str]:
 
   If 'workload' is specified, only sub-benchmarks with the given workload are
   returned.
+
+  Args:
+    benchmark_path: Path to the benchmark binary.
+    workload: The workload to filter for. If empty, all sub-benchmarks are
+      returned.
+
+  Returns:
+    A list of sub-benchmark names that match the workload filter.
   """
   cmd = [benchmark_path, "--benchmark_list_tests"]
 
@@ -84,11 +92,11 @@ def GetWorkloads(benchmark_path: str):
   benchmarks = GetSubBenchmarks(benchmark_path, "all")
   workload_pattern = r"BM_(?P<workload>[^_]+)"
 
-  def extract_workload(benchmark):
+  def _ExtractWorkload(benchmark):
     match = re.search(workload_pattern, benchmark)
     return match.group("workload") if match else None
 
-  return list(set(filter(None, map(extract_workload, benchmarks))))
+  return list(set(filter(None, map(_ExtractWorkload, benchmarks))))
 
 
 class Benchmark:
@@ -116,3 +124,101 @@ def Path(self):
 
   def __str__(self):
     return self.Name()
+
+
+def _CreateBenchmarks(bm_target: str, names: list[str]) -> dict[str, Benchmark]:
+  """Creates benchmark dictionary with the benchmark name as the key."""
+  benchmarks = {}
+  for name in names:
+    benchmark = Benchmark(bm_target, name)
+    benchmarks[benchmark.Name()] = benchmark
+  return benchmarks
+
+
+def _CreateMatchingBenchmarks(
+    bm_target: str, bm_filter: str, bm_candidates: list[str]
+) -> dict[str, Benchmark]:
+  """Creates benchmarks that match the given filter."""
+  matching_bm_names = [name for name in bm_candidates if bm_filter in name]
+  if not matching_bm_names:
+    raise ValueError(f"Can't find benchmarks matching {bm_filter}.")
+  return _CreateBenchmarks(bm_target, matching_bm_names)
+
+
+def GetDefaultBenchmarks(
+    benchmark_target: str, benchmark_filters: list[str]
+) -> dict[str, Benchmark]:
+  """Get a list of benchmarks from the default target.
+
+    Filtering options:
+  - Empty list: Returns all default benchmarks.
+  - Keyword list: Returns benchmarks from the default list matching the provided
+                  keyword (e.g., "Cold Hot").
+
+  Args:
+    benchmark_target: Path to the benchmark binary.
+    benchmark_filters: List of filters to apply to the benchmarks to run.
+
+  Returns:
+    A map of benchmark names to Benchmark objects.
+  """
+  benchmarks = {}
+  sub_benchmarks = GetSubBenchmarks(benchmark_target)
+
+  # Gets default benchmark sets
+  if not benchmark_filters:
+    return _CreateBenchmarks(benchmark_target, sub_benchmarks)
+
+  # Gets benchmark sets from filters
+  for bm_filter in benchmark_filters:
+    benchmarks.update(
+        _CreateMatchingBenchmarks(benchmark_target, bm_filter, sub_benchmarks)
+    )
+  return benchmarks
+
+
+def GetWorkloadBenchmarks(
+    benchmark_target: str, workload_filters: list[str]
+) -> dict[str, Benchmark]:
+  """Get a list of benchmarks from the given workload that match the filters.
+
+  Filtering options:
+    - Workload name + keyword(s): Returns benchmarks associated with the
+        specified workload, further filtered by keywords (e.g.,
+        "libc,Memcpy,Memcmp").
+    - Workload name + "all": Returns all benchmarks associated with the
+        specified workload (e.g., "proto,all").
+  Args:
+    benchmark_target: Path to the benchmark binary.
+    workload_filters: List of filters to apply to the benchmarks to run.
+
+  Returns:
+    A map of benchmark names to Benchmark objects.
+  """
+  benchmarks = {}
+
+  # Get all unique workloads
+  workloads = GetWorkloads(benchmark_target)
+
+  def _GetWorkloadAndFilter(bm_filter: str) -> tuple[str, list[str]]:
+    parts = bm_filter.split(",")
+    if parts[0].upper() not in workloads:
+      raise ValueError(f"Workload {parts[0]} not supported in Fleetbench.")
+    return parts[0], parts[1:]
+
+  for workload_filter in workload_filters:
+    workload, bm_filters = _GetWorkloadAndFilter(workload_filter)
+    workload_bms = GetSubBenchmarks(benchmark_target, workload)
+    if bm_filters == ["all"]:
+      benchmarks.update(
+          _CreateMatchingBenchmarks(
+              benchmark_target, workload.upper(), workload_bms
+          )
+      )
+    else:
+      for bm_filter in bm_filters:
+        benchmarks.update(
+            _CreateMatchingBenchmarks(benchmark_target, bm_filter, workload_bms)
+        )
+
+  return benchmarks
diff --git a/fleetbench/parallel/benchmark_test.py b/fleetbench/parallel/benchmark_test.py
@@ -147,6 +147,97 @@ def testGetSubBenchmarksWorkloadWithUnmatchedBM(self, mock_run):
         env=mock.ANY,
     )
 
+  @mock.patch.object(benchmark, "GetSubBenchmarks", autospec=True)
+  @flagsaver.flagsaver(
+      benchmark_dir=absltest.get_default_test_tmpdir(),
+  )
+  def test_getbenchmark_without_filter(self, mock_get_subbenchmarks):
+    mock_get_subbenchmarks.return_value = ["BM_Test1", "BM_Test2"]
+    self.create_tempfile(
+        os.path.join(absltest.get_default_test_tmpdir(), "fake_bench")
+    )
+
+    benchmarks = benchmark.GetDefaultBenchmarks("fake_bench", [])
+    self.assertLen(benchmarks, 2)
+    self.assertCountEqual(
+        benchmarks.keys(),
+        ["fake_bench (BM_Test1)", "fake_bench (BM_Test2)"],
+    )
+
+  @mock.patch.object(benchmark, "GetSubBenchmarks", autospec=True)
+  @flagsaver.flagsaver(
+      benchmark_dir=absltest.get_default_test_tmpdir(),
+  )
+  def test_getbenchmark_with_filter_partial_match(self, mock_get_subbenchmarks):
+    mock_get_subbenchmarks.return_value = ["BM_Test1", "BM_Test2"]
+    self.create_tempfile(
+        os.path.join(absltest.get_default_test_tmpdir(), "fake_bench")
+    )
+    benchmarks = benchmark.GetDefaultBenchmarks("fake_bench", ["Test1"])
+    self.assertLen(benchmarks, 1)
+    self.assertCountEqual(
+        benchmarks.keys(),
+        ["fake_bench (BM_Test1)"],
+    )
+
+  @mock.patch.object(benchmark, "GetWorkloads", autospec=True)
+  @mock.patch.object(benchmark, "GetSubBenchmarks", autospec=True)
+  @flagsaver.flagsaver(
+      benchmark_dir=absltest.get_default_test_tmpdir(),
+  )
+  def test_getworkloadbenchmark_subset(
+      self, mock_get_subbenchmarks, mock_get_workloads
+  ):
+    mock_get_workloads.return_value = ["PROTO", "CORD"]
+    mock_get_subbenchmarks.side_effect = [
+        ["BM_PROTO_Test1", "BM_PROTO_Test2", "BM_PROTO_Test3"],
+        ["BM_CORD_Test1"],
+    ]
+    self.create_tempfile(
+        os.path.join(absltest.get_default_test_tmpdir(), "fake_bench")
+    )
+
+    benchmarks = benchmark.GetWorkloadBenchmarks(
+        "fake_bench", ["proto,1,2", "cord,all"]
+    )
+    self.assertLen(benchmarks, 3)
+    self.assertCountEqual(
+        benchmarks.keys(),
+        [
+            "fake_bench (BM_PROTO_Test1)",
+            "fake_bench (BM_PROTO_Test2)",
+            "fake_bench (BM_CORD_Test1)",
+        ],
+    )
+
+  @mock.patch.object(benchmark, "GetWorkloads", autospec=True)
+  @mock.patch.object(benchmark, "GetSubBenchmarks", autospec=True)
+  @flagsaver.flagsaver(
+      benchmark_dir=absltest.get_default_test_tmpdir(),
+  )
+  def test_getworkloadbenchmark_all(
+      self, mock_get_subbenchmarks, mock_get_workloads
+  ):
+    mock_get_workloads.return_value = ["PROTO"]
+    mock_get_subbenchmarks.return_value = [
+        "BM_PROTO_Test1",
+        "BM_PROTO_Test2",
+        "BM_PROTO_Test3",
+    ]
+    self.create_tempfile(
+        os.path.join(absltest.get_default_test_tmpdir(), "fake_bench")
+    )
+    benchmarks = benchmark.GetWorkloadBenchmarks("fake_bench", ["proto,all"])
+    self.assertLen(benchmarks, 3)
+    self.assertCountEqual(
+        benchmarks.keys(),
+        [
+            "fake_bench (BM_PROTO_Test1)",
+            "fake_bench (BM_PROTO_Test2)",
+            "fake_bench (BM_PROTO_Test3)",
+        ],
+    )
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/fleetbench/parallel/parallel_bench_lib.py b/fleetbench/parallel/parallel_bench_lib.py
@@ -26,7 +26,6 @@
 
 from fleetbench.parallel import benchmark as bm
 from fleetbench.parallel import cpu
-from fleetbench.parallel import result
 from fleetbench.parallel import run
 from fleetbench.parallel import worker
 
@@ -41,7 +40,7 @@ def ParseBenchmarkWeights(
   filter should be in ALL CAPS to ensure case-insensitive matching.
 
   Args:
-    benchmark_list: A list of strings to parse.
+    benchmark_weights_list: A list of strings to parse.
 
   Returns:
     A dictionary of {capitalized string: float} or None if the list is empty.
@@ -56,114 +55,14 @@ def ParseBenchmarkWeights(
       benchmark_weights[key.upper()] = float(value_str)
     except ValueError:
       logging.warning(
-          f"Invalid benchmark string: %s. The format should be"
-          f" <benchmark_name|benchmark_filter>:<weight>. Skipping...",
+          "Invalid benchmark string: %s. The format should be"
+          " <benchmark_name|benchmark_filter>:<weight>. Skipping...",
           weights,
       )
 
   return benchmark_weights
 
 
-def _CreateBenchmarks(
-    bm_target: str, names: list[str]
-) -> dict[str, bm.Benchmark]:
-  """Creates benchmark dictionary with the benchmark name as the key."""
-  benchmarks = {}
-  for name in names:
-    benchmark = bm.Benchmark(bm_target, name)
-    benchmarks[benchmark.Name()] = benchmark
-  return benchmarks
-
-
-def _CreateMatchingBenchmarks(
-    bm_target: str, bm_filter: str, bm_candidates: list[str]
-) -> dict[str, bm.Benchmark]:
-  """Creates benchmarks that match the given filter."""
-  matching_bm_names = [name for name in bm_candidates if bm_filter in name]
-  if not matching_bm_names:
-    raise ValueError(f"Can't find benchmarks matching {bm_filter}.")
-  return _CreateBenchmarks(bm_target, matching_bm_names)
-
-
-def _GetDefaultBenchmarks(
-    benchmark_target: str, benchmark_filters: list[str]
-) -> dict[str, bm.Benchmark]:
-  """Get a list of benchmarks from the default target.
-
-    Filtering options:
-  - Empty list: Returns all default benchmarks.
-  - Keyword list: Returns benchmarks from the default list matching the provided
-                  keyword (e.g., "Cold Hot").
-
-  Args:
-    benchmark_target: Path to the benchmark binary.
-    benchmark_filters: List of filters to apply to the benchmarks to run.
-
-  Returns:
-    A map of benchmark names to Benchmark objects.
-  """
-  benchmarks = {}
-  sub_benchmarks = bm.GetSubBenchmarks(benchmark_target)
-
-  # Gets default benchmark sets
-  if not benchmark_filters:
-    return _CreateBenchmarks(benchmark_target, sub_benchmarks)
-
-  # Gets benchmark sets from filters
-  for bm_filter in benchmark_filters:
-    benchmarks.update(
-        _CreateMatchingBenchmarks(benchmark_target, bm_filter, sub_benchmarks)
-    )
-  return benchmarks
-
-
-def _GetWorkloadBenchmarks(
-    benchmark_target: str, workload_filters: list[str]
-) -> dict[str, bm.Benchmark]:
-  """Get a list of benchmarks from the given workload that match the filters.
-
-  Filtering options:
-    - Workload name + keyword(s): Returns benchmarks associated with the
-        specified workload, further filtered by keywords (e.g.,
-        "libc,Memcpy,Memcmp").
-    - Workload name + "all": Returns all benchmarks associated with the
-        specified workload (e.g., "proto,all").
-  Args:
-    benchmark_target: Path to the benchmark binary.
-    workload_filters: List of filters to apply to the benchmarks to run.
-
-  Returns:
-    A map of benchmark names to Benchmark objects.
-  """
-  benchmarks = {}
-
-  # Get all unique workloads
-  workloads = bm.GetWorkloads(benchmark_target)
-
-  def _GetWorkloadAndFilter(bm_filter: str) -> tuple[str, list[str]]:
-    parts = bm_filter.split(",")
-    if parts[0].upper() not in workloads:
-      raise ValueError(f"Workload {parts[0]} not supported in Fleetbench.")
-    return parts[0], parts[1:]
-
-  for workload_filter in workload_filters:
-    workload, bm_filters = _GetWorkloadAndFilter(workload_filter)
-    workload_bms = bm.GetSubBenchmarks(benchmark_target, workload)
-    if bm_filters == ["all"]:
-      benchmarks.update(
-          _CreateMatchingBenchmarks(
-              benchmark_target, workload.upper(), workload_bms
-          )
-      )
-    else:
-      for bm_filter in bm_filters:
-        benchmarks.update(
-            _CreateMatchingBenchmarks(benchmark_target, bm_filter, workload_bms)
-        )
-
-  return benchmarks
-
-
 def _SetExtraBenchmarkFlags(
     benchmark_perf_counters: str,
     benchmark_repetitions: int,
@@ -254,11 +153,11 @@ def _PreRun(
     logging.info("Initializing benchmarks and worker threads...")
 
     if workload_filters:
-      self.benchmarks = _GetWorkloadBenchmarks(
+      self.benchmarks = bm.GetWorkloadBenchmarks(
           benchmark_target, workload_filters
       )
     else:
-      self.benchmarks = _GetDefaultBenchmarks(
+      self.benchmarks = bm.GetDefaultBenchmarks(
           benchmark_target, benchmark_filters
       )
 
@@ -271,8 +170,8 @@ def _PreRun(
       for benchmark in self.benchmarks.values():
         benchmark.AddCommandFlags(benchmark_flags)
 
-    # Initialize the runtimes with a fake wall time of 1. This causes all benchmarks
-    # to be equally likely at first.
+    # Initialize the runtimes with a fake wall time of 1. This causes all
+    # benchmarks to be equally likely at first.
     self.runtimes = {
         benchmark: [
             BenchmarkMetrics(
diff --git a/fleetbench/parallel/parallel_bench_lib_test.py b/fleetbench/parallel/parallel_bench_lib_test.py
