bvandeusen
80f8eb4756
The siglip/ccip backfills skip images that already have current-version regions, so adding crop detectors only affected NEW images — the back-catalogue would never be re-cropped. Add a reprocess trigger that resets every done/error job of a task back to pending, so the agent re-runs the FULL pipeline (figure detection + CCIP + concept/panel crops) over the whole library under the current detectors. - reprocess_gpu_jobs(task='ccip') task + POST /api/gpu/reprocess. - gpu store reprocess() + GpuAgentCard "Re-process library (re-detect + re-crop)" button with a confirm (it's heavy). - Test: a done job resets to pending (attempts cleared). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
702 lines
27 KiB
Python
702 lines
27 KiB
Python
"""ML Celery tasks: per-image embedding, backfill discovery, head training,
|
||
model self-heal.
|
||
|
||
All run on the ml-worker (queue 'ml'). 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
|
||
|
||
|
||
@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:
|
||
"""Compute + store one image's SigLIP embedding.
|
||
|
||
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. On no-frames returns status='no_frames' (not an
|
||
error). (Camie tagging was retired #1189 — heads + CCIP are the tag source.)
|
||
"""
|
||
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("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"
|
||
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(
|
||
"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"
|
||
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(
|
||
"tag_and_embed 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(
|
||
"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: %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 tag_and_embed (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.
|
||
"""
|
||
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:
|
||
tag_and_embed.delay(image_id)
|
||
enqueued += 1
|
||
last_id = rows[-1]
|
||
return enqueued
|
||
|
||
|
||
@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"
|
||
|
||
|
||
@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.enqueue_gpu_backfill")
|
||
def enqueue_gpu_backfill(task_name: str) -> int:
|
||
"""Enqueue a gpu_job for every image that still needs `task_name` (one
|
||
INSERT…SELECT, so it scales to a full library). The desktop agent drains the
|
||
queue over HTTP. Returns the number enqueued.
|
||
|
||
'siglip' gates on the RESULT (no concept region yet) rather than on a prior
|
||
job, so it picks up the back-catalogue of images that were CCIP-embedded
|
||
before concept crops existed, and retries images whose concept embed failed —
|
||
without re-touching their figure/CCIP regions."""
|
||
from sqlalchemy import exists, insert, literal, or_
|
||
from sqlalchemy import select as sa_select
|
||
|
||
from ..models import GpuJob, ImageRecord, ImageRegion, MLSettings
|
||
|
||
SessionLocal = _sync_session_factory()
|
||
with SessionLocal() as session:
|
||
cur_version = session.execute(
|
||
select(MLSettings.embedder_model_version).where(MLSettings.id == 1)
|
||
).scalar_one()
|
||
if task_name == "embed":
|
||
# Whole-image GPU re-embed (#1190): images with no embedding, or one
|
||
# stamped under a DIFFERENT model version (an operator model swap).
|
||
stale = or_(
|
||
ImageRecord.siglip_embedding.is_(None),
|
||
ImageRecord.siglip_model_version.is_(None),
|
||
ImageRecord.siglip_model_version != cur_version,
|
||
)
|
||
queued = exists().where(
|
||
GpuJob.image_record_id == ImageRecord.id,
|
||
GpuJob.task == "embed",
|
||
GpuJob.status.in_(["pending", "leased"]),
|
||
)
|
||
sel = sa_select(
|
||
ImageRecord.id, literal("embed"), literal("pending")
|
||
).where(stale).where(~queued)
|
||
elif task_name == "siglip":
|
||
# Concept-crop re-embed: enqueue when there's no concept region AT THE
|
||
# CURRENT model version — so a model swap re-triggers crops too, not
|
||
# only the never-embedded back-catalogue.
|
||
has_current_concept = exists().where(
|
||
ImageRegion.image_record_id == ImageRecord.id,
|
||
ImageRegion.kind == "concept",
|
||
ImageRegion.embedding_version == cur_version,
|
||
)
|
||
queued = exists().where(
|
||
GpuJob.image_record_id == ImageRecord.id,
|
||
GpuJob.task == "siglip",
|
||
GpuJob.status.in_(["pending", "leased"]),
|
||
)
|
||
sel = sa_select(
|
||
ImageRecord.id, literal("siglip"), literal("pending")
|
||
).where(~has_current_concept).where(~queued)
|
||
else:
|
||
already = exists().where(
|
||
GpuJob.image_record_id == ImageRecord.id,
|
||
GpuJob.task == task_name,
|
||
GpuJob.status.in_(["pending", "leased", "done"]),
|
||
)
|
||
sel = sa_select(
|
||
ImageRecord.id, literal(task_name), literal("pending")
|
||
).where(~already)
|
||
# RETURNING + count: result.rowcount is unreliable for INSERT…SELECT.
|
||
rows = session.execute(
|
||
insert(GpuJob)
|
||
.from_select(["image_record_id", "task", "status"], sel)
|
||
.returning(GpuJob.id)
|
||
).fetchall()
|
||
session.commit()
|
||
return len(rows)
|
||
|
||
|
||
@celery.task(name="backend.app.tasks.ml.recover_orphaned_gpu_jobs")
|
||
def recover_orphaned_gpu_jobs() -> int:
|
||
"""Reset expired GPU-job leases back to pending — recovers work orphaned by an
|
||
agent that died mid-job (no graceful release). Short beat cadence so orphans
|
||
get picked back up quickly + the queue counts read honestly. Returns the
|
||
number recovered."""
|
||
from datetime import UTC, datetime
|
||
|
||
from sqlalchemy import update
|
||
|
||
from ..models import GpuJob
|
||
|
||
SessionLocal = _sync_session_factory()
|
||
with SessionLocal() as session:
|
||
now = datetime.now(UTC)
|
||
res = session.execute(
|
||
update(GpuJob)
|
||
.where(GpuJob.status == "leased", GpuJob.lease_expires_at < now)
|
||
.values(
|
||
status="pending", lease_token=None, leased_at=None,
|
||
lease_expires_at=None, updated_at=now,
|
||
)
|
||
)
|
||
session.commit()
|
||
return res.rowcount or 0
|
||
|
||
|
||
@celery.task(name="backend.app.tasks.ml.reprocess_gpu_jobs")
|
||
def reprocess_gpu_jobs(task_name: str = "ccip") -> int:
|
||
"""Reset every done/error job of `task_name` back to pending so the agent
|
||
re-runs the WHOLE library under the CURRENT pipeline — e.g. after adding crop
|
||
detectors (#1202), re-cropping existing images. Heavy + operator-triggered;
|
||
the back-catalogue won't otherwise re-process (the backfills skip images that
|
||
already have current-version regions). Returns the number reset."""
|
||
from datetime import UTC, datetime
|
||
|
||
from sqlalchemy import update
|
||
|
||
from ..models import GpuJob
|
||
|
||
SessionLocal = _sync_session_factory()
|
||
with SessionLocal() as session:
|
||
now = datetime.now(UTC)
|
||
res = session.execute(
|
||
update(GpuJob)
|
||
.where(
|
||
GpuJob.task == task_name,
|
||
GpuJob.status.in_(["done", "error"]),
|
||
)
|
||
.values(
|
||
status="pending", attempts=0, lease_token=None, leased_at=None,
|
||
lease_expires_at=None, updated_at=now,
|
||
)
|
||
)
|
||
session.commit()
|
||
return res.rowcount or 0
|
||
|
||
|
||
@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}"
|