feat(worker): detect stuck RUNNING jobs via progress heartbeat

alembic/versions/a7f3c2e9b104_add_progress_updated_at_to_job_runs.py

@@ -0,0 +1,44 @@
+"""Add progress_updated_at heartbeat column to job_runs
+
+Revision ID: a7f3c2e9b104
+Revises: 4dbf24ed1857
+Create Date: 2026-07-09 00:00:00.000000
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "a7f3c2e9b104"
+down_revision = "4dbf24ed1857"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "job_runs",
+        sa.Column(
+            "progress_updated_at",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=False,
+        ),
+    )
+    op.create_index(
+        "ix_job_runs_status_progress_updated",
+        "job_runs",
+        ["status", "progress_updated_at"],
+        unique=False,
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index("ix_job_runs_status_progress_updated", table_name="job_runs")
+    op.drop_column("job_runs", "progress_updated_at")
+    # ### end Alembic commands ###

src/mavedb/models/job_run.py

@@ -70,6 +70,14 @@ class JobRun(Base):
     progress_total: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
     progress_message: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)
 
+    # Liveness heartbeat. `onupdate=func.now()` bumps this on *any* commit that dirties the row,
+    # so every progress checkpoint (and every lifecycle transition) refreshes it with zero per-job
+    # code. The stalled-job sweeper treats a RUNNING job whose heartbeat has gone stale as wedged,
+    # which distinguishes a healthy long-running job (still committing progress) from a dead one.
+    progress_updated_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now()
+    )
+
     # Correlation for tracing
     correlation_id: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
 
@@ -98,6 +106,8 @@ class JobRun(Base):
         Index("ix_job_runs_created_at", "created_at"),
         Index("ix_job_runs_correlation_id", "correlation_id"),
         Index("ix_job_runs_status_scheduled", "status", "scheduled_at"),
+        # Supports the stalled-job sweeper's RUNNING heartbeat query (status + heartbeat staleness).
+        Index("ix_job_runs_status_progress_updated", "status", "progress_updated_at"),
         CheckConstraint(
             "status IN ('pending', 'queued', 'running', 'succeeded', 'failed', 'errored', 'cancelled', 'skipped')",
             name="ck_job_runs_status_valid",

src/mavedb/worker/jobs/registry.py

@@ -71,8 +71,10 @@
     cron(
         cleanup_stalled_jobs,
         name="cleanup_stalled_jobs_cron",
-        minute={15, 45},  # Run at :15 and :45 past each hour (every 30 minutes)
-        keep_result=timedelta(minutes=25).total_seconds(),
+        # Every 10 minutes. Paired with the 20-min progress-stall window this gives an effective
+        # stall-detection latency of ~20-30 min (vs. up to 50 min on the old 30-min cadence).
+        minute={5, 15, 25, 35, 45, 55},
+        keep_result=timedelta(minutes=9).total_seconds(),
     ),
 ]
 

src/mavedb/worker/jobs/system/cleanup.py

@@ -13,6 +13,7 @@
 The cleanup job acts as a safety net to ensure jobs don't remain in limbo forever.
 """
 
+import asyncio
 import logging
 from datetime import datetime, timedelta, timezone
 
@@ -37,14 +38,30 @@
 logger = logging.getLogger(__name__)
 
 # Timeout thresholds for detecting stalled jobs (in minutes).
-# RUNNING_TIMEOUT_MINUTES must stay below ArqWorkerSettings.job_timeout (currently 2 hours)
-# to avoid marking legitimately running jobs as stalled.
-RUNNING_TIMEOUT_MINUTES = 150  # RUNNING jobs should complete within 150 min (30 min buffer under ARQ timeout)
+#
+# RUNNING jobs are evaluated on two signals:
+#   1. PROGRESS_STALL_MINUTES — the primary signal. A healthy job checkpoints progress frequently
+#      (JobRun.progress_updated_at is bumped on every committing progress update via onupdate=now()).
+#      If the heartbeat has not advanced in this window the job is considered wedged, regardless of
+#      how long it has legitimately been running. This is what lets a multi-hour VEP fan-out run to
+#      completion without being killed, while still catching a job that has genuinely hung.
+#   2. RUNNING_TIMEOUT_MINUTES — a wall-clock backstop for jobs that never emit progress (or whose
+#      heartbeat column is somehow stuck). Sits ~30 min under ArqWorkerSettings.job_timeout (8h) so
+#      the DB-driven sweeper, not ARQ's hard kill, is what terminates a job — keeping DB state and
+#      actual work in sync.
+PROGRESS_STALL_MINUTES = 20  # RUNNING jobs should checkpoint progress at least this often
+RUNNING_TIMEOUT_MINUTES = 450  # Wall-clock backstop (7.5h), 30 min under the 8h ARQ ceiling
 PENDING_TIMEOUT_MINUTES = 5  # PENDING jobs which are actionable within pipelines should be enqueued within 5 minutes
 PIPELINE_STUCK_TIMEOUT_MINUTES = (
     5  # Pipelines in non-terminal states with no active jobs should resolve within 5 minutes
 )
 
+# How long the sweeper waits for an aborted RUNNING job to actually die before giving up. A job
+# cancelled at an await point dies within a poll cycle (sub-second); a job wedged in a blocking sync
+# call never honors the abort and will hit this timeout — in which case we do NOT re-drive it (see
+# _handle_stalled_running_job) to avoid running it twice.
+ABORT_CONFIRM_TIMEOUT_SECONDS = 15
+
 
 async def _handle_stalled_job_retry(
     job: JobRun,
@@ -193,19 +210,90 @@ async def _handle_stalled_job_retry(
         return False
 
 
+async def _running_job_confirmed_dead(arq_id: str, redis: ArqRedis) -> bool:
+    """Abort a RUNNING job's current ARQ attempt and report whether it is safe to re-drive.
+
+    RUNNING recovery differs from QUEUED/PENDING recovery: a RUNNING job may still have a live
+    coroutine holding a worker slot. Re-enqueuing without cancelling it first may mean two
+    coroutines are running the same job simultaneously, which is unsafe.
+
+    We therefore issue Job.abort() on the current attempt and interpret the outcome:
+
+    - abort() returns True  -> the coroutine was alive and has now been cancelled + recorded a
+      terminal result. Safe to re-drive.
+    - abort() returns False -> the job is absent from ARQ (crashed worker / never enqueued / already
+      finished). Nothing is running. Safe to re-drive — this is the common crash-recovery path.
+    - abort() raises TimeoutError -> the job is still present in ARQ and did not die within the
+      window: a coroutine wedged in a blocking sync call that cannot honor cooperative cancellation.
+      NOT safe to re-drive; leave it for the next sweep or ARQ's own ceiling.
+    - any other error -> treat as not-confirmably-alive and allow re-drive; a genuinely live job
+      would have either confirmed the abort or timed out, not errored here.
+
+    Returns True when it is safe to re-drive the job, False when the job may still be running.
+    """
+    try:
+        await ArqJob(arq_id, redis).abort(timeout=ABORT_CONFIRM_TIMEOUT_SECONDS)
+        return True
+    except asyncio.TimeoutError:
+        logger.warning(
+            f"Abort of RUNNING job attempt {arq_id} did not confirm within "
+            f"{ABORT_CONFIRM_TIMEOUT_SECONDS}s; job may be wedged in a blocking call — not re-driving"
+        )
+        return False
+    except Exception as e:
+        logger.warning(f"Abort of RUNNING job attempt {arq_id} raised {e!r}; treating as not running and re-driving")
+        return True
+
+
+async def _handle_stalled_running_job(
+    job: JobRun,
+    manager: JobManager,
+    redis: ArqRedis,
+    stall_reason: str,
+    db: Session,
+) -> bool:
+    """Recover a stalled RUNNING job via abort-first ordering.
+
+    Aborts the current attempt (urn#N) BEFORE any retry bookkeeping, and only proceeds to re-drive
+    (which increments retry_count and enqueues urn#N+1) once the abort confirms the original is dead
+    or absent. If the original refuses to die, we leave it alone rather than risk a double-run.
+
+    Returns True if the sweeper acted on the job (re-drove or permanently failed it), False if it was
+    left in place because the abort could not be confirmed (so it should not be counted as cleaned).
+    """
+    arq_id = arq_job_id(job)
+
+    # Could not confirm the coroutine is dead. Leave the job RUNNING; the next sweep will retry
+    # the abort, and ARQ's job_timeout ceiling remains as a final backstop.
+    if not await _running_job_confirmed_dead(arq_id, redis):
+        logger.warning(
+            f"Leaving stalled RUNNING job {job.urn} in place — abort unconfirmed, re-driving would risk a double-run",
+            extra=manager.logging_context(),
+        )
+        return False
+
+    # Abort confirmed dead (or the job was absent from ARQ). Refresh our view — if the aborted
+    # coroutine's own lifecycle handler already marked the row terminal, we want to see that — then
+    # re-drive through the shared retry handler, which preserves pipeline dependency semantics.
+    db.refresh(job)
+    await _handle_stalled_job_retry(job, manager, redis, stall_reason, db)
+    return True
+
+
 @with_guaranteed_job_run_record("cron_job")
 @with_job_management
 async def cleanup_stalled_jobs(ctx: dict, job_id: int, job_manager: JobManager) -> JobExecutionOutcome:
     """Detect and handle jobs that have stalled in intermediate states.
 
-    This job runs periodically (every 15 minutes) to find jobs that have been
-    stuck in QUEUED, RUNNING, or PENDING states beyond reasonable timeouts
-    and handles them appropriately.
+    This job runs periodically to find jobs that have been stuck in QUEUED, RUNNING, or PENDING
+    states beyond reasonable timeouts and handles them appropriately.
 
     Stalled job detection criteria:
     - QUEUED: Present in DB as QUEUED but absent from ARQ's Redis queue
       (process crashed between prepare_queue and redis.enqueue_job)
-    - RUNNING: Started > 60 minutes ago but not finished (worker likely crashed)
+    - RUNNING: Progress heartbeat idle > PROGRESS_STALL_MINUTES, or wall-clock runtime past the
+      RUNNING_TIMEOUT_MINUTES backstop, and not finished. Recovered abort-first (see
+      _handle_stalled_running_job) so a wedged attempt is never re-driven alongside its own retry.
     - PENDING: Created > 5 minutes ago in a pipeline and currently runnable
       (coordination failure)
     - Pipeline stuck: Non-terminal pipeline with no active jobs older than 5 minutes
@@ -230,7 +318,8 @@ async def cleanup_stalled_jobs(ctx: dict, job_id: int, job_manager: JobManager)
 
         Job stalled in RUNNING (worker crash):
         - Worker started job, marked it RUNNING, then crashed
-        - After 60 minutes (longer than ARQ timeout), janitor detects and retries
+        - Once the progress heartbeat goes stale (or the wall-clock backstop is hit), the janitor
+          aborts the dead attempt and re-drives it
     """
     job_manager.save_to_context(
         {
@@ -305,14 +394,19 @@ async def cleanup_stalled_jobs(ctx: dict, job_id: int, job_manager: JobManager)
     job_manager.save_to_context({"cleaned_queued_jobs": queued_jobs})
     logger.debug("Completed cleaning stalled QUEUED jobs.", extra=job_manager.logging_context())
 
-    # Find RUNNING jobs that have been running too long OR have missing started_at
-    # These likely indicate worker crashes (worker died mid-execution) or data inconsistencies
+    # Find RUNNING jobs that are wedged. A job is a candidate when EITHER its progress heartbeat has
+    # gone stale (the primary, runtime-agnostic signal) OR it has blown past the wall-clock backstop
+    # (or is missing started_at entirely — a data inconsistency). Healthy long-running jobs keep
+    # bumping progress_updated_at and so are never selected until the backstop.
+    progress_stall_threshold = now - timedelta(minutes=PROGRESS_STALL_MINUTES)
     running_threshold = now - timedelta(minutes=RUNNING_TIMEOUT_MINUTES)
     running_jobs = job_manager.db.scalars(
         select(JobRun).where(
             JobRun.status == JobStatus.RUNNING,
-            (JobRun.started_at < running_threshold) | (JobRun.started_at.is_(None)),
             JobRun.finished_at.is_(None),
+            (JobRun.progress_updated_at < progress_stall_threshold)
+            | (JobRun.started_at < running_threshold)
+            | (JobRun.started_at.is_(None)),
         )
     ).all()
 
@@ -333,18 +427,26 @@ async def cleanup_stalled_jobs(ctx: dict, job_id: int, job_manager: JobManager)
             continue
 
         elapsed_minutes = (now - job.started_at).total_seconds() / 60
+        heartbeat_stale = job.progress_updated_at < progress_stall_threshold
+        if heartbeat_stale:
+            heartbeat_age = (now - job.progress_updated_at).total_seconds() / 60
+            stall_reason = (
+                f"Job stalled in RUNNING state: progress heartbeat idle for {heartbeat_age:.1f} minutes "
+                f"(running for {elapsed_minutes:.1f} minutes)"
+            )
+        else:
+            stall_reason = (
+                f"Job stalled in RUNNING state for {elapsed_minutes:.1f} minutes "
+                f"(exceeded {RUNNING_TIMEOUT_MINUTES} min backstop; likely worker crash)"
+            )
 
-        logger.warning(
-            f"Detected stalled RUNNING job {job.urn} "
-            f"(started {job.started_at}, running for {elapsed_minutes:.1f} minutes)",
-            extra=manager.logging_context(),
-        )
+        logger.warning(f"Detected stalled RUNNING job {job.urn}: {stall_reason}", extra=manager.logging_context())
 
-        stall_reason = f"Job stalled in RUNNING state for {elapsed_minutes:.1f} minutes (likely worker crash)"
-        await _handle_stalled_job_retry(job, manager, job_manager.redis, stall_reason, job_manager.db)
+        acted = await _handle_stalled_running_job(job, manager, job_manager.redis, stall_reason, job_manager.db)
 
         manager.db.commit()
-        cleaned_jobs["running"].append(job.urn)
+        if acted:
+            cleaned_jobs["running"].append(job.urn)
 
     job_manager.save_to_context({"cleaned_running_jobs": running_jobs})
     logger.debug("Completed cleaning stalled RUNNING jobs.", extra=job_manager.logging_context())

src/mavedb/worker/lib/decorators/job_management.py

@@ -5,6 +5,7 @@
 Includes JobManager injection for advanced operations and robust error handling.
 """
 
+import asyncio
 import functools
 import inspect
 import logging
@@ -15,7 +16,7 @@
 
 from mavedb.lib.slack import send_slack_error, send_slack_job_error, send_slack_job_failure
 from mavedb.lib.types.workflow import JobExecutionOutcome
-from mavedb.models.enums.job_pipeline import JobStatus
+from mavedb.models.enums.job_pipeline import FailureCategory, JobStatus
 from mavedb.worker.lib.decorators.utils import ensure_ctx, ensure_job_id, ensure_session_ctx, is_test_mode
 from mavedb.worker.lib.managers import JobManager
 from mavedb.worker.lib.managers.constants import TERMINAL_JOB_STATUSES
@@ -143,6 +144,43 @@ async def _execute_managed_job(func: Callable[..., Awaitable[JobExecutionOutcome
 
         return result
 
+    # The coroutine is being cancelled — either ARQ hit job_timeout, or cleanup_stalled_jobs
+    # aborted this attempt (Job.abort()) as part of stall recovery. CancelledError is a
+    # BaseException, so the `except Exception` handler below never sees it; without this branch
+    # the JobRun row would be orphaned as RUNNING forever while the coroutine dies.
+    #
+    # We mark the row terminally (FAILED / TIMEOUT) so DB state matches reality, then RE-RAISE.
+    # Re-raising is essential: it lets ARQ record the cancellation as the job's result, which is
+    # what makes a concurrent Job.abort() confirm the job actually died.
+    except asyncio.CancelledError:
+        logger.warning(f"Job {job_id} cancelled (timeout or stall-recovery abort); marking FAILED before re-raising")
+        try:
+            db_session.rollback()
+            if job_manager is not None:
+                job_manager.fail_job(
+                    result=JobExecutionOutcome.failed(
+                        reason="Job cancelled due to timeout, stall-recovery, or internal ARQ abort",
+                        data={"reason": "cancelled"},
+                        failure_category=FailureCategory.TIMEOUT,
+                    )
+                )
+                db_session.commit()
+                job = job_manager.get_job()
+                send_slack_job_failure(
+                    job_urn=job.urn,
+                    job_function=job.job_function,
+                    reason="Job cancelled due to timeout.",
+                    failure_category=str(FailureCategory.TIMEOUT),
+                    retry_count=job.retry_count,
+                    max_retries=job.max_retries,
+                    will_retry=False,
+                )
+
+        except Exception as inner_e:
+            logger.critical(f"Failed to mark cancelled job {job_id} as failed: {inner_e}")
+            send_slack_error(inner_e)
+        raise
+
     except Exception as e:
         # Prioritize salvaging lifecycle state
         will_retry = False

src/mavedb/worker/settings/worker.py

@@ -25,7 +25,13 @@
 # driver, enabling incremental migration of job functions without touching
 # the FastAPI layer. Once all jobs use async sessions, raise MAX_JOBS to 10+.
 MAX_JOBS = 2
-JOB_TIMEOUT_SECONDS = 3 * 60 * 60  # 3 hours — matches RUNNING_TIMEOUT_MINUTES (150 min) with buffer
+# ARQ's hard coroutine-kill ceiling. Kept deliberately high (8h) so ARQ is the *last* resort:
+# our own DB-driven sweeper (cleanup_stalled_jobs) should detect and recover stalls long before
+# this fires, keeping the JobRun state machine in sync with reality. VEP fan-out over large allele
+# sets can legitimately run for hours, so a low ceiling would kill healthy work.
+# cleanup's RUNNING_TIMEOUT_MINUTES backstop sits ~30 min under this to try and sweep these jobs before
+# ARQ's hard timeout fires.
+JOB_TIMEOUT_SECONDS = 8 * 60 * 60  # 8 hours
 
 
 class ArqWorkerSettings:
@@ -44,3 +50,7 @@ class ArqWorkerSettings:
 
     max_jobs = MAX_JOBS
     job_timeout = JOB_TIMEOUT_SECONDS
+    # Required for cleanup_stalled_jobs to abort a wedged RUNNING job's coroutine (via Job.abort())
+    # before re-driving it. Without this, ARQ ignores abort requests and the sweeper cannot safely
+    # confirm a stalled job is dead before retrying it.
+    allow_abort_jobs = True