From 2cb04278685ca5e45c30455423a73cced3d3baa7 Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Mon, 29 Jun 2026 19:07:40 -0400 Subject: [PATCH] =?UTF-8?q?feat(gpu):=20fast=20orphan=20recovery=20?= =?UTF-8?q?=E2=80=94=20graceful=20release=20+=2060s=20sweep=20(#114)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit So work an agent orphaned gets picked back up quickly, three layers: - GpuJobService.release(): a graceful agent stop hands its still-leased jobs back to pending instantly (POST /api/gpu/jobs/release), no waiting out the lease. - GpuJobService.recover_orphaned() + recover_orphaned_gpu_jobs Celery task on a 60s beat: resets expired leases (a hard-crashed agent) to pending and keeps the queue counts honest even when nothing is leasing. - Lease TTL 300→180s: still well above any single job (a capped-frame video embed is tens of seconds, and a live worker heartbeats), but a hard crash recovers faster once the sweep fires. Tests: release returns-to-pending (token-scoped), recover_orphaned resets only expired leases, release API round-trip. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa --- backend/app/api/gpu.py | 15 ++++++++ backend/app/celery_app.py | 4 +++ backend/app/services/ml/gpu_jobs.py | 55 +++++++++++++++++++++++++---- backend/app/tasks/ml.py | 27 ++++++++++++++ tests/test_api_gpu.py | 19 ++++++++++ tests/test_gpu_jobs.py | 35 ++++++++++++++++++ 6 files changed, 149 insertions(+), 6 deletions(-) diff --git a/backend/app/api/gpu.py b/backend/app/api/gpu.py index 0a9b46a..e463d81 100644 --- a/backend/app/api/gpu.py +++ b/backend/app/api/gpu.py @@ -196,3 +196,18 @@ async def fail(): ) await session.commit() return jsonify({"ok": ok}) + + +@gpu_bp.route("/jobs/release", methods=["POST"]) +async def release(): + """Graceful stop: the agent hands its still-leased jobs back to pending so + they're picked up immediately instead of waiting out the lease.""" + body = await request.get_json(silent=True) or {} + agent_id = str(body.get("agent_id") or "agent") + job_ids = [int(x) for x in (body.get("job_ids") or [])] + async with get_session() as session: + if not await _agent_authed(session): + return jsonify({"error": "unauthorized"}), 401 + n = await GpuJobService(session).release(agent_id, job_ids) + await session.commit() + return jsonify({"released": n}) diff --git a/backend/app/celery_app.py b/backend/app/celery_app.py index 9617d6c..fc98ae5 100644 --- a/backend/app/celery_app.py +++ b/backend/app/celery_app.py @@ -117,6 +117,10 @@ def make_celery() -> Celery: "task": "backend.app.tasks.ml.scheduled_apply_head_tags", "schedule": 86400.0, # no-op unless head_auto_apply_enabled }, + "recover-orphaned-gpu-jobs": { + "task": "backend.app.tasks.ml.recover_orphaned_gpu_jobs", + "schedule": 60.0, # quick pickup of work a dead agent orphaned + }, "snapshot-head-metrics-daily": { "task": "backend.app.tasks.maintenance.snapshot_head_metrics", "schedule": 86400.0, diff --git a/backend/app/services/ml/gpu_jobs.py b/backend/app/services/ml/gpu_jobs.py index b760da9..3312b5e 100644 --- a/backend/app/services/ml/gpu_jobs.py +++ b/backend/app/services/ml/gpu_jobs.py @@ -1,10 +1,13 @@ -"""GPU-job queue engine (#114): enqueue / lease / heartbeat / complete / fail. +"""GPU-job queue engine (#114): enqueue / lease / heartbeat / complete / fail +/ release / recover_orphaned. Backs the HTTP API the desktop agent pulls work from. The lease claims pending -OR expired-leased jobs with FOR UPDATE SKIP LOCKED, so concurrent agents (or a -retry after an agent died) never grab the same job and the queue self-heals -without a separate recovery sweep. Result-writing (regions) is done by the API -handler via RegionService; complete() just closes the job. +OR expired-leased jobs with FOR UPDATE SKIP LOCKED, so concurrent agents/workers +never grab the same job. Orphan recovery is three-layered: a graceful agent stop +calls release() to hand its in-flight jobs back instantly; a hard crash is caught +by recover_orphaned() (a 60s beat sweep) which resets expired leases to pending; +and the lease itself reclaims expired leases as a final backstop. Result-writing +(regions) is done by the API handler via RegionService; complete() just closes. """ from datetime import UTC, datetime, timedelta @@ -14,7 +17,10 @@ from sqlalchemy.ext.asyncio import AsyncSession from ...models import GpuJob -DEFAULT_LEASE_TTL = 300 # seconds an agent holds a job before it can be re-leased +# Lease window. Kept comfortably above any single job (a capped-frame video embed +# is tens of seconds) so a live, heartbeating worker is never falsely expired, +# but short enough that a hard crash recovers fast once the sweep fires. +DEFAULT_LEASE_TTL = 180 # seconds an agent holds a job before it can be re-leased DEFAULT_BATCH = 8 MAX_ATTEMPTS = 3 @@ -132,3 +138,40 @@ class GpuJobService: job.error = (error or "")[:1000] job.updated_at = datetime.now(UTC) return True + + async def release(self, token: str, job_ids: list[int]) -> int: + """Hand the agent's still-leased jobs back to pending NOW (graceful stop), + so another worker picks them up immediately instead of waiting out the + lease. Scoped to the token's own leases. Returns rows released.""" + if not job_ids: + return 0 + now = datetime.now(UTC) + res = await self.session.execute( + update(GpuJob) + .where( + GpuJob.id.in_(job_ids), + GpuJob.lease_token == token, + GpuJob.status == "leased", + ) + .values( + status="pending", lease_token=None, leased_at=None, + lease_expires_at=None, updated_at=now, + ) + ) + return res.rowcount or 0 + + async def recover_orphaned(self) -> int: + """Reset every expired lease back to pending — catches agents that died + mid-job (no graceful release). Run on a short beat so the queue recovers + + reads honestly even when no worker is actively leasing. Returns rows + recovered.""" + now = datetime.now(UTC) + res = await self.session.execute( + update(GpuJob) + .where(GpuJob.status == "leased", GpuJob.lease_expires_at < now) + .values( + status="pending", lease_token=None, leased_at=None, + lease_expires_at=None, updated_at=now, + ) + ) + return res.rowcount or 0 diff --git a/backend/app/tasks/ml.py b/backend/app/tasks/ml.py index ecdfcbd..bcc1903 100644 --- a/backend/app/tasks/ml.py +++ b/backend/app/tasks/ml.py @@ -768,3 +768,30 @@ def enqueue_gpu_backfill(task_name: str) -> int: ).fetchall() session.commit() return len(rows) + + +@celery.task(name="backend.app.tasks.ml.recover_orphaned_gpu_jobs") +def recover_orphaned_gpu_jobs() -> int: + """Reset expired GPU-job leases back to pending — recovers work orphaned by an + agent that died mid-job (no graceful release). Short beat cadence so orphans + get picked back up quickly + the queue counts read honestly. Returns the + number recovered.""" + from datetime import UTC, datetime + + from sqlalchemy import update + + from ..models import GpuJob + + SessionLocal = _sync_session_factory() + with SessionLocal() as session: + now = datetime.now(UTC) + res = session.execute( + update(GpuJob) + .where(GpuJob.status == "leased", GpuJob.lease_expires_at < now) + .values( + status="pending", lease_token=None, leased_at=None, + lease_expires_at=None, updated_at=now, + ) + ) + session.commit() + return res.rowcount or 0 diff --git a/tests/test_api_gpu.py b/tests/test_api_gpu.py index 4993f75..40345d8 100644 --- a/tests/test_api_gpu.py +++ b/tests/test_api_gpu.py @@ -96,3 +96,22 @@ async def test_backfill_enqueues_then_is_idempotent(db): n = enqueue_gpu_backfill("ccip") # sync task, own session assert n >= 2 assert enqueue_gpu_backfill("ccip") == 0 # all already pending + + +@pytest.mark.asyncio +async def test_release_hands_job_back_to_pending(client, db): + img = await _img(db, "e" * 64) + await GpuJobService(db).enqueue(img.id, "ccip") + await db.commit() + token = (await (await client.post("/api/gpu/token/rotate")).get_json())["token"] + hdr = {"Authorization": f"Bearer {token}"} + j = (await (await client.post( + "/api/gpu/jobs/lease", json={"agent_id": "a1"}, headers=hdr, + )).get_json())["jobs"][0] + + resp = await client.post("/api/gpu/jobs/release", json={ + "agent_id": "a1", "job_ids": [j["job_id"]], + }, headers=hdr) + assert resp.status_code == 200 and (await resp.get_json())["released"] == 1 + st = await (await client.get("/api/gpu/status")).get_json() + assert st["pending"] == 1 and st["leased"] == 0 diff --git a/tests/test_gpu_jobs.py b/tests/test_gpu_jobs.py index d2a6ea7..9d1f885 100644 --- a/tests/test_gpu_jobs.py +++ b/tests/test_gpu_jobs.py @@ -123,3 +123,38 @@ async def test_fail_requeues_until_cap(db): assert await svc.fail("agent-1", job.id, "boom again") is True await db.commit() assert (await db.get(GpuJob, job.id)).status == "error" + + +@pytest.mark.asyncio +async def test_release_returns_to_pending(db): + img = await _img(db, "01" + "a" * 62) + svc = GpuJobService(db) + await svc.enqueue(img.id, "ccip") + await db.commit() + job = (await svc.lease("agent-1"))[0] + await db.commit() + + assert await svc.release("other", [job.id]) == 0 # not this token's lease + assert await svc.release("agent-1", [job.id]) == 1 # graceful hand-back + await db.commit() + fresh = await db.get(GpuJob, job.id) + assert fresh.status == "pending" and fresh.lease_token is None + + +@pytest.mark.asyncio +async def test_recover_orphaned_resets_only_expired(db): + img1 = await _img(db, "02" + "a" * 62) + img2 = await _img(db, "03" + "a" * 62) + svc = GpuJobService(db) + await svc.enqueue(img1.id, "ccip") + await svc.enqueue(img2.id, "ccip") + await db.commit() + expired, fresh = await svc.lease("dead", batch_size=2) + # One lease is in the past (orphaned), the other still valid. + expired.lease_expires_at = datetime.now(UTC) - timedelta(minutes=10) + await db.commit() + + assert await svc.recover_orphaned() == 1 + await db.commit() + assert (await db.get(GpuJob, expired.id)).status == "pending" + assert (await db.get(GpuJob, fresh.id)).status == "leased" # untouched