diff --git a/alembic/versions/0070_gpu_job_lease_indexes.py b/alembic/versions/0070_gpu_job_lease_indexes.py new file mode 100644 index 0000000..10ec3f9 --- /dev/null +++ b/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") diff --git a/backend/app/models/gpu_job.py b/backend/app/models/gpu_job.py index a8ed5c1..5e14e2d 100644 --- a/backend/app/models/gpu_job.py +++ b/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 diff --git a/backend/app/services/ml/gpu_jobs.py b/backend/app/services/ml/gpu_jobs.py index 3312b5e..f41086b 100644 --- a/backend/app/services/ml/gpu_jobs.py +++ b/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( diff --git a/tests/test_gpu_jobs.py b/tests/test_gpu_jobs.py index e1c103e..0678dc5 100644 --- a/tests/test_gpu_jobs.py +++ b/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)