Redesign stuck-RUNNING-job detection: progress-heartbeat + abort-first recovery

[bencap]

Summary

Replace the wall-clock-only detection of stuck RUNNING jobs in cleanup_stalled_jobs with a progress-staleness heartbeat, make the DB job machinery authoritative over job termination via ARQ abort, and demote the ARQ job_timeout to a genuine last resort. This also closes two latent correctness bugs (a job double-run on retry, and a JobRun row orphaned as RUNNING when ARQ cancels a timed-out job).

Problem

cleanup_stalled_jobs flags a RUNNING job as stalled purely on total wall-clock since started_at (RUNNING_TIMEOUT_MINUTES = 150, sitting under JOB_TIMEOUT_SECONDS = 3h). A wall-clock timeout cannot distinguish a healthy job that is simply long-running from a wedged job. One threshold, two failure modes: set it low and healthy long jobs get killed/retried; set it high and genuine stalls sit for hours.

This now bites in practice: VEP annotation (populate_vep_for_score_set, which batches HGVS strings to Ensembl VEP and Variant Recoder) fans out with allele count and can legitimately run past the 150-minute timeout.

Two bugs fall out of the current behavior and are resolved by this work:

Bug 1 — Job double-run on stall retry

When a RUNNING job crosses RUNNING_TIMEOUT_MINUTES, _handle_stalled_job_retry re-enqueues the job without cancelling the still-running coroutine — it assumes the worker has already crashed. arq_job_id embeds retry_count (f"{urn}#{retry_count}"), so the retry (urn#1) does not dedupe against the still-live original (urn#0); a fresh attempt starts from scratch while the original keeps running to ARQ's job_timeout. With max_jobs = 2, a single logical job can occupy both worker slots running itself twice.

• Steps to reproduce: enqueue a job that stays healthy but runs longer than RUNNING_TIMEOUT_MINUTES (e.g. a large VEP score set). Let cleanup_stalled_jobs run against it.
• Expected: the job is either left alone (still making progress) or terminated-then-restarted exactly once.
• Observed: a second attempt starts ("job starts over") while the first continues executing until the ARQ job_timeout cancels it; both run concurrently for up to the gap between the two timeouts.
• Low severity (only triggers on slow-but-not-crashed jobs) but incorrect.

Bug 2 — JobRun orphaned as RUNNING on ARQ timeout

ARQ enforces job_timeout by cancelling the job coroutine, which raises asyncio.CancelledError — a BaseException, not an Exception. The with_job_management recovery handler catches except Exception, so it misses the cancellation entirely and never writes a terminal status. The JobRun row is left orphaned in RUNNING.

• Expected: a job cancelled by the ARQ ceiling ends in a terminal DB state (FAILED, category TIMEOUT).
• Observed: the row stays RUNNING indefinitely.

Proposed behavior

Adopt progress-staleness detection with abort-first recovery:

• A RUNNING job is considered alive while it keeps committing to its row; the existing per-batch progress checkpoints supply that heartbeat for free.
• Detection keys off heartbeat staleness (fast) plus a high wall-clock backstop (catch-all), never on total runtime alone.
• Before re-driving a suspected-stalled job, cleanup aborts the current attempt and only re-enqueues once the abort is confirmed, so the DB machinery — not ARQ's ceiling — owns termination.
• ARQ job_timeout remains only as a last-resort ceiling, and when it does fire the DB row is now updated rather than orphaned.

Acceptance criteria

• with_job_management catches asyncio.CancelledError (or BaseException cancellation), marks the JobRun FAILED with FailureCategory.TIMEOUT, commits, and re-raises. A job cancelled by the ARQ ceiling ends in a terminal DB state, never orphaned as RUNNING.
• JobRun has a progress_updated_at column that is automatically bumped on any row-mutating commit (no per-job code changes), backed by an index on (status, progress_updated_at), with an Alembic migration.
• Calling update_progress (and the other progress-checkpoint paths) advances progress_updated_at.
• Worker settings: JOB_TIMEOUT_SECONDS raised to 8 hours; allow_abort_jobs = True; stale comments referencing the old 150-min / 3h relationship corrected.
• cleanup_stalled_jobs uses PROGRESS_STALL_MINUTES = 20 for heartbeat staleness and RUNNING_TIMEOUT_MINUTES = 450 as the wall-clock backstop (30-minute buffer under the 8h ceiling).
• RUNNING-branch recovery aborts the current attempt (arq_job_id(job) at urn#N) before prepare_retry mutates retry_count, confirms the abort, then calls prepare_retry and enqueues urn#(N+1).
• If the abort is not confirmed, cleanup does not re-drive the job; it leaves it for the wall-clock backstop / ARQ ceiling (avoiding the double-run).
• A healthy long-running job whose heartbeat is fresh but whose started_at is old (the VEP case) is left untouched.
• A stalled job with a confirmed abort is re-driven exactly once (no concurrent second attempt).
• cleanup_stalled_jobs docstring updated to describe heartbeat + abort-first semantics.
• API tests cover: decorator marks FAILED on cancellation; progress_updated_at bumps on update_progress; stale heartbeat + confirmed abort → single re-drive; fresh heartbeat + old started_at → left alone; job past the 450-min backstop → handled.

Implementation notes

Affected files in mavedb-api:

• src/mavedb/worker/lib/decorators/job_management.py — add the CancelledError handler (mark FAILED/TIMEOUT, commit, re-raise).
• src/mavedb/models/job_run.py — add progress_updated_at (onupdate=func.now(), server_default=func.now()) and the composite index.
• src/mavedb/worker/settings/worker.py — JOB_TIMEOUT_SECONDS = 8h, allow_abort_jobs = True, comment fixes.
• src/mavedb/worker/jobs/system/cleanup.py — new thresholds, rewritten RUNNING branch with abort-first ordering.
• src/mavedb/worker/lib/managers/utils.py — arq_job_id is the reason a naive re-enqueue double-runs; the abort must target the pre-prepare_retry id.
• src/mavedb/worker/jobs/external_services/vep.py — context for the motivating workload (the per-batch progress checkpoints are the heartbeat source).

Tradeoff — cooperative cancellation: CancelledError only lands at an await. VEP work is awaited, so aborts land promptly; a coroutine blocked in a synchronous call will not honor the abort (nor can ARQ's own timeout), which is why we re-drive only on a confirmed abort. Worker throughput is not a concern (low upload volume), so a rare sync-wedged job holding a slot until the ceiling is acceptable.

Checks to resolve during implementation:

• Exact arq Job.abort() semantics and confirmation/return behavior; verify allow_abort_jobs = True has no unwanted side effects on normal completion.
• Whether VEP / Variant Recoder / variant-creation writes are idempotent (determines whether the historical double-run merely wasted compute or could have raced into duplicate/partial rows).
• Cron cadence: heartbeat window is 20 min but the cleanup cron runs every 30 min → effective detection latency ~30–50 min; tighten cadence if faster recovery is required.
• Open value: worst-case legitimate VEP runtime, which sizes both the 8h ceiling and the 20-min stall window.