refactor: use static registry for post-processing task lookup

Author: mohamedelabbas1996

ami/jobs/models.py

@@ -18,6 +18,7 @@
 from ami.jobs.tasks import run_job
 from ami.main.models import Deployment, Project, SourceImage, SourceImageCollection
 from ami.ml.models import Pipeline
+from ami.ml.post_processing.registry import get_postprocessing_task
 from ami.utils.schemas import OrderedEnum
 
 logger = logging.getLogger(__name__)
@@ -651,9 +652,6 @@ class PostProcessingJob(JobType):
 
     @classmethod
     def run(cls, job: "Job"):
-        import ami.ml.post_processing  # noqa F401
-        from ami.ml.post_processing.base import get_postprocessing_task
-
         job.progress.add_stage(cls.name, key=cls.key)
         job.update_status(JobState.STARTED)
         job.started_at = datetime.datetime.now()
@@ -664,7 +662,7 @@ def run(cls, job: "Job"):
         config = params.get("config", {})
         job.logger.info(f"Post-processing task: {task_key} with params: {job.params}")
 
-        task_cls = get_postprocessing_task(task_key)
+        task_cls = get_postprocessing_task(key=task_key)
         if not task_cls:
             raise ValueError(f"Unknown post-processing task '{task_key}'")
 

ami/ml/post_processing/base.py

@@ -1,35 +1,13 @@
 import abc
 import logging
-from typing import Any
+import typing
+from typing import Any, Optional
 
-from ami.jobs.models import Job
 from ami.ml.models import Algorithm
 from ami.ml.models.algorithm import AlgorithmTaskType
 
-# Registry of available post-processing tasks
-POSTPROCESSING_TASKS: dict[str, type["BasePostProcessingTask"]] = {}
-
-
-def register_postprocessing_task(task_cls: type["BasePostProcessingTask"]):
-    """
-    Decorator to register a post-processing task in the global registry.
-    Each task must define a unique `key`.
-    Ensures an Algorithm entry exists for this task.
-    """
-    if not hasattr(task_cls, "key") or not task_cls.key:
-        raise ValueError(f"Task {task_cls.__name__} missing required 'key' attribute")
-
-    # Register the task
-    POSTPROCESSING_TASKS[task_cls.key] = task_cls
-    return task_cls
-
-
-def get_postprocessing_task(name: str) -> type["BasePostProcessingTask"] | None:
-    """
-    Get a task class by its registry key.
-    Returns None if not found.
-    """
-    return POSTPROCESSING_TASKS.get(name)
+if typing.TYPE_CHECKING:
+    from ami.jobs.models import Job
 
 
 class BasePostProcessingTask(abc.ABC):
@@ -43,7 +21,7 @@ class BasePostProcessingTask(abc.ABC):
 
     def __init__(
         self,
-        job: Job | None = None,
+        job: Optional["Job"] = None,
         logger: logging.Logger | None = None,
         **config: Any,
     ):

ami/ml/post_processing/registry.py

@@ -0,0 +1,11 @@
+# Registry of available post-processing tasks
+from ami.ml.post_processing.small_size_filter import SmallSizeFilterTask
+
+POSTPROCESSING_TASKS = {
+    SmallSizeFilterTask.key: SmallSizeFilterTask,
+}
+
+
+def get_postprocessing_task(key: str):
+    """Return a post-processing task class by key."""
+    return POSTPROCESSING_TASKS.get(key)

ami/ml/post_processing/small_size_filter.py

@@ -2,10 +2,9 @@
 from django.utils import timezone
 
 from ami.main.models import Detection, SourceImageCollection, Taxon, TaxonRank
-from ami.ml.post_processing.base import BasePostProcessingTask, register_postprocessing_task
+from ami.ml.post_processing.base import BasePostProcessingTask
 
 
-@register_postprocessing_task
 class SmallSizeFilterTask(BasePostProcessingTask):
     key = "small_size_filter"
     name = "Small size filter"