Merge pull request #1248 from NASA-IMPACT/integrate_classification_queue

Integrate classification queue

Author: CarsonDavis

config/settings/base.py

@@ -349,3 +349,4 @@
 LRM_DEV_TOKEN = env("LRM_DEV_TOKEN")
 XLI_TOKEN = env("XLI_TOKEN")
 INFERENCE_API_URL = env("INFERENCE_API_URL")
+TDAMM_CLASSIFICATION_THRESHOLD = env("TDAMM_CLASSIFICATION_THRESHOLD", default="0.5")

inference/models/inference.py

@@ -9,6 +9,7 @@
     InferenceJobStatus,
 )
 from inference.utils.batch import BatchProcessor
+from inference.utils.classification_utils import update_url_with_classification_results
 from inference.utils.inference_api_client import InferenceAPIClient
 
 
@@ -167,6 +168,11 @@ def initiate(self, inference_api_url=settings.INFERENCE_API_URL) -> None:
 
             if not created_batch:
                 self.log_error_and_set_status_failed("No external jobs created")
+                self.status = InferenceJobStatus.FAILED
+                self.updated_at = timezone.now()
+                self.completed_at = timezone.now()
+                self.save()
+                return
 
             self.status = InferenceJobStatus.PENDING
             self.save()
@@ -184,17 +190,37 @@ def refresh_external_jobs_status_and_store_results(self) -> None:
     def reevaluate_progress_and_update_status(self) -> None:
         """Evaluate overall job status and handle completion"""
 
+        if self.status == InferenceJobStatus.QUEUED:
+            return
+
+        if not self.external_jobs.exists() and self.status == InferenceJobStatus.PENDING:
+            self.status = InferenceJobStatus.FAILED
+            self.error_message = "No external jobs created for pending job"
+            self.completed_at = timezone.now()
+            self.save()
+            return
+
         if self.get_ongoing_external_jobs().exists():
             self.status = InferenceJobStatus.PENDING
+            self.updated_at = timezone.now()
         else:
             if self.get_failed_external_jobs().exists():
                 self.status = InferenceJobStatus.FAILED
+                self.updated_at = timezone.now()
             else:
                 self.status = InferenceJobStatus.COMPLETED
+                self.updated_at = timezone.now()
             self.completed_at = timezone.now()
             self.unload_model()
         self.save()
 
+        # If job is completed or failed, check if all classifications are done
+        # if self.status in [InferenceJobStatus.COMPLETED, InferenceJobStatus.FAILED]:
+        #     self.collection.check_classifications_complete_and_finish_migration()
+
+        if self.status in [InferenceJobStatus.COMPLETED]:
+            self.collection.check_classifications_complete_and_finish_migration()
+
     def unload_model(self) -> None:
         """
         Check that no other jobs are using the loaded model
@@ -247,6 +273,25 @@ def store_results(self, results) -> None:
         """Store results and mark as completed"""
         try:
             self.results = results
+            if results:
+                collection = self.inference_job.collection
+
+                for idx, url_id in enumerate(self.url_ids):
+                    if idx < len(results):
+                        try:
+                            dump_url = collection.dump_urls.get(id=url_id)
+                            result = results[idx]
+                            # print(f"Processing result {idx}: {result}")
+                            if isinstance(result, dict) and "confidence" in result:
+                                # Ensure confidence is float
+                                result["confidence"] = float(result["confidence"])
+
+                            update_url_with_classification_results(dump_url, results[idx])
+                            # tdamm_tags = update_url_with_classification_results(dump_url, results[idx])
+                            # print(f"tdamm_tags added: {tdamm_tags}")
+                        except collection.dump_urls.model.DoesNotExist:
+                            continue
+
             self.mark_completed()
 
         except Exception as e:
@@ -256,7 +301,8 @@ def refresh_status_and_store_results(self) -> None:
         """Process this external job and update status/results"""
         try:
             api_client = InferenceAPIClient()
-            model_version = ModelVersion.objects.get(classification_type=self.inference_job.classification_type)
+            # model_version = ModelVersion.objects.get(classification_type=self.inference_job.classification_type)
+            model_version = self.inference_job.model_version
 
             response = api_client.get_job_status(model_version.api_identifier, self.external_job_id)
 
@@ -268,7 +314,7 @@ def refresh_status_and_store_results(self) -> None:
             # Handle completion or failure
             if new_status == ExternalJobStatus.COMPLETED:
                 self.store_results(response.get("results"))
-                self.completed_at = timezone.now()
+                # self.completed_at = timezone.now() # completed in mark_completed called in store_results
             self.save()
 
         except Exception as e:

inference/tasks.py

@@ -5,13 +5,6 @@
 from inference.utils.advisory_lock import AdvisoryLock
 
 
-def generate_inference_job(collection, classification_type):
-    """Creates a new inference job for a collection."""
-    return InferenceJob.objects.create(
-        collection=collection, classification_type=classification_type, status=InferenceJobStatus.QUEUED
-    )
-
-
 @shared_task
 def process_inference_job_queue():
     """
@@ -25,13 +18,18 @@ def process_inference_job_queue():
             return "Queue processing already in progress"
 
         try:
+            # Reevaluate progress and update status of all inference jobs that are not currently queued
+            # for job in InferenceJob.objects.exclude(status=InferenceJobStatus.QUEUED):
+            #     job.reevaluate_progress_and_update_status()
+
             # Look for pending jobs first
             pending_jobs = InferenceJob.objects.filter(status=InferenceJobStatus.PENDING)
 
             if pending_jobs.exists():
-                # Process pending jobs
+                # Refresh and process pending jobs
                 for job in pending_jobs:
                     job.refresh_external_jobs_status_and_store_results()
+                    job.reevaluate_progress_and_update_status()
             else:
                 # If no pending jobs, try to initiate a queued job
                 queued_job = (

inference/tests/local_test_inference_api_client.py

@@ -1,4 +1,5 @@
-# inference/tests/test_inference_api_client.py
+# inference/tests/local_test_inference_api_client.py
+# docker-compose -f local.yml run --rm django pytest inference/tests/local_test_inference_api_client.py
 
 """
 This is a test designed to be run on a local machine which has the inference pipeline running

inference/tests/test_batch.py

@@ -1,4 +1,5 @@
 # inference/tests/test_batch.py
+# docker-compose -f local.yml run --rm django pytest inference/tests/test_batch.py
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest