Add save extractor job to model storage

Author: gabriel-piles

src/trainable_entity_extractor/adapters/LocalModelStorage.py

@@ -2,64 +2,24 @@
 import json
 from typing import Optional
 
-from trainable_entity_extractor.config import EXTRACTOR_JOB_PATH, CACHE_PATH
+from trainable_entity_extractor.config import EXTRACTOR_JOB_PATH
 from trainable_entity_extractor.domain.ExtractionIdentifier import ExtractionIdentifier
-from trainable_entity_extractor.domain.Option import Option
 from trainable_entity_extractor.domain.TrainableEntityExtractorJob import TrainableEntityExtractorJob
 from trainable_entity_extractor.ports.ModelStorage import ModelStorage
 
 
 class LocalModelStorage(ModelStorage):
-    def __init__(self):
-        self.completion_signals = {}
 
     def upload_model(self, extraction_identifier: ExtractionIdentifier, extractor_job: TrainableEntityExtractorJob) -> bool:
-        try:
-            model_path = extraction_identifier.get_path()
-            if not os.path.exists(model_path):
-                os.makedirs(model_path, exist_ok=True)
-
-            extractor_job_dir = os.path.join(model_path, EXTRACTOR_JOB_PATH.parent)
-            if not os.path.exists(extractor_job_dir):
-                os.makedirs(extractor_job_dir, exist_ok=True)
-
-            job_file_path = os.path.join(model_path, EXTRACTOR_JOB_PATH)
-            job_data = self.serialize_job_to_dict(extractor_job)
-
-            with open(job_file_path, "w", encoding="utf-8") as f:
-                json.dump(job_data, f, indent=2, ensure_ascii=False)
-
-            return True
-        except Exception as e:
-            print(f"Error saving job: {e}")
-            return False
+        return self.save_extractor_job(extraction_identifier, extractor_job)
 
     def download_model(self, extraction_identifier: ExtractionIdentifier) -> bool:
-        """Download/load model locally"""
         try:
             model_path = extraction_identifier.get_path()
             return os.path.exists(model_path)
         except Exception:
             return False
 
-    def check_model_completion_signal(self, extraction_identifier: ExtractionIdentifier) -> bool:
-        key = f"{extraction_identifier.run_name}_{extraction_identifier.extraction_name}"
-        return self.completion_signals.get(key, False)
-
-    def create_model_completion_signal(self, extraction_identifier: ExtractionIdentifier) -> bool:
-        try:
-            key = f"{extraction_identifier.run_name}_{extraction_identifier.extraction_name}"
-            self.completion_signals[key] = True
-
-            # Also create a physical completion signal file
-            completion_file = os.path.join(extraction_identifier.get_path(), "training_complete.signal")
-            with open(completion_file, "w") as f:
-                f.write("Training completed successfully")
-
-            return True
-        except Exception:
-            return False
-
     def get_extractor_job(self, extraction_identifier: ExtractionIdentifier) -> Optional[TrainableEntityExtractorJob]:
         try:
             model_path = extraction_identifier.get_path()

src/trainable_entity_extractor/ports/JobExecutor.py

@@ -62,27 +62,7 @@ def upload_model(self, extraction_identifier: ExtractionIdentifier, extractor_jo
         try:
             extraction_identifier.clean_extractor_folder(extractor_job.method_name)
             shutil.rmtree(CACHE_PATH / extraction_identifier.run_name, ignore_errors=True)
-            upload_success = self.model_storage.upload_model(extraction_identifier, extractor_job)
-            if upload_success:
-                signal_success = self.model_storage.create_model_completion_signal(extraction_identifier)
-                if signal_success:
-                    self.logger.log(
-                        extraction_identifier, f"Model and completion signal uploaded for method {extractor_job.method_name}"
-                    )
-                    return True
-                else:
-                    self.logger.log(
-                        extraction_identifier,
-                        f"Model uploaded but completion signal creation failed for method {extractor_job.method_name}",
-                        LogSeverity.error,
-                    )
-                    return False
-            else:
-                self.logger.log(
-                    extraction_identifier, f"Model upload failed for method {extractor_job.method_name}", LogSeverity.error
-                )
-                return False
-
+            return self.model_storage.upload_model(extraction_identifier, extractor_job)
         except Exception as e:
             self.logger.log(extraction_identifier, f"Model upload failed with exception: {e}", LogSeverity.error, e)
             return False

src/trainable_entity_extractor/ports/ModelStorage.py

@@ -1,5 +1,9 @@
+import json
+import os
 from abc import ABC, abstractmethod
 from typing import Optional
+
+from trainable_entity_extractor.config import EXTRACTOR_JOB_PATH
 from trainable_entity_extractor.domain.ExtractionIdentifier import ExtractionIdentifier
 from trainable_entity_extractor.domain.Option import Option
 from trainable_entity_extractor.domain.TrainableEntityExtractorJob import TrainableEntityExtractorJob
@@ -16,16 +20,31 @@ def download_model(self, extraction_identifier: ExtractionIdentifier) -> bool:
         pass
 
     @abstractmethod
-    def check_model_completion_signal(self, extraction_identifier: ExtractionIdentifier) -> bool:
+    def get_extractor_job(self, extraction_identifier: ExtractionIdentifier) -> Optional[TrainableEntityExtractorJob]:
         pass
 
-    @abstractmethod
-    def create_model_completion_signal(self, extraction_identifier: ExtractionIdentifier) -> bool:
-        pass
+    def save_extractor_job(
+        self, extraction_identifier: ExtractionIdentifier, extractor_job: TrainableEntityExtractorJob
+    ) -> bool:
+        try:
+            model_path = extraction_identifier.get_path()
+            if not os.path.exists(model_path):
+                os.makedirs(model_path, exist_ok=True)
 
-    @abstractmethod
-    def get_extractor_job(self, extraction_identifier: ExtractionIdentifier) -> Optional[TrainableEntityExtractorJob]:
-        pass
+            extractor_job_dir = os.path.join(model_path, EXTRACTOR_JOB_PATH.parent)
+            if not os.path.exists(extractor_job_dir):
+                os.makedirs(extractor_job_dir, exist_ok=True)
+
+            job_file_path = os.path.join(model_path, EXTRACTOR_JOB_PATH)
+            job_data = self.serialize_job_to_dict(extractor_job)
+
+            with open(job_file_path, "w", encoding="utf-8") as f:
+                json.dump(job_data, f, indent=2, ensure_ascii=False)
+
+            return True
+        except Exception as e:
+            print(f"Error saving job: {e}")
+            return False
 
     @staticmethod
     def serialize_job_to_dict(job: TrainableEntityExtractorJob) -> dict:
@@ -56,7 +75,6 @@ def deserialize_job_from_dict(job_data: dict) -> TrainableEntityExtractorJob:
         gpu_needed = job_data.get("gpu_needed", False)
         timeout = job_data.get("timeout", 3600)
 
-        additional_fields = {}
         if version != "1.0":
             pass
 

src/trainable_entity_extractor/use_cases/OrchestratorUseCase.py

@@ -113,9 +113,6 @@ def _process_prediction_job(self, distributed_job: DistributedJob) -> JobProcess
             )
 
     def _process_performance_job(self, distributed_job: DistributedJob) -> JobProcessingResult:
-        if len(distributed_job.sub_jobs) == [x for x in distributed_job.sub_jobs if x.status == JobStatus.WAITING]:
-            self.job_executor.recreate_model_folder(distributed_job.extraction_identifier)
-
         self._start_pending_performance_evaluations(distributed_job)
 
         if self._has_perfect_score_job(distributed_job):
@@ -129,8 +126,8 @@ def _process_performance_job(self, distributed_job: DistributedJob) -> JobProces
                 gpu_needed=any(getattr(job.extractor_job, "requires_gpu", False) for job in distributed_job.sub_jobs),
             )
 
-        self._log_performance_summary(distributed_job)
         self._remove_job_from_queue(distributed_job)
+        self._log_performance_summary(distributed_job)
 
         return self._handle_performance_results(distributed_job)
 
@@ -194,18 +191,6 @@ def _handle_performance_results(self, distributed_job: DistributedJob) -> JobPro
             gpu_needed=getattr(best_job.extractor_job, "requires_gpu", False),
         )
 
-    def _finalize_best_model(self, distributed_job: DistributedJob, best_job: DistributedSubJob) -> JobProcessingResult:
-        if self.job_executor.upload_model(distributed_job.extraction_identifier, best_job.extractor_job):
-            performance_score = self._extract_performance_score(best_job)
-            return JobProcessingResult(
-                finished=True,
-                success=True,
-                error_message=f"Best model selected: {best_job.extractor_job.method_name} with performance {performance_score}",
-                gpu_needed=getattr(best_job.extractor_job, "requires_gpu", False),
-            )
-        else:
-            return JobProcessingResult(finished=True, success=False, error_message="Best model selected but upload failed")
-
     @staticmethod
     def _extract_performance_score(best_job: DistributedSubJob) -> str:
         if best_job.result and hasattr(best_job.result, "performance_score"):