Add option to wait until jobs finish before exiting

Signed-off-by: Colin Wilk <colin.wilk@tum.de>

src/benchmarking_tool/coordinator.py

@@ -89,10 +89,11 @@ def __schedule_every_2_hours(self) -> None:
     def __init__(self) -> None:
         pass
 
-    def trigger_benchmark(self) -> None:
+    def trigger_benchmark(self) -> list:
         trigger_time: datetime.datetime = datetime.datetime.now()
         wave_id: str = str(uuid.uuid4())
         l.info("Triggering Benchmark")
+        all_results: list = []
         for filename in self.filenames:
             groups: list = []
             for worker_group in self.worker_groups:
@@ -106,7 +107,9 @@ def trigger_benchmark(self) -> None:
                 groups.append(grouped_tasks)
 
             chained_tasks = chain(*groups)
-            chained_tasks.apply_async()
+            result = chained_tasks.apply_async()
+            all_results.append(result)
+        return all_results
 
     def run(self) -> NoReturn:
         self.__schedule_every_2_hours()

src/benchmarking_tool/main.py

@@ -153,6 +153,12 @@
     help="Trigger a Job immediately and then exit. Useful for testing",
     action="store_true",
 )
+parser_coordinator.add_argument(
+    "-w",
+    "--wait",
+    help="Wait for all Jobs to finish before exiting.",
+    action="store_true",
+)
 
 if __name__ == "__main__":
     args = parser.parse_args()
@@ -179,6 +185,9 @@
             Coordinator().set_worker_groups(args.groups).set_filenames(args.filenames)
         )
         if args.trigger:
-            coordinator.trigger_benchmark()
+            all_results: list = coordinator.trigger_benchmark()
+            if args.wait:
+                for result in all_results:
+                    result.get()
         else:
             coordinator.run()