"""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)