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'), which under B3 (2026-07-02) is an
OPTIONAL container: its only processing role is the CPU whole-image embed
fallback (gated by ml_settings.cpu_embed_enabled) for stacks without a GPU
agent — plus head training / auto-apply, which need sklearn/numpy and so
live on this image. GPU-queue coordination (backfill enqueues, orphan
recovery, reprocess) deliberately does NOT live here — see tasks/gpu_queue.py
(maintenance lane), so the agent pipeline works with no ml-worker at all.
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


def cpu_embed_enabled() -> bool:
    """Dispatch gate for the CPU embed fallback (B3, operator 2026-07-02):
    stacks that run a GPU agent and DROP the (optional) ml-worker container
    turn ml_settings.cpu_embed_enabled off, so the import hooks stop queueing
    embed work into a queue nothing consumes — the daily GPU 'embed' backfill
    covers those images instead. Opens its own short session because the four
    dispatch sites sit in different session scopes; defaults ON when the
    settings row is missing (a fresh install must work agent-less)."""
    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        val = session.execute(
            select(MLSettings.cpu_embed_enabled).where(MLSettings.id == 1)
        ).scalar_one_or_none()
    return True if val is None else bool(val)


@celery.task(
    name="backend.app.tasks.ml.embed_image",
    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 embed_image(self, image_id: int) -> dict:
    """Compute + store one image's whole-image SigLIP embedding — the CPU
    fallback path (B3, operator 2026-07-02): this is the ml-worker's ONLY
    processing role, keeping search/similarity/head-suggestions alive on
    deployments without a GPU agent. Detection, cropping and CCIP are
    deliberately agent-only, and their backfill predicates read image_region /
    gpu_job state — never image_record.siglip_embedding — so a CPU whole-image
    embed can NEVER mark crop work as done. (Renamed from tag_and_embed —
    Camie tagging was retired #1189; the old name kept implying a tagging step
    that no longer exists.)

    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 — the same shape as the GPU agent's video
    handling. On no-frames returns status='no_frames' (not an error).
    """
    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("embed_image start: %s", ctx)
            if not src.is_file():
                log.warning("embed_image 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(
                        "embed_image 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(
                    "embed_image 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(
            "embed_image 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(
            "embed_image FAILED in phase=%s after %.0fs: %s",
            phase, _elapsed(), ctx,
        )
        raise

    log.info("embed_image 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 embed_image (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.
    """
    if not cpu_embed_enabled():
        log.info("cpu backfill skipped: cpu_embed_enabled is off (B3 — the "
                 "GPU 'embed' backfill owns whole-image embeds on this stack)")
        return 0
    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:
                embed_image.delay(image_id)
                enqueued += 1
            last_id = rows[-1]
    return enqueued


@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.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}"