Enable async and distributed processing for the ML backend

.envs/.ci/.django

@@ -19,3 +19,7 @@ MINIO_BROWSER_REDIRECT_URL=http://minio:9001
 
 DEFAULT_PROCESSING_SERVICE_NAME=Test Processing Service
 DEFAULT_PROCESSING_SERVICE_ENDPOINT=http://ml_backend:2000
+
+CELERY_BROKER_URL=amqp://rabbituser:rabbitpass@rabbitmq:5672//
+RABBITMQ_DEFAULT_USER=rabbituser
+RABBITMQ_DEFAULT_PASS=rabbitpass

.envs/.ci/.postgres

@@ -3,3 +3,4 @@ POSTGRES_PORT=5432
 POSTGRES_DB=ami-ci
 POSTGRES_USER=4JXkOnTAeDmDyIapSRrGEE
 POSTGRES_PASSWORD=d4xojpnJU3OzPQ0apSCLP1oHR1TYvyMzAlF5KpE9HFL6MPlnbDibwI
+DATABASE_URL=postgres://xekSryPnqczJXkOnTAeDmDyIapSRrGEE:iMRQjJEGflj5xojpnJU3OzPQ0apSCLP1oHR1TYvyMzAlF5KpE9HFL6MPlnbDibwI@postgres:5432/ami

.envs/.local/.django

@@ -42,3 +42,7 @@ MINIO_BROWSER_REDIRECT_URL=http://minio:9001
 DEFAULT_PROCESSING_SERVICE_NAME=Local Processing Service
 DEFAULT_PROCESSING_SERVICE_ENDPOINT=http://ml_backend:2000
 # DEFAULT_PIPELINES_ENABLED=random,constant  # When set to None, all pipelines will be enabled.
+
+CELERY_BROKER_URL=amqp://rabbituser:rabbitpass@rabbitmq:5672//
+RABBITMQ_DEFAULT_USER=rabbituser
+RABBITMQ_DEFAULT_PASS=rabbitpass

.envs/.local/.postgres

@@ -3,3 +3,4 @@ POSTGRES_PORT=5432
 POSTGRES_DB=ami
 POSTGRES_USER=xekSryPnqczJXkOnTAeDmDyIapSRrGEE
 POSTGRES_PASSWORD=iMRQjJEGflj5xojpnJU3OzPQ0apSCLP1oHR1TYvyMzAlF5KpE9HFL6MPlnbDibwI
+DATABASE_URL=postgres://xekSryPnqczJXkOnTAeDmDyIapSRrGEE:iMRQjJEGflj5xojpnJU3OzPQ0apSCLP1oHR1TYvyMzAlF5KpE9HFL6MPlnbDibwI@postgres:5432/ami

ami/jobs/admin.py

@@ -4,7 +4,7 @@
 
 from ami.main.admin import AdminBase
 
-from .models import Job, get_job_type_by_inferred_key
+from .models import Job, MLTaskRecord, get_job_type_by_inferred_key
 
 
 @admin.register(Job)
@@ -54,3 +54,23 @@ def inferred_job_type(self, obj: Job) -> str:
         "progress",
         "result",
     )
+
+
+@admin.register(MLTaskRecord)
+class MLTaskRecordAdmin(AdminBase):
+    """Admin panel example for ``MLTaskRecord`` model."""
+
+    list_display = (
+        "job",
+        "task_id",
+        "task_name",
+        "status",
+    )
+
+    @admin.action()
+    def kill_task(self, request: HttpRequest, queryset: QuerySet[MLTaskRecord]) -> None:
+        for ml_task_record in queryset:
+            ml_task_record.kill_task()
+        self.message_user(request, f"Killed {queryset.count()} ML task(s).")
+
+    actions = [kill_task]

ami/jobs/migrations/0019_mltaskrecord.py

@@ -0,0 +1,72 @@
+# Generated by Django 4.2.10 on 2025-08-10 22:17
+
+import ami.ml.schemas
+from django.db import migrations, models
+import django.db.models.deletion
+import django_pydantic_field.fields
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("main", "0060_alter_sourceimagecollection_method"),
+        ("jobs", "0018_alter_job_job_type_key"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="MLTaskRecord",
+            fields=[
+                ("id", models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name="ID")),
+                ("created_at", models.DateTimeField(auto_now_add=True)),
+                ("updated_at", models.DateTimeField(auto_now=True)),
+                ("task_id", models.CharField(max_length=255)),
+                (
+                    "task_name",
+                    models.CharField(
+                        choices=[
+                            ("process_pipeline_request", "process_pipeline_request"),
+                            ("save_results", "save_results"),
+                        ],
+                        default="process_pipeline_request",
+                        max_length=255,
+                    ),
+                ),
+                (
+                    "status",
+                    models.CharField(
+                        choices=[("STARTED", "STARTED"), ("SUCCESS", "SUCCESS"), ("FAIL", "FAIL")],
+                        default="STARTED",
+                        max_length=255,
+                    ),
+                ),
+                ("raw_results", models.JSONField(blank=True, default=dict, null=True)),
+                ("raw_traceback", models.TextField(blank=True, null=True)),
+                (
+                    "pipeline_request",
+                    django_pydantic_field.fields.PydanticSchemaField(
+                        blank=True, config=None, null=True, schema=ami.ml.schemas.PipelineRequest
+                    ),
+                ),
+                (
+                    "pipeline_response",
+                    django_pydantic_field.fields.PydanticSchemaField(
+                        blank=True, config=None, null=True, schema=ami.ml.schemas.PipelineResultsResponse
+                    ),
+                ),
+                ("num_captures", models.IntegerField(default=0, help_text="Same as number of source_images")),
+                ("num_detections", models.IntegerField(default=0)),
+                ("num_classifications", models.IntegerField(default=0)),
+                ("subtask_id", models.CharField(blank=True, max_length=255, null=True)),
+                (
+                    "job",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE, related_name="ml_task_records", to="jobs.job"
+                    ),
+                ),
+                ("source_images", models.ManyToManyField(related_name="ml_task_records", to="main.sourceimage")),
+            ],
+            options={
+                "abstract": False,
+            },
+        ),
+    ]

ami/jobs/migrations/0020_alter_job_logs_alter_job_progress.py

@@ -0,0 +1,28 @@
+# Generated by Django 4.2.10 on 2025-09-04 10:42
+
+import ami.jobs.models
+from django.db import migrations
+import django_pydantic_field.fields
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("jobs", "0019_mltaskrecord"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="job",
+            name="logs",
+            field=django_pydantic_field.fields.PydanticSchemaField(
+                config=None, default=ami.jobs.models.JobLogs, schema=ami.jobs.models.JobLogs
+            ),
+        ),
+        migrations.AlterField(
+            model_name="job",
+            name="progress",
+            field=django_pydantic_field.fields.PydanticSchemaField(
+                config=None, default=ami.jobs.models.default_job_progress, schema=ami.jobs.models.JobProgress
+            ),
+        ),
+    ]

ami/jobs/migrations/0021_remove_mltaskrecord_subtask_id_and_more.py

@@ -0,0 +1,30 @@
+# Generated by Django 4.2.10 on 2025-10-16 19:31
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("jobs", "0020_alter_job_logs_alter_job_progress"),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name="mltaskrecord",
+            name="subtask_id",
+        ),
+        migrations.AlterField(
+            model_name="mltaskrecord",
+            name="status",
+            field=models.CharField(
+                choices=[("PENDING", "PENDING"), ("STARTED", "STARTED"), ("SUCCESS", "SUCCESS"), ("FAIL", "FAIL")],
+                default="STARTED",
+                max_length=255,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="mltaskrecord",
+            name="task_id",
+            field=models.CharField(blank=True, max_length=255, null=True),
+        ),
+    ]

ami/jobs/migrations/0022_job_last_checked.py

@@ -0,0 +1,17 @@
+# Generated by Django 4.2.10 on 2025-10-17 01:27
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("jobs", "0021_remove_mltaskrecord_subtask_id_and_more"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="job",
+            name="last_checked",
+            field=models.DateTimeField(blank=True, null=True),
+        ),
+    ]

ami/jobs/migrations/0023_alter_job_last_checked_alter_mltaskrecord_status.py

@@ -0,0 +1,33 @@
+# Generated by Django 4.2.10 on 2025-11-04 11:44
+
+import datetime
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("jobs", "0022_job_last_checked"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="job",
+            name="last_checked",
+            field=models.DateTimeField(blank=True, default=datetime.datetime.now, null=True),
+        ),
+        migrations.AlterField(
+            model_name="mltaskrecord",
+            name="status",
+            field=models.CharField(
+                choices=[
+                    ("PENDING", "PENDING"),
+                    ("STARTED", "STARTED"),
+                    ("SUCCESS", "SUCCESS"),
+                    ("FAIL", "FAIL"),
+                    ("REVOKED", "REVOKED"),
+                ],
+                default="STARTED",
+                max_length=255,
+            ),
+        ),
+    ]