FabledCurator/backend/app/api/gpu.py

"""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.embedder import MODEL_NAME as EMBED_MODEL_NAME
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 = dict(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,
                # The embedding model the agent must use for concept crops, so
                # its region vectors land in the SAME space the heads trained in.
                # Server-announced → the agent stays model-agnostic; a swap is a
                # server setting + a re-embed migration, never an agent change.
                "embed_model_name": EMBED_MODEL_NAME,
                "embed_version": ml.embedder_model_version,
            })
    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})


@gpu_bp.route("/jobs/release", methods=["POST"])
async def release():
    """Graceful stop: the agent hands its still-leased jobs back to pending so
    they're picked up immediately instead of waiting out the lease."""
    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).release(agent_id, job_ids)
        await session.commit()
    return jsonify({"released": n})