FabledCurator/backend/app/tasks/ml.py

"""ML Celery tasks: per-image embedding, backfill discovery, head training,
model self-heal.

All run on the ml-worker (queue 'ml'). Sync sessions (Celery workers are sync
processes), same pattern as FC-2a tasks.
"""

import logging
from pathlib import Path

from celery.exceptions import SoftTimeLimitExceeded
from sqlalchemy import select
from sqlalchemy.exc import DBAPIError, OperationalError

from ..celery_app import celery
from ..models import ImageRecord, MLSettings
from ._sync_engine import sync_session_factory as _sync_session_factory

log = logging.getLogger(__name__)

IMAGES_ROOT = Path("/images")
VIDEO_EXTS = {".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v", ".wmv", ".flv"}


def _is_video(path: Path) -> bool:
    return path.suffix.lower() in VIDEO_EXTS


@celery.task(
    name="backend.app.tasks.ml.tag_and_embed",
    bind=True,
    autoretry_for=(OperationalError, DBAPIError, OSError),
    retry_backoff=5,
    retry_backoff_max=60,
    retry_jitter=True,
    max_retries=3,
    # Sized for the video branch: sample 6 frames, run tagger +
    # embedder on each (≈12 GPU ops vs 2 for an image). A loaded
    # ml-worker can take 5-10 min on a long video; bumped from
    # 5min/7min on 2026-05-28 after operator-flagged image 6288 (a
    # .mp4) hit the recovery sweep at 5 min while still legitimately
    # processing. Image runs return in seconds; the bump doesn't
    # affect their UX.
    soft_time_limit=900,   # 15 min
    time_limit=1200,       # 20 min hard
)
def tag_and_embed(self, image_id: int) -> dict:
    """Compute + store one image's SigLIP embedding.

    Video (#747): sample frames at a fixed cadence (ml_settings
    video_frame_interval_seconds, capped at video_max_frames) and mean-pool the
    per-frame SigLIP embeddings. On no-frames returns status='no_frames' (not an
    error). (Camie tagging was retired #1189 — heads + CCIP are the tag source.)
    """
    import time

    from ..services.ml.embedder import get_embedder

    # Phase + file context, so a timeout/crash names WHICH file and WHERE it
    # died instead of a bare SoftTimeLimitExceeded() (operator-flagged 2026-06-08:
    # the activity told them nothing about the file or why). `ctx` is enriched
    # once the record is loaded; both feed the worker log AND the re-raised
    # exception message (which becomes the activity's error_message).
    started = time.monotonic()
    phase = "open_session"
    ctx = f"image_id={image_id}"

    def _elapsed() -> float:
        return time.monotonic() - started

    try:
        SessionLocal = _sync_session_factory()
        with SessionLocal() as session:
            record = session.get(ImageRecord, image_id)
            if record is None:
                return {"status": "missing", "image_id": image_id}
            settings = session.execute(
                select(MLSettings).where(MLSettings.id == 1)
            ).scalar_one()

            src = Path(record.path)
            is_vid = _is_video(src)
            ctx = (
                f"image_id={image_id} path={record.path} mime={record.mime} "
                f"bytes={record.size_bytes} video={is_vid}"
            )
            log.info("tag_and_embed start: %s", ctx)
            if not src.is_file():
                log.warning("tag_and_embed file missing on disk: %s", ctx)
                return {"status": "file_missing", "image_id": image_id}

            phase = "load_models"
            embedder = get_embedder(settings.embedder_model_name)

            if is_vid:
                # Layer-3 isolation: ffprobe (a separate process) validates
                # the container before we burn GPU ops sampling frames from it.
                # A corrupt video that would crash the frame decoder is rejected
                # cleanly here instead of taking down the ml-worker.
                phase = "video_probe"
                from ..utils import safe_probe
                vprobe = safe_probe.probe_video(src)
                if not vprobe.ok:
                    log.warning(
                        "tag_and_embed bad video (%s): %s", vprobe.reason, ctx
                    )
                    return {
                        "status": "bad_video", "image_id": image_id,
                        "reason": vprobe.reason,
                    }
                phase = "video_sample_frames"
                frames = _sample_video_frames(
                    src,
                    interval=settings.video_frame_interval_seconds,
                    max_frames=settings.video_max_frames,
                )
                if not frames:
                    return {"status": "no_frames", "image_id": image_id}
                phase = "video_embed"
                import numpy as np

                # Mean-pool the per-frame SigLIP embeddings into one vector.
                embedding = np.mean(
                    [embedder.infer(f) for f in frames], axis=0
                ).astype("float32")
                for f in frames:
                    f.unlink(missing_ok=True)
            else:
                phase = "embed"
                t0 = time.monotonic()
                embedding = embedder.infer(src)
                log.info(
                    "tag_and_embed embedded in %.1fs: %s",
                    time.monotonic() - t0, ctx,
                )

            phase = "persist"
            record.siglip_embedding = embedding.tolist()
            record.siglip_model_version = settings.embedder_model_version
            session.add(record)
            session.commit()
    except SoftTimeLimitExceeded:
        log.error(
            "tag_and_embed TIMED OUT after %.0fs in phase=%s: %s",
            _elapsed(), phase, ctx,
        )
        # Re-raise as SoftTimeLimitExceeded (preserves the 'timeout' status in
        # the task_run signal) but WITH context, so the activity error_message
        # names the file + phase instead of being empty.
        raise SoftTimeLimitExceeded(
            f"timed out in phase={phase} after {_elapsed():.0f}s ({ctx})"
        ) from None
    except Exception:
        # OSError/DBAPIError/OperationalError are autoretried — re-raise the
        # ORIGINAL so the type is preserved; just make sure it's logged with
        # context first.
        log.exception(
            "tag_and_embed FAILED in phase=%s after %.0fs: %s",
            phase, _elapsed(), ctx,
        )
        raise

    log.info("tag_and_embed ok in %.1fs: %s", _elapsed(), ctx)
    return {"status": "ok", "image_id": image_id}


def _sample_video_frames(
    src: Path, *, interval: float, max_frames: int,
) -> list[Path]:
    """Sample frames at a fixed CADENCE — one every `interval` seconds — so a
    tag's frame-presence reflects real screen time regardless of video length
    (#747). The count is capped at `max_frames`: a video longer than
    interval×max_frames stretches the spacing instead of exploding the frame
    count (keeps cost bounded so a long video can't hog the single ml-worker).
    Frames are taken across the 5%–95% window (skip intro/outro black/cards) via
    per-frame fast-seek. Returns temp file paths (caller deletes); [] on failure.
    """
    import json
    import subprocess
    import tempfile

    try:
        probe = subprocess.run(
            [
                "ffprobe", "-v", "quiet", "-print_format", "json",
                "-show_format", str(src),
            ],
            check=True, capture_output=True, timeout=30,
        )
        duration = float(json.loads(probe.stdout)["format"]["duration"])
    except Exception:
        return []
    if duration <= 0:
        return []

    start, end = duration * 0.05, duration * 0.95
    span = max(end - start, 0.0)
    # Cadence count, clamped to [1, max_frames]. int(span/interval)+1 ≈ one frame
    # per `interval` seconds across the window; the cap stretches spacing on very
    # long videos.
    n = max(1, min(int(span / interval) + 1, max(1, max_frames)))
    step = span / max(n - 1, 1)
    out: list[Path] = []
    tmpdir = Path(tempfile.mkdtemp(prefix="fc_vid_"))
    for i in range(n):
        ts = start + i * step
        dest = tmpdir / f"frame_{i:04d}.jpg"
        try:
            subprocess.run(
                [
                    "ffmpeg", "-ss", f"{ts:.2f}", "-i", str(src),
                    "-frames:v", "1", "-q:v", "3", "-y", str(dest),
                ],
                check=True, capture_output=True, timeout=30,
            )
            if dest.is_file():
                out.append(dest)
        except Exception:
            continue
    return out


@celery.task(name="backend.app.tasks.ml.backfill", bind=True)
def backfill(self) -> int:
    """Enqueue tag_and_embed (embed-only) for images with no SigLIP embedding.
    Keyset pagination by id ASC (restart-safe).

    NB: a siglip MODEL-VERSION mismatch (an operator model swap, #1190) is NOT
    re-embedded here — the CPU ml-worker can't churn the library at 384/512px;
    the GPU agent owns version re-embeds via the 'embed' job.
    """
    SessionLocal = _sync_session_factory()
    enqueued = 0
    last_id = 0
    with SessionLocal() as session:
        while True:
            rows = session.execute(
                select(ImageRecord.id)
                .where(ImageRecord.id > last_id)
                .where(ImageRecord.siglip_embedding.is_(None))
                .order_by(ImageRecord.id.asc())
                .limit(500)
            ).scalars().all()
            if not rows:
                break
            for image_id in rows:
                tag_and_embed.delay(image_id)
                enqueued += 1
            last_id = rows[-1]
    return enqueued


@celery.task(
    name="backend.app.tasks.ml.tag_eval_run",
    bind=True,
    # The head-vs-centroid eval (#1130) loads embeddings + fits sklearn heads
    # for several concepts — minutes, not seconds. Runs on the ml queue because
    # only that worker has numpy/scikit-learn.
    soft_time_limit=1800, time_limit=2100,
)
def tag_eval_run(self, run_id: int) -> str:
    """Compute the eval report into the persisted TagEvalRun row so it survives
    navigation (the admin card rehydrates from the row, not transient state)."""
    from datetime import UTC, datetime

    from ..models import TagEvalRun
    from ..services.ml.tag_eval import run_eval

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        run = session.get(TagEvalRun, run_id)
        if run is None:
            return "missing"
        run.last_progress_at = datetime.now(UTC)
        session.commit()
        try:
            report = run_eval(session, run.params)
        except SoftTimeLimitExceeded:
            run.status = "error"
            run.error = "timed out"
            run.finished_at = datetime.now(UTC)
            session.commit()
            raise
        except Exception as exc:
            log.exception("tag_eval_run %d failed", run_id)
            run.status = "error"
            run.error = str(exc)
            run.finished_at = datetime.now(UTC)
            session.commit()
            return "error"
        run.report = report
        run.status = "ready"
        run.finished_at = datetime.now(UTC)
        session.commit()
        return "ready"


@celery.task(
    name="backend.app.tasks.ml.train_heads",
    bind=True,
    # Trains a logistic-regression head per eligible concept over stored SigLIP
    # embeddings — minutes for a full library. Runs on the ml queue (only that
    # worker has scikit-learn). Commits per head so a kill leaves progress.
    soft_time_limit=3600, time_limit=3900,
)
def train_heads(self, run_id: int) -> str:
    """(Re)train all eligible concept heads into tag_head, tracked by the
    HeadTrainingRun row so the admin card shows live + historical status."""
    from datetime import UTC, datetime

    from ..models import HeadTrainingRun
    from ..services.ml.heads import train_all_heads

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        run = session.get(HeadTrainingRun, run_id)
        if run is None:
            return "missing"
        run.last_progress_at = datetime.now(UTC)
        session.commit()
        try:
            result = train_all_heads(session, run.params, run)
        except SoftTimeLimitExceeded:
            run.status = "error"
            run.error = "timed out"
            run.finished_at = datetime.now(UTC)
            session.commit()
            raise
        except Exception as exc:
            log.exception("train_heads %d failed", run_id)
            run.status = "error"
            run.error = str(exc)
            run.finished_at = datetime.now(UTC)
            session.commit()
            return "error"
        run.n_trained = result["n_trained"]
        run.n_skipped = result["n_skipped"]
        run.status = "ready"
        run.finished_at = datetime.now(UTC)
        session.commit()
        return "ready"


@celery.task(name="backend.app.tasks.ml.scheduled_train_heads")
def scheduled_train_heads() -> str:
    """Nightly passive retrain (#114): fold the day's accepts/rejects + any
    newly-eligible concepts into the heads without the operator clicking. Skips
    if a run is already in flight (one at a time). Creates + COMMITS the run row
    before dispatching so the ml-queue worker can always find it."""
    from datetime import UTC, datetime

    from sqlalchemy import select as sa_select

    from ..models import HeadTrainingRun

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        running = session.execute(
            sa_select(HeadTrainingRun.id).where(HeadTrainingRun.status == "running")
        ).scalar_one_or_none()
        if running is not None:
            return "already running"
        run = HeadTrainingRun(
            params={"source": "scheduled"}, status="running",
            last_progress_at=datetime.now(UTC),
        )
        session.add(run)
        session.commit()
        run_id = run.id
    train_heads.delay(run_id)
    return "dispatched"


@celery.task(
    name="backend.app.tasks.ml.apply_head_tags",
    bind=True,
    # Scores the whole library against the graduated heads and applies their
    # tags (or, dry_run, just counts). Streams embeddings in chunks; numpy only,
    # but ml queue keeps it off the API workers. Commits per chunk.
    soft_time_limit=3600, time_limit=3900,
)
def apply_head_tags(self, run_id: int) -> str:
    """Run an earned-auto-apply sweep into the persisted HeadAutoApplyRun row."""
    from datetime import UTC, datetime

    from ..models import HeadAutoApplyRun
    from ..services.ml.heads import auto_apply_sweep

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        run = session.get(HeadAutoApplyRun, run_id)
        if run is None:
            return "missing"
        run.last_progress_at = datetime.now(UTC)
        session.commit()
        try:
            result = auto_apply_sweep(session, run, run.dry_run)
        except SoftTimeLimitExceeded:
            run.status = "error"
            run.error = "timed out"
            run.finished_at = datetime.now(UTC)
            session.commit()
            raise
        except Exception as exc:
            log.exception("apply_head_tags %d failed", run_id)
            run.status = "error"
            run.error = str(exc)
            run.finished_at = datetime.now(UTC)
            session.commit()
            return "error"
        run.n_applied = result["n_applied"]
        run.report = {"concepts": result["concepts"]}
        run.status = "ready"
        run.finished_at = datetime.now(UTC)
        session.commit()
        return "ready"


@celery.task(name="backend.app.tasks.ml.scheduled_apply_head_tags")
def scheduled_apply_head_tags() -> str:
    """Daily passive auto-apply sweep (#114) — only when the master switch is on.
    Skips if a sweep is already in flight. Creates + COMMITS the run before
    dispatching so the worker always finds it."""
    from datetime import UTC, datetime

    from sqlalchemy import select as sa_select

    from ..models import HeadAutoApplyRun, MLSettings

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        enabled = session.execute(
            sa_select(MLSettings.head_auto_apply_enabled).where(MLSettings.id == 1)
        ).scalar_one_or_none()
        if not enabled:
            return "disabled"
        running = session.execute(
            sa_select(HeadAutoApplyRun.id).where(HeadAutoApplyRun.status == "running")
        ).scalar_one_or_none()
        if running is not None:
            return "already running"
        run = HeadAutoApplyRun(
            dry_run=False, params={"dry_run": False, "source": "scheduled"},
            status="running", last_progress_at=datetime.now(UTC),
        )
        session.add(run)
        session.commit()
        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 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.

    'siglip' gates on the RESULT (no concept region yet) rather than on a prior
    job, so it picks up the back-catalogue of images that were CCIP-embedded
    before concept crops existed, and retries images whose concept embed failed —
    without re-touching their figure/CCIP regions."""
    from sqlalchemy import exists, insert, literal, or_
    from sqlalchemy import select as sa_select

    from ..models import GpuJob, ImageRecord, ImageRegion, MLSettings

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        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,
            )
            queued = exists().where(
                GpuJob.image_record_id == ImageRecord.id,
                GpuJob.task == "embed",
                GpuJob.status.in_(["pending", "leased"]),
            )
            sel = sa_select(
                ImageRecord.id, literal("embed"), literal("pending")
            ).where(stale).where(~queued)
        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,
            )
            queued = exists().where(
                GpuJob.image_record_id == ImageRecord.id,
                GpuJob.task == "siglip",
                GpuJob.status.in_(["pending", "leased"]),
            )
            sel = sa_select(
                ImageRecord.id, literal("siglip"), literal("pending")
            ).where(~has_current_concept).where(~queued)
        else:
            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)
        # 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.ml.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). 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 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.status == "leased", GpuJob.lease_expires_at < now)
            .values(
                status="pending", lease_token=None, leased_at=None,
                lease_expires_at=None, updated_at=now,
            )
        )
        session.commit()
        return res.rowcount or 0


@celery.task(
    name="backend.app.tasks.ml.scheduled_ccip_auto_apply",
    soft_time_limit=1800, time_limit=2100,
)
def scheduled_ccip_auto_apply() -> str:
    """Auto-tag confident CCIP character matches (source='ccip_auto') so identity
    tags keep flowing without a button. No-op unless ccip_auto_apply_enabled.
    References come only from single-character images (unambiguous); a tag is
    applied where any figure's best cosine to a character's prototypes clears
    ccip_auto_apply_threshold and it isn't already applied/rejected. Reversible."""
    import numpy as np
    from sqlalchemy import func
    from sqlalchemy import select as sa_select
    from sqlalchemy.dialects.postgresql import insert as pg_insert

    from ..models import ImageRegion, MLSettings, Tag, TagKind, TagSuggestionRejection
    from ..models.tag import image_tag

    fig = ("face", "figure")

    def _l2(m):
        n = np.linalg.norm(m, axis=1, keepdims=True)
        n[n == 0] = 1.0
        return m / n

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        s = session.get(MLSettings, 1)
        if s is None or not s.ccip_auto_apply_enabled:
            return "disabled"
        thr = float(s.ccip_auto_apply_threshold)

        single = (
            sa_select(image_tag.c.image_record_id)
            .join(Tag, Tag.id == image_tag.c.tag_id)
            .where(Tag.kind == TagKind.character)
            .group_by(image_tag.c.image_record_id)
            .having(func.count() == 1)
        )
        ref_rows = session.execute(
            sa_select(image_tag.c.tag_id, ImageRegion.ccip_embedding)
            .select_from(ImageRegion)
            .join(
                image_tag,
                image_tag.c.image_record_id == ImageRegion.image_record_id,
            )
            .join(Tag, Tag.id == image_tag.c.tag_id)
            .where(Tag.kind == TagKind.character)
            .where(ImageRegion.kind.in_(fig))
            .where(ImageRegion.ccip_embedding.is_not(None))
            .where(ImageRegion.image_record_id.in_(single))
        ).all()
        if not ref_rows:
            return "no-references"

        by_char: dict[int, list] = {}
        for tid, vec in ref_rows:
            by_char.setdefault(tid, []).append(vec)
        ref_tags = list(by_char)
        mats = [_l2(np.asarray(by_char[t], dtype=np.float32)) for t in ref_tags]
        allref = np.vstack(mats)                                  # (total, 768)
        seg = np.cumsum([0] + [len(m) for m in mats])[:-1]        # per-char start

        # Per character: images that already carry OR rejected the tag — skip.
        skip = {t: set() for t in ref_tags}
        for t in ref_tags:
            for (iid,) in session.execute(
                sa_select(image_tag.c.image_record_id).where(
                    image_tag.c.tag_id == t
                )
            ):
                skip[t].add(iid)
            for (iid,) in session.execute(
                sa_select(TagSuggestionRejection.image_record_id).where(
                    TagSuggestionRejection.tag_id == t
                )
            ):
                skip[t].add(iid)

        img_ids = list(session.execute(
            sa_select(ImageRegion.image_record_id)
            .where(ImageRegion.kind.in_(fig), ImageRegion.ccip_embedding.is_not(None))
            .distinct()
        ).scalars())

        applied = 0
        chunk_n = 500
        for start in range(0, len(img_ids), chunk_n):
            chunk = img_ids[start:start + chunk_n]
            rows = session.execute(
                sa_select(ImageRegion.image_record_id, ImageRegion.ccip_embedding)
                .where(
                    ImageRegion.image_record_id.in_(chunk),
                    ImageRegion.kind.in_(fig),
                    ImageRegion.ccip_embedding.is_not(None),
                )
            ).all()
            by_img: dict[int, list] = {}
            for iid, vec in rows:
                by_img.setdefault(iid, []).append(vec)
            for iid, vecs in by_img.items():
                q = _l2(np.asarray(vecs, dtype=np.float32))      # (nq, 768)
                colmax = (q @ allref.T).max(axis=0)              # (total,)
                charmax = np.maximum.reduceat(colmax, seg)       # (n_chars,)
                for ci in np.where(charmax >= thr)[0]:
                    t = ref_tags[int(ci)]
                    if iid in skip[t]:
                        continue
                    skip[t].add(iid)
                    session.execute(
                        pg_insert(image_tag)
                        .values(
                            image_record_id=iid, tag_id=t, source="ccip_auto",
                        )
                        .on_conflict_do_nothing()
                    )
                    applied += 1
            session.commit()
        return f"applied={applied}"