Add option to delete files after benchmark

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

src/benchmarking_tool/celery_app.py

@@ -21,6 +21,7 @@
             registration, initialization, and task execution.
 """
 
+import glob
 import os
 import socket
 from datetime import datetime
@@ -81,6 +82,7 @@ def __init__(self) -> None:
         self.worker_group: str | None = None
         self.worker_id: str | None = None
         self.worker_directory: str | None = None
+        self.delete: bool | None = None
 
         self.app: Celery = Celery(
             "fio",
@@ -94,11 +96,16 @@ def __init__(self) -> None:
         )
 
     def register_worker(
-        self, worker_group: str, worker_id: str, worker_directory: str
+        self,
+        worker_group: str,
+        worker_id: str,
+        worker_directory: str,
+        delete: bool = True,
     ) -> Self:
         self.worker_id = worker_id
         self.worker_group = worker_group
         self.worker_directory = worker_directory
+        self.delete = delete
 
         self.r.sadd(f"workers_{self.worker_group}", self.worker_id)
 
@@ -114,6 +121,11 @@ def start_worker(self) -> None:
         )
         self.__del__()  # pylint: disable=unnecessary-dunder-call
 
+    def delete_fio_files(self) -> None:
+        r = glob.glob(f"{self.worker_directory}*")
+        for i in r:
+            os.remove(i)
+
     def __del__(self) -> None:
         if self.worker_group is not None and self.worker_id is not None:
             print(
@@ -136,6 +148,8 @@ def run_benchmark(filename: str, wave_id: str, timestamp: datetime) -> dict[Any,
     fio = Fio()
     config: FioConfig = FioConfig(filename)
     res: FioResult = fio.run(config, worker.worker_directory or "/tmp")
+    if worker.delete:
+        worker.delete_fio_files()
     return {
         "status": "Passed",
         "worker_id": worker.worker_id,

src/benchmarking_tool/main.py

@@ -123,6 +123,9 @@
     type=str,
     required=True,
 )
+parser_worker.add_argument(
+    "--delete", help="Delete files after benchmark", action="store_true"
+)
 
 ################################################################################
 # Coordinator Parameters
@@ -189,7 +192,9 @@
     if args.role == "worker":
         if args.hostname is None:
             args.hostname = hostname
-        worker.register_worker(args.group, args.hostname, args.directory).start_worker()
+        worker.register_worker(
+            args.group, args.hostname, args.directory, args.delete
+        ).start_worker()
 
     if args.role == "coordinator":
         coordinator = (