feat(gpu): video-ready regions + the HTTP GPU-job queue engine (#114 slice 3)
Answers "how are videos/all media handled by the GPU worker": a job is per ITEM, but the agent fans a VIDEO into per-frame instances (ffmpeg in the agent, the existing cadence), each stored with a timestamp — so a video becomes a BAG of frame embeddings (fixes the mean-embedding muddle) instead of one washed-out vector. Stills → frame_time NULL; animated GIF/WebP treated like short video. - image_region.frame_time (migration 0061, not yet deployed so folded in): the source frame's seconds for video/animated media; NULL for stills. RegionService passes it through. A whole frame is just kind='frame'. - gpu_job + GpuJobService (migration 0062): the durable work list that keeps the desktop agent HTTP-only — enqueue (dedupes (image,task)) / lease (FOR UPDATE SKIP LOCKED, re-claims expired leases so the queue self-heals) / heartbeat / complete / fail (re-queues until MAX_ATTEMPTS then 'error'). The server enqueues; the agent leases+submits over the web API; Redis/Postgres stay private. Tests: enqueue dedupe, lease-then-skip-when-held, expired-lease reclaim, scoped heartbeat, complete, fail-requeue-then-error. region test now covers frame_time. NEXT: the thin HTTP API (lease/submit/heartbeat) + bearer-token auth, then the agent container + control UI. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
This commit is contained in:
@@ -0,0 +1,125 @@
|
||||
"""GPU-job queue engine (#114): enqueue dedupe + lease/heartbeat/complete/fail."""
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import select
|
||||
|
||||
from backend.app.models import GpuJob, ImageRecord
|
||||
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_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_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"
|
||||
@@ -23,7 +23,7 @@ async def test_replace_and_get_regions(db):
|
||||
svc = RegionService(db)
|
||||
n = await svc.replace_regions(img.id, ["figure"], [
|
||||
{"kind": "figure", "bbox": (0.1, 0.1, 0.3, 0.4),
|
||||
"score": 0.9, "detector_version": "det-v1"},
|
||||
"score": 0.9, "detector_version": "det-v1", "frame_time": 42.5},
|
||||
])
|
||||
await db.commit()
|
||||
assert n == 1
|
||||
@@ -33,6 +33,7 @@ async def test_replace_and_get_regions(db):
|
||||
assert r.kind == "figure"
|
||||
assert r.rw == pytest.approx(0.3) and r.rh == pytest.approx(0.4)
|
||||
assert r.score == pytest.approx(0.9)
|
||||
assert r.frame_time == pytest.approx(42.5) # video frame timestamp
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
Reference in New Issue
Block a user