bvandeusen
77baee49fd
Two cadences for keeping heads in sync with your tagging: - PASSIVE: a nightly `scheduled_train_heads` beat (skips if a run is already in flight; creates+commits the run row before dispatching train_heads so the ml worker always finds it). Folds the day's accepts/rejects + newly-eligible concepts into the heads without anyone clicking. - ACTIVE: a "Retrain heads" button in the Explore trail bar — bank the +/- feedback you just gave while walking content, without a trip to Settings. Shared logic in a new useHeadTraining composable (trigger + poll + start/finish toasts), used by the Explore button; reflects an already-running run (incl. the nightly one) on mount. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
662 lines
25 KiB
Python
662 lines
25 KiB
Python
"""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"
|
||
|
||
|
||
@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"
|