perf(gpu-queue): partial indexes + two-phase lease so leasing stays O(batch)

The throughput bottleneck was curator-side, not the network. lease() claimed the
lowest-id pending/expired jobs with `... ORDER BY id LIMIT n`, but with only a
plain `status` index Postgres walked the primary key from id=1, skipping the
entire prefix of already done/error rows before reaching pending ones. As `done`
grew (69k+), every lease became an O(done) scan — leasing crawled, the DB
saturated, and even /status (the queue GROUP BY count) stalled the agent.

- Migration 0070 adds two partial indexes over just the live slice: pending rows
  indexed by id (hot path), and leased rows by lease_expires_at (crash-recovery
  + orphan sweep). They stay tiny no matter how large the done/error history.
- lease() split into two phases so each uses a partial index: claim pending
  first (id-ordered, O(batch)); reclaim expired leases only when pending can't
  fill the batch. Same semantics (SKIP LOCKED, attempts++, expired reclaim).
- Model __table_args__ declares the indexes so ORM and schema agree.
- Test: a done-prefix at low ids must not stop the lease reaching pending.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa

backend/app/models/gpu_job.py

@@ -14,7 +14,16 @@ pending for another agent).
 
 from datetime import datetime
 
-from sqlalchemy import DateTime, ForeignKey, Integer, String, Text, func
+from sqlalchemy import (
+    DateTime,
+    ForeignKey,
+    Index,
+    Integer,
+    String,
+    Text,
+    func,
+    text,
+)
 from sqlalchemy.orm import Mapped, mapped_column
 
 from .base import Base
@@ -23,6 +32,17 @@ from .base import Base
 class GpuJob(Base):
     __tablename__ = "gpu_job"
 
+    # Partial indexes over just the live slice (see migration 0070): the lease
+    # reads the lowest-id pending jobs on the hot path, and reclaims expired
+    # leases as a backstop — both stay O(batch) as done/error history grows.
+    __table_args__ = (
+        Index("ix_gpu_job_pending", "id", postgresql_where=text("status = 'pending'")),
+        Index(
+            "ix_gpu_job_leased_expires", "lease_expires_at",
+            postgresql_where=text("status = 'leased'"),
+        ),
+    )
+
     id: Mapped[int] = mapped_column(Integer, primary_key=True)
     image_record_id: Mapped[int] = mapped_column(
         ForeignKey("image_record.id", ondelete="CASCADE"), index=True

backend/app/services/ml/gpu_jobs.py

@@ -12,7 +12,7 @@ and the lease itself reclaims expired leases as a final backstop. Result-writing
 
 from datetime import UTC, datetime, timedelta
 
-from sqlalchemy import and_, or_, select, update
+from sqlalchemy import and_, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from ...models import GpuJob
@@ -51,25 +51,33 @@ class GpuJobService:
     async def lease(
         self, token: str, batch_size: int = DEFAULT_BATCH, ttl: int = DEFAULT_LEASE_TTL
     ) -> list[GpuJob]:
-        """Claim up to batch_size pending (or expired-leased) jobs for `token`."""
+        """Claim up to batch_size pending (or expired-leased) jobs for `token`.
+
+        Two phases so each hits a partial index (0070) and stays O(batch) no
+        matter how many done/error rows have accumulated: the pending pool is the
+        hot path; expired leases are reclaimed only when pending can't fill the
+        batch (a crashed agent's work — rare). The old single OR-query walked the
+        primary key past the whole done-prefix in id order → O(done), which is
+        why leasing crawled — and the DB saturated — as the run progressed."""
         now = datetime.now(UTC)
-        picked = (
-            await self.session.execute(
-                select(GpuJob.id)
-                .where(
-                    or_(
-                        GpuJob.status == "pending",
-                        and_(
-                            GpuJob.status == "leased",
-                            GpuJob.lease_expires_at < now,
-                        ),
+
+        async def _claim(condition, limit: int) -> list[int]:
+            return list(
+                (
+                    await self.session.execute(
+                        select(GpuJob.id).where(condition)
+                        .order_by(GpuJob.id).limit(limit)
+                        .with_for_update(skip_locked=True)
                     )
-                )
-                .order_by(GpuJob.id)
-                .limit(batch_size)
-                .with_for_update(skip_locked=True)
+                ).scalars().all()
+            )
+
+        picked = await _claim(GpuJob.status == "pending", batch_size)
+        if len(picked) < batch_size:     # pending exhausted → reclaim expired leases
+            picked += await _claim(
+                and_(GpuJob.status == "leased", GpuJob.lease_expires_at < now),
+                batch_size - len(picked),
             )
-        ).scalars().all()
         if not picked:
             return []
         await self.session.execute(