remove tmp dirs after submit handling

Author: Nigma-Ks

backend/config/paths.py

@@ -1,7 +1,11 @@
 import os
 
 
-def get_process_filepath(uid, filepath):
+def get_tmp_thread_files(uid):
+    return [get_thread_filepath(uid, UPLOAD_DIR), get_thread_filepath(uid, RESULTS_DIR)]
+
+
+def get_thread_filepath(uid, filepath):
     base_prefix = TMP_FILE_DIR + "/"
     new_prefix = f"{TMP_FILE_DIR}/{uid}_"
     return filepath.replace(base_prefix, new_prefix)

backend/file_utils/files.py

@@ -10,7 +10,7 @@ def save_upload_file(file: UploadFile, dst: str) -> None:
             f.write(chunk)
 
 
-def reset_dir(path: str) -> None:
-    if os.path.exists(path):
-        shutil.rmtree(path)
-    os.makedirs(path)
+def reset_dirs(paths: list) -> None:
+    for path in paths:
+        if os.path.exists(path):
+            shutil.rmtree(path)

backend/main.py

@@ -2,7 +2,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 import shortuuid
 
-from backend.file_utils.files import reset_dir
+from backend.file_utils.files import reset_dirs
 from backend.utils.docker_runner import run_pipeline
 from backend.utils.data_uploader import handle_upload
 from backend.utils.methods_handler import Methods
@@ -11,7 +11,8 @@
     RESULTS_DIR,
     DATASET_FILE,
     METHODS_TS_FILE,
-    get_process_filepath, BASE_DIR,
+    get_thread_filepath,
+    get_tmp_thread_files
 )
 
 dataset_methods = Methods(data_filepath=DATASET_FILE, output_filepath=METHODS_TS_FILE)
@@ -41,12 +42,14 @@ async def handle_submit(
 
     send_folder_by_email(
         email,
-        get_process_filepath(task_uid, RESULTS_DIR),
+        get_thread_filepath(task_uid, RESULTS_DIR),
         experiment,
         file.filename,
     )
 
-    reset_dir(BASE_DIR)
+    tmp_thread_dirs = get_tmp_thread_files(task_uid)
+    reset_dirs(tmp_thread_dirs)
+
     return {
         "experiment": experiment,
         "filename": file.filename,

backend/utils/data_uploader.py

@@ -3,7 +3,7 @@
 from backend.file_utils.csv_methods_writer import write_launch_info_csv
 from backend.file_utils.files import save_upload_file
 from backend.utils.methods_handler import Methods
-from backend.config.paths import LAUNCH_INFO_FILE, MODEL_ONNX_FILE, get_process_filepath
+from backend.config.paths import LAUNCH_INFO_FILE, MODEL_ONNX_FILE, get_thread_filepath
 
 
 def handle_upload(
@@ -14,10 +14,10 @@ def handle_upload(
 ) -> None:
     selection_methods = json.loads(methods)
 
-    save_upload_file(file, get_process_filepath(uid, MODEL_ONNX_FILE))
+    save_upload_file(file, get_thread_filepath(uid, MODEL_ONNX_FILE))
 
     write_launch_info_csv(
         methods=selection_methods,
         dataset_methods=dataset_methods,
-        output_file=get_process_filepath(uid, LAUNCH_INFO_FILE),
+        output_file=get_thread_filepath(uid, LAUNCH_INFO_FILE),
     )

backend/utils/docker_runner.py

@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 
 from backend.config.paths import (
-    get_process_filepath,
+    get_thread_filepath,
     RESULTS_DIR,
     LAUNCH_INFO_FILE,
     MODEL_ONNX_FILE,
@@ -40,9 +40,9 @@ class RunstratBaseline(DockerRunner):
     def _volumes(self, uid) -> list[str]:
         return [
             "-v",
-            f"{get_process_filepath(uid, RESULTS_DIR)}:/workspace/PySymGym/tools/runstrat/results",
+            f"{get_thread_filepath(uid, RESULTS_DIR)}:/workspace/PySymGym/tools/runstrat/results",
             "-v",
-            f"{get_process_filepath(uid, LAUNCH_INFO_FILE)}"
+            f"{get_thread_filepath(uid, LAUNCH_INFO_FILE)}"
             f":/workspace/PySymGym/tools/runstrat/resources/launch_info.csv",
         ]
 
@@ -67,12 +67,12 @@ class RunstratAI(DockerRunner):
     def _volumes(self, uid) -> list[str]:
         return [
             "-v",
-            f"{get_process_filepath(uid, RESULTS_DIR)}:/workspace/PySymGym/tools/runstrat/results",
+            f"{get_thread_filepath(uid, RESULTS_DIR)}:/workspace/PySymGym/tools/runstrat/results",
             "-v",
-            f"{get_process_filepath(uid, LAUNCH_INFO_FILE)}"
+            f"{get_thread_filepath(uid, LAUNCH_INFO_FILE)}"
             f":/workspace/PySymGym/tools/runstrat/resources/launch_info.csv",
             "-v",
-            f"{get_process_filepath(uid, MODEL_ONNX_FILE)}"
+            f"{get_thread_filepath(uid, MODEL_ONNX_FILE)}"
             f":/workspace/PySymGym/tools/runstrat/resources/model.onnx",
         ]
 
@@ -99,13 +99,13 @@ class Compstrat(DockerRunner):
     def _volumes(self, uid) -> list[str]:
         return [
             "-v",
-            f"{get_process_filepath(uid, ARTIFACTS_AI_CSV_FILE)}"
+            f"{get_thread_filepath(uid, ARTIFACTS_AI_CSV_FILE)}"
             ":/workspace/PySymGym/tools/compstrat/strat/ai_strat",
             "-v",
-            f"{get_process_filepath(uid, ARTIFACTS_BASELINE_CSV_FILE)}"
+            f"{get_thread_filepath(uid, ARTIFACTS_BASELINE_CSV_FILE)}"
             ":/workspace/PySymGym/tools/compstrat/strat/baseline_strat",
             "-v",
-            f"{get_process_filepath(uid, COMPSTRAT_RESULTS_DIR)}"
+            f"{get_thread_filepath(uid, COMPSTRAT_RESULTS_DIR)}"
             ":/workspace/PySymGym/tools/compstrat/results",
         ]