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 pathlib import Path

from quart import Blueprint, jsonify, request
from sqlalchemy import func, or_, select, update
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, error_dedupe_statements
from ..services.ml.gpu_triage import classify_reason, recover_defective_image
from ..services.ml.regions import RegionService

gpu_bp = Blueprint("gpu", __name__, url_prefix="/api/gpu")

# Same container mount the maintenance tasks use (tasks/admin.py) — recovery
# deletes the defective original + thumbnail under it.
_IMAGES_ROOT = Path("/images")

_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


@gpu_bp.route("/reprocess", methods=["POST"])
async def reprocess():
    """Reset every done/error job of `task` back to pending so the agent re-runs
    the WHOLE library under the current pipeline (e.g. after adding crop
    detectors). Heavy — the back-catalogue is otherwise skipped by the backfills."""
    body = await request.get_json(silent=True) or {}
    task = str(body.get("task") or "ccip")
    from ..tasks.ml import reprocess_gpu_jobs

    r = reprocess_gpu_jobs.delay(task)
    return jsonify({"celery_task_id": r.id, "task": task}), 202


@gpu_bp.route("/retry_errors", methods=["POST"])
async def retry_errors():
    """Requeue every ERRORED job (all task types) back to pending — the scoped
    recovery after an agent-side fix (e.g. the short-video sampler), where
    /reprocess would needlessly re-run the whole done library too. Attempts and
    the stored error reset so each job gets its full retry budget under the
    fixed pipeline. Stale tombstones are pruned FIRST (loop-era duplicates and
    rows a later success made moot — the same statements the backfills run), so
    one failing file requeues as ONE job, never a fan-out of duplicates. Small
    row count (errors only) → inline statements; the response carries the
    counts for the UI toast. Triage-confirmed defects are NOT requeued (see
    the WHERE below) — they stay on the recovery surface."""
    async with get_session() as session:
        pruned = 0
        for stmt in error_dedupe_statements():
            pruned += (await session.execute(stmt)).rowcount or 0
        res = await session.execute(
            update(GpuJob)
            .where(
                GpuJob.status == "error",
                # Triage-confirmed DEFECTS stay errored: the integrity probe
                # already proved the FILE is bad, so re-running the job just
                # burns agent time re-minting the same tombstone — those go
                # through /errors/<id>/recover instead.
                or_(GpuJob.triage_status.is_(None),
                    GpuJob.triage_status != "defect"),
            )
            .values(
                status="pending", attempts=0, error=None, lease_token=None,
                leased_at=None, lease_expires_at=None, triage_status=None,
                updated_at=func.now(),
            )
        )
        kept = (
            await session.execute(
                select(func.count()).select_from(GpuJob)
                .where(GpuJob.status == "error")
            )
        ).scalar_one()
        await session.commit()
    return jsonify({
        "requeued": res.rowcount or 0, "pruned": pruned, "defects_kept": kept,
    })


# --- Failure triage + recovery (#125) ------------------------------------

@gpu_bp.route("/errors", methods=["GET"])
async def errors():
    """The triage view of the error tombstones: every errored job joined with
    its image's integrity verdict, bucketed by reason for the overview. The
    probe sweep (triage_gpu_errors, 15-min beat) fills triage_status; 'defect'
    rows are the recovery surface's list."""
    async with get_session() as session:
        rows = (
            await session.execute(
                select(
                    GpuJob.id, GpuJob.image_record_id, GpuJob.task,
                    GpuJob.error, GpuJob.triage_status, GpuJob.updated_at,
                    ImageRecord.integrity_status, ImageRecord.mime,
                    ImageRecord.path, ImageRecord.thumbnail_path,
                )
                .join(ImageRecord, ImageRecord.id == GpuJob.image_record_id)
                .where(GpuJob.status == "error")
                .order_by(GpuJob.updated_at.desc())
                .limit(500)
            )
        ).all()
        total = (
            await session.execute(
                select(func.count()).select_from(GpuJob)
                .where(GpuJob.status == "error")
            )
        ).scalar_one()
    by_class: dict[str, int] = {}
    triage = {"defect": 0, "file_ok": 0, "unclassified": 0}
    items = []
    for r in rows:
        cls = classify_reason(r.error)
        by_class[cls] = by_class.get(cls, 0) + 1
        bucket = r.triage_status or "unclassified"
        triage[bucket] = triage.get(bucket, 0) + 1
        items.append({
            "job_id": r.id,
            "image_id": r.image_record_id,
            "task": r.task,
            "error": r.error,
            "reason_class": cls,
            "triage_status": r.triage_status,
            "integrity_status": r.integrity_status,
            "mime": r.mime,
            "image_url": image_url(r.path),
            "thumbnail_url": (
                image_url(r.thumbnail_path) if r.thumbnail_path else None
            ),
            "updated_at": r.updated_at.isoformat() if r.updated_at else None,
        })
    return jsonify({
        "total": total, "by_class": by_class, "triage": triage, "items": items,
    })


@gpu_bp.route("/errors/triage", methods=["POST"])
async def errors_triage():
    """Run the probe sweep NOW (the card's button) instead of waiting out the
    15-minute beat cadence."""
    from ..tasks.maintenance import triage_gpu_errors

    r = triage_gpu_errors.delay()
    return jsonify({"celery_task_id": r.id}), 202


@gpu_bp.route("/errors/<int:image_id>/recover", methods=["POST"])
async def errors_recover(image_id: int):
    """Recover a defect-triaged original: delete the bad copy + record and
    re-poll its subscription Source (a fresh fetch re-imports the file, which
    re-enters the GPU pipeline). Returns status 'no_source' when nothing
    pollable resolves — the file needs manual replacement there."""
    async with get_session() as session:
        result = await session.run_sync(
            lambda s: recover_defective_image(
                s, image_id, images_root=_IMAGES_ROOT,
            )
        )
    return jsonify(result)


# --- 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 + the
                # whole-image 'embed' task, so its vectors land in the SAME space
                # the heads trained in. Server-announced FROM THE SETTING → the
                # agent stays model-agnostic; an operator swap is a setting + a
                # re-embed, never an agent change.
                "embed_model_name": ml.embedder_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/submit_embedding", methods=["POST"])
async def submit_embedding():
    """Store a whole-image SigLIP embedding (the 'embed' task) on image_record +
    close the job. Body: {agent_id, job_id, embedding:[...], embedding_version}.
    This is how the GPU agent re-embeds the library under a new model (#1190) —
    much faster than the CPU ml-worker at higher resolutions."""
    body = await request.get_json(silent=True) or {}
    agent_id = str(body.get("agent_id") or "agent")
    job_id = body.get("job_id")
    embedding = body.get("embedding")
    version = body.get("embedding_version")
    if job_id is None or not embedding or not version:
        return jsonify({"error": "job_id, embedding, embedding_version required"}), 400
    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
        img = await session.get(ImageRecord, job.image_record_id)
        if img is not None:
            img.siglip_embedding = embedding
            img.siglip_model_version = version
        await GpuJobService(session).complete(agent_id, int(job_id))
        await session.commit()
    return jsonify({"ok": True})


@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})