"""ML Celery tasks: per-image inference, backfill discovery, centroid recompute, allowlist auto-apply, model self-heal. All run on the ml-worker (queue 'ml') except recompute_centroids and apply_allowlist_tags sweeps which are 'maintenance' lane. 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 delete, select from sqlalchemy.exc import DBAPIError, OperationalError from ..celery_app import celery from ..models import ImagePrediction, 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: """Run Camie + SigLIP on one image; store predictions + embedding; then enqueue per-image allowlist application. Video (#747): sample frames at a fixed cadence (ml_settings video_frame_interval_seconds, capped at video_max_frames), keep a tag only if it appears in >= video_min_tag_frames frames and average its confidence over those frames (mean-pool, not max — kills one-frame noise); mean-pool the SigLIP embeddings. On no-frames returns status='no_frames' (not an error). """ import time from ..services.ml.embedder import get_embedder from ..services.ml.tagger import get_tagger # 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" tagger = get_tagger() embedder = get_embedder() if is_vid: # Layer-3 isolation: ffprobe (a separate process) validates # the container before we burn ~20 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. Operator-flagged 2026-05-28. 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" t0 = time.monotonic() frames = _sample_video_frames( src, interval=settings.video_frame_interval_seconds, max_frames=settings.video_max_frames, ) log.info( "tag_and_embed sampled %d frame(s) in %.1fs: %s", len(frames), time.monotonic() - t0, ctx, ) if not frames: return {"status": "no_frames", "image_id": image_id} phase = "video_infer" import numpy as np preds = _aggregate_video_predictions( [tagger.infer(f, store_floor=settings.tagger_store_floor) for f in frames], min_frames=settings.video_min_tag_frames, ) embedding = np.mean( [embedder.infer(f) for f in frames], axis=0 ).astype("float32") log.info( "tag_and_embed video aggregated %d tag(s) from %d frame(s) " "(min_frames=%d): %s", len(preds), len(frames), settings.video_min_tag_frames, ctx, ) for f in frames: f.unlink(missing_ok=True) else: phase = "tag" t0 = time.monotonic() raw = tagger.infer(src, store_floor=settings.tagger_store_floor) log.info( "tag_and_embed tagged in %.1fs (%d tags): %s", time.monotonic() - t0, len(raw), ctx, ) preds = { name: {"category": p.category, "confidence": p.confidence} for name, p in raw.items() } 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.tagger_model_version = settings.tagger_model_version record.siglip_embedding = embedding.tolist() record.siglip_model_version = settings.embedder_model_version session.add(record) # Write the normalized image_prediction rows (#768) — the sole home # for predictions now (image_record.tagger_predictions was dropped in # migration 0046). Delete-then-insert keeps a re-tag idempotent; # tagger_store_floor was already applied in tagger.infer, so preds is # the >=floor set. session.execute( delete(ImagePrediction).where( ImagePrediction.image_record_id == image_id ) ) session.add_all([ ImagePrediction( image_record_id=image_id, raw_name=name, category=p.get("category", "general"), score=float(p.get("confidence", 0.0)), ) for name, p in preds.items() ]) 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 (%d tags): %s", _elapsed(), len(preds), ctx ) apply_allowlist_tags.delay(image_id=image_id) return {"status": "ok", "image_id": image_id, "tags": len(preds)} 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 def _aggregate_video_predictions(per_frame: list[dict], *, min_frames: int) -> dict: """Aggregate per-frame {name: TagPrediction} into one prediction set (#747). A tag is kept only if it appears (≥ the tagger store floor, already applied) in at least `min_frames` of the sampled frames — because sampling is at a fixed cadence, that means it was on screen for roughly min_frames×interval seconds, so a single-frame flicker / scene-transition artifact is dropped while a genuine scene-local tag in a long video survives. Confidence is the MEAN over the frames where the tag appears (not max — max re-inflated the one-frame noise this whole change exists to remove). `min_frames` is clamped to the number of frames actually sampled so a very short video (1–2 frames) still tags instead of dropping everything. """ n = len(per_frame) if n == 0: return {} threshold = max(1, min(min_frames, n)) agg: dict[str, dict] = {} for frame_preds in per_frame: for name, p in frame_preds.items(): cur = agg.get(name) if cur is None: agg[name] = {"category": p.category, "sum": p.confidence, "count": 1} else: cur["sum"] += p.confidence cur["count"] += 1 return { name: {"category": v["category"], "confidence": v["sum"] / v["count"]} for name, v in agg.items() if v["count"] >= threshold } @celery.task(name="backend.app.tasks.ml.backfill", bind=True) def backfill(self) -> int: """Enqueue tag_and_embed for images missing predictions/embeddings for the current model versions. Keyset pagination by id ASC (restart-safe). """ SessionLocal = _sync_session_factory() enqueued = 0 last_id = 0 with SessionLocal() as session: settings = session.execute( select(MLSettings).where(MLSettings.id == 1) ).scalar_one() while True: rows = session.execute( select(ImageRecord.id) .where(ImageRecord.id > last_id) .where( (ImageRecord.tagger_model_version.is_(None)) | ( ImageRecord.tagger_model_version != settings.tagger_model_version ) | (ImageRecord.siglip_embedding.is_(None)) | ( ImageRecord.siglip_model_version != settings.embedder_model_version ) ) .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.apply_allowlist_tags", bind=True, # Audit 2026-06-02 — the full-sweep mode (neither tag_id nor image_id) # is O(images × allowlist) and legitimately runs >5 min on large # libraries. Cap matches the maintenance queue's recovery threshold. soft_time_limit=1800, time_limit=2100, ) def apply_allowlist_tags(self, tag_id: int | None = None, image_id: int | None = None) -> int: """Retroactively apply allowlisted tags. Modes: - tag_id only : scan all images for this tag. - image_id only : scan all allowlisted tags for this image. - both : just the (image, tag) pair. - neither : full sweep (daily beat). Skips: already-applied, rejected (tag_suggestion_rejection), or confidence below the tag's allowlist min_confidence. Applied with source='ml_auto'. """ from sqlalchemy import and_ from sqlalchemy import select as sa_select from sqlalchemy.dialects.postgresql import insert as pg_insert from ..models import TagAllowlist, TagSuggestionRejection from ..models.tag import image_tag SessionLocal = _sync_session_factory() applied = 0 with SessionLocal() as session: allow_rows = session.execute( sa_select(TagAllowlist.tag_id, TagAllowlist.min_confidence) if tag_id is None else sa_select( TagAllowlist.tag_id, TagAllowlist.min_confidence ).where(TagAllowlist.tag_id == tag_id) ).all() allow = {r[0]: r[1] for r in allow_rows} if not allow: return 0 # Images that have any predictions (#768: from image_prediction, not # the old JSON column), optionally narrowed to one image. img_ids_query = sa_select(ImagePrediction.image_record_id).distinct() if image_id is not None: img_ids_query = img_ids_query.where( ImagePrediction.image_record_id == image_id ) for (img_id,) in session.execute(img_ids_query).all(): preds = _load_predictions_sync(session, img_id) for a_tag_id, min_conf in allow.items(): exists = session.execute( sa_select(image_tag.c.tag_id).where( and_( image_tag.c.image_record_id == img_id, image_tag.c.tag_id == a_tag_id, ) ) ).scalar_one_or_none() if exists is not None: continue rej = session.get( TagSuggestionRejection, (img_id, a_tag_id) ) if rej is not None: continue from ..models import Tag tag = session.get(Tag, a_tag_id) if tag is None: continue conf = _confidence_for_tag(session, tag, preds) if conf is None or conf < min_conf: continue stmt = pg_insert(image_tag).values( image_record_id=img_id, tag_id=a_tag_id, source="ml_auto", ) stmt = stmt.on_conflict_do_nothing( index_elements=["image_record_id", "tag_id"] ) session.execute(stmt) applied += 1 session.commit() return applied def _load_predictions_sync(session, image_id: int) -> dict: """Predictions for one image from image_prediction (#768), in the {raw_name: {category, confidence}} shape _confidence_for_tag consumes — keeps the allowlist resolution logic unchanged.""" from sqlalchemy import select as sa_select rows = session.execute( sa_select( ImagePrediction.raw_name, ImagePrediction.category, ImagePrediction.score, ).where(ImagePrediction.image_record_id == image_id) ).all() return { r.raw_name: {"category": r.category, "confidence": r.score} for r in rows } def _confidence_for_tag(session, tag, preds: dict) -> float | None: """Highest confidence among predictions that resolve to `tag` — either the prediction name equals the tag name, or an alias maps (prediction name, category) -> tag.id. """ from sqlalchemy import select as sa_select from ..models import TagAlias best: float | None = None direct = preds.get(tag.name) if direct is not None: best = float(direct.get("confidence", 0.0)) alias_rows = session.execute( sa_select(TagAlias.alias_string, TagAlias.alias_category).where( TagAlias.canonical_tag_id == tag.id ) ).all() for alias_string, alias_category in alias_rows: p = preds.get(alias_string) if p is None: continue if p.get("category") != alias_category: continue c = float(p.get("confidence", 0.0)) if best is None or c > best: best = c return best @celery.task(name="backend.app.tasks.ml.recompute_centroid", bind=True) def recompute_centroid(self, tag_id: int) -> bool: import asyncio from ..services.ml.centroids import CentroidService from ._async_session import async_session_factory async def _run() -> bool: # Per-task NullPool engine bound to THIS asyncio.run loop — the shared # process-wide engine reuses connections across loops and raises # "Future attached to a different loop" on every call after the first. async_factory, async_engine = async_session_factory() try: async with async_factory() as session: svc = CentroidService(session) result = await svc.recompute_for_tag(tag_id) await session.commit() return result finally: await async_engine.dispose() return asyncio.run(_run()) @celery.task( name="backend.app.tasks.ml.recompute_centroids", bind=True, # Audit 2026-06-02 — drifted-centroid rebuild over potentially # hundreds of tags. soft_time_limit=1800, time_limit=2100, ) def recompute_centroids(self) -> int: """Daily: find drifted centroids, enqueue recompute_centroid for each.""" import asyncio from ..services.ml.centroids import CentroidService from ._async_session import async_session_factory async def _list() -> list[int]: # Per-task NullPool engine bound to this loop (see recompute_centroid). async_factory, async_engine = async_session_factory() try: async with async_factory() as session: return await CentroidService(session).list_drifted() finally: await async_engine.dispose() drifted = asyncio.run(_list()) for tid in drifted: recompute_centroid.delay(tid) return len(drifted) @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"