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:
@@ -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),
|
||||
|
||||
@@ -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")
|
||||
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
)
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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"),
|
||||
|
||||
@@ -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