bvandeusen
181f1c6a27
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
266 lines
9.1 KiB
Python
266 lines
9.1 KiB
Python
"""GPU-job queue engine (#114): enqueue dedupe + lease/heartbeat/complete/fail."""
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import pytest
|
|
from sqlalchemy import func, select
|
|
|
|
from backend.app.models import GpuJob, ImageRecord, ImageRegion
|
|
from backend.app.services.ml.gpu_jobs import GpuJobService
|
|
|
|
pytestmark = pytest.mark.integration
|
|
|
|
|
|
async def _img(db, sha) -> ImageRecord:
|
|
img = ImageRecord(
|
|
path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg",
|
|
width=1, height=1, origin="imported_filesystem", integrity_status="unknown",
|
|
)
|
|
db.add(img)
|
|
await db.flush()
|
|
return img
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enqueue_siglip_backfill_gates_on_concept_region(db):
|
|
# 'siglip' backfill enqueues images that lack a concept region (the
|
|
# back-catalogue) and skips ones that already have one — and never double-
|
|
# enqueues an image that already has a pending siglip job.
|
|
from backend.app.models import MLSettings
|
|
from backend.app.tasks.ml import enqueue_gpu_backfill
|
|
|
|
cur = (await db.execute(
|
|
select(MLSettings.embedder_model_version).where(MLSettings.id == 1)
|
|
)).scalar_one()
|
|
need = await _img(db, "e1" * 32) # no concept region → wants one
|
|
have = await _img(db, "e2" * 32) # concept @ current version → skip
|
|
db.add(ImageRegion(
|
|
image_record_id=have.id, kind="concept", rx=0.0, ry=0.0, rw=1.0, rh=1.0,
|
|
siglip_embedding=[0.0] * 1152, embedding_version=cur,
|
|
))
|
|
await db.commit()
|
|
|
|
assert enqueue_gpu_backfill("siglip") >= 1
|
|
|
|
queued = {
|
|
j.image_record_id for j in (
|
|
await db.execute(select(GpuJob).where(GpuJob.task == "siglip"))
|
|
).scalars()
|
|
}
|
|
assert need.id in queued
|
|
assert have.id not in queued
|
|
|
|
# Idempotent: the now-pending job means a second run doesn't re-enqueue it.
|
|
enqueue_gpu_backfill("siglip")
|
|
n_for_need = (
|
|
await db.execute(
|
|
select(func.count()).select_from(GpuJob).where(
|
|
GpuJob.task == "siglip", GpuJob.image_record_id == need.id
|
|
)
|
|
)
|
|
).scalar_one()
|
|
assert n_for_need == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enqueue_embed_backfill_selects_stale_and_unembedded(db):
|
|
# Whole-image GPU re-embed (#1190): enqueue images with no embedding or one
|
|
# stamped under a DIFFERENT model version (an operator swap); skip ones
|
|
# already at the current version.
|
|
from backend.app.models import MLSettings
|
|
from backend.app.tasks.ml import enqueue_gpu_backfill
|
|
|
|
cur = (await db.execute(
|
|
select(MLSettings.embedder_model_version).where(MLSettings.id == 1)
|
|
)).scalar_one()
|
|
current = await _img(db, "f1" * 32)
|
|
current.siglip_embedding = [0.0] * 1152
|
|
current.siglip_model_version = cur # up to date → skip
|
|
stale = await _img(db, "f2" * 32)
|
|
stale.siglip_embedding = [0.0] * 1152
|
|
stale.siglip_model_version = "old-embedder-v0" # wrong space → re-embed
|
|
unembedded = await _img(db, "f3" * 32) # never embedded → embed
|
|
await db.commit()
|
|
|
|
assert enqueue_gpu_backfill("embed") >= 2
|
|
queued = {
|
|
j.image_record_id for j in (
|
|
await db.execute(select(GpuJob).where(GpuJob.task == "embed"))
|
|
).scalars()
|
|
}
|
|
assert stale.id in queued and unembedded.id in queued
|
|
assert current.id not in queued
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reprocess_resets_done_jobs_to_pending(db):
|
|
# Re-process (#1202): done/error jobs of a task go back to pending so the
|
|
# agent re-runs the whole library under the current pipeline.
|
|
from backend.app.tasks.ml import reprocess_gpu_jobs
|
|
|
|
img = await _img(db, "r1" * 32)
|
|
job = await GpuJobService(db).enqueue(img.id, "ccip")
|
|
job.status = "done"
|
|
job.attempts = 2
|
|
await db.commit()
|
|
|
|
assert reprocess_gpu_jobs("ccip") >= 1
|
|
await db.refresh(job)
|
|
assert job.status == "pending"
|
|
assert job.attempts == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_enqueue_dedupes_same_pair(db):
|
|
img = await _img(db, "a" * 64)
|
|
svc = GpuJobService(db)
|
|
first = await svc.enqueue(img.id, "ccip")
|
|
dup = await svc.enqueue(img.id, "ccip")
|
|
other = await svc.enqueue(img.id, "siglip_region")
|
|
await db.commit()
|
|
assert first is not None
|
|
assert dup is None # same (image, task) already queued
|
|
assert other is not None # different task is fine
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lease_claims_then_skips_when_held(db):
|
|
img = await _img(db, "b" * 64)
|
|
svc = GpuJobService(db)
|
|
await svc.enqueue(img.id, "ccip")
|
|
await db.commit()
|
|
|
|
leased = await svc.lease("agent-1", batch_size=8)
|
|
await db.commit()
|
|
assert len(leased) == 1
|
|
assert leased[0].status == "leased" and leased[0].lease_token == "agent-1"
|
|
assert leased[0].attempts == 1
|
|
|
|
# Already leased + not expired → a second agent gets nothing.
|
|
again = await svc.lease("agent-2", batch_size=8)
|
|
await db.commit()
|
|
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)
|
|
svc = GpuJobService(db)
|
|
job = await svc.enqueue(img.id, "ccip")
|
|
await db.commit()
|
|
# Force the lease into the past.
|
|
job.status = "leased"
|
|
job.lease_token = "dead-agent"
|
|
job.lease_expires_at = datetime.now(UTC) - timedelta(minutes=10)
|
|
await db.commit()
|
|
|
|
leased = await svc.lease("agent-2", batch_size=8)
|
|
await db.commit()
|
|
assert len(leased) == 1
|
|
assert leased[0].lease_token == "agent-2"
|
|
assert leased[0].attempts == 1 # re-lease incremented from 0 (was set directly)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_heartbeat_extends_only_own_lease(db):
|
|
img = await _img(db, "d" * 64)
|
|
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.heartbeat("agent-1", [job.id]) == 1
|
|
assert await svc.heartbeat("someone-else", [job.id]) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_complete_closes_job(db):
|
|
img = await _img(db, "e" * 64)
|
|
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.complete("wrong-token", job.id) is False
|
|
assert await svc.complete("agent-1", job.id) is True
|
|
await db.commit()
|
|
fresh = await db.get(GpuJob, job.id)
|
|
assert fresh.status == "done" and fresh.lease_token is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fail_requeues_until_cap(db):
|
|
img = await _img(db, "f" * 64)
|
|
svc = GpuJobService(db)
|
|
await svc.enqueue(img.id, "ccip")
|
|
await db.commit()
|
|
job = (await svc.lease("agent-1"))[0] # attempts -> 1
|
|
await db.commit()
|
|
# Under the cap → back to pending for a retry.
|
|
assert await svc.fail("agent-1", job.id, "boom") is True
|
|
await db.commit()
|
|
assert (await db.get(GpuJob, job.id)).status == "pending"
|
|
|
|
# At the attempt cap → terminal 'error'.
|
|
j = await db.get(GpuJob, job.id)
|
|
j.attempts = 3
|
|
j.status = "leased"
|
|
j.lease_token = "agent-1"
|
|
j.lease_expires_at = datetime.now(UTC) + timedelta(minutes=5)
|
|
await db.commit()
|
|
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
|