Fix per-cgroup bench script

Author: tzussman

bench/bench_lib.py

@@ -22,8 +22,13 @@
 
 
 class CacheExtPolicy:
+
+    def set_cgroup(self, cgroup: str):
+        """Set the cgroup path for the policy."""
+        self.cgroup_path = f"/sys/fs/cgroup/{cgroup}"
+
     def __init__(self, cgroup: str, loader_path: str, watch_dir: str):
-        self.cgroup_path = "/sys/fs/cgroup/%s" % cgroup
+        self.set_cgroup(cgroup)
         self.loader_path = loader_path
         self.watch_dir = watch_dir
         self.has_started = False

bench/bench_per_cgroup.py

@@ -307,10 +307,11 @@ def before_benchmark(self, config):
                     cgroup=f"{DEFAULT_CACHE_EXT_CGROUP}_2",
                     limit_in_bytes=config["cgroup_config"].policy2_size,
                 )
-                self.cache_ext_policy.start(cgroup=f"{DEFAULT_CACHE_EXT_CGROUP}_1")
-                self.second_cache_ext_policy.start(
-                    cgroup=f"{DEFAULT_CACHE_EXT_CGROUP}_2"
-                )
+                self.cache_ext_policy.set_cgroup(f"{DEFAULT_CACHE_EXT_CGROUP}_1")
+                self.second_cache_ext_policy.set_cgroup(f"{DEFAULT_CACHE_EXT_CGROUP}_2")
+
+                self.cache_ext_policy.start()
+                self.second_cache_ext_policy.start()
             else:
                 size = (
                     config["cgroup_config"].policy1_size
@@ -320,14 +321,14 @@ def before_benchmark(self, config):
                     recreate_cache_ext_cgroup(
                         cgroup=f"{DEFAULT_CACHE_EXT_CGROUP}_1", limit_in_bytes=size
                     )
-                    self.cache_ext_policy.start(cgroup=f"{DEFAULT_CACHE_EXT_CGROUP}_1")
+                    self.cache_ext_policy.set_cgroup(f"{DEFAULT_CACHE_EXT_CGROUP}_1")
+                    self.cache_ext_policy.start()
                 else:
                     recreate_cache_ext_cgroup(
                         cgroup=f"{DEFAULT_CACHE_EXT_CGROUP}_2", limit_in_bytes=size
                     )
-                    self.second_cache_ext_policy.start(
-                        cgroup=f"{DEFAULT_CACHE_EXT_CGROUP}_2"
-                    )
+                    self.second_cache_ext_policy.set_cgroup(f"{DEFAULT_CACHE_EXT_CGROUP}_2")
+                    self.second_cache_ext_policy.start()
         else:
             if config["cgroup_config"].split_cgroups:
                 recreate_baseline_cgroup(