From c87f8a1bb35a44549d43781ff8507aa28b0c968f Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Mon, 13 Jul 2026 21:41:34 -0400 Subject: [PATCH] refactor(maintenance): DRY the recover_stalled head-run twins (#161) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit recover_stalled_head_training_runs and recover_stalled_head_auto_apply_runs were near-exact copies (coalesce-flip + keep-last-N prune, differing only in model + two constants). Extract _recover_stalled_runs(model, stall_minutes, keep_runs, label); the two tasks become thin wrappers. The other two recover tasks are deliberately NOT folded in (library-audit has no prune tail; backup uses a single started_at cutoff). test_recover_stalled_head_runs.py covers both wrappers (stalled→error, fresh survives) — previously untested. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01NsmJSQxnNxGgtM5Yz4GAqi --- backend/app/tasks/maintenance.py | 117 +++++++++--------------- tests/test_recover_stalled_head_runs.py | 63 +++++++++++++ 2 files changed, 108 insertions(+), 72 deletions(-) create mode 100644 tests/test_recover_stalled_head_runs.py diff --git a/backend/app/tasks/maintenance.py b/backend/app/tasks/maintenance.py index 2ad0b3f..67432c2 100644 --- a/backend/app/tasks/maintenance.py +++ b/backend/app/tasks/maintenance.py @@ -776,89 +776,62 @@ def recover_stalled_library_audit_runs() -> int: return recovered +def _recover_stalled_runs(model, *, stall_minutes: int, keep_runs: int, label: str) -> int: + """Shared recovery + retention sweep for the head run-tracking tables + (HeadTrainingRun / HeadAutoApplyRun, which share the + status/last_progress_at/started_at/finished_at/error/id columns): flip 'running' + rows with no progress past `stall_minutes` to 'error', then prune to the last + `keep_runs` (rule 89). Returns the number recovered. NOTE the two other recover + tasks are deliberately NOT folded in — library-audit has no prune tail and + backup uses a single started_at cutoff.""" + SessionLocal = _sync_session_factory() + now = datetime.now(UTC) + cutoff = now - timedelta(minutes=stall_minutes) + with SessionLocal() as session: + result = session.execute( + update(model) + .where(model.status == "running") + .where(func.coalesce(model.last_progress_at, model.started_at) < cutoff) + .values( + status="error", finished_at=now, + error=f"stranded by recovery sweep (no progress for {stall_minutes} min)", + ) + ) + keep = session.execute( + select(model.id).order_by(model.id.desc()).limit(keep_runs) + ).scalars().all() + if keep: + session.execute(delete(model).where(model.id.not_in(keep))) + session.commit() + recovered = result.rowcount or 0 + if recovered: + log.info("%s: recovered %d rows", label, recovered) + return recovered + + @celery.task(name="backend.app.tasks.maintenance.recover_stalled_head_training_runs") def recover_stalled_head_training_runs() -> int: """Flip HeadTrainingRun rows stuck in 'running' past the stall threshold to 'error', and prune old runs to the last HEAD_TRAINING_KEEP_RUNS (retention, rule 89). Runs every 5 min on the maintenance lane; no-op when idle.""" - SessionLocal = _sync_session_factory() - now = datetime.now(UTC) - cutoff = now - timedelta(minutes=HEAD_TRAINING_STALL_THRESHOLD_MINUTES) - with SessionLocal() as session: - result = session.execute( - update(HeadTrainingRun) - .where(HeadTrainingRun.status == "running") - .where( - func.coalesce( - HeadTrainingRun.last_progress_at, HeadTrainingRun.started_at - ) - < cutoff - ) - .values( - status="error", finished_at=now, - error=( - f"stranded by recovery sweep (no progress for " - f"{HEAD_TRAINING_STALL_THRESHOLD_MINUTES} min)" - ), - ) - ) - keep = session.execute( - select(HeadTrainingRun.id).order_by(HeadTrainingRun.id.desc()) - .limit(HEAD_TRAINING_KEEP_RUNS) - ).scalars().all() - if keep: - session.execute( - delete(HeadTrainingRun).where(HeadTrainingRun.id.not_in(keep)) - ) - session.commit() - recovered = result.rowcount or 0 - if recovered: - log.info( - "recover_stalled_head_training_runs: recovered %d rows", recovered - ) - return recovered + return _recover_stalled_runs( + HeadTrainingRun, + stall_minutes=HEAD_TRAINING_STALL_THRESHOLD_MINUTES, + keep_runs=HEAD_TRAINING_KEEP_RUNS, + label="recover_stalled_head_training_runs", + ) @celery.task(name="backend.app.tasks.maintenance.recover_stalled_head_auto_apply_runs") def recover_stalled_head_auto_apply_runs() -> int: """Flip stalled HeadAutoApplyRun 'running' rows to 'error' + prune to the last HEAD_AUTO_APPLY_KEEP_RUNS (retention, rule 89). 5-min maintenance lane.""" - SessionLocal = _sync_session_factory() - now = datetime.now(UTC) - cutoff = now - timedelta(minutes=HEAD_AUTO_APPLY_STALL_THRESHOLD_MINUTES) - with SessionLocal() as session: - result = session.execute( - update(HeadAutoApplyRun) - .where(HeadAutoApplyRun.status == "running") - .where( - func.coalesce( - HeadAutoApplyRun.last_progress_at, HeadAutoApplyRun.started_at - ) - < cutoff - ) - .values( - status="error", finished_at=now, - error=( - f"stranded by recovery sweep (no progress for " - f"{HEAD_AUTO_APPLY_STALL_THRESHOLD_MINUTES} min)" - ), - ) - ) - keep = session.execute( - select(HeadAutoApplyRun.id).order_by(HeadAutoApplyRun.id.desc()) - .limit(HEAD_AUTO_APPLY_KEEP_RUNS) - ).scalars().all() - if keep: - session.execute( - delete(HeadAutoApplyRun).where(HeadAutoApplyRun.id.not_in(keep)) - ) - session.commit() - recovered = result.rowcount or 0 - if recovered: - log.info( - "recover_stalled_head_auto_apply_runs: recovered %d rows", recovered - ) - return recovered + return _recover_stalled_runs( + HeadAutoApplyRun, + stall_minutes=HEAD_AUTO_APPLY_STALL_THRESHOLD_MINUTES, + keep_runs=HEAD_AUTO_APPLY_KEEP_RUNS, + label="recover_stalled_head_auto_apply_runs", + ) # Keep ~6 months of daily head-metric snapshots (enough to see tuning trends). diff --git a/tests/test_recover_stalled_head_runs.py b/tests/test_recover_stalled_head_runs.py new file mode 100644 index 0000000..778e491 --- /dev/null +++ b/tests/test_recover_stalled_head_runs.py @@ -0,0 +1,63 @@ +"""recover_stalled_head_training_runs + recover_stalled_head_auto_apply_runs share +one helper (_recover_stalled_runs, DRY pass #161). These pin BOTH wrappers so the +shared source stays correct: a 'running' row with no progress past the stall +threshold flips to 'error'; a fresh 'running' row is left alone.""" +from datetime import UTC, datetime, timedelta + +import pytest +from sqlalchemy import select + +from backend.app.models import HeadAutoApplyRun, HeadTrainingRun + +pytestmark = pytest.mark.integration + + +def test_recover_stalled_head_training_runs_flips_stalled_keeps_fresh(db_sync): + from backend.app.tasks.maintenance import recover_stalled_head_training_runs + + stale = HeadTrainingRun( + params={}, status="running", + last_progress_at=datetime.now(UTC) - timedelta(days=1), + ) + fresh = HeadTrainingRun( + params={}, status="running", last_progress_at=datetime.now(UTC), + ) + db_sync.add_all([stale, fresh]) + db_sync.commit() + stale_id, fresh_id = stale.id, fresh.id + + assert recover_stalled_head_training_runs.apply().get() == 1 + + db_sync.expire_all() + assert db_sync.execute( + select(HeadTrainingRun.status).where(HeadTrainingRun.id == stale_id) + ).scalar_one() == "error" + assert db_sync.execute( + select(HeadTrainingRun.status).where(HeadTrainingRun.id == fresh_id) + ).scalar_one() == "running" + + +def test_recover_stalled_head_auto_apply_runs_flips_stalled_keeps_fresh(db_sync): + from backend.app.tasks.maintenance import recover_stalled_head_auto_apply_runs + + stale = HeadAutoApplyRun( + dry_run=False, params={}, status="running", + last_progress_at=datetime.now(UTC) - timedelta(days=1), + ) + fresh = HeadAutoApplyRun( + dry_run=False, params={}, status="running", + last_progress_at=datetime.now(UTC), + ) + db_sync.add_all([stale, fresh]) + db_sync.commit() + stale_id, fresh_id = stale.id, fresh.id + + assert recover_stalled_head_auto_apply_runs.apply().get() == 1 + + db_sync.expire_all() + assert db_sync.execute( + select(HeadAutoApplyRun.status).where(HeadAutoApplyRun.id == stale_id) + ).scalar_one() == "error" + assert db_sync.execute( + select(HeadAutoApplyRun.status).where(HeadAutoApplyRun.id == fresh_id) + ).scalar_one() == "running"