Support cancelling ML async jobs

ami/jobs/models.py

@@ -15,7 +15,7 @@
 
 from ami.base.models import BaseModel
 from ami.base.schemas import ConfigurableStage, ConfigurableStageParam
-from ami.jobs.tasks import run_job
+from ami.jobs.tasks import cleanup_async_job_if_needed, 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
@@ -970,14 +970,16 @@ def cancel(self):
         """
         self.status = JobState.CANCELING
         self.save()
+
+        cleanup_async_job_if_needed(self)
         if self.task_id:
             task = run_job.AsyncResult(self.task_id)
             if task:
                 task.revoke(terminate=True)
-                self.save()
-        else:
-            self.status = JobState.REVOKED
-            self.save()
+        # For sync jobs the task revoke will update the job status. However, for async jobs we need to set the status
+        # to revoked here since the task already finished (it only queues the images)
+        self.status = JobState.REVOKED
+        self.save()
 
     def update_status(self, status=None, save=True):
         """

ami/jobs/tasks.py

@@ -16,7 +16,7 @@
 from config import celery_app
 
 if TYPE_CHECKING:
-    from ami.jobs.models import JobState
+    from ami.jobs.models import Job, JobState
 
 logger = logging.getLogger(__name__)
 # Minimum success rate. Jobs with fewer than this fraction of images
@@ -94,6 +94,13 @@ def process_nats_pipeline_result(self, job_id: int, result_data: dict, reply_sub
         )
         raise self.retry(countdown=5, max_retries=10)
 
+    if progress_info.unknown:
+        logger.warning(
+            f"Progress info is unknown for job {job_id} when processing results. Job may be cancelled."
+            f"Or this could be a transient Redis error and the NATS task will be retried."
+        )
+        return
+
     try:
         complete_state = JobState.SUCCESS
         if progress_info.total > 0 and (progress_info.failed / progress_info.total) > FAILURE_THRESHOLD:
@@ -272,10 +279,10 @@ def _update_job_progress(
     # Clean up async resources for completed jobs that use NATS/Redis
     if job.progress.is_complete():
         job = Job.objects.get(pk=job_id)  # Re-fetch outside transaction
-        _cleanup_job_if_needed(job)
+        cleanup_async_job_if_needed(job)
 
 
-def _cleanup_job_if_needed(job) -> None:
+def cleanup_async_job_if_needed(job: "Job") -> None:
     """
     Clean up async resources (NATS/Redis) if this job uses them.
 
@@ -330,7 +337,7 @@ def update_job_status(sender, task_id, task, state: str, retval=None, **kwargs):
 
     # Clean up async resources for revoked jobs
     if state == JobState.REVOKED:
-        _cleanup_job_if_needed(job)
+        cleanup_async_job_if_needed(job)
 
 
 @task_failure.connect(sender=run_job, retry=False)
@@ -345,7 +352,7 @@ def update_job_failure(sender, task_id, exception, *args, **kwargs):
     job.save()
 
     # Clean up async resources for failed jobs
-    _cleanup_job_if_needed(job)
+    cleanup_async_job_if_needed(job)
 
 
 def log_time(start: float = 0, msg: str | None = None) -> tuple[float, Callable]:

ami/ml/orchestration/async_job_state.py

@@ -43,6 +43,7 @@ class JobStateProgress:
     may need to be generalized.
     """
 
+    unknown: bool = False  # True if progress cannot be determined (e.g. missing Redis keys)
     remaining: int = 0  # source images not yet processed in this stage
     total: int = 0  # total source images in the job
     processed: int = 0  # source images completed (success + failed)
@@ -169,7 +170,8 @@ def _commit_update(
         pending_images = cache.get(self._get_pending_key(stage))
         total_images = cache.get(self._total_key)
         if pending_images is None or total_images is None:
-            return None
+            # important that this is distinguishable from not getting the lock
+            return JobStateProgress(unknown=True)
         remaining_images = [img_id for img_id in pending_images if img_id not in processed_image_ids]
         assert len(pending_images) >= len(remaining_images)
         cache.set(self._get_pending_key(stage), remaining_images, timeout=self.TIMEOUT)
@@ -210,3 +212,4 @@ def cleanup(self) -> None:
             cache.delete(self._get_pending_key(stage))
         cache.delete(self._failed_key)
         cache.delete(self._total_key)
+        cache.delete(_lock_key(self.job_id))

ami/ml/orchestration/nats_queue.py

@@ -95,21 +95,31 @@ def _get_consumer_name(self, job_id: int) -> str:
         """Get consumer name from job_id."""
         return f"job-{job_id}-consumer"
 
-    async def _ensure_stream(self, job_id: int):
-        """Ensure stream exists for the given job."""
+    async def _stream_exists(self, job_id: int) -> bool:
+        """Check if stream exists for the given job."""
         if self.js is None:
             raise RuntimeError("Connection is not open. Use TaskQueueManager as an async context manager.")
 
         stream_name = self._get_stream_name(job_id)
-        subject = self._get_subject(job_id)
-
         try:
             await asyncio.wait_for(self.js.stream_info(stream_name), timeout=NATS_JETSTREAM_TIMEOUT)
             logger.debug(f"Stream {stream_name} already exists")
+            return True
         except asyncio.TimeoutError:
             raise  # NATS unreachable — let caller handle it rather than creating a stream blindly
-        except Exception as e:
-            logger.warning(f"Stream {stream_name} does not exist: {e}")
+        except Exception:
+            return False
+
+    async def _ensure_stream(self, job_id: int):
+        """Ensure stream exists for the given job."""
+        if self.js is None:
+            raise RuntimeError("Connection is not open. Use TaskQueueManager as an async context manager.")
+
+        if not await self._stream_exists(job_id):
+            stream_name = self._get_stream_name(job_id)
+            subject = self._get_subject(job_id)
+
+            logger.warning(f"Stream {stream_name} does not exist")
             # Stream doesn't exist, create it
             await asyncio.wait_for(
                 self.js.add_stream(
@@ -207,7 +217,10 @@ async def reserve_tasks(self, job_id: int, count: int, timeout: float = 5) -> li
             raise RuntimeError("Connection is not open. Use TaskQueueManager as an async context manager.")
 
         try:
-            await self._ensure_stream(job_id)
+            if not await self._stream_exists(job_id):
+                logger.debug(f"Stream for job '{job_id}' does not exist when reserving task")
+                return []
+
             await self._ensure_consumer(job_id)
 
             consumer_name = self._get_consumer_name(job_id)

ami/ml/tests.py

@@ -902,7 +902,7 @@ def test_small_size_filter_assigns_not_identifiable(self):
             )
 
 
-class TestTaskStateManager(TestCase):
+class TestAsyncJobStateManager(TestCase):
     """Test TaskStateManager for job progress tracking."""
 
     def setUp(self):
@@ -1025,7 +1025,7 @@ def test_cleanup(self):
 
         # Verify keys are gone
         progress = self.manager._commit_update(set(), "process")
-        self.assertIsNone(progress)
+        self.assertTrue(progress is None or progress.unknown)
 
     def test_failed_image_tracking(self):
         """Test basic failed image tracking with no double-counting on retries."""