feat(gpu): video-ready regions + the HTTP GPU-job queue engine (#114 slice 3)
CI / lint (push) Successful in 2s
CI / frontend-build (push) Successful in 19s
CI / backend-lint-and-test (push) Successful in 26s
CI / integration (push) Successful in 3m30s

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:
2026-06-29 11:18:28 -04:00
parent 0ea7ecdea5
commit b735432d02
9 changed files with 377 additions and 2 deletions
+2
View File
@@ -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),
+55
View File
@@ -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")
+2
View File
@@ -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",
+50
View File
@@ -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()
)
+6 -1
View File
@@ -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)
+134
View File
@@ -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
+1
View File
@@ -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"),
+125
View File
@@ -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"
+2 -1
View File
@@ -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