Add support for --workload_filter in parallel benchmark.

PiperOrigin-RevId: 663807393
Change-Id: I3d83d9584cadaf32f890e9e7b70256b9bc2e5890

Author: liyuying0000

fleetbench/parallel/benchmark.py

@@ -15,6 +15,7 @@
 """Represent a Fleetbench benchmark."""
 
 import os
+import re
 import subprocess
 
 from absl import flags
@@ -44,12 +45,36 @@ def _FindBenchmarkPath(benchmark: str) -> str:
   raise FileNotFoundError(f"Benchmark not found: {benchmark}")
 
 
-def GetSubBenchmarks(benchmark_path: str):
+def GetSubBenchmarks(benchmark_path: str, workload: str = "") -> list[str]:
+  """Retrieves a list of sub-benchmarks from a benchmark executable.
+
+  If 'workload' is specified, only sub-benchmarks with the given workload are
+  returned.
+  """
   cmd = [benchmark_path, "--benchmark_list_tests"]
+
+  if workload == "all":
+    cmd += ["--benchmark_filter=all"]
+  elif workload:
+    cmd += [
+        f"--benchmark_filter=BM_{workload.upper()}",
+    ]
   p = subprocess.run(cmd, capture_output=True, text=True, check=True)
   return p.stdout.split("\n")[:-1]
 
 
+def GetWorkloads(benchmark_path: str):
+  """Retrieves a list of unique workloads from a benchmark executable."""
+  benchmarks = GetSubBenchmarks(benchmark_path, "all")
+  workload_pattern = r"BM_(?P<workload>[^_]+)"
+
+  def extract_workload(benchmark):
+    match = re.search(workload_pattern, benchmark)
+    return match.group("workload") if match else None
+
+  return list(set(filter(None, map(extract_workload, benchmarks))))
+
+
 class Benchmark:
   """Represents a benchmark binary and filter."""
 

fleetbench/parallel/benchmark_test.py

@@ -90,6 +90,62 @@ def testCommandLine(self):
         ],
     )
 
+  @mock.patch.object(subprocess, "run", autospec=True)
+  def testGetWorkloads(self, mock_run):
+    mock_run.return_value = subprocess.CompletedProcess(
+        args=[],
+        returncode=0,
+        stdout=(
+            "BM_LIBC_Test1\nBM_LIBC_Test2\nBM_PROTO_Test1\nBM_TCMALLOC_Test1\n"
+        ),
+        stderr="",
+    )
+    self.assertCountEqual(
+        benchmark.GetWorkloads("/path/to/benchmark"),
+        ["LIBC", "PROTO", "TCMALLOC"],
+    )
+
+  @mock.patch.object(subprocess, "run", autospec=True)
+  def testGetSubBenchmarksWorkload(self, mock_run):
+    mock_run.return_value = subprocess.CompletedProcess(
+        args=[],
+        returncode=0,
+        stdout="BM_PROTO_Arena\nBM_PROTO_NoArena\n",
+        stderr="",
+    )
+    self.assertEqual(
+        benchmark.GetSubBenchmarks("/path/to/benchmark", "proto"),
+        ["BM_PROTO_Arena", "BM_PROTO_NoArena"],
+    )
+
+  @mock.patch.object(subprocess, "run", autospec=True)
+  def testGetSubBenchmarksWorkloadWithUnmatchedBM(self, mock_run):
+    # Simulate the full list of benchmarks
+    full_benchmarks = ["BM_PROTO_Arena", "BM_PROTO_NoArena", "BM_CORD_Fleet"]
+
+    # Simulate the subprocess output
+    mock_run.return_value = subprocess.CompletedProcess(
+        args=[],
+        returncode=0,
+        stdout="\n".join(full_benchmarks),
+        stderr="",
+    )
+
+    sub_benchmarks = benchmark.GetSubBenchmarks("/path/to/benchmark", "proto")
+
+    # Assert the expected behavior
+    self.assertEqual(sub_benchmarks, ["BM_PROTO_Arena", "BM_PROTO_NoArena"])
+    mock_run.assert_called_once_with(
+        [
+            "/path/to/benchmark",
+            "--benchmark_list_tests",
+            "--benchmark_filter=BM_PROTO",
+        ],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+
 
 if __name__ == "__main__":
   absltest.main()

fleetbench/parallel/parallel_bench.py

@@ -49,7 +49,32 @@
 )
 
 _BENCHMARK_FILTER = flags.DEFINE_multi_string(
-    "benchmark_filter", [], "Specifies subset of benchmarks to run."
+    "benchmark_filter",
+    [],
+    """Specifies subset of benchmarks to run.
+
+    Filtering options:
+    - Empty list: Selects all default benchmarks.
+    - Keyword list: Selects benchmarks from the default list matching any \n
+        provided keyword, one keyword per filter \n
+        (e.g., "--benchmark_filter=Cold --benchmark_filter=Hot").""",
+)
+
+
+_WORKLOAD_FILTER = flags.DEFINE_multi_string(
+    "workload_filter",
+    [],
+    """Selects benchmarks associated with specified workloads. This will \n
+       overwrite the `--benchmark_filter` flag.
+
+    Filtering options:
+    - Workload name + keyword(s): Selects benchmarks associated with the \n
+        specified workload, further filtered by keywords 
+        (e.g., "--workload_filter=libc,Memcpy,Memcmp").
+    - Workload name + "all": Selects all benchmarks associated with the \n
+        specified workload 
+        (e.g., "--workload_filter=proto,all")
+    """,
 )
 
 _BENCHMARK_PERF_COUNTERS = flags.DEFINE_string(
@@ -112,6 +137,7 @@ def main(argv: Sequence[str]) -> None:
   results = bench.Run(
       benchmark_target=_BENCHMARK_TARGET.value,
       benchmark_filter=_BENCHMARK_FILTER.value,
+      workload_filter=_WORKLOAD_FILTER.value,
       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

@@ -32,27 +32,88 @@
 from fleetbench.parallel import worker
 
 
-def _GetBenchmarks(
+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 given target."""
+  """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").
+  """
   benchmarks = {}
   sub_benchmarks = bm.GetSubBenchmarks(benchmark_target)
 
+  # Gets default benchmark sets
   if not benchmark_filters:
-    for name in sub_benchmarks:
-      benchmark = bm.Benchmark(benchmark_target, name)
-      benchmarks[benchmark.Name()] = benchmark
-  else:
-    for bm_filter in benchmark_filters:
-      matching_bm_names = [name for name in sub_benchmarks if bm_filter in name]
-      if matching_bm_names:
-        for name in matching_bm_names:
-          benchmark = bm.Benchmark(benchmark_target, name)
-          benchmarks[benchmark.Name()] = benchmark
-      else:
-        raise ValueError(
-            f"Benchmark {bm_filter} not found in {benchmark_target}."
+    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").
+  """
+  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
@@ -131,6 +192,7 @@ def _PreRun(
       self,
       benchmark_target: str,
       benchmark_filters: list[str],
+      workload_filters: list[str],
       benchmark_perf_counters: str,
       benchmark_repetitions: int,
       benchmark_min_time: str,
@@ -139,7 +201,14 @@ def _PreRun(
 
     logging.info("Initializing benchmarks and worker threads...")
 
-    self.benchmarks = _GetBenchmarks(benchmark_target, benchmark_filters)
+    if workload_filters:
+      self.benchmarks = _GetWorkloadBenchmarks(
+          benchmark_target, workload_filters
+      )
+    else:
+      self.benchmarks = _GetDefaultBenchmarks(
+          benchmark_target, benchmark_filters
+      )
 
     benchmark_flags = _SetExtraBenchmarkFlags(
         benchmark_perf_counters, benchmark_repetitions, benchmark_min_time
@@ -331,14 +400,25 @@ def Run(
       self,
       benchmark_target: str,
       benchmark_filter: list[str] = [],
+      workload_filter: list[str] = [],
       benchmark_perf_counters: str = "",
       benchmark_repetitions: int = 0,
       benchmark_min_time: str = "",
   ) -> list[result.Result]:
     """Run benchmarks in parallel."""
+    logging.info("Running with benchmark_filter: %s", benchmark_filter)
+    logging.info("Running with workload_filter: %s", workload_filter)
+
+    if benchmark_filter and workload_filter:
+      logging.warning(
+          "Both benchmark_filter and workload_filter specified. "
+          "benchmark_filter will be ignored."
+      )
+
     self._PreRun(
         benchmark_target,
         benchmark_filter,
+        workload_filter,
         benchmark_perf_counters,
         benchmark_repetitions,
         benchmark_min_time,