diff --git a/alembic/versions/a7f3c2e9b104_add_progress_updated_at_to_job_runs.py b/alembic/versions/a7f3c2e9b104_add_progress_updated_at_to_job_runs.py new file mode 100644 index 00000000..d04ceab9 --- /dev/null +++ b/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 ### diff --git a/src/mavedb/models/job_run.py b/src/mavedb/models/job_run.py index 877eeab0..b079a175 100644 --- a/src/mavedb/models/job_run.py +++ b/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", diff --git a/src/mavedb/worker/jobs/registry.py b/src/mavedb/worker/jobs/registry.py index 94cd76e3..ba9754a0 100644 --- a/src/mavedb/worker/jobs/registry.py +++ b/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(), ), ] diff --git a/src/mavedb/worker/jobs/system/cleanup.py b/src/mavedb/worker/jobs/system/cleanup.py index 89089828..4a2dc975 100644 --- a/src/mavedb/worker/jobs/system/cleanup.py +++ b/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()) diff --git a/src/mavedb/worker/lib/decorators/job_management.py b/src/mavedb/worker/lib/decorators/job_management.py index 58aad427..0dcb463a 100644 --- a/src/mavedb/worker/lib/decorators/job_management.py +++ b/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 diff --git a/src/mavedb/worker/settings/worker.py b/src/mavedb/worker/settings/worker.py index 6a558e6b..3b0ab4e8 100644 --- a/src/mavedb/worker/settings/worker.py +++ b/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 diff --git a/tests/worker/jobs/system/test_cleanup.py b/tests/worker/jobs/system/test_cleanup.py index 522b6167..104c3af6 100644 --- a/tests/worker/jobs/system/test_cleanup.py +++ b/tests/worker/jobs/system/test_cleanup.py @@ -12,6 +12,7 @@ pytest.importorskip("arq") # Skip tests if arq is not installed +import asyncio from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, patch @@ -29,6 +30,7 @@ from mavedb.worker.jobs.system.cleanup import ( PENDING_TIMEOUT_MINUTES, PIPELINE_STUCK_TIMEOUT_MINUTES, + PROGRESS_STALL_MINUTES, RUNNING_TIMEOUT_MINUTES, cleanup_stalled_jobs, ) @@ -40,17 +42,55 @@ @pytest.fixture def mock_arq_job_not_found(): - """Mock ArqJob.status() to return not_found, simulating a crashed-enqueue QUEUED job.""" + """Mock ArqJob for a crashed-enqueue QUEUED job: status() is not_found, and abort() confirms dead.""" with patch("mavedb.worker.jobs.system.cleanup.ArqJob") as mock_arq_job: mock_arq_job.return_value.status = AsyncMock(return_value=ArqJobStatus.not_found) + mock_arq_job.return_value.abort = AsyncMock(return_value=True) yield mock_arq_job @pytest.fixture def mock_arq_job_in_redis(): - """Mock ArqJob.status() to return queued, simulating a legitimately queued job in ARQ Redis.""" + """Mock ArqJob for a legitimately queued job: status() is queued, and abort() confirms dead.""" with patch("mavedb.worker.jobs.system.cleanup.ArqJob") as mock_arq_job: mock_arq_job.return_value.status = AsyncMock(return_value=ArqJobStatus.queued) + mock_arq_job.return_value.abort = AsyncMock(return_value=True) + yield mock_arq_job + + +@pytest.fixture +def mock_arq_abort_confirmed(): + """Mock ArqJob.abort() to return True — the stalled RUNNING attempt was alive and is now dead. + + This is the signal that makes it safe for the sweeper to re-drive a stalled RUNNING job. + """ + with patch("mavedb.worker.jobs.system.cleanup.ArqJob") as mock_arq_job: + mock_arq_job.return_value.status = AsyncMock(return_value=ArqJobStatus.not_found) + mock_arq_job.return_value.abort = AsyncMock(return_value=True) + yield mock_arq_job + + +@pytest.fixture +def mock_arq_abort_absent(): + """Mock ArqJob.abort() to return False — the job is absent from ARQ (crashed / never enqueued). + + Nothing is running, so it is still safe to re-drive. + """ + with patch("mavedb.worker.jobs.system.cleanup.ArqJob") as mock_arq_job: + mock_arq_job.return_value.status = AsyncMock(return_value=ArqJobStatus.not_found) + mock_arq_job.return_value.abort = AsyncMock(return_value=False) + yield mock_arq_job + + +@pytest.fixture +def mock_arq_abort_wedged(): + """Mock ArqJob.abort() to time out — the coroutine is wedged in a blocking call and won't die. + + The sweeper must NOT re-drive such a job (it could still be running), to avoid a double-run. + """ + with patch("mavedb.worker.jobs.system.cleanup.ArqJob") as mock_arq_job: + mock_arq_job.return_value.status = AsyncMock(return_value=ArqJobStatus.in_progress) + mock_arq_job.return_value.abort = AsyncMock(side_effect=asyncio.TimeoutError()) yield mock_arq_job @@ -148,7 +188,7 @@ async def test_cleanup_stalled_queued_job_max_retries_reached( assert "stalled" in stalled_job.error_message.lower() async def test_cleanup_stalled_running_job_with_retries( - self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job + self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job, mock_arq_abort_confirmed ): """Test cleanup of a stalled RUNNING job with retries remaining.""" # Create a stalled RUNNING job in the database @@ -184,7 +224,7 @@ async def test_cleanup_stalled_running_job_with_retries( assert stalled_job.finished_at is None async def test_cleanup_stalled_running_job_max_retries_reached( - self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job + self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job, mock_arq_abort_confirmed ): """Test cleanup of a stalled RUNNING job with max retries reached.""" # Create a stalled RUNNING job with max retries @@ -257,6 +297,115 @@ async def test_cleanup_stalled_running_job_missing_started_at( assert stalled_job.status == JobStatus.RUNNING assert stalled_job.retry_count == 0 + async def test_cleanup_detects_running_job_by_stale_heartbeat( + self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job, mock_arq_abort_confirmed + ): + """A RUNNING job with a recent started_at but a stale progress heartbeat is detected and re-driven. + + This is the primary staleness signal: staleness is judged by progress_updated_at, not total runtime. + """ + stalled_job = JobRun( + job_type="test_job", + job_function="test_function", + status=JobStatus.RUNNING, + created_at=datetime.now(timezone.utc) - timedelta(minutes=40), + started_at=datetime.now(timezone.utc) - timedelta(minutes=30), # well within the wall-clock backstop + progress_updated_at=datetime.now(timezone.utc) - timedelta(minutes=PROGRESS_STALL_MINUTES + 5), + finished_at=None, + max_retries=3, + retry_count=0, + job_params={}, + ) + session.add(stalled_job) + session.commit() + + result = await cleanup_stalled_jobs( + mock_worker_ctx, None, JobManager(session, mock_worker_ctx["redis"], sample_cleanup_job_run.id) + ) + + # Abort of the wedged attempt was confirmed, so the job was re-driven. + mock_arq_abort_confirmed.return_value.abort.assert_awaited_once() + mock_worker_ctx["redis"].enqueue_job.assert_called_once() + + assert result.data["total_cleaned"] == 1 + assert stalled_job.urn in result.data["running_jobs"] + + session.refresh(stalled_job) + assert stalled_job.status == JobStatus.QUEUED + assert stalled_job.retry_count == 1 + + async def test_cleanup_leaves_running_job_with_fresh_heartbeat( + self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job + ): + """A RUNNING job that has been running a long time but is still checkpointing progress is left alone. + + This is the VEP fan-out case: legitimate multi-hour runtime, well past the old 150-min limit, + but under the wall-clock backstop and with a fresh heartbeat — must NOT be treated as stalled. + """ + healthy_job = JobRun( + job_type="test_job", + job_function="test_function", + status=JobStatus.RUNNING, + created_at=datetime.now(timezone.utc) - timedelta(minutes=310), + started_at=datetime.now(timezone.utc) - timedelta(minutes=300), # 5h — past old limit, under backstop + progress_updated_at=datetime.now(timezone.utc) - timedelta(minutes=2), # fresh heartbeat + finished_at=None, + max_retries=3, + retry_count=0, + job_params={}, + ) + session.add(healthy_job) + session.commit() + + result = await cleanup_stalled_jobs( + mock_worker_ctx, None, JobManager(session, mock_worker_ctx["redis"], sample_cleanup_job_run.id) + ) + + assert result.data["total_cleaned"] == 0 + assert healthy_job.urn not in result.data["running_jobs"] + + session.refresh(healthy_job) + assert healthy_job.status == JobStatus.RUNNING + assert healthy_job.retry_count == 0 + + async def test_cleanup_does_not_redrive_wedged_running_job( + self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job, mock_arq_abort_wedged + ): + """A stalled RUNNING job whose abort does not confirm is left in place, not re-driven. + + Re-driving a job that may still be executing (wedged in a blocking sync call) would recreate + the historical double-run, so the sweeper takes no action and leaves it for a later sweep. + """ + wedged_job = JobRun( + job_type="test_job", + job_function="test_function", + status=JobStatus.RUNNING, + created_at=datetime.now(timezone.utc) - timedelta(minutes=40), + started_at=datetime.now(timezone.utc) - timedelta(minutes=30), + progress_updated_at=datetime.now(timezone.utc) - timedelta(minutes=PROGRESS_STALL_MINUTES + 5), + finished_at=None, + max_retries=3, + retry_count=0, + job_params={}, + ) + session.add(wedged_job) + session.commit() + + result = await cleanup_stalled_jobs( + mock_worker_ctx, None, JobManager(session, mock_worker_ctx["redis"], sample_cleanup_job_run.id) + ) + + # Abort was attempted but never confirmed, so no retry was enqueued. + mock_arq_abort_wedged.return_value.abort.assert_awaited_once() + mock_worker_ctx["redis"].enqueue_job.assert_not_called() + + assert result.data["total_cleaned"] == 0 + assert wedged_job.urn not in result.data["running_jobs"] + + session.refresh(wedged_job) + assert wedged_job.status == JobStatus.RUNNING + assert wedged_job.retry_count == 0 + async def test_cleanup_stalled_pending_job_with_retries( self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job ): @@ -477,7 +626,7 @@ async def test_cleanup_stalled_queued_standalone_job_enqueue_failure( assert "Failed to enqueue after stall recovery" in stalled_job.error_message async def test_cleanup_stalled_running_standalone_job_enqueue_failure( - self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job + self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job, mock_arq_abort_confirmed ): """Test that stalled standalone RUNNING job is marked FAILED if ARQ enqueue fails.""" @@ -559,7 +708,7 @@ async def test_cleanup_stalled_queued_pipeline_job_dependencies_satisfied( assert stalled_job.retry_count == 1 async def test_cleanup_stalled_running_pipeline_job_dependencies_satisfied( - self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job + self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job, mock_arq_abort_confirmed ): """Test that stalled pipeline RUNNING job with satisfied dependencies is enqueued.""" # Create a pipeline with all dependencies satisfied @@ -735,7 +884,7 @@ async def test_cleanup_stalled_queued_pipeline_job_dependencies_not_ready( assert stalled_job.retry_count == 1 async def test_cleanup_stalled_running_pipeline_job_dependencies_failed( - self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job + self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job, mock_arq_abort_confirmed ): """Test that stalled pipeline RUNNING job with failed dependencies is skipped.""" # Create a pipeline @@ -867,7 +1016,7 @@ async def test_cleanup_stalled_pending_pipeline_job_dependencies_failed( assert stalled_job.retry_count == 0 async def test_cleanup_stalled_running_pipeline_job_dependencies_not_ready( - self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job + self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job, mock_arq_abort_confirmed ): """Test that stalled pipeline RUNNING job with dependencies not ready is skipped.""" # Create a pipeline @@ -1249,7 +1398,7 @@ async def test_cleanup_sends_slack_when_max_retries_reached_queued_job( assert call_kwargs["max_retries"] == 3 async def test_cleanup_sends_slack_when_max_retries_reached_running_job( - self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job + self, session, mock_worker_ctx, sample_cleanup_job_run, with_cleanup_job, mock_arq_abort_confirmed ): """Test that send_slack_job_failure is called when sweeper permanently fails a stalled RUNNING job.""" stalled_job = JobRun( diff --git a/tests/worker/lib/decorators/test_job_management.py b/tests/worker/lib/decorators/test_job_management.py index dd2e10ef..fbfb370d 100644 --- a/tests/worker/lib/decorators/test_job_management.py +++ b/tests/worker/lib/decorators/test_job_management.py @@ -16,7 +16,7 @@ from sqlalchemy import select 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.models.job_run import JobRun from mavedb.worker.lib.decorators.job_management import with_job_management from mavedb.worker.lib.managers.exceptions import JobStateError @@ -133,6 +133,67 @@ async def sample_error(ctx: dict, job_id: int, job_manager: JobManager): mock_start_job.assert_called_once() mock_error_job.assert_called_once() + async def test_decorator_marks_failed_and_reraises_on_cancellation( + self, session, mock_worker_ctx, mock_job_manager + ): + """A CancelledError (ARQ timeout or stall-recovery abort) must mark the row FAILED then re-raise. + + Re-raising is what lets ARQ record the cancellation as the job result so a concurrent + Job.abort() can confirm the job actually died. Without this handler the row would be + orphaned as RUNNING (CancelledError is a BaseException, not caught by `except Exception`). + """ + + @with_job_management + async def sample_cancel(ctx: dict, job_id: int, job_manager: JobManager): + raise asyncio.CancelledError() + + with ( + patch("mavedb.worker.lib.decorators.job_management.JobManager") as mock_job_manager_class, + patch("mavedb.worker.lib.decorators.job_management.send_slack_job_failure") as mock_send_slack_job_failure, + patch.object(mock_job_manager, "start_job", return_value=None) as mock_start_job, + patch.object(mock_job_manager, "fail_job", return_value=None) as mock_fail_job, + ): + mock_job_manager_class.return_value = mock_job_manager + + with pytest.raises(asyncio.CancelledError): + await sample_cancel(mock_worker_ctx, 999) + + mock_start_job.assert_called_once() + mock_fail_job.assert_called_once() + # The row is failed with a TIMEOUT category. + outcome = mock_fail_job.call_args.kwargs["result"] + assert outcome.status == JobStatus.FAILED + assert outcome.failure_category == FailureCategory.TIMEOUT + # A permanent-failure alert is emitted for the cancellation. + mock_send_slack_job_failure.assert_called_once() + + async def test_decorator_reraises_cancellation_even_if_marking_failed_errors( + self, session, mock_worker_ctx, mock_job_manager + ): + """If marking the cancelled job FAILED itself raises, the decorator still re-raises CancelledError. + + The salvage path must never swallow the cancellation — ARQ needs it to record the result. + """ + + @with_job_management + async def sample_cancel(ctx: dict, job_id: int, job_manager: JobManager): + raise asyncio.CancelledError() + + with ( + patch("mavedb.worker.lib.decorators.job_management.JobManager") as mock_job_manager_class, + patch("mavedb.worker.lib.decorators.job_management.send_slack_error") as mock_send_slack_error, + patch("mavedb.worker.lib.decorators.job_management.send_slack_job_failure"), + patch.object(mock_job_manager, "start_job", return_value=None), + patch.object(mock_job_manager, "fail_job", side_effect=RuntimeError("db down")), + ): + mock_job_manager_class.return_value = mock_job_manager + + with pytest.raises(asyncio.CancelledError): + await sample_cancel(mock_worker_ctx, 999) + + # The inner failure is reported, and the original cancellation still propagates. + mock_send_slack_error.assert_called_once() + async def test_decorator_calls_start_job_and_skip_job_when_wrapped_function_returns_skipped_status( self, session, mock_worker_ctx, mock_job_manager ): diff --git a/tests/worker/lib/managers/test_job_manager.py b/tests/worker/lib/managers/test_job_manager.py index 980a30f0..84f709bf 100644 --- a/tests/worker/lib/managers/test_job_manager.py +++ b/tests/worker/lib/managers/test_job_manager.py @@ -11,6 +11,7 @@ pytest.importorskip("arq") import re +from datetime import datetime, timedelta, timezone from unittest.mock import Mock, PropertyMock, patch from arq import ArqRedis @@ -1348,6 +1349,28 @@ def test_update_progress_success(self, session, arq_redis, with_populated_job_da assert job.progress_total == 100 assert job.progress_message == "Halfway done" + def test_update_progress_bumps_progress_updated_at_heartbeat( + self, session, arq_redis, with_populated_job_data, sample_job_run + ): + """Committing a progress update refreshes progress_updated_at via onupdate=now(). + + This is the liveness heartbeat the stalled-job sweeper reads: it must advance on every + committing progress checkpoint with no explicit code in the job. + """ + + manager = JobManager(session, arq_redis, sample_job_run.id) + + # Pin the heartbeat to a stale value so the bump is unambiguous. + stale = datetime.now(timezone.utc) - timedelta(hours=1) + job = session.execute(select(JobRun).where(JobRun.id == sample_job_run.id)).scalar_one() + job.progress_updated_at = stale + session.commit() + + manager.update_progress(10, 100, "tick", commit=True) # commit triggers the UPDATE + onupdate + + job = session.execute(select(JobRun).where(JobRun.id == sample_job_run.id)).scalar_one() + assert job.progress_updated_at > stale + def test_update_progress_success_does_not_overwrite_old_message_when_no_new_message_is_provided( self, session, arq_redis, with_populated_job_data, sample_job_run ):