refactor Methods class, move build and fetch functions to app_setup

Author: Nigma-Ks

backend/file_utils/csv_methods_writer.py

@@ -1,21 +1,16 @@
 import csv
 
-from backend.utils.methods_handler import Methods
 
-
-def write_launch_info_csv(
+def write_launch_info_to_csv(
     *,
-    methods: list[str],
-    dataset_methods: Methods,
+    parsed_methods: list[str],
     output_file: str,
 ) -> None:
-    parsed = dataset_methods.get_methods_from_selection(methods)
-
     with open(output_file, "w") as csvfile:
         writer = csv.writer(csvfile)
         writer.writerow(["dll", "method"])
 
-        for item in parsed:
+        for item in parsed_methods:
             if "," in item:
                 dll, method = item.split(",", 1)
                 writer.writerow([dll, method])

backend/file_utils/files.py

@@ -3,9 +3,9 @@
 from fastapi import UploadFile
 
 
-def save_upload_file(file: UploadFile, dst: str) -> None:
-    os.makedirs(os.path.dirname(dst), exist_ok=True)
-    with open(dst, "wb") as f:
+def save_upload_file(file: UploadFile, dest_filepath: str) -> None:
+    os.makedirs(os.path.dirname(dest_filepath), exist_ok=True)
+    with open(dest_filepath, "wb") as f:
         while chunk := file.file.read(1024 * 1024):
             f.write(chunk)
 

backend/launch_service/__init__.py

@@ -1,7 +1,13 @@
 import os
 import subprocess
 import requests
-from backend.config.paths import RESOURCES_DIR, DATASET_FILE, DOCKER_DIR
+from backend.config.paths import (
+    RESOURCES_DIR,
+    DATASET_FILE,
+    DOCKER_DIR,
+    METHODS_TS_FILE,
+)
+from backend.utils.methods_handler import Methods
 
 IMAGE_NAME = "pysymgym-test"
 
@@ -21,12 +27,21 @@ def fetch_dataset():
 
 def build_container():
     print(f"Building Docker image '{IMAGE_NAME}'...")
-    subprocess.run(["docker", "build", "--no-cache", "-t", IMAGE_NAME, DOCKER_DIR], check=True)
+    subprocess.run(
+        ["docker", "build", "--no-cache", "-t", IMAGE_NAME, DOCKER_DIR], check=True
+    )
     print("Docker build completed.\n")
 
     os.makedirs(RESOURCES_DIR, exist_ok=True)
 
 
+def update_frontend_selection_options():
+    print("Updating frontend selection options...")
+    selection_tree = Methods.parse_dataset_file_for_front_selection(DATASET_FILE)
+    Methods.write_selection_dataset_to_front_file(selection_tree, METHODS_TS_FILE)
+
+
 if __name__ == "__main__":
     build_container()
     fetch_dataset()
+    update_frontend_selection_options()

docker/build_container.py backend/launch_service/app_setup.pydocker/build_container.py renamed to backend/launch_service/app_setup.py

@@ -2,15 +2,14 @@
 from fastapi import FastAPI, UploadFile, File, Form
 from fastapi.middleware.cors import CORSMiddleware
 
-from backend.config.paths import (
-    DATASET_FILE,
-    METHODS_TS_FILE,
-)
+from backend.config.paths import METHODS_TS_FILE
 from backend.utils.data_uploader import handle_upload
 from backend.utils.methods_handler import Methods
 from backend.utils.task import process_and_cleanup_task
 
-dataset_methods = Methods(data_filepath=DATASET_FILE, output_filepath=METHODS_TS_FILE)
+DATASET_DLLS_AND_METHODS = (
+    Methods.get_dlls_with_all_methods_dict_from_front_options_file(METHODS_TS_FILE)
+)
 
 app = FastAPI()
 
@@ -24,15 +23,14 @@
 
 @app.post("/api/upload")
 async def handle_submit(
-        file: UploadFile = File(...),
-        email: str = Form(...),
-        methods: str = Form(...),
-        experiment: str = Form(...),
+    file: UploadFile = File(...),
+    email: str = Form(...),
+    methods: str = Form(...),
+    experiment: str = Form(...),
 ):
     task_uid = str(shortuuid.uuid())
 
-    handle_upload(task_uid, file, methods, dataset_methods)
-
+    handle_upload(task_uid, file, methods, DATASET_DLLS_AND_METHODS)
     process_and_cleanup_task.delay(task_uid, email, experiment, file.filename)
 
     return {

backend/main.py

@@ -1,23 +1,27 @@
 import json
+from collections import defaultdict
 
-from backend.file_utils.csv_methods_writer import write_launch_info_csv
+from fastapi import UploadFile
+
+from backend.file_utils.csv_methods_writer import write_launch_info_to_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_thread_filepath
 
 
 def handle_upload(
-    uid,
-    file,
+    uid: str,
+    file: UploadFile,
     methods: str,
-    dataset_methods: Methods,
+    dataset_dll_and_methods: defaultdict,
 ) -> None:
-    selection_methods = json.loads(methods)
+    launch_methods = Methods.get_launch_info_list_from_selected(
+        dataset_dll_and_methods, json.loads(methods)
+    )
 
     save_upload_file(file, get_thread_filepath(uid, MODEL_ONNX_FILE))
 
-    write_launch_info_csv(
-        methods=selection_methods,
-        dataset_methods=dataset_methods,
+    write_launch_info_to_csv(
+        parsed_methods=launch_methods,
         output_file=get_thread_filepath(uid, LAUNCH_INFO_FILE),
     )

backend/utils/data_uploader.py

@@ -10,7 +10,7 @@
     ARTIFACTS_BASELINE_CSV_FILE,
     COMPSTRAT_RESULTS_DIR,
 )
-from docker.build_container import IMAGE_NAME
+from backend.launch_service.app_setup import IMAGE_NAME
 
 
 class DockerRunner(ABC):

backend/utils/docker_runner.py

@@ -3,17 +3,42 @@
 
 
 class Methods:
-    def __init__(self, data_filepath, output_filepath):
+    @staticmethod
+    def write_selection_dataset_to_front_file(
+        selection_dataset, front_selection_resource_path
+    ):
+        with open(front_selection_resource_path, "w") as f:
+            f.write("export const METHODS = ")
+            json.dump(selection_dataset, f, indent=4)
+
+    @staticmethod
+    def get_dlls_with_all_methods_dict_from_front_options_file(
+        front_selection_resource_path,
+    ):
+        with open(front_selection_resource_path, "r") as f:
+            content = f.read()
+        prefix = "export const METHODS = "
+        json_str = content[len(prefix) :]
+        selection_tree = json.loads(json_str)
+
+        dll_methods = defaultdict(list)
+        for node in selection_tree:
+            dll_name = node["title"] + ".dll"
+            dll_methods[dll_name] = []
+            for child in node["children"]:
+                dll_methods[dll_name].append(child["value"])
+        return dll_methods
+
+    @staticmethod
+    def parse_dataset_file_for_front_selection(data_filepath):
         with open(data_filepath, "r") as f:
             data = json.load(f)
 
-        self.dll_methods = defaultdict(list)
         groups = defaultdict(list)
         for item in data:
             groups[item["AssemblyFullName"]].append(item)
-            self.dll_methods[item["AssemblyFullName"]] = []
 
-        tree = []
+        dataset_tree = []
         for assembly, items in groups.items():
             node = {"title": assembly[:-4], "value": assembly, "children": []}
             seen_children = set()
@@ -22,21 +47,18 @@ def __init__(self, data_filepath, output_filepath):
                 if name in seen_children:
                     continue
                 seen_children.add(name)
-                dll_method = f"{assembly},{name}"
-                self.dll_methods[assembly].append(dll_method)
-                child_node = {"title": name, "value": dll_method}
+                dll_and_method = f"{assembly},{name}"
+                child_node = {"title": name, "value": dll_and_method}
                 node["children"].append(child_node)
-            tree.append(node)
-
-        with open(output_filepath, "w") as f:
-            f.write("export const METHODS = ")
-            json.dump(tree, f, indent=4)
+            dataset_tree.append(node)
+        return dataset_tree
 
-    def get_methods_from_selection(self, selection_list):
-        selected_methods = []
-        for item in selection_list:
-            if item in self.dll_methods.keys():
-                selected_methods.extend(self.dll_methods[item])
+    @staticmethod
+    def get_launch_info_list_from_selected(dll_methods: defaultdict, selected):
+        launch_info_methods = []
+        for item in selected:
+            if item in dll_methods.keys():
+                launch_info_methods.extend(dll_methods[item])
             else:
-                selected_methods.append(item)
-        return selected_methods
+                launch_info_methods.append(item)
+        return launch_info_methods

backend/utils/methods_handler.py

@@ -1,4 +1,5 @@
 from celery import Celery
+
 from backend.utils.docker_runner import run_pipeline
 from backend.utils.results_sender import send_folder_by_email
 from backend.config.paths import RESULTS_DIR, get_thread_filepath, get_tmp_thread_files
@@ -10,7 +11,12 @@
 
 
 @celery_app.task
-def process_and_cleanup_task(task_uid: str, email: str, experiment: str, filename: str):
+def process_and_cleanup_task(
+    task_uid: str,
+    email: str,
+    experiment: str,
+    filename: str,
+):
     run_pipeline(task_uid)
 
     send_folder_by_email(

backend/utils/task.py

@@ -37,7 +37,7 @@ const ComparisonForm: React.FC = () => {
             });
             if (!res.ok) throw new Error('Upload failed');
             const data = await res.json();
-            message.success('Data successfully submitted');
+            message.success(data.message);
             console.log(data);
         } catch (err) {
             console.error(err);