Use transaction.on_commit with all async celery tasks

Author: vanessavmac

ami/jobs/models.py

@@ -579,10 +579,8 @@ def update_job_progress(cls, job: "Job"):
                 failed=num_failed_save_tasks,
             )
 
-            # The ML job is completed, log general job stags
-            if job.status != JobState.FAILURE:
-                # the job might've already been marked as failed because of unsent process pipeline request tasks
-                job.update_status(JobState.FAILURE if any_failed_tasks else JobState.SUCCESS, save=False)
+            # The ML job is completed, log general job stats
+            job.update_status(JobState.FAILURE if any_failed_tasks else JobState.SUCCESS, save=False)
 
             if any_failed_tasks:
                 failed_save_task_ids = [
@@ -708,23 +706,11 @@ def run(cls, job: "Job"):
 
         except Exception as e:
             job.logger.error(f"Failed to submit all images: {e}")
-            # mark the job as failed
-            job.progress.update_stage(
-                "process",
-                status=JobState.FAILURE,
-                progress=1,
-                failed=image_count,
-                processed=0,
-                remaining=image_count,
-            )
             job.update_status(JobState.FAILURE)
             job.save()
-        finally:
-            # Handle the successfully submitted tasks
-            subtasks = job.ml_task_records.all()
-            if subtasks:
-                check_ml_job_status.apply_async([job.pk])
-            else:
+        else:
+            subtasks = job.ml_task_records.filter(created_at__gte=job.started_at)
+            if not subtasks:
                 # No tasks were scheduled, mark the job as done
                 job.logger.info("No subtasks were scheduled, ending the job.")
                 job.progress.update_stage(
@@ -740,6 +726,16 @@ def run(cls, job: "Job"):
                 job.update_status(JobState.SUCCESS, save=False)
                 job.finished_at = datetime.datetime.now()
                 job.save()
+            else:
+                job.logger.info(
+                    f"Continue processing the remaining {subtasks.count()} process image request subtasks."
+                )
+                from django.db import transaction
+
+                transaction.on_commit(lambda: check_ml_job_status.apply_async([job.pk]))
+        finally:
+            # TODO: clean up?
+            pass
 
 
 class DataStorageSyncJob(JobType):

ami/jobs/views.py

@@ -168,7 +168,9 @@ def check_inprogress_subtasks(self, request, pk=None):
         has_inprogress_tasks = job.check_inprogress_subtasks()
         if has_inprogress_tasks:
             # Schedule task to update the job status
+            from django.db import transaction
+
             from ami.ml.tasks import check_ml_job_status
 
-            check_ml_job_status.apply_async((job.pk,))
+            transaction.on_commit(lambda: check_ml_job_status.apply_async((job.pk,)))
         return Response({"inprogress_subtasks": has_inprogress_tasks})

ami/ml/models/pipeline.py

@@ -16,7 +16,7 @@
 
 import requests
 from celery.result import AsyncResult
-from django.db import models
+from django.db import models, transaction
 from django.utils.text import slugify
 from django.utils.timezone import now
 from django_pydantic_field import SchemaField
@@ -185,8 +185,8 @@ def process_pipeline_request(pipeline_request: dict, project_id: int):
 def process_images(
     pipeline: Pipeline,
     images: typing.Iterable[SourceImage],
+    project_id: int,
     job_id: int | None = None,
-    project_id: int | None = None,
     process_sync: bool = False,
 ) -> PipelineResultsResponse | None:
     """
@@ -207,13 +207,12 @@ def process_images(
         job = Job.objects.get(pk=job_id)
         task_logger = job.logger
 
-    if project_id:
-        project = Project.objects.get(pk=project_id)
-    else:
-        task_logger.warning(f"Pipeline {pipeline} is not associated with a project")
-        project = None
+    # Pipelines must be associated with a project in order to select a processing service
+    # A processing service is required to send requests to the /process endpoint
+    project = Project.objects.get(pk=project_id)
+    task_logger.info(f"Using project: {project}")
 
-    pipeline_config = pipeline.get_config(project_id=project_id)
+    pipeline_config = pipeline.get_config(project_id=project.pk)
     task_logger.info(f"Using pipeline config: {pipeline_config}")
 
     prefiltered_images = list(images)
@@ -265,15 +264,16 @@ def process_images(
     task_logger.info(f"Found {len(detection_requests)} existing detections.")
 
     if not process_sync:
+        assert job_id is not None, "job_id is required to process images using async tasks."
         handle_async_process_images(
             pipeline.slug,
             source_image_requests,
             images,
             pipeline_config,
             detection_requests,
+            project_id,
             job_id,
             task_logger,
-            project_id,
         )
         return
     else:
@@ -289,12 +289,11 @@ def handle_async_process_images(
     source_images: list[SourceImage],
     pipeline_config: PipelineRequestConfigParameters,
     detection_requests: list[DetectionRequest],
-    job_id: int | None = None,
+    project_id: int,
+    job_id: int,
     task_logger: logging.Logger = logger,
-    project_id: int | None = None,
 ):
     """Handle asynchronous processing by submitting tasks to the appropriate pipeline queue."""
-    task_ids = []
     batch_size = pipeline_config.get("batch_size", 1)
 
     # Group source images into batches
@@ -333,14 +332,18 @@ def handle_async_process_images(
             detections=detections_batch,
             config=pipeline_config,
         )
-        task_result = process_pipeline_request.apply_async(
-            args=[prediction_request.dict(), project_id],
-            # TODO: make ml-pipeline an environment variable (i.e. PIPELINE_QUEUE_PREFIX)?
-            queue=f"ml-pipeline-{pipeline}",
-            # all pipelines have their own queue beginning with "ml-pipeline-"
-            # the antenna celeryworker should subscribe to all pipeline queues
+
+        task_id = str(uuid.uuid4())
+        transaction.on_commit(
+            lambda: process_pipeline_request.apply_async(
+                args=[prediction_request.dict(), project_id],
+                task_id=task_id,
+                # TODO: make ml-pipeline an environment variable (i.e. PIPELINE_QUEUE_PREFIX)?
+                queue=f"ml-pipeline-{pipeline}",
+                # all pipelines have their own queue beginning with "ml-pipeline-"
+                # the antenna celeryworker should subscribe to all pipeline queues
+            )
         )
-        task_ids.append(task_result.id)
 
         if job_id:
             from ami.jobs.models import Job, MLTaskRecord
@@ -349,21 +352,17 @@ def handle_async_process_images(
             # Create a new MLTaskRecord for this task
             ml_task_record = MLTaskRecord.objects.create(
                 job=job,
-                task_id=task_result.id,
+                task_id=task_id,
                 task_name="process_pipeline_request",
                 pipeline_request=prediction_request,
-                num_captures=len(source_image_batches[idx]),
+                num_captures=len(source_image_batches[i]),
             )
-            ml_task_record.source_images.set(source_image_batches[idx])
+            ml_task_record.source_images.set(source_image_batches[i])
             ml_task_record.save()
-            # job.logger.info(
-            #     f"Created MLTaskRecord for job {job_id} with task ID {task_result.id}"
-            #     " and task name process_pipeline_request"
-            # )
         else:
             task_logger.warning("No job ID provided, MLTaskRecord will not be created.")
 
-    task_logger.info(f"Submitted {len(task_ids)} batch image processing task(s).")
+    task_logger.info(f"Submitted {len(source_image_request_batches)} batch image processing task(s).")
 
 
 def handle_sync_process_images(
@@ -377,9 +376,6 @@ def handle_sync_process_images(
     job: Job | None,
 ) -> PipelineResultsResponse:
     """Handle synchronous processing by sending HTTP requests to the processing service."""
-    if project_id is None:
-        raise ValueError("Project ID must be provided when syncronously processing images.")
-
     processing_service = pipeline.choose_processing_service_for_pipeline(job_id, pipeline.name, project_id)
     if not processing_service.endpoint_url:
         raise ValueError(f"No endpoint URL configured for this pipeline's processing service ({processing_service})")
@@ -1207,8 +1203,8 @@ def get_config(self, project_id: int | None = None) -> PipelineRequestConfigPara
                 )
             except self.project_pipeline_configs.model.DoesNotExist as e:
                 logger.warning(f"No project-pipeline config for Pipeline {self} " f"and Project #{project_id}: {e}")
-
-        logger.warning("No project_id, no pipeline config is used.")
+        else:
+            logger.warning("No project_id. No pipeline config is used. Using default empty config instead.")
 
         return config
 

ami/ml/tasks.py

@@ -26,12 +26,10 @@ def process_source_images_async(pipeline_choice: str, image_ids: list[int], job_
 
     images = SourceImage.objects.filter(pk__in=image_ids)
     pipeline = Pipeline.objects.get(slug=pipeline_choice)
+    project = pipeline.projects.first()
+    assert project, f"Pipeline {pipeline} must be associated with a project."
 
-    results = process_images(
-        pipeline=pipeline,
-        images=images,
-        job_id=job_id,
-    )
+    results = process_images(pipeline=pipeline, images=images, job_id=job_id, project_id=project.pk)
 
     try:
         save_results(results=results, job_id=job_id)
@@ -136,5 +134,9 @@ def check_ml_job_status(ml_job_id: int):
         logger.info(f"ML Job {ml_job_id} is complete.")
         job.logger.info(f"ML Job {ml_job_id} is complete.")
     else:
+        from django.db import transaction
+
         logger.info(f"ML Job {ml_job_id} still in progress. Checking again for completed tasks.")
-        check_ml_job_status.apply_async([ml_job_id], countdown=10)  # check again in 10 seconds
+        transaction.on_commit(
+            lambda: check_ml_job_status.apply_async([ml_job_id], countdown=10)
+        )  # check again in 10 seconds