Clean up task queue when job is revoked or re-started, fix duplicate tasks

[mihow]

Bug

NATS task queue (stream + consumer) is not cleaned up when an async_api job is revoked or re-started. Two distinct lifecycle gaps:

1. Job marked REVOKED → stream + consumer left behind

The job-watchdog / healthcheck flips a stuck job to REVOKED, but the per-job NATS stream and consumer are untouched. Messages sit in the stream indefinitely.

2. User clicks "Start" on a non-running job → tasks re-published, stream not purged

Re-starting (from REVOKED, FAILURE, or CANCELLED) re-runs the collect stage and publishes a fresh round of tasks on top of whatever is still in the stream. The existing consumer's delivered_seq continues from where it was, so:

• Stream message count roughly doubles (or worse on multiple restarts).
• Already-ACKed tasks from the previous attempt that were re-published become duplicates that workers will process again, producing duplicate detections / classifications.
• Stage params are overwritten with fresh values, which manifests in the UI as the "Processed" counter visibly resetting to 0 while the user is watching.

Concrete observation

A real job hit both gaps in sequence today:

1. Job started, collect completed with 44,286 tasks → stream job_<N> published 44,286 messages.
2. Workers were starved by an unrelated /next filter bug (separate ticket: CANCELLED jobs leak through /next filter, starve newer async_api jobs #1282), never delivered. Healthcheck flipped the job to REVOKED after ~18 min. Stream + consumer still present, 44,286 pending.
3. After the upstream filter bug was unblocked manually, user clicked "Start". Second collect ran, found 43,597 tasks (689 source images had soft-deleted in between), and published them on top of the existing 44,286.
4. Resulting NATS state: messages = 87,883, delivered_seq = wherever the consumer had reached, ~half the workload now duplicated across the two rounds.
5. UI showed the process stage counter snap from ~3,000 back to 0 and start counting up again, while detection / classification counts ran ahead of the processed-image count (reflecting re-detection of already-processed images).

Suggested cleanup hooks

• On status transition into REVOKED / FAILURE / CANCELLED (whether by user, watchdog, or admin tooling): delete the consumer, purge or delete the stream. After that point the job is terminal; nothing should still be drainable.
• On Start action for a non-running job: idempotently delete any leftover stream + consumer for that job ID before publishing the new round, so the new run begins from a known-empty queue. Reset stage params to fresh defaults at the same point so the UI counter starts cleanly at 0 rather than briefly reflecting stale state.
• Both code paths should tolerate "stream/consumer doesn't exist" without raising — the cleanup must be safe to run when there's nothing to clean.

Why this matters

• Duplicate task processing wastes GPU time and produces duplicate Detection / Classification rows that have to be reconciled downstream.
• The visible counter reset confuses operators about whether work was lost.
• Long-lived REVOKED streams + consumers accumulate on the NATS server and obscure real triage signal (num_pending counts from dead jobs).

Out of scope

A separate "resume" action that does keep the existing stream + consumer would be a useful follow-up, but the current "Start" button has restart-from-scratch semantics and the cleanup-then-publish sequence matches that intent.

[mihow]

Update — correction to Claude's earlier comment

I previously claimed (in this same comment, before edit) that run_job for an async_api job is long-running and that the "split-out" would be to make it publish-and-exit. That claim was wrong — run_job for async_api already does collect → publish → exit per MLPipelineJobType.run in ami/jobs/models.py:535–543:

if job.dispatch_mode == JobDispatchMode.ASYNC_API:
    queued = queue_images_to_nats(job, images)
    if not queued:
        # ... fail path
        return
    # When all stages are already complete (e.g. 0 images to process),
    # finalize the job now since no async results will arrive to trigger completion.
    if job.progress.is_complete():
        # ... immediate-finalize path
else:
    cls.process_images(job, images)   # ← long-lived synchronous path, not async_api

After queue_images_to_nats, MLPipelineJobType.run returns. run_job then logs Waiting for workers to pick up tasks for pipeline X (Y/Z online recently) via _log_worker_availability (ami/jobs/tasks.py:159–162) and exits. Verified live: a recent ~1.8k-image async_api collect saw the run_job Celery task succeed in ~146 s — that's the full collect + publish window, not the lifetime of the job. After that the slot is freed.

What I likely actually saw

The "run_job from 7+ hours earlier was still active in Celery (inspect showed it for hours after the cancel)" observation is most likely inspect.active() cache staleness, not a real long-lived task. Observed again today on a separate job: inspect.active() showed run_job as ACTIVE several minutes after the worker log line Task ami.jobs.tasks.run_job[…] succeeded in 146.12 s had already printed. So inspect.active() was lying, not run_job leaking.

What is still real (and still in scope for #1283)

• Acks for a deleted consumer kept arriving. Even after Cancel correctly tore down the NATS stream + consumer, queued process_nats_pipeline_result / create_detection_images Celery tasks from before the cancel kept executing and posting to $JS.ACK.job_<N>.job-<N>-consumer.…. Those land on a consumer/stream that no longer exists. The cleanup hook should drain or revoke any pending result-callback tasks tied to the cancelled job, not just the upstream run_job.
• Defensive run_job revoke on cancel is still worth doing in case the cancel arrives while run_job is mid-collect or mid-publish. It's a short window in steady state, but a big collect (S3 listing per deployment + filter_processed_images() DB scan) can take a couple minutes — long enough to want to interrupt cleanly on user cancel.

Retracted

The "Architectural concern" section in the previous version of this comment proposed splitting out a refactor to make run_job publish-and-exit. That work is not needed — it already does that. Apologies for the noise; leaving this correction here so the next reader doesn't chase a non-bug.

[mihow]

Also need to clean up if a Job is deleted from Django (via api or admin). Should only be allowed by superadmins

[mihow]

Claude says: Some implementation notes from looking at the existing cleanup paths. Posting these here so #1283 has a starting point.

Existing infra we can build on

There are already two pieces that do most of the work:

• cleanup_async_job_resources(job_id) in ami/ml/orchestration/jobs.py — clears Redis state via AsyncJobStateManager.cleanup() and deletes the NATS stream + consumer via TaskQueueManager.cleanup_job_resources(). Already idempotent: catches exceptions, logs, returns a success bool. Tolerates "stream/consumer doesn't exist".
• cleanup_async_job_if_needed(job) in ami/jobs/tasks.py — thin wrapper that checks dispatch_mode == ASYNC_API and calls the resource cleanup. Safe to call for any job.

So the cleanup primitive is solid. What's missing is a single canonical entry point on the Job model that bundles Celery revoke + status transition + cleanup, and consistent invocation of that entry point across every code path that ends or restarts a job.

Proposed helper

Add Job.terminate(status: str = JobState.REVOKED, reason: str | None = None) to ami/jobs/models.py. Idempotent. Roughly:

def terminate(self, status: str = JobState.REVOKED, reason: str | None = None) -> None:
    """
    Terminate a job: revoke the Celery task, transition to a final status,
    and clean up async (NATS/Redis) resources. Safe to call multiple times
    and safe to call on jobs that never used async resources.
    """
    if reason:
        self.logger.info(f"Terminating job {self.pk} ({status}): {reason}")

    if self.task_id:
        task = run_job.AsyncResult(self.task_id)
        if task:
            task.revoke(terminate=True)

    if self.status not in JobState.final_states():
        self.status = status
        self.finished_at = self.finished_at or timezone.now()
        self.save()

    cleanup_async_job_if_needed(self)

Job.cancel() then becomes a thin wrapper that does the CANCELING → REVOKED UI transition and delegates the actual termination to terminate().

Call sites that should funnel through it

Gap 1 in the issue (REVOKED via watchdog, no cleanup) and gap 2 (Start re-publishes on top of stale stream) both come from cleanup not being uniformly wired. After adding terminate(), every site below should call it:

Path	File	Currently
User clicks Cancel	ami/jobs/views.py cancel action → Job.cancel()	calls cleanup ✅
User deletes Job (superuser-only after #1284)	JobViewSet.perform_destroy	does not cancel/cleanup ❌
Watchdog reaps stalled job → REVOKED	ami/jobs/tasks.py _run_stale_jobs_check	calls cleanup ✅ but via separate code
task_failure signal	ami/jobs/tasks.py update_job_failure	calls cleanup ✅
task_postrun signal on REVOKED	ami/jobs/tasks.py update_job_status	calls cleanup ✅
_fail_job (FAILURE_THRESHOLD breach, lost-images reconciler)	ami/jobs/tasks.py:417	calls cleanup_async_job_resources directly
Admin tooling / management commands	various	inconsistent

The win from the helper is that adding a new way to end a job (a future "force fail this stuck job" admin action, a new watchdog rule, etc.) is one line.

Gap 2: cleanup before re-Start

For the "Start on a non-running job" path, the existing cleanup_async_job_if_needed helper does what we need — it's just not called before re-publish. Suggested location: at the top of the collect stage, when the job is being rehydrated from a non-running state. Roughly:

def enqueue(self):
    # If this job was previously REVOKED / FAILURE / CANCELLED, scrub any
    # leftover NATS stream + consumer + Redis state before re-publishing.
    # Idempotent: a no-op when there's nothing to clean.
    if self.status in JobState.failed_states() or self.progress.summary.status in JobState.failed_states():
        cleanup_async_job_if_needed(self)
    # ...existing enqueue logic...

The collect stage's stage params (the total / processed counters that visibly snap back to 0 in the UI) should also be reset to fresh defaults at the same point. Currently they're overwritten by the new collect, but the brief inconsistency is what produces the visible reset; resetting them explicitly first makes the transition deterministic.

Out of scope but worth flagging

The zombie_streams sub-check (_run_zombie_streams_check in jobs_health_check) is a janitor for streams whose owning job was never properly terminated. Once terminate() is wired in, the zombie-stream rate should drop. Worth keeping the janitor as a backstop, but its unfixable count is a useful signal of how many cleanup paths we still missed.

Suggested PR sequencing

1. Urgent permission-removal PR — fix(perms): restrict job deletion to superusers #1284: drop delete_job from MLDataManager so non-superusers can no longer delete. Stops the immediate bleed.
2. This issue, part 1: introduce Job.terminate(), refactor Job.cancel() and the watchdog path to use it, wire it into perform_destroy so superuser deletions cancel in-flight work first.
3. This issue, part 2: idempotent cleanup at the top of enqueue() for restart-from-non-running, reset stage params before re-publish.

(2) and (3) can be one PR or two — they're independent but both small.

[mihow]

Claude says: Partial fix has landed for gap 1 — flagging so the scope can narrow.

The zombie_streams sub-check inside jobs_health_check (ami/jobs/tasks.py:1210-1211) does delete the consumer + stream once a job is in final_states() (SUCCESS / FAILURE / REVOKED). So a watchdog-flipped REVOKED job will eventually have its NATS stream torn down without operator action.

Caveats that keep this issue alive:

1. Eventually, not synchronously. Beat tick is every 15 min (migration 0020_schedule_job_monitoring_beat_tasks.py) and the age guard is 60 min (Job.ZOMBIE_STREAMS_MAX_AGE_MINUTES = STALLED_JOBS_MAX_MINUTES * 6). So a stream can sit for up to ~75 min after the job goes terminal before it's drained.
2. Start-on-non-running race window is unaddressed. If a user clicks Start before the zombie drain fires (i.e. within that 60-min age window), collect re-publishes on top of the existing stream and the duplicate-task / counter-snap-to-0 behaviour described in the original report still happens. This is the more user-visible half of the bug.
3. REVOKED-via-manual-cancel path is the same. The drain only depends on the job being in final_states(), so a manual revoke is covered the same way as the watchdog path — eventually drained, but not synchronously, and not before a Start click can race it.

Suggested narrowed scope for this issue:

• On Start action for any non-running job: idempotently delete consumer + stream for that job ID before the new collect publishes. This closes the duplicate-tasks / counter-reset path independent of the beat schedule.
• On status transition into a final_states() value (any code path — watchdog, manual revoke, exception handler): trigger consumer + stream cleanup synchronously instead of waiting for the next beat tick. The async drain remains as the safety net.

The 60-min eventually-drained behaviour is fine as a backstop. The Start race + the synchronous-on-transition cleanup are the parts that still need code changes.

Observed today on prod (job 2543, 2026-04-30): a watchdog FAILURE → stream was drained on the next beat tick, no manual cleanup needed. So gap 1 is closed in practice for the FAILURE path; the Start-action race + counter reset are what's left.