Files
FabledCurator/backend/app/tasks/ml.py
T
bvandeusen 6e3c5f697f
CI / lint (push) Successful in 3s
CI / frontend-build (push) Successful in 20s
CI / backend-lint-and-test (push) Successful in 26s
CI / integration (push) Successful in 3m23s
feat(ml): tag-eval backend — head-vs-centroid learning-curve eval (persisted)
Slice 1 of milestone #114 (tagging v2). Proves the frozen-embedding + trained-
head spine on the operator's own data, reusing the SigLIP embeddings already
stored on image_record — no re-embedding, no GPU.

Per concept: train a logistic-regression HEAD (positives + negatives = explicit
rejections + sampled unlabeled) vs the old single-CENTROID baseline; report
cross-validated precision/recall/AP for both, a LEARNING CURVE (AP/F1 as tagged
positives grow 10→30→100→300), and example image ids (head-would-suggest /
head-doubts-positive) to eyeball.

Persisted so the report SURVIVES navigation (operator-flagged): the run + full
report live in a new tag_eval_run row (mirrors library_audit_run); the admin
card will rehydrate from GET on mount, not transient state.

- models.TagEvalRun + migration 0056; runs on the ml queue (only worker with
  numpy/sklearn) — numpy/sklearn lazy-imported so the API can still enqueue.
- services/ml/tag_eval (compute + start helper, one-running guard), tasks.ml
  .tag_eval_run, api/tag-eval (POST create, GET history light / detail w/ report).
- recover_stalled_tag_eval_runs sweep + retention (keep last 20) + 5-min beat
  (rule 89). scikit-learn added to requirements-ml.
- tests: param normalization + the rehydrate read-path + create/conflict.

Frontend admin card (trigger + render persisted report) follows next.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-27 22:49:10 -04:00

586 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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 (12 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"