Create a weights.py file for all functionalities to manipulate benchmark weights.

PiperOrigin-RevId: 736648945
Change-Id: If970a454623bbe9e47e59db8f6c780a2c075294d

Author: liyuying0000

fleetbench/parallel/BUILD

@@ -58,6 +58,7 @@ py_library(
         ":cpu",
         ":reporter",
         ":run",
+        ":weights",
         ":worker",
         "@com_google_absl_py//absl/logging",
         requirement("numpy"),
@@ -74,6 +75,25 @@ py_library(
     ],
 )
 
+py_library(
+    name = "weights",
+    srcs = ["weights.py"],
+    deps = [
+        ":benchmark",
+        "@com_google_absl_py//absl/logging",
+    ],
+)
+
+py_test(
+    name = "weights_test",
+    srcs = ["weights_test.py"],
+    deps = [
+        ":benchmark",
+        ":weights",
+        "@com_google_absl_py//absl/testing:absltest",
+    ],
+)
+
 py_binary(
     name = "parallel_bench",
     testonly = True,

fleetbench/parallel/benchmark.py

@@ -122,6 +122,9 @@ def Name(self):
   def Path(self):
     return f"{self._path}"
 
+  def BenchmarkName(self):
+    return self._benchmark_filter
+
   def __str__(self):
     return self.Name()
 

fleetbench/parallel/parallel_bench.py

@@ -105,7 +105,7 @@
 )
 
 
-_BENCHMARK_WEIGHTS = flags.DEFINE_multi_string(
+_CUSTOM_BENCHMARK_WEIGHTS = flags.DEFINE_multi_string(
     "benchmark_weights",
     [],
     "Weights for selected benchmarks. Default weight of 1.0 is used if not"
@@ -154,25 +154,22 @@ def main(argv: Sequence[str]) -> None:
   )
   cpus, target_utilization = scheduling_mode.SelectCPURangeAndSetUtilization()
 
-  # Parse benchmark weights.
-  benchmark_weights = parallel_bench_lib.ParseBenchmarkWeights(
-      _BENCHMARK_WEIGHTS.value
-  )
-  logging.info("Running with selected benchmark weights: %s", benchmark_weights)
-
   bench = parallel_bench_lib.ParallelBench(
       cpus=cpus,
       cpu_affinity=_CPU_AFFINITY.value,
-      benchmark_weights=benchmark_weights,
       utilization=target_utilization,
       duration=_DURATION.value,
       temp_root=_TEMP_ROOT.value,
   )
 
+  bench.SetWeights(
+      _BENCHMARK_TARGET.value,
+      _BENCHMARK_FILTER.value,
+      _WORKLOAD_FILTER.value,
+      _CUSTOM_BENCHMARK_WEIGHTS.value,
+  )
+
   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

@@ -28,42 +28,10 @@
 from fleetbench.parallel import cpu
 from fleetbench.parallel import reporter
 from fleetbench.parallel import run
+from fleetbench.parallel import weights
 from fleetbench.parallel import worker
 
 
-def ParseBenchmarkWeights(
-    benchmark_weights_list: list[str],
-) -> dict[str, float] | None:
-  """Parses a list of benchmark weights into a dictionary.
-
-  The string element in the list should be in the format:
-  <benchmark_name|benchmark_filter>:<weight>. Note that the benchmark name or
-  filter should be in ALL CAPS to ensure case-insensitive matching.
-
-  Args:
-    benchmark_weights_list: A list of strings to parse.
-
-  Returns:
-    A dictionary of {capitalized string: float} or None if the list is empty.
-  """
-  if not benchmark_weights_list:
-    return None
-
-  benchmark_weights = {}
-  for weights in benchmark_weights_list:
-    try:
-      key, value_str = weights.split(":")
-      benchmark_weights[key.upper()] = float(value_str)
-    except ValueError:
-      logging.warning(
-          "Invalid benchmark string: %s. The format should be"
-          " <benchmark_name|benchmark_filter>:<weight>. Skipping...",
-          weights,
-      )
-
-  return benchmark_weights
-
-
 def _SetExtraBenchmarkFlags(
     benchmark_perf_counters: str,
     benchmark_repetitions: int,
@@ -119,7 +87,6 @@ def __init__(
       self,
       cpus: list[int],
       cpu_affinity: bool,
-      benchmark_weights: dict[str, float] | None,
       utilization: float,
       duration: float,
       temp_root: str,
@@ -133,7 +100,7 @@ def __init__(
     self.controller_cpu = cpus[0]
     self.cpus = cpus[1:]
     self.cpu_affinity = cpu_affinity
-    self.benchmark_weights = benchmark_weights
+    self.benchmark_weights: dict[str, float] = {}
     self.benchmarks: dict[str, bm.Benchmark] = {}
     self.target_utilization = utilization * 100
     self.duration = duration
@@ -142,27 +109,44 @@ def __init__(
     self.workers: dict[int, worker.Worker] = {}
     self.utilization_samples: list[tuple[pd.Timestamp, float]] = []
 
-  def _PreRun(
+  def SetWeights(
       self,
       benchmark_target: str,
-      benchmark_filters: list[str] | None,
-      workload_filters: list[str] | None,
-      benchmark_perf_counters: str,
-      benchmark_repetitions: int,
-      benchmark_min_time: str,
+      benchmark_filter: list[str] | None,
+      workload_filter: list[str] | None,
+      custom_benchmark_weights: list[str] | None,
   ) -> None:
-    """Initial configuration steps."""
-
-    logging.info("Initializing benchmarks and worker threads...")
+    """Sets the benchmark weights."""
+    logging.info("Running with benchmark_filter: %s", benchmark_filter)
+    logging.info("Running with workload_filter: %s", workload_filter)
 
-    if workload_filters:
+    if benchmark_filter and workload_filter:
+      logging.warning(
+          "Both benchmark_filter and workload_filter specified. "
+          "benchmark_filter will be ignored."
+      )
+    if workload_filter:
       self.benchmarks = bm.GetWorkloadBenchmarks(
-          benchmark_target, workload_filters
+          benchmark_target, workload_filter
       )
     else:
       self.benchmarks = bm.GetDefaultBenchmarks(
-          benchmark_target, benchmark_filters
+          benchmark_target, benchmark_filter
       )
+    # Gets the number of workloads and num of benchmark for each workload
+    self.benchmark_weights = weights.GetBenchmarkWeights(
+        self.benchmarks, custom_benchmark_weights
+    )
+
+  def _PreRun(
+      self,
+      benchmark_perf_counters: str,
+      benchmark_repetitions: int,
+      benchmark_min_time: str,
+  ) -> None:
+    """Initial configuration steps."""
+
+    logging.info("Initializing benchmarks and worker threads...")
 
     benchmark_flags = _SetExtraBenchmarkFlags(
         benchmark_perf_counters, benchmark_repetitions, benchmark_min_time
@@ -402,27 +386,13 @@ def PostProcessBenchmarkResults(self, benchmark_perf_counters: str) -> None:
 
   def Run(
       self,
-      benchmark_target: str,
-      benchmark_filter: list[str] | None = None,
-      workload_filter: list[str] | None = None,
       benchmark_perf_counters: str = "",
       benchmark_repetitions: int = 0,
       benchmark_min_time: str = "",
   ):
     """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,

fleetbench/parallel/parallel_bench_lib_test.py

@@ -35,7 +35,6 @@ def setUp(self):
     self.pb = parallel_bench_lib.ParallelBench(
         cpus=[0, 1],
         cpu_affinity=False,
-        benchmark_weights=None,
         utilization=0.5,
         duration=0.1,
         temp_root=absltest.get_default_test_tmpdir(),
@@ -89,12 +88,17 @@ def fake_utilization(unused_cpus):
     self.pb = parallel_bench_lib.ParallelBench(
         cpus=[0, 1],
         cpu_affinity=False,
-        benchmark_weights=None,
         utilization=0.5,
         duration=0.1,
         temp_root=absltest.get_default_test_tmpdir(),
     )
-    self.pb.Run("fake_bench", [])
+    self.pb.SetWeights(
+        benchmark_target="fake_bench",
+        benchmark_filter=None,
+        workload_filter=None,
+        custom_benchmark_weights=None,
+    )
+    self.pb.Run()
     mock_execute.assert_called_once()
 
   def test_set_extra_benchmark_flags(self):
@@ -241,43 +245,5 @@ def test_post_processing_benchmark_results(
     mock_save_benchmark_results.assert_called_once()
 
 
-class ParseBenchmarkWeightsTest(absltest.TestCase):
-
-  def test_empty_list(self):
-    self.assertIsNone(parallel_bench_lib.ParseBenchmarkWeights([]))
-
-  def test_valid_list(self):
-    benchmark_list = [
-        "cold:0.5",
-        "PROTO:0.3",
-        "Cord:0.2",
-    ]
-    expected_output = {
-        "COLD": 0.5,
-        "PROTO": 0.3,
-        "CORD": 0.2,
-    }
-    self.assertEqual(
-        parallel_bench_lib.ParseBenchmarkWeights(benchmark_list),
-        expected_output,
-        "Should return correct dictionary for valid list",
-    )
-
-  def test_invalid_string(self):
-    benchmark_list = [
-        "cold:0.5",
-        "PROTO:invalid",
-        "TCMALLOC",
-        "Cord:0.2",
-    ]
-    # Even with an invalid string, the function should still process the valid
-    # ones.
-    expected_output = {"COLD": 0.5, "CORD": 0.2}
-    self.assertEqual(
-        parallel_bench_lib.ParseBenchmarkWeights(benchmark_list),
-        expected_output,
-    )
-
-
 if __name__ == "__main__":
   absltest.main()

fleetbench/parallel/weights.py

@@ -0,0 +1,83 @@
+# Copyright 2025 The Fleetbench Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""This file contains functions for handling benchmark weights."""
+
+from absl import logging
+from fleetbench.parallel import benchmark as bm
+
+
+def _ParseBenchmarkWeights(
+    benchmark_weights_list: list[str],
+) -> dict[str, float] | None:
+  """Parses a list of benchmark weights into a dictionary.
+
+  The string element in the list should be in the format:
+  <benchmark_name|benchmark_filter>:<weight>. Note that the benchmark name or
+  filter should be in ALL CAPS to ensure case-insensitive matching.
+
+  Args:
+    benchmark_weights_list: A list of strings to parse.
+
+  Returns:
+    A dictionary of {capitalized string: float} or None if the list is empty.
+  """
+  if not benchmark_weights_list:
+    return None
+
+  benchmark_weights = {}
+  for weights in benchmark_weights_list:
+    try:
+      key, value_str = weights.split(":")
+      benchmark_weights[key.upper()] = float(value_str)
+    except ValueError:
+      logging.warning(
+          "Invalid benchmark string: %s. The format should be"
+          " <benchmark_name|benchmark_filter>:<weight>. Skipping...",
+          weights,
+      )
+
+  return benchmark_weights
+
+
+def GetBenchmarkWeights(
+    benchmarks: dict[str, bm.Benchmark],
+    custom_weight_strings: list[str] | None = None,
+) -> dict[str, float]:
+  """Get the weights for each benchmark.
+
+  The default weight for each benchmark is 1.0. If custom_weight_strings is
+  provided, the weights for the benchmarks in the list will be updated.
+
+  Args:
+    benchmarks: A dictionary of {benchmark name: Benchmark object}.
+    custom_weight_strings: A list of strings to parse for custom weights.
+
+  Returns:
+    A dictionary of {benchmark name: weight}.
+  """
+  benchmark_weights = {
+      benchmark.BenchmarkName(): 1.0 for benchmark in benchmarks.values()
+  }
+
+  # Update the benchmark weights with the custom weights.
+  if custom_weight_strings:
+    logging.info(
+        "Running with selected benchmark weights: %s", custom_weight_strings
+    )
+    custom_weights = _ParseBenchmarkWeights(custom_weight_strings)
+    if custom_weights:
+      benchmark_weights.update(custom_weights)
+  return benchmark_weights