feat(gpu): fast orphan recovery — graceful release + 60s sweep (#114)

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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa

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})

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,

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
+OR expired-leased jobs with FOR UPDATE SKIP LOCKED, so concurrent agents/workers
-retry after an agent died) never grab the same job and the queue self-heals
+never grab the same job. Orphan recovery is three-layered: a graceful agent stop
-without a separate recovery sweep. Result-writing (regions) is done by the API
+calls release() to hand its in-flight jobs back instantly; a hard crash is caught
-handler via RegionService; complete() just closes the job.
+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

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

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

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