feat(gpu): HTTP job API + token auth + backfill — the agent's server side (#114 slice 3b)
The thin HTTP surface over the queue so the desktop agent stays HTTP-only: - Agent endpoints (Authorization: Bearer <token>): POST /api/gpu/jobs/lease (returns jobs + image_url + mime + video frame cadence), /submit (stores regions via RegionService + closes the job; 409 on a stale lease), /heartbeat, /fail. Token validated against AppSetting (mirrors the extension-key pattern, constant-time compare). - Admin (browser): GET/POST /api/gpu/token[/rotate] (generate + show the agent token), GET /api/gpu/status (queue counts), POST /api/gpu/backfill → dispatches enqueue_gpu_backfill. - enqueue_gpu_backfill(task): one INSERT…SELECT enqueues a job per image lacking one for the task (scales to the full library; idempotent). Agent flow: lease over HTTP → fetch pixels via the normal FC image URL → compute on the GPU → submit. Redis/Postgres never exposed. Tests: bearer required (+ wrong-token 401), lease→submit round-trip (region+CCIP vector stored, job done via /status), stale-lease 409, backfill enqueue + idempotency. NEXT: the agent container + control UI, then the CCIP detector/embedder + matcher. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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})
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user