diff --git a/backend/app/api/__init__.py b/backend/app/api/__init__.py index f39a966..4e51bb7 100644 --- a/backend/app/api/__init__.py +++ b/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, diff --git a/backend/app/api/gpu.py b/backend/app/api/gpu.py new file mode 100644 index 0000000..6eb4d39 --- /dev/null +++ b/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 ) 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}) diff --git a/backend/app/tasks/ml.py b/backend/app/tasks/ml.py index bc22c5b..72c2f80 100644 --- a/backend/app/tasks/ml.py +++ b/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 diff --git a/tests/test_api_gpu.py b/tests/test_api_gpu.py new file mode 100644 index 0000000..4993f75 --- /dev/null +++ b/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