refactor: refactor job runner to initialize post-processing tasks with job context

Author: mohamedelabbas1996

ami/jobs/models.py

@@ -662,13 +662,16 @@ def run(cls, job: "Job"):
         params = job.params or {}
         task_key: str = params.get("task", "")
         config = params.get("config", {})
-        job.logger.info(f"Post-processing task: {task_key} with params: {job.params}")
-
+        job.logger.info(f"Post-processing task: {task_key} with params: {config}")
+        # Get the registered task class
         task_cls = get_postprocessing_task(task_key)
         if not task_cls:
             raise ValueError(f"Unknown post-processing task '{task_key}'")
 
+        # Instantiate the task with job context and config
         task = task_cls(job=job, **config)
+
+        # Run the task
         task.run()
         job.progress.update_stage(cls.key, status=JobState.SUCCESS, progress=1)
         job.finished_at = datetime.datetime.now()

ami/ml/post_processing/base.py

@@ -2,7 +2,7 @@
 
 import abc
 import logging
-from typing import Any
+from typing import Any, Optional
 
 from ami.jobs.models import Job
 from ami.ml.models import Algorithm
@@ -39,54 +39,34 @@ class BasePostProcessingTask(abc.ABC):
     Abstract base class for all post-processing tasks.
     """
 
-    # Each task must override these
     key: str = ""
     name: str = ""
 
     def __init__(
         self,
-        job: Job | None = None,
-        logger: logging.Logger | None = None,
+        job: Optional["Job"] = None,
+        task_logger: logging.Logger | None = None,
         **config: Any,
     ):
-        self.job = job
-        self.config = config
-        # Choose the right logger
-        if logger is not None:
-            self.logger = logger
-        elif job is not None:
-            self.logger = job.logger
-        else:
-            self.logger = logging.getLogger(f"ami.post_processing.{self.key}")
-
-        algorithm, _ = Algorithm.objects.get_or_create(
-            name=self.__class__.__name__,
-            defaults={
-                "description": f"Post-processing task: {self.key}",
-                "task_type": AlgorithmTaskType.POST_PROCESSING.value,
-            },
-        )
-        self.algorithm: Algorithm = algorithm
-
-        self.logger.info(f"Initialized {self.__class__.__name__} with config={self.config}, job={job}")
-
-    def update_progress(self, progress: float):
         """
-        Update progress if job is present, otherwise just log.
+        Initialize task with optional job and logger context.
         """
+        self.job = job
+        self.config: dict[str, Any] = config
 
-        if self.job:
-            self.job.progress.update_stage(self.job.job_type_key, progress=progress)
-            self.job.save(update_fields=["progress"])
-
+        if job:
+            self.logger = job.logger
+        elif task_logger:
+            self.logger = task_logger
         else:
-            # No job object — fallback to plain logging
-            self.logger.info(f"[{self.name}] Progress {progress:.0%}")
+            self.logger = logging.getLogger(f"ami.post_processing.{self.key}")
+        self.log_config()
 
     @abc.abstractmethod
     def run(self) -> None:
-        """
-        Run the task logic.
-        Must be implemented by subclasses.
-        """
-        raise NotImplementedError("BasePostProcessingTask subclasses must implement run()")
+        """Run the task logic. Must be implemented by subclasses."""
+        raise NotImplementedError("Subclasses must implement run()")
+
+    def log_config(self):
+        """Helper to log the task configuration at start."""
+        self.logger.info(f"Running task {self.name} ({self.key}) with config: {self.config}")

ami/ml/post_processing/small_size_filter.py

@@ -31,7 +31,7 @@ def run(self) -> None:
 
         try:
             collection = SourceImageCollection.objects.get(pk=collection_id)
-            self.logger.info(f"Loaded SourceImageCollection {collection_id} (Project={collection.project})")
+            self.logger.info(f"Loaded SourceImageCollection {collection_id} " f"(Project={collection.project})")
         except SourceImageCollection.DoesNotExist:
             msg = f"SourceImageCollection {collection_id} not found"
             self.logger.error(msg)
@@ -85,11 +85,6 @@ def run(self) -> None:
                             comment=f"Auto-set by {self.name} post-processing task",
                         )
                 modified += 1
-                self.logger.info(f"Detection {det.pk}: marked as 'Not identifiable'")
-
-            # Update progress every 10 detections
-            if i % 10 == 0 or i == total:
-                progress = i / total if total > 0 else 1.0
-                self.update_progress(progress)
+                self.logger.debug(f"Detection {det.pk}: marked as 'Not identifiable'")
 
         self.logger.info(f"=== Completed {self.name}: {modified}/{total} detections modified ===")