FabledCurator/backend/app/tasks/gpu_queue.py

"""GPU-job queue coordination: backfill enqueues, orphan recovery, reprocess.

These are pure-DB sweeps (INSERT…SELECT / UPDATE) — no torch, no sklearn —
that keep the desktop GPU agent's work queue fed and self-healing. They lived
in tasks/ml.py (routed to the 'ml' queue) purely by colocation, which made the
ml-worker container a hard dependency of the GPU pipeline; under B3 the
ml-worker is OPTIONAL (its only processing role is the CPU embed fallback), so
these moved here and route to the 'maintenance' quick lane with the other
recovery sweeps. A stack with no ml-worker keeps a fully-working GPU pipeline.
"""

import logging

from sqlalchemy import select

from ..celery_app import celery
from ._sync_engine import sync_session_factory as _sync_session_factory

log = logging.getLogger(__name__)


@celery.task(name="backend.app.tasks.gpu_queue.enqueue_gpu_backfill")
def enqueue_gpu_backfill(task_name: str) -> int:
    """Enqueue a gpu_job for every image that still needs `task_name` (one
    INSERT…SELECT, so it scales to a full library). The desktop agent drains the
    queue over HTTP. Returns the number enqueued.

    Completion is judged PER PIPELINE, never across them (B3, operator
    2026-07-02): 'ccip' by prior gpu_job rows, 'siglip' by concept regions at
    the current model version, and only 'embed' by image_record's whole-image
    embedding — the one artifact the CPU fallback also produces. A CPU embed
    therefore never closes crop/detect work for the agent.

    An ERRORED job is a tombstone for its (image, task): no variant re-enqueues
    it. Retry is deliberate-only (/retry_errors), which also means an errored
    back-catalogue needs one "Retry errored jobs" press after a model swap.
    Before the tombstone rule, this loop re-minted a fresh doomed job for every
    permanently-bad file each run — ~24 duplicate error rows/day per file (the
    2026-07-02 "unprocessable" flood)."""
    from sqlalchemy import exists, insert, literal, or_
    from sqlalchemy import select as sa_select

    from ..models import GpuJob, ImageRecord, ImageRegion, MLSettings
    from ..services.ml.gpu_jobs import error_dedupe_statements

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        # Prune stale tombstones first (loop-era duplicates + rows made moot by
        # a later success), so 'error' reads as one row per distinct failing
        # file and the skip-guards below see a clean picture.
        pruned = sum(
            session.execute(s).rowcount or 0 for s in error_dedupe_statements()
        )
        if pruned:
            log.info("gpu backfill: pruned %d stale/duplicate error rows", pruned)
        cur_version = session.execute(
            select(MLSettings.embedder_model_version).where(MLSettings.id == 1)
        ).scalar_one()
        if task_name == "embed":
            # Whole-image GPU re-embed (#1190): images with no embedding, or one
            # stamped under a DIFFERENT model version (an operator model swap).
            stale = or_(
                ImageRecord.siglip_embedding.is_(None),
                ImageRecord.siglip_model_version.is_(None),
                ImageRecord.siglip_model_version != cur_version,
            )
            # 'error' blocks too — tombstone rule, see docstring.
            blocked = exists().where(
                GpuJob.image_record_id == ImageRecord.id,
                GpuJob.task == "embed",
                GpuJob.status.in_(["pending", "leased", "error"]),
            )
            sel = sa_select(
                ImageRecord.id, literal("embed"), literal("pending")
            ).where(stale).where(~blocked)
        elif task_name == "siglip":
            # Concept-crop re-embed: enqueue when there's no concept region AT THE
            # CURRENT model version — so a model swap re-triggers crops too, not
            # only the never-embedded back-catalogue.
            has_current_concept = exists().where(
                ImageRegion.image_record_id == ImageRecord.id,
                ImageRegion.kind == "concept",
                ImageRegion.embedding_version == cur_version,
            )
            # 'error' blocks too — tombstone rule, see docstring.
            blocked = exists().where(
                GpuJob.image_record_id == ImageRecord.id,
                GpuJob.task == "siglip",
                GpuJob.status.in_(["pending", "leased", "error"]),
            )
            sel = sa_select(
                ImageRecord.id, literal("siglip"), literal("pending")
            ).where(~has_current_concept).where(~blocked)
        else:
            # ANY prior row blocks — including 'error' (tombstone rule, see
            # docstring): pre-fix this branch ran HOURLY and was the loop.
            already = exists().where(
                GpuJob.image_record_id == ImageRecord.id,
                GpuJob.task == task_name,
                GpuJob.status.in_(["pending", "leased", "done", "error"]),
            )
            sel = sa_select(
                ImageRecord.id, literal(task_name), literal("pending")
            ).where(~already)
        # RETURNING + count: result.rowcount is unreliable for INSERT…SELECT.
        rows = session.execute(
            insert(GpuJob)
            .from_select(["image_record_id", "task", "status"], sel)
            .returning(GpuJob.id)
        ).fetchall()
        session.commit()
        return len(rows)


@celery.task(name="backend.app.tasks.gpu_queue.recover_orphaned_gpu_jobs")
def recover_orphaned_gpu_jobs() -> int:
    """Reset expired GPU-job leases back to pending — recovers work orphaned by an
    agent that died mid-job (no graceful release) — and convert poison-loopers
    (release/expiry cycles that never reach fail()'s attempt cap) to 'error'.
    Statements are shared with GpuJobService.recover_orphaned so the sweep and
    the service can't drift. Short beat cadence so orphans get picked back up
    quickly + the queue counts read honestly. Returns the number recovered."""
    from datetime import UTC, datetime

    from ..services.ml.gpu_jobs import recover_statements

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        counts = {
            name: session.execute(stmt).rowcount or 0
            for name, stmt in recover_statements(datetime.now(UTC)).items()
        }
        session.commit()
        if counts["poison_expired"] or counts["poison_pending"]:
            log.warning(
                "gpu jobs poisoned -> error: %d crash-loop (expired lease), "
                "%d never-complete (pending)",
                counts["poison_expired"], counts["poison_pending"],
            )
        return counts["recovered"]


@celery.task(name="backend.app.tasks.gpu_queue.reprocess_gpu_jobs")
def reprocess_gpu_jobs(task_name: str = "ccip") -> int:
    """Reset every done/error job of `task_name` back to pending so the agent
    re-runs the WHOLE library under the CURRENT pipeline — e.g. after adding crop
    detectors (#1202), re-cropping existing images. Heavy + operator-triggered;
    the back-catalogue won't otherwise re-process (the backfills skip images that
    already have current-version regions). Returns the number reset."""
    from datetime import UTC, datetime

    from sqlalchemy import update

    from ..models import GpuJob

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        now = datetime.now(UTC)
        res = session.execute(
            update(GpuJob)
            .where(
                GpuJob.task == task_name,
                GpuJob.status.in_(["done", "error"]),
            )
            .values(
                status="pending", attempts=0, lease_token=None, leased_at=None,
                lease_expires_at=None, updated_at=now,
            )
        )
        session.commit()
        return res.rowcount or 0