From b735432d0217c6ecd1b1b2f8ce8ea61bd200d935 Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Mon, 29 Jun 2026 11:18:28 -0400 Subject: [PATCH] feat(gpu): video-ready regions + the HTTP GPU-job queue engine (#114 slice 3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa --- alembic/versions/0061_image_region.py | 2 + alembic/versions/0062_gpu_job.py | 55 +++++++++++ backend/app/models/__init__.py | 2 + backend/app/models/gpu_job.py | 50 ++++++++++ backend/app/models/image_region.py | 7 +- backend/app/services/ml/gpu_jobs.py | 134 ++++++++++++++++++++++++++ backend/app/services/ml/regions.py | 1 + tests/test_gpu_jobs.py | 125 ++++++++++++++++++++++++ tests/test_regions.py | 3 +- 9 files changed, 377 insertions(+), 2 deletions(-) create mode 100644 alembic/versions/0062_gpu_job.py create mode 100644 backend/app/models/gpu_job.py create mode 100644 backend/app/services/ml/gpu_jobs.py create mode 100644 tests/test_gpu_jobs.py diff --git a/alembic/versions/0061_image_region.py b/alembic/versions/0061_image_region.py index 892951e..b3af8a9 100644 --- a/alembic/versions/0061_image_region.py +++ b/alembic/versions/0061_image_region.py @@ -32,6 +32,8 @@ def upgrade() -> None: sa.ForeignKey("image_record.id", ondelete="CASCADE"), nullable=False, ), sa.Column("kind", sa.String(length=16), nullable=False), + # Video/animated: source frame timestamp (seconds); NULL for stills. + sa.Column("frame_time", sa.Float(), nullable=True), sa.Column("rx", sa.Float(), nullable=False), sa.Column("ry", sa.Float(), nullable=False), sa.Column("rw", sa.Float(), nullable=False), diff --git a/alembic/versions/0062_gpu_job.py b/alembic/versions/0062_gpu_job.py new file mode 100644 index 0000000..a044995 --- /dev/null +++ b/alembic/versions/0062_gpu_job.py @@ -0,0 +1,55 @@ +"""gpu_job: the HTTP-leased GPU work queue for the desktop agent (#114) + +The agent stays HTTP-only — the server enqueues per-(image, task) jobs here and +the agent leases/submits over the web API; Redis/Postgres stay private. + +Revision ID: 0062 +Revises: 0061 +Create Date: 2026-06-29 +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +revision: str = "0062" +down_revision: Union[str, None] = "0061" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.create_table( + "gpu_job", + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column( + "image_record_id", sa.Integer(), + sa.ForeignKey("image_record.id", ondelete="CASCADE"), nullable=False, + ), + sa.Column("task", sa.String(length=32), nullable=False), + sa.Column( + "status", sa.String(length=16), nullable=False, + server_default="pending", + ), + sa.Column("lease_token", sa.String(length=64), nullable=True), + sa.Column("leased_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("attempts", sa.Integer(), nullable=False, server_default="0"), + sa.Column("error", sa.Text(), nullable=True), + sa.Column( + "created_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + sa.Column( + "updated_at", sa.DateTime(timezone=True), nullable=False, + server_default=sa.func.now(), + ), + ) + op.create_index("ix_gpu_job_image_record_id", "gpu_job", ["image_record_id"]) + op.create_index("ix_gpu_job_status", "gpu_job", ["status"]) + + +def downgrade() -> None: + op.drop_index("ix_gpu_job_status", table_name="gpu_job") + op.drop_index("ix_gpu_job_image_record_id", table_name="gpu_job") + op.drop_table("gpu_job") diff --git a/backend/app/models/__init__.py b/backend/app/models/__init__.py index 9df332a..087ff0b 100644 --- a/backend/app/models/__init__.py +++ b/backend/app/models/__init__.py @@ -8,6 +8,7 @@ from .base import Base from .credential import Credential from .download_event import DownloadEvent from .external_link import ExternalLink +from .gpu_job import GpuJob from .head_auto_apply_run import HeadAutoApplyRun from .head_metric import HeadMetric from .head_metrics_snapshot import HeadMetricsSnapshot @@ -67,6 +68,7 @@ __all__ = [ "image_tag", "DownloadEvent", "ExternalLink", + "GpuJob", "ImportBatch", "ImportTask", "ImportSettings", diff --git a/backend/app/models/gpu_job.py b/backend/app/models/gpu_job.py new file mode 100644 index 0000000..a8ed5c1 --- /dev/null +++ b/backend/app/models/gpu_job.py @@ -0,0 +1,50 @@ +"""GpuJob — a unit of GPU work the desktop agent pulls over HTTP (#114). + +The durable work list that lets the agent stay HTTP-only: the server enqueues a +job per (image, task) — e.g. detect figures + CCIP-embed — and the agent LEASES a +batch, computes on its GPU, then SUBMITS results, all over the already-exposed web +API. Redis/Postgres stay private. A lease has an expiry; the lease query itself +re-claims expired leases (agent died / stopped mid-batch), so the queue is +self-healing without a separate sweep. One job is per ITEM; the agent fans a +VIDEO out into per-frame instances internally (see image_region.frame_time). + +State: pending → leased → done | error (a failure under the attempt cap returns to +pending for another agent). +""" + +from datetime import datetime + +from sqlalchemy import DateTime, ForeignKey, Integer, String, Text, func +from sqlalchemy.orm import Mapped, mapped_column + +from .base import Base + + +class GpuJob(Base): + __tablename__ = "gpu_job" + + id: Mapped[int] = mapped_column(Integer, primary_key=True) + image_record_id: Mapped[int] = mapped_column( + ForeignKey("image_record.id", ondelete="CASCADE"), index=True + ) + # What to compute, e.g. 'ccip' (detect figures + CCIP-embed) or 'siglip_region'. + task: Mapped[str] = mapped_column(String(32), nullable=False) + status: Mapped[str] = mapped_column( + String(16), nullable=False, default="pending", index=True + ) + # pending | leased | done | error + lease_token: Mapped[str | None] = mapped_column(String(64), nullable=True) + leased_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + lease_expires_at: Mapped[datetime | None] = mapped_column( + DateTime(timezone=True), nullable=True + ) + attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + error: Mapped[str | None] = mapped_column(Text, nullable=True) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) diff --git a/backend/app/models/image_region.py b/backend/app/models/image_region.py index 5312926..58cfa8c 100644 --- a/backend/app/models/image_region.py +++ b/backend/app/models/image_region.py @@ -30,8 +30,13 @@ class ImageRegion(Base): image_record_id: Mapped[int] = mapped_column( ForeignKey("image_record.id", ondelete="CASCADE"), index=True ) - # 'face' | 'figure' (→ CCIP character id) | 'concept' (→ SigLIP head bag). + # 'frame' (a whole video frame → SigLIP bag) | 'face' | 'figure' (→ CCIP + # character id) | 'concept' (→ SigLIP head bag). kind: Mapped[str] = mapped_column(String(16), nullable=False) + # For video/animated media: the source frame's timestamp in SECONDS. NULL for + # static images. Lets a video be a BAG of per-frame instances (fixes the + # mean-embedding muddle) + grounds a tag to "appears at 0:42". + frame_time: Mapped[float | None] = mapped_column(Float, nullable=True) # Normalized bbox in [0,1]: top-left (rx, ry) + size (rw, rh). Named rx/ry/… # rather than x/y/by to dodge SQL keyword ambiguity ('by'). rx: Mapped[float] = mapped_column(Float, nullable=False) diff --git a/backend/app/services/ml/gpu_jobs.py b/backend/app/services/ml/gpu_jobs.py new file mode 100644 index 0000000..b760da9 --- /dev/null +++ b/backend/app/services/ml/gpu_jobs.py @@ -0,0 +1,134 @@ +"""GPU-job queue engine (#114): enqueue / lease / heartbeat / complete / fail. + +Backs the HTTP API the desktop agent pulls work from. The lease claims pending +OR expired-leased jobs with FOR UPDATE SKIP LOCKED, so concurrent agents (or a +retry after an agent died) never grab the same job and the queue self-heals +without a separate recovery sweep. Result-writing (regions) is done by the API +handler via RegionService; complete() just closes the job. +""" + +from datetime import UTC, datetime, timedelta + +from sqlalchemy import and_, or_, select, update +from sqlalchemy.ext.asyncio import AsyncSession + +from ...models import GpuJob + +DEFAULT_LEASE_TTL = 300 # seconds an agent holds a job before it can be re-leased +DEFAULT_BATCH = 8 +MAX_ATTEMPTS = 3 + + +class GpuJobService: + def __init__(self, session: AsyncSession): + self.session = session + + async def enqueue(self, image_id: int, task: str) -> GpuJob | None: + """Queue a (image, task) job. Idempotent: returns None if one is already + pending/leased for the same pair (no duplicate work).""" + dup = ( + await self.session.execute( + select(GpuJob.id).where( + GpuJob.image_record_id == image_id, + GpuJob.task == task, + GpuJob.status.in_(["pending", "leased"]), + ) + ) + ).first() + if dup: + return None + job = GpuJob(image_record_id=image_id, task=task, status="pending") + self.session.add(job) + await self.session.flush() + return job + + 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`.""" + 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, + ), + ) + ) + .order_by(GpuJob.id) + .limit(batch_size) + .with_for_update(skip_locked=True) + ) + ).scalars().all() + if not picked: + return [] + await self.session.execute( + update(GpuJob) + .where(GpuJob.id.in_(picked)) + .values( + status="leased", lease_token=token, leased_at=now, + lease_expires_at=now + timedelta(seconds=ttl), + attempts=GpuJob.attempts + 1, updated_at=now, + ) + ) + # populate_existing: overwrite identity-map copies with the post-UPDATE + # values so the returned jobs reflect the new lease/attempts, not stale + # pre-lease state. + return list( + ( + await self.session.execute( + select(GpuJob) + .where(GpuJob.id.in_(picked)) + .order_by(GpuJob.id) + .execution_options(populate_existing=True) + ) + ).scalars() + ) + + async def heartbeat( + self, token: str, job_ids: list[int], ttl: int = DEFAULT_LEASE_TTL + ) -> int: + """Extend the lease on the agent's in-flight jobs. Returns rows touched.""" + now = datetime.now(UTC) + res = await self.session.execute( + update(GpuJob) + .where( + GpuJob.id.in_(job_ids), + GpuJob.lease_token == token, + GpuJob.status == "leased", + ) + .values(lease_expires_at=now + timedelta(seconds=ttl), updated_at=now) + ) + return res.rowcount or 0 + + async def complete(self, token: str, job_id: int) -> bool: + """Close a leased job (after its results were stored). False if the job + isn't leased by this token (a stale/expired submit).""" + job = await self.session.get(GpuJob, job_id) + if job is None or job.status != "leased" or job.lease_token != token: + return False + job.status = "done" + job.lease_token = None + job.lease_expires_at = None + job.error = None + job.updated_at = datetime.now(UTC) + return True + + async def fail(self, token: str, job_id: int, error: str) -> bool: + """Report a failure: re-queue (pending) until MAX_ATTEMPTS, then 'error'.""" + job = await self.session.get(GpuJob, job_id) + if job is None or job.lease_token != token: + return False + if job.attempts >= MAX_ATTEMPTS: + job.status = "error" + else: + job.status = "pending" + job.lease_token = None + job.lease_expires_at = None + job.error = (error or "")[:1000] + job.updated_at = datetime.now(UTC) + return True diff --git a/backend/app/services/ml/regions.py b/backend/app/services/ml/regions.py index 8c3d8ab..7008536 100644 --- a/backend/app/services/ml/regions.py +++ b/backend/app/services/ml/regions.py @@ -46,6 +46,7 @@ class RegionService: rx, ry, rw, rh = r["bbox"] self.session.add(ImageRegion( image_record_id=image_id, kind=r["kind"], + frame_time=r.get("frame_time"), rx=rx, ry=ry, rw=rw, rh=rh, score=r.get("score"), detector_version=r.get("detector_version"), diff --git a/tests/test_gpu_jobs.py b/tests/test_gpu_jobs.py new file mode 100644 index 0000000..d2a6ea7 --- /dev/null +++ b/tests/test_gpu_jobs.py @@ -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" diff --git a/tests/test_regions.py b/tests/test_regions.py index 63c8ca5..f060d77 100644 --- a/tests/test_regions.py +++ b/tests/test_regions.py @@ -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