fix(jobs): close prerun-clobber + orphan-NATS-dispatch races

Review on #1324 surfaced two races that left the early-guard non-functional
in production:

1. ``task_prerun`` (``pre_update_job_status``) wrote PENDING to the row
   before the ``run_job`` body inspected status. A canceled or redelivered
   message therefore had its REVOKED/CANCELING overwritten with PENDING,
   and the early-guard added in the parent commit never tripped. The
   existing tests passed only because they invoked ``run_job.apply(args=[…])``
   while production uses ``kwargs={"job_id": …}`` — under args, the prerun
   handler raised ``KeyError`` and exited silently. Switching the tests to
   ``kwargs=`` reproduces the production code path; the prerun handler now
   short-circuits when ``Job.is_settled()`` is true, preserving the status
   the early-guard reads next.

2. For ASYNC_API jobs ``Job.cancel()`` revokes without ``terminate=True``,
   marks the row REVOKED, and tears down the NATS stream + Redis state.
   ``MLJob.run`` running in a worker that's still inside ``collect_images``
   (slow for large collections) would then proceed to ``queue_images_to_nats``
   and recreate the stream the cancel just deleted, dispatching real GPU
   work to ADC for a revoked job; the results came back to no Redis state
   and ``_fail_job`` silently overwrote REVOKED with FAILURE. The bootstrap
   now checks ``Job.status`` (via a values-only read so the in-memory
   ``progress`` mutations don't clobber the cancel's REVOKED) right after
   the collect stage and bails out before any dispatch.

Adds ``Job.is_settled()`` to centralize the "terminal or being torn down"
predicate that ``run_job``'s early-guard, the prerun handler, ``_fail_job``,
and the bootstrap guard all needed. Adds two regression tests: one for the
prerun-then-guard chain, one for the cancel-during-bootstrap race.

Co-Authored-By: Claude <noreply@anthropic.com>

Author: mihow

ami/jobs/models.py

@@ -529,6 +529,28 @@ def run(cls, job: "Job"):
             progress=1,
         )
 
+        # Mid-bootstrap cancel guard. ``collect_images`` above can run for many
+        # minutes on large collections (S3 list + DB joins), and the user may
+        # cancel during that window. ``Job.cancel()`` for ASYNC_API does
+        # ``revoke(terminate=False)`` to avoid SIGKILL'ing this worker, then
+        # writes REVOKED + tears down the NATS stream / Redis state. Without
+        # this check we would (a) clobber the cancel's REVOKED via the next
+        # full ``job.save()`` and (b) proceed to ``queue_images_to_nats``,
+        # recreating the stream the cancel just deleted and dispatching real
+        # GPU work to ADC for a revoked job. Refresh is read-only against the
+        # ``status`` column; the in-memory ``progress`` mutations from the
+        # collect stage are intentionally dropped on the bail path because the
+        # job is settled — no further progress writes make sense. Covers
+        # ASYNC_API (NATS dispatch) and SYNC paths (Celery sub-tasks in
+        # ``process_images``); INTERNAL jobs benefit too. See
+        # RolnickLab/antenna#1323.
+        db_status = Job.objects.values_list("status", flat=True).get(pk=job.pk)
+        if db_status in JobState.final_states() or db_status == JobState.CANCELING:
+            job.logger.info(
+                f"Job {job.pk} settled to {db_status} during bootstrap; " f"skipping dispatch of {len(images)} images"
+            )
+            return
+
         # End image collection stage
         job.save()
 
@@ -1041,6 +1063,21 @@ def setup(self, save=True):
         if save:
             self.save()
 
+    def is_settled(self) -> bool:
+        """Return True when the job is in a terminal state or being cancelled.
+
+        Used by every code path that must not start (or continue) work for a
+        job whose lifecycle has effectively ended: the ``run_job`` early-guard
+        (after acks_late redelivery), the ``MLJob.run`` mid-bootstrap cancel
+        check (before ``queue_images_to_nats`` dispatches GPU work to ADC),
+        the prerun signal handler (so a redelivered or canceled message does
+        not get its status reset to PENDING), and ``_fail_job``. Centralized
+        so the predicate stays in one place — adding a new "do not resume"
+        status only requires touching :meth:`JobState.final_states` or this
+        method.
+        """
+        return self.status in JobState.final_states() or self.status == JobState.CANCELING
+
     def run(self):
         """
         Run the job.

ami/jobs/tasks.py

@@ -150,7 +150,7 @@ def update_async_services_seen_for_project(project_id: int) -> None:
     reject_on_worker_lost=True,
 )
 def run_job(self, job_id: int) -> None:
-    from ami.jobs.models import Job, JobState
+    from ami.jobs.models import Job
 
     try:
         job = Job.objects.get(pk=job_id)
@@ -161,8 +161,10 @@ def run_job(self, job_id: int) -> None:
     # Early-guard: under acks_late, the broker may redeliver this message after a
     # worker SIGKILL/OOM, and Job.cancel() may also flip status to CANCELING /
     # REVOKED while the message sits in the prefetch buffer. Don't re-run a job
-    # that's already settled or being torn down.
-    if job.status in JobState.final_states() or job.status == JobState.CANCELING:
+    # that's already settled or being torn down. The companion guard in
+    # pre_update_job_status above prevents the task_prerun signal from
+    # overwriting that status with PENDING before we get here.
+    if job.is_settled():
         job.logger.info(
             f"Skipping run_job for job {job.pk}: already in status {job.status} "
             f"(redelivery or cancellation in flight)"
@@ -444,7 +446,7 @@ def _fail_job(job_id: int, reason: str) -> None:
     try:
         with transaction.atomic():
             job = Job.objects.select_for_update().get(pk=job_id)
-            if job.status in (JobState.CANCELING, *JobState.final_states()):
+            if job.is_settled():
                 return
             job.update_status(JobState.FAILURE, save=False)
             job.finished_at = datetime.datetime.now()
@@ -1327,7 +1329,34 @@ def cleanup_async_job_if_needed(job) -> None:
 
 @task_prerun.connect(sender=run_job)
 def pre_update_job_status(sender, task_id, task, **kwargs):
-    # in the prerun signal, set the job status to PENDING
+    """Bump the job to PENDING when a worker picks the message up.
+
+    Skipped when the job is already settled (terminal state) or being
+    cancelled. Without that guard, a broker redelivery (acks_late + worker
+    crash) or a cancel that arrived while the message was still in the
+    prefetch buffer would have its REVOKED/CANCELING status silently
+    overwritten with PENDING here, and the ``run_job`` early-guard
+    (which reads ``Job.status`` after this signal fires) would then fail
+    to short-circuit and re-run the job. See RolnickLab/antenna#1323.
+    """
+    from ami.jobs.models import Job
+
+    job_id = task.request.kwargs.get("job_id") if task.request.kwargs else None
+    if job_id is None and task.request.args:
+        job_id = task.request.args[0]
+    if job_id is not None:
+        try:
+            job = Job.objects.only("status").get(pk=job_id)
+        except Job.DoesNotExist:
+            pass
+        else:
+            if job.is_settled():
+                logger.info(
+                    "task_prerun: skipping PENDING write for job %s in status %s " "(redelivery or cancel in flight)",
+                    job_id,
+                    job.status,
+                )
+                return
     update_job_status(sender, task_id, task, "PENDING", **kwargs)
 
 

ami/jobs/tests/test_jobs.py

@@ -439,6 +439,50 @@ def test_cancel_sync_api_job_terminates_celery_task(self):
         job.refresh_from_db()
         self.assertEqual(job.status, JobState.REVOKED)
 
+    def test_mljob_run_bails_when_cancelled_during_bootstrap(self):
+        """Regression for the cancel-during-bootstrap race in ASYNC_API jobs.
+
+        ``Job.cancel()`` revokes without terminate=True (so the local worker is
+        not SIGTERM'd), marks the row REVOKED, and tears down NATS/Redis state.
+        If the worker is still inside ``MLJob.run`` (typically blocked in the
+        slow ``collect_images`` step), it must refresh the row and bail BEFORE
+        calling ``queue_images_to_nats`` — otherwise it would recreate the
+        stream and dispatch real GPU work to ADC for a revoked job.
+        """
+        from unittest.mock import patch
+
+        from ami.jobs.models import MLJob
+
+        pipeline = Pipeline.objects.create(name="Cancel-race pipeline", slug="cancel-race-pipeline")
+        pipeline.projects.add(self.project)
+        collection = SourceImageCollection.objects.create(name="Cancel-race collection", project=self.project)
+        job = Job.objects.create(
+            project=self.project,
+            name="Cancel-race",
+            pipeline=pipeline,
+            source_image_collection=collection,
+            status=JobState.STARTED,
+            dispatch_mode=JobDispatchMode.ASYNC_API,
+        )
+        job.setup()
+
+        def cancel_mid_collect(*_args, **_kwargs):
+            # Simulate the user clicking cancel while collect_images is still
+            # running: rewrite the DB row out from under this in-flight task.
+            Job.objects.filter(pk=job.pk).update(status=JobState.REVOKED)
+            return []
+
+        with patch.object(
+            pipeline,
+            "collect_images",
+            side_effect=cancel_mid_collect,
+        ), patch("ami.ml.orchestration.jobs.queue_images_to_nats") as mock_queue:
+            MLJob.run(job)
+
+        mock_queue.assert_not_called()
+        job.refresh_from_db()
+        self.assertEqual(job.status, JobState.REVOKED)
+
     def test_cancel_job_without_task_id_still_revokes(self):
         """A job that never made it to enqueue (no task_id) still transitions
         to REVOKED and triggers async-cleanup (a no-op for non-ASYNC_API)."""

ami/jobs/tests/test_tasks.py

@@ -667,7 +667,7 @@ def test_skips_when_job_already_revoked(self):
         job = self._make_job(JobState.REVOKED)
 
         with patch.object(Job, "run") as mock_run:
-            result = run_job.apply(args=[job.pk])
+            result = run_job.apply(kwargs={"job_id": job.pk})
 
         self.assertTrue(result.successful(), msg=f"task should succeed, got {result.state}: {result.traceback}")
         mock_run.assert_not_called()
@@ -679,7 +679,7 @@ def test_skips_when_job_canceling(self):
         job = self._make_job(JobState.CANCELING)
 
         with patch.object(Job, "run") as mock_run:
-            result = run_job.apply(args=[job.pk])
+            result = run_job.apply(kwargs={"job_id": job.pk})
 
         self.assertTrue(result.successful(), msg=f"task should succeed, got {result.state}: {result.traceback}")
         mock_run.assert_not_called()
@@ -691,7 +691,7 @@ def test_skips_when_job_already_success(self):
         job = self._make_job(JobState.SUCCESS)
 
         with patch.object(Job, "run") as mock_run:
-            result = run_job.apply(args=[job.pk])
+            result = run_job.apply(kwargs={"job_id": job.pk})
 
         self.assertTrue(result.successful(), msg=f"task should succeed, got {result.state}: {result.traceback}")
         mock_run.assert_not_called()
@@ -703,10 +703,29 @@ def test_runs_when_job_pending(self):
         job = self._make_job(JobState.PENDING)
 
         with patch.object(Job, "run") as mock_run:
-            run_job.apply(args=[job.pk])
+            run_job.apply(kwargs={"job_id": job.pk})
 
         mock_run.assert_called_once()
 
+    def test_prerun_signal_does_not_clobber_revoked_status(self):
+        """
+        Regression: the ``task_prerun`` signal would otherwise call
+        ``update_job_status(state="PENDING")`` and overwrite a REVOKED/CANCELING
+        status before the ``run_job`` early-guard reads it. With the prerun
+        guard in place, the status survives the signal, the early-guard fires,
+        and ``Job.run()`` is not called.
+        """
+        from ami.jobs.tasks import run_job
+
+        job = self._make_job(JobState.REVOKED)
+
+        with patch.object(Job, "run") as mock_run:
+            run_job.apply(kwargs={"job_id": job.pk})
+
+        job.refresh_from_db()
+        self.assertEqual(job.status, JobState.REVOKED)
+        mock_run.assert_not_called()
+
 
 class TestResultEndpointWithError(APITestCase):
     """Integration test for the result API endpoint with error results."""