Distributed Worker Architecture for ML Processing (Processing Service V2) pt. 1

.dockerignore

@@ -1,20 +1,72 @@
 .editorconfig
 .gitattributes
 .github
-.gitignore
 .gitlab-ci.yml
 .idea
 .pre-commit-config.yaml
 .readthedocs.yml
 .travis.yml
-.git
 ui
 ami/media
 backups
-venv
-.venv
-.env
-.envs
+venv/
+.venv/
+.env/
+.envs/
 .envs/*
 node_modules
 data
+
+# Python cache / bytecode
+__pycache__/
+*.py[cod]
+*.pyo
+*.pyd
+*.pdb
+*.egg-info/
+*.egg
+*.whl
+
+
+# Django / runtime artifacts
+*.log
+*.pot
+*.pyc
+db.sqlite3
+media/
+staticfiles/     # collected static files (use collectstatic inside container)
+
+# Node / UI dependencies (if using React/Vue in your UI service)
+node_modules/
+npm-debug.log
+yarn-error.log
+.pnpm-debug.log
+
+# Docs build artifacts
+/docs/_build/
+
+# Git / VCS
+.git/
+.gitignore
+.gitattributes
+*.swp
+*.swo
+
+# IDE / editor
+.vscode/
+.idea/
+*.iml
+
+# OS cruft
+.DS_Store
+Thumbs.db
+
+# Docker itself
+.dockerignore
+Dockerfile
+docker-compose*.yml
+
+# Build / dist
+build/
+dist/
+.eggs/

.envs/.local/.django

@@ -12,6 +12,9 @@ DJANGO_SUPERUSER_PASSWORD=localadmin
 # Redis
 REDIS_URL=redis://redis:6379/0
 
+# NATS
+NATS_URL=nats://nats:4222
+
 # Celery / Flower
 CELERY_FLOWER_USER=QSocnxapfMvzLqJXSsXtnEZqRkBtsmKT
 CELERY_FLOWER_PASSWORD=BEQgmCtgyrFieKNoGTsux9YIye0I7P5Q7vEgfJD2C4jxmtHDetFaE2jhS7K7rxaf

.gitignore

@@ -276,3 +276,6 @@ sandbox/
 
 # Other
 flower
+
+# huggingface cache
+huggingface_cache/

.vscode/launch.json

@@ -2,7 +2,29 @@
     "version": "0.2.0",
     "configurations": [
         {
-            "name": "Python Debugger: Remote Attach",
+            "name": "Current File",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${file}",
+            "console": "integratedTerminal"
+        },
+        {
+            "name": "Django attach",
+            "type": "debugpy",
+            "request": "attach",
+            "connect": {
+                "host": "localhost",
+                "port": 5679
+            },
+            "pathMappings": [
+                {
+                    "localRoot": "${workspaceFolder}",
+                    "remoteRoot": "."
+                }
+            ]
+        },
+        {
+            "name": "Celery worker attach",
             "type": "debugpy",
             "request": "attach",
             "connect": {

README.md

@@ -40,6 +40,7 @@ docker compose -f processing_services/example/docker-compose.yml up -d
 - Django admin: http://localhost:8000/admin/
 - OpenAPI / Swagger documentation: http://localhost:8000/api/v2/docs/
 - Minio UI: http://minio:9001, Minio service: http://minio:9000
+- NATS dashboard: https://natsdashboard.com/ (Add localhost)
 
 NOTE: If one of these services is not working properly, it could be due another process is using the port. You can check for this with `lsof -i :<PORT_NUMBER>`.
 

ami/base/views.py

@@ -38,7 +38,9 @@ def get_active_project(
     if not project_id:
         # Look for project_id in GET query parameters or POST data
         # POST data returns a list of ints, but QueryDict.get() returns a single value
-        project_id = request.query_params.get(param) or request.data.get(param)
+        project_id = request.query_params.get(param) or (request.data if isinstance(request.data, dict) else {}).get(
+            param
+        )
 
         project_id = SingleParamSerializer[int].clean(
             param_name=param,

ami/jobs/models.py

@@ -15,10 +15,13 @@
 
 from ami.base.models import BaseModel
 from ami.base.schemas import ConfigurableStage, ConfigurableStageParam
+from ami.jobs.task_state import TaskStateManager
 from ami.jobs.tasks import run_job
+from ami.jobs.utils import _run_in_async_loop
 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.nats_queue import TaskQueueManager
 from ami.utils.schemas import OrderedEnum
 
 logger = logging.getLogger(__name__)
@@ -72,7 +75,7 @@ def get_status_label(status: JobState, progress: float) -> str:
     if status in [JobState.CREATED, JobState.PENDING, JobState.RECEIVED]:
         return "Waiting to start"
     elif status in [JobState.STARTED, JobState.RETRY, JobState.SUCCESS]:
-        return f"{progress:.0%} complete"
+        return f"{progress:.0%} complete"  # noqa E231
     else:
         return f"{status.name}"
 
@@ -133,14 +136,14 @@ def get_stage(self, stage_key: str) -> JobProgressStageDetail:
         for stage in self.stages:
             if stage.key == stage_key:
                 return stage
-        raise ValueError(f"Job stage with key '{stage_key}' not found in progress")
+        raise ValueError(f"Job stage with key '{stage_key}' not found in progress")  # noqa E713
 
     def get_stage_param(self, stage_key: str, param_key: str) -> ConfigurableStageParam:
         stage = self.get_stage(stage_key)
         for param in stage.params:
             if param.key == param_key:
                 return param
-        raise ValueError(f"Job stage parameter with key '{param_key}' not found in stage '{stage_key}'")
+        raise ValueError(f"Job stage parameter with key '{param_key}' not found in stage '{stage_key}'")  # noqa E713
 
     def add_stage_param(self, stage_key: str, param_name: str, value: typing.Any = None) -> ConfigurableStageParam:
         stage = self.get_stage(stage_key)
@@ -327,10 +330,6 @@ def run(cls, job: "Job"):
         job.finished_at = None
         job.save()
 
-        # Keep track of sub-tasks for saving results, pair with batch number
-        save_tasks: list[tuple[int, AsyncResult]] = []
-        save_tasks_completed: list[tuple[int, AsyncResult]] = []
-
         if job.delay:
             update_interval_seconds = 2
             last_update = time.time()
@@ -365,7 +364,7 @@ def run(cls, job: "Job"):
             progress=0,
         )
 
-        images = list(
+        images: list[SourceImage] = list(
             # @TODO return generator plus image count
             # @TODO pass to celery group chain?
             job.pipeline.collect_images(
@@ -389,8 +388,6 @@ def run(cls, job: "Job"):
             images = images[: job.limit]
             image_count = len(images)
             job.progress.add_stage_param("collect", "Limit", image_count)
-        else:
-            image_count = source_image_count
 
         job.progress.update_stage(
             "collect",
@@ -401,6 +398,16 @@ def run(cls, job: "Job"):
         # End image collection stage
         job.save()
 
+        cls.process_images(job, images)
+        # TODO CGJS: do this conditionally based on the type of processing service this job is using
+        # cls.queue_images_to_nats(job, images)
+
+    @classmethod
+    def process_images(cls, job, images):
+        image_count = len(images)
+        # Keep track of sub-tasks for saving results, pair with batch number
+        save_tasks: list[tuple[int, AsyncResult]] = []
+        save_tasks_completed: list[tuple[int, AsyncResult]] = []
         total_captures = 0
         total_detections = 0
         total_classifications = 0
@@ -420,7 +427,7 @@ def run(cls, job: "Job"):
                     job_id=job.pk,
                     project_id=job.project.pk,
                 )
-                job.logger.info(f"Processed image batch {i+1} in {time.time() - request_sent:.2f}s")
+                job.logger.info(f"Processed image batch {i+1} in {time.time() - request_sent:.2f}s")  # noqa E231
             except Exception as e:
                 # Log error about image batch and continue
                 job.logger.error(f"Failed to process image batch {i+1}: {e}")
@@ -472,7 +479,7 @@ def run(cls, job: "Job"):
 
         if image_count:
             percent_successful = 1 - len(request_failed_images) / image_count if image_count else 0
-            job.logger.info(f"Processed {percent_successful:.0%} of images successfully.")
+            job.logger.info(f"Processed {percent_successful:.0%} of images successfully.")  # noqa E231
 
         # Check all Celery sub-tasks if they have completed saving results
         save_tasks_remaining = set(save_tasks) - set(save_tasks_completed)
@@ -512,6 +519,110 @@ def run(cls, job: "Job"):
         job.finished_at = datetime.datetime.now()
         job.save()
 
+    # TODO: This needs to happen once a job is done
+    @classmethod
+    def cleanup_nats_resources(cls, job: "Job"):
+        """
+        Clean up NATS JetStream resources (stream and consumer) for a completed job.
+
+        Args:
+            job: The Job instance
+        """
+        job_id = f"job{job.pk}"
+
+        async def cleanup():
+            async with TaskQueueManager() as manager:
+                success = await manager.cleanup_job_resources(job_id)
+                return success
+
+        _run_in_async_loop(cleanup, f"cleaning up NATS resources for job '{job_id}'")
+
+    @classmethod
+    def queue_images_to_nats(cls, job: "Job", images: list[SourceImage]):
+        """
+        Queue all images for a job to a NATS JetStream stream for the job.
+
+        Args:
+            job: The Job instance
+            images: List of SourceImage instances to queue
+
+        Returns:
+            bool: True if all images were successfully queued, False otherwise
+        """
+        job_id = f"job{job.pk}"
+        job.logger.info(f"Queuing {len(images)} images to NATS stream for job '{job_id}'")
+
+        # Prepare all messages outside of async context to avoid Django ORM issues
+        messages = []
+        image_ids = []
+        for i, image in enumerate(images):
+            image_id = str(image.pk)
+            image_ids.append(image_id)
+            message = {
+                "job_id": job.pk,
+                "image_id": image_id,
+                "image_url": image.url() if hasattr(image, "url") else None,
+                "timestamp": (
+                    image.timestamp.isoformat() if hasattr(image, "timestamp") and image.timestamp else None
+                ),
+                "batch_index": i,
+                "total_images": len(images),
+                "queue_timestamp": datetime.datetime.now().isoformat(),
+            }
+            messages.append((image.pk, message))
+
+        # Store all image IDs in Redis for progress tracking
+        state_manager = TaskStateManager(job.pk)
+        state_manager.initialize_job(image_ids)
+        job.logger.info(f"Initialized task state tracking for {len(image_ids)} images")
+
+        async def queue_all_images():
+            successful_queues = 0
+            failed_queues = 0
+
+            async with TaskQueueManager() as manager:
+                for i, (image_pk, message) in enumerate(messages):
+                    try:
+                        logger.info(f"Queueing image {image_pk} to stream for job '{job_id}': {message}")
+                        # Use TTR of 300 seconds (5 minutes) for image processing
+                        success = await manager.publish_job(
+                            job_id=job_id,
+                            data=message,
+                            ttr=120,  # visibility timeout in seconds
+                        )
+                    except Exception as e:
+                        logger.error(f"Failed to queue image {image_pk} to stream for job '{job_id}': {e}")
+                        success = False
+
+                    if success:
+                        successful_queues += 1
+                    else:
+                        failed_queues += 1
+
+            return successful_queues, failed_queues
+
+        result = _run_in_async_loop(queue_all_images, f"queuing images to NATS for job '{job_id}'")
+        if result is None:
+            job.logger.error(f"Failed to queue images to NATS for job '{job_id}'")
+            return False
+        successful_queues, failed_queues = result
+
+        if not images:
+            job.progress.update_stage("results", status=JobState.SUCCESS, progress=1.0)
+            job.save()
+
+        # Log results (back in sync context)
+        if successful_queues > 0:
+            job.logger.info(
+                f"Successfully queued {successful_queues}/{len(images)} images to stream for job '{job_id}'"
+            )
+
+        if failed_queues > 0:
+            job.logger.warning(f"Failed to queue {failed_queues}/{len(images)} images to stream for job '{job_id}'")
+            return False
+
+        return True
+
 
 class DataStorageSyncJob(JobType):
     name = "Data storage sync"