alembic/versions/0070_gpu_job_lease_indexes.py

@@ -0,0 +1,44 @@
+"""partial indexes so GPU-job leasing stays O(batch), not O(completed)
+
+The lease claims the lowest-id pending (or expired-leased) jobs. With only a
+plain `status` index, `... ORDER BY id LIMIT n` walked the primary-key index from
+the start, skipping the entire prefix of already-done/error rows before reaching
+pending ones — so leasing slowed to a crawl as `done` piled up (the whole reason
+throughput fell off a cliff mid-run and /status stalled). Two partial indexes fix
+it: the pending one is id-ordered so the hot path reads just the first n entries,
+and the leased-expiry one keeps the crash-recovery reclaim + the orphan sweep
+cheap. They cover only the small live slice of the table, so they stay tiny even
+as the done/error history grows to millions.
+
+Revision ID: 0070
+Revises: 0069
+Create Date: 2026-06-30
+"""
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+revision: str = "0070"
+down_revision: Union[str, None] = "0069"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # Hot path: lowest-id pending jobs. Index on id, restricted to pending, so
+    # `WHERE status='pending' ORDER BY id LIMIT n` is a short index-order scan.
+    op.create_index(
+        "ix_gpu_job_pending", "gpu_job", ["id"],
+        postgresql_where=sa.text("status = 'pending'"),
+    )
+    # Crash-recovery: expired leases, for the lease backstop + recover_orphaned.
+    op.create_index(
+        "ix_gpu_job_leased_expires", "gpu_job", ["lease_expires_at"],
+        postgresql_where=sa.text("status = 'leased'"),
+    )
+
+
+def downgrade() -> None:
+    op.drop_index("ix_gpu_job_leased_expires", table_name="gpu_job")
+    op.drop_index("ix_gpu_job_pending", table_name="gpu_job")

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(
+        async def _claim(condition, limit: int) -> list[int]:
-                select(GpuJob.id)
+            return list(
-                .where(
+                (
-                    or_(
+                    await self.session.execute(
-                        GpuJob.status == "pending",
+                        select(GpuJob.id).where(condition)
-                        and_(
+                        .order_by(GpuJob.id).limit(limit)
-                            GpuJob.status == "leased",
+                        .with_for_update(skip_locked=True)
-                            GpuJob.lease_expires_at < now,
-                        ),
                     )
-                )
+                ).scalars().all()
-                .order_by(GpuJob.id)
+            )
-                .limit(batch_size)
+
-                .with_for_update(skip_locked=True)
+        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(

tests/test_gpu_jobs.py

@@ -141,6 +141,22 @@ async def test_lease_claims_then_skips_when_held(db):
     assert again == []
 
 
+@pytest.mark.asyncio
+async def test_lease_skips_done_prefix(db):
+    # A long prefix of already-done jobs at the LOW ids must not stop the lease
+    # from reaching the pending ones (the O(done) scan bug this indexing fixes).
+    svc = GpuJobService(db)
+    for i in range(5):
+        j = await svc.enqueue((await _img(db, f"d{i}" * 32)).id, "ccip")
+        j.status = "done"
+    target = await svc.enqueue((await _img(db, "dz" * 32)).id, "ccip")  # highest id
+    await db.commit()
+
+    leased = await svc.lease("agent-1", batch_size=8)
+    await db.commit()
+    assert [j.id for j in leased] == [target.id]
+
+
 @pytest.mark.asyncio
 async def test_expired_lease_is_reclaimed(db):
     img = await _img(db, "c" * 64)