backend/app/api/__init__.py

@@ -25,6 +25,7 @@ def all_blueprints() -> list[Blueprint]:
     from .downloads import downloads_bp
     from .extension import extension_bp
     from .gallery import gallery_bp
+    from .gpu import gpu_bp
     from .heads import heads_bp
     from .import_admin import import_admin_bp
     from .ml_admin import ml_admin_bp
@@ -60,6 +61,7 @@ def all_blueprints() -> list[Blueprint]:
         aliases_bp,
         tag_eval_bp,
         heads_bp,
+        gpu_bp,
         ml_admin_bp,
         thumbnails_bp,
         sources_bp,

backend/app/api/gpu.py

@@ -0,0 +1,198 @@
+"""GPU-job API (#114): the HTTP surface the desktop agent pulls work from.
+
+The agent stays HTTP-only — it leases jobs, fetches image pixels via the normal
+FC image URLs, and submits embeddings/regions back, all over this API. Redis and
+Postgres are never exposed. The agent endpoints are gated by a bearer token
+(Authorization: Bearer <token>) stored in AppSetting; the admin endpoints
+(token / backfill / status) ride the browser session like the rest of FC's
+homelab admin.
+"""
+
+import secrets
+
+from quart import Blueprint, jsonify, request
+from sqlalchemy import func, select
+from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+from ..extensions import get_session
+from ..models import AppSetting, GpuJob, ImageRecord, MLSettings
+from ..services.gallery_service import image_url
+from ..services.ml.gpu_jobs import GpuJobService
+from ..services.ml.regions import RegionService
+
+gpu_bp = Blueprint("gpu", __name__, url_prefix="/api/gpu")
+
+_TOKEN_KEY = "gpu_agent_token"
+
+
+def _bearer() -> str | None:
+    h = request.headers.get("Authorization", "")
+    return h[7:].strip() if h.startswith("Bearer ") else None
+
+
+async def _agent_authed(session) -> bool:
+    supplied = _bearer()
+    if not supplied:
+        return False
+    stored = (
+        await session.execute(
+            select(AppSetting.value).where(AppSetting.key == _TOKEN_KEY)
+        )
+    ).scalar_one_or_none()
+    return stored is not None and secrets.compare_digest(supplied, stored)
+
+
+# --- Admin (browser): token + backfill + status -------------------------
+
+@gpu_bp.route("/token", methods=["GET"])
+async def get_token():
+    async with get_session() as session:
+        tok = (
+            await session.execute(
+                select(AppSetting.value).where(AppSetting.key == _TOKEN_KEY)
+            )
+        ).scalar_one_or_none()
+    return jsonify({"token": tok, "configured": tok is not None})
+
+
+@gpu_bp.route("/token/rotate", methods=["POST"])
+async def rotate_token():
+    token = secrets.token_urlsafe(32)
+    async with get_session() as session:
+        await session.execute(
+            pg_insert(AppSetting)
+            .values(key=_TOKEN_KEY, value=token)
+            .on_conflict_do_update(index_elements=["key"], set_={"value": token})
+        )
+        await session.commit()
+    return jsonify({"token": token})
+
+
+@gpu_bp.route("/status", methods=["GET"])
+async def status():
+    async with get_session() as session:
+        rows = (
+            await session.execute(
+                select(GpuJob.status, func.count()).group_by(GpuJob.status)
+            )
+        ).all()
+    counts = {s: c for s, c in rows}
+    return jsonify({
+        "pending": counts.get("pending", 0),
+        "leased": counts.get("leased", 0),
+        "done": counts.get("done", 0),
+        "error": counts.get("error", 0),
+    })
+
+
+@gpu_bp.route("/backfill", methods=["POST"])
+async def backfill():
+    """Enqueue a job for every image that doesn't already have one for `task`."""
+    body = await request.get_json(silent=True) or {}
+    task = str(body.get("task") or "ccip")
+    from ..tasks.ml import enqueue_gpu_backfill
+
+    r = enqueue_gpu_backfill.delay(task)
+    return jsonify({"celery_task_id": r.id, "task": task}), 202
+
+
+# --- Agent (bearer token): lease / submit / heartbeat / fail ------------
+
+@gpu_bp.route("/jobs/lease", methods=["POST"])
+async def lease():
+    body = await request.get_json(silent=True) or {}
+    agent_id = str(body.get("agent_id") or "agent")
+    try:
+        batch = min(max(int(body.get("batch_size", 8)), 1), 64)
+    except (TypeError, ValueError):
+        batch = 8
+    async with get_session() as session:
+        if not await _agent_authed(session):
+            return jsonify({"error": "unauthorized"}), 401
+        jobs = await GpuJobService(session).lease(agent_id, batch_size=batch)
+        ml = (
+            await session.execute(select(MLSettings).where(MLSettings.id == 1))
+        ).scalar_one()
+        # image rows for url/mime in one shot
+        ids = [j.image_record_id for j in jobs]
+        imgs = {
+            i.id: i for i in (
+                await session.execute(
+                    select(ImageRecord).where(ImageRecord.id.in_(ids))
+                )
+            ).scalars()
+        } if ids else {}
+        await session.commit()
+        out = []
+        for j in jobs:
+            img = imgs.get(j.image_record_id)
+            if img is None:
+                continue
+            out.append({
+                "job_id": j.id,
+                "image_id": j.image_record_id,
+                "task": j.task,
+                "mime": img.mime,
+                "image_url": image_url(img.path),
+                # For video/animated: the agent samples at this cadence.
+                "frame_interval_seconds": ml.video_frame_interval_seconds,
+                "max_frames": ml.video_max_frames,
+            })
+    return jsonify({"jobs": out})
+
+
+@gpu_bp.route("/jobs/heartbeat", methods=["POST"])
+async def heartbeat():
+    body = await request.get_json(silent=True) or {}
+    agent_id = str(body.get("agent_id") or "agent")
+    job_ids = [int(x) for x in (body.get("job_ids") or [])]
+    async with get_session() as session:
+        if not await _agent_authed(session):
+            return jsonify({"error": "unauthorized"}), 401
+        n = await GpuJobService(session).heartbeat(agent_id, job_ids)
+        await session.commit()
+    return jsonify({"extended": n})
+
+
+@gpu_bp.route("/jobs/submit", methods=["POST"])
+async def submit():
+    """Store a job's regions + close it. regions: [{kind, bbox:[x,y,w,h],
+    frame_time?, score?, *_version?, ccip_embedding?, siglip_embedding?}].
+    replace_kinds defaults to the kinds present in the submitted regions."""
+    body = await request.get_json(silent=True) or {}
+    agent_id = str(body.get("agent_id") or "agent")
+    job_id = body.get("job_id")
+    regions = body.get("regions") or []
+    if job_id is None:
+        return jsonify({"error": "job_id required"}), 400
+    kinds = body.get("replace_kinds") or sorted({r["kind"] for r in regions})
+    async with get_session() as session:
+        if not await _agent_authed(session):
+            return jsonify({"error": "unauthorized"}), 401
+        job = await session.get(GpuJob, int(job_id))
+        if job is None or job.status != "leased" or job.lease_token != agent_id:
+            return jsonify({"error": "lease_invalid"}), 409
+        if kinds:
+            await RegionService(session).replace_regions(
+                job.image_record_id, kinds, regions
+            )
+        await GpuJobService(session).complete(agent_id, int(job_id))
+        await session.commit()
+    return jsonify({"ok": True, "stored": len(regions)})
+
+
+@gpu_bp.route("/jobs/fail", methods=["POST"])
+async def fail():
+    body = await request.get_json(silent=True) or {}
+    agent_id = str(body.get("agent_id") or "agent")
+    job_id = body.get("job_id")
+    if job_id is None:
+        return jsonify({"error": "job_id required"}), 400
+    async with get_session() as session:
+        if not await _agent_authed(session):
+            return jsonify({"error": "unauthorized"}), 401
+        ok = await GpuJobService(session).fail(
+            agent_id, int(job_id), str(body.get("error") or "")
+        )
+        await session.commit()
+    return jsonify({"ok": ok})

backend/app/tasks/ml.py

@@ -738,3 +738,31 @@ def scheduled_apply_head_tags() -> str:
         run_id = run.id
     apply_head_tags.delay(run_id)
     return "dispatched"
+
+
+@celery.task(name="backend.app.tasks.ml.enqueue_gpu_backfill")
+def enqueue_gpu_backfill(task_name: str) -> int:
+    """Enqueue a gpu_job for every image that doesn't already have one for
+    `task_name` (one INSERT…SELECT, so it scales to a full library). The desktop
+    agent drains the queue over HTTP. Returns the number enqueued."""
+    from sqlalchemy import exists, insert, literal, select as sa_select
+
+    from ..models import GpuJob, ImageRecord
+
+    SessionLocal = _sync_session_factory()
+    with SessionLocal() as session:
+        already = exists().where(
+            GpuJob.image_record_id == ImageRecord.id,
+            GpuJob.task == task_name,
+            GpuJob.status.in_(["pending", "leased", "done"]),
+        )
+        sel = sa_select(
+            ImageRecord.id, literal(task_name), literal("pending")
+        ).where(~already)
+        result = session.execute(
+            insert(GpuJob).from_select(
+                ["image_record_id", "task", "status"], sel
+            )
+        )
+        session.commit()
+        return result.rowcount or 0

tests/test_api_gpu.py

@@ -0,0 +1,98 @@
+"""GPU-job HTTP API (#114): bearer auth + lease/submit round-trip + backfill."""
+import pytest
+
+from backend.app.models import ImageRecord
+from backend.app.services.ml.gpu_jobs import GpuJobService
+from backend.app.services.ml.regions import RegionService
+
+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_agent_endpoints_require_bearer(client, db):
+    resp = await client.post("/api/gpu/jobs/lease", json={"agent_id": "a1"})
+    assert resp.status_code == 401
+    # A wrong token is also rejected.
+    await (await client.post("/api/gpu/token/rotate")).get_json()
+    bad = await client.post(
+        "/api/gpu/jobs/lease", json={"agent_id": "a1"},
+        headers={"Authorization": "Bearer nope"},
+    )
+    assert bad.status_code == 401
+
+
+@pytest.mark.asyncio
+async def test_lease_submit_round_trip(client, db):
+    img = await _img(db, "a" * 64)
+    await GpuJobService(db).enqueue(img.id, "ccip")
+    await db.commit()
+
+    token = (await (await client.post("/api/gpu/token/rotate")).get_json())["token"]
+    hdr = {"Authorization": f"Bearer {token}"}
+
+    leased = await client.post(
+        "/api/gpu/jobs/lease", json={"agent_id": "a1", "batch_size": 5}, headers=hdr,
+    )
+    assert leased.status_code == 200
+    jobs = (await leased.get_json())["jobs"]
+    assert len(jobs) == 1
+    j = jobs[0]
+    assert j["image_id"] == img.id and j["task"] == "ccip"
+    assert j["image_url"].startswith("/images/")
+
+    submitted = await client.post("/api/gpu/jobs/submit", json={
+        "agent_id": "a1", "job_id": j["job_id"],
+        "regions": [{
+            "kind": "figure", "bbox": [0.1, 0.1, 0.4, 0.4],
+            "ccip_embedding": [0.1] * 768, "embedding_version": "ccip-test",
+        }],
+    }, headers=hdr)
+    assert submitted.status_code == 200
+    assert (await submitted.get_json())["stored"] == 1
+
+    # Job closed (read on the app's own connection via the status endpoint).
+    st = await (await client.get("/api/gpu/status")).get_json()
+    assert st["done"] == 1 and st["pending"] == 0 and st["leased"] == 0
+
+    # Region persisted with its CCIP vector.
+    regs = await RegionService(db).get_regions(img.id, kinds=["figure"])
+    assert len(regs) == 1 and len(list(regs[0].ccip_embedding)) == 768
+
+
+@pytest.mark.asyncio
+async def test_submit_with_stale_lease_is_409(client, db):
+    img = await _img(db, "b" * 64)
+    await GpuJobService(db).enqueue(img.id, "ccip")
+    await db.commit()
+    token = (await (await client.post("/api/gpu/token/rotate")).get_json())["token"]
+    hdr = {"Authorization": f"Bearer {token}"}
+    j = (await (await client.post(
+        "/api/gpu/jobs/lease", json={"agent_id": "a1"}, headers=hdr,
+    )).get_json())["jobs"][0]
+    # A different agent can't submit someone else's lease.
+    resp = await client.post("/api/gpu/jobs/submit", json={
+        "agent_id": "other", "job_id": j["job_id"], "regions": [],
+    }, headers=hdr)
+    assert resp.status_code == 409
+
+
+@pytest.mark.asyncio
+async def test_backfill_enqueues_then_is_idempotent(db):
+    await _img(db, "c" * 64)
+    await _img(db, "d" * 64)
+    await db.commit()
+    from backend.app.tasks.ml import enqueue_gpu_backfill
+
+    n = enqueue_gpu_backfill("ccip")     # sync task, own session
+    assert n >= 2
+    assert enqueue_gpu_backfill("ccip") == 0   # all already pending