Add random strategy for group scheduling

Signed-off-by: Colin Wilk <[email protected]>

src/benchmarking_tool/coordinator.py

@@ -22,6 +22,7 @@
 """
 
 import datetime
+import random as r
 import time
 import uuid
 from itertools import product
@@ -69,7 +70,7 @@ class Coordinator:
             Retrieves the filename used in benchmark tasks.
     """
 
-    def __schedule_every_2_hours(self) -> None:
+    def __schedule_every_2_hours(self, random: bool = False) -> None:
         for t in [
             "00",
             "02",
@@ -84,17 +85,21 @@ def __schedule_every_2_hours(self) -> None:
             "20",
             "22",
         ]:
-            schedule.every().day.at(f"{t}:00").do(self.trigger_benchmark)
+            schedule.every().day.at(f"{t}:00").do(
+                lambda: self.trigger_benchmark(random)
+            )
 
     def __init__(self) -> None:
         pass
 
-    def trigger_benchmark(self) -> None:
+    def trigger_benchmark(self, random: bool = False) -> None:
         trigger_time: datetime.datetime = datetime.datetime.now()
         wave_id: str = str(uuid.uuid4())
         l.info("Triggering Benchmark")
         for worker_group, filename in product(self.worker_groups, self.filenames):
-            workers: set[Any] = worker.get_workers(worker_group)
+            workers: list[Any] = list(worker.get_workers(worker_group))
+            if random:  # Pick one element if we use random strategy
+                workers = [r.choice(workers)]
             for worker_instance in workers:
                 queue_name = worker_instance.decode()
                 task = (
@@ -110,8 +115,8 @@ def trigger_benchmark(self) -> None:
                 )
                 task.get()  # Wait until the task is done
 
-    def run(self) -> NoReturn:
-        self.__schedule_every_2_hours()
+    def run(self, random: bool = False) -> NoReturn:
+        self.__schedule_every_2_hours(random)
         l.info("Coordinator Started")
         while True:
             schedule.run_pending()

src/benchmarking_tool/main.py

@@ -153,6 +153,11 @@
     help="Trigger a Job immediately and then exit. Useful for testing",
     action="store_true",
 )
+parser_coordinator.add_argument(
+    "--random",
+    help="Chooses a node from the group at random instead of running on all nodes",
+    action="store_true",
+)
 
 if __name__ == "__main__":
     args = parser.parse_args()
@@ -179,7 +184,7 @@
             Coordinator().set_worker_groups(args.groups).set_filenames(args.filenames)
         )
         if args.trigger:
-            coordinator.trigger_benchmark()
+            coordinator.trigger_benchmark(args.random)
 
         else:
-            coordinator.run()
+            coordinator.run(args.random)