feat(ml): operator model swap — GPU re-embed + embedder as a setting (#1190)
Make the SigLIP embedder an operator choice (drop-in to SigLIP 2:
google/siglip2-so400m-patch16-512 is a verified 1152-d model at 512px → no
schema change, better small-cue fidelity). A swap = set model + re-embed +
retrain, all operator-driven; the GPU agent does the re-embed so it's fast.
- settings: embedder_model_name is now a setting (migration 0065) alongside the
existing embedder_model_version; both editable + validated (non-empty) in the
ml admin API. The server embedder loads by HF name (AutoImageProcessor/Model,
model-agnostic), preferring the pre-downloaded local dir for the default so
existing deploys don't re-download; rebuilds on a name change.
- agent: new 'embed' job = whole-image SigLIP embedding (mean-pool video frames)
under the lease-announced model → POST /jobs/submit_embedding writes
image_record.siglip_embedding + siglip_model_version. The lease now announces
the model FROM THE SETTING (not a constant).
- re-embed routing: enqueue_gpu_backfill('embed') selects unembedded + stale-
version images; 'siglip' now re-embeds concept crops whose version != current
(so a swap re-triggers crops, not just the never-embedded back-catalogue). The
CPU ml-worker backfill no longer re-embeds on a version mismatch (it can't
churn the library at 512px) — the GPU agent owns version re-embeds. Daily
'embed' + 'siglip' beats self-heal.
- scoring: score_image only bags embeddings in the CURRENT model's space (whole-
image gated by siglip_model_version, concept regions by embedding_version) so a
mid-swap stale vector isn't scored by new-space heads; legacy NULL = current.
- UI: GpuAgentCard "Embedding model (advanced)" — edit name/version, Save, and
"Re-embed library (GPU)" (queues embed + siglip); points at SigLIP 2.
Tests: lease announces model + submit_embedding round-trip; enqueue 'embed'
selects stale/unembedded; stale-version excluded from scoring; embedder model
settable + empty rejected; siglip gate updated to current-version concept.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
This commit is contained in:
+34
-6
@@ -17,7 +17,6 @@ from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
from ..extensions import get_session
|
||||
from ..models import AppSetting, GpuJob, ImageRecord, MLSettings
|
||||
from ..services.gallery_service import image_url
|
||||
from ..services.ml.embedder import MODEL_NAME as EMBED_MODEL_NAME
|
||||
from ..services.ml.gpu_jobs import GpuJobService
|
||||
from ..services.ml.regions import RegionService
|
||||
|
||||
@@ -138,11 +137,12 @@ async def lease():
|
||||
# For video/animated: the agent samples at this cadence.
|
||||
"frame_interval_seconds": ml.video_frame_interval_seconds,
|
||||
"max_frames": ml.video_max_frames,
|
||||
# The embedding model the agent must use for concept crops, so
|
||||
# its region vectors land in the SAME space the heads trained in.
|
||||
# Server-announced → the agent stays model-agnostic; a swap is a
|
||||
# server setting + a re-embed migration, never an agent change.
|
||||
"embed_model_name": EMBED_MODEL_NAME,
|
||||
# The embedding model the agent must use for concept crops + the
|
||||
# whole-image 'embed' task, so its vectors land in the SAME space
|
||||
# the heads trained in. Server-announced FROM THE SETTING → the
|
||||
# agent stays model-agnostic; an operator swap is a setting + a
|
||||
# re-embed, never an agent change.
|
||||
"embed_model_name": ml.embedder_model_name,
|
||||
"embed_version": ml.embedder_model_version,
|
||||
})
|
||||
return jsonify({"jobs": out})
|
||||
@@ -188,6 +188,34 @@ async def submit():
|
||||
return jsonify({"ok": True, "stored": len(regions)})
|
||||
|
||||
|
||||
@gpu_bp.route("/jobs/submit_embedding", methods=["POST"])
|
||||
async def submit_embedding():
|
||||
"""Store a whole-image SigLIP embedding (the 'embed' task) on image_record +
|
||||
close the job. Body: {agent_id, job_id, embedding:[...], embedding_version}.
|
||||
This is how the GPU agent re-embeds the library under a new model (#1190) —
|
||||
much faster than the CPU ml-worker at higher resolutions."""
|
||||
body = await request.get_json(silent=True) or {}
|
||||
agent_id = str(body.get("agent_id") or "agent")
|
||||
job_id = body.get("job_id")
|
||||
embedding = body.get("embedding")
|
||||
version = body.get("embedding_version")
|
||||
if job_id is None or not embedding or not version:
|
||||
return jsonify({"error": "job_id, embedding, embedding_version required"}), 400
|
||||
async with get_session() as session:
|
||||
if not await _agent_authed(session):
|
||||
return jsonify({"error": "unauthorized"}), 401
|
||||
job = await session.get(GpuJob, int(job_id))
|
||||
if job is None or job.status != "leased" or job.lease_token != agent_id:
|
||||
return jsonify({"error": "lease_invalid"}), 409
|
||||
img = await session.get(ImageRecord, job.image_record_id)
|
||||
if img is not None:
|
||||
img.siglip_embedding = embedding
|
||||
img.siglip_model_version = version
|
||||
await GpuJobService(session).complete(agent_id, int(job_id))
|
||||
await session.commit()
|
||||
return jsonify({"ok": True})
|
||||
|
||||
|
||||
@gpu_bp.route("/jobs/fail", methods=["POST"])
|
||||
async def fail():
|
||||
body = await request.get_json(silent=True) or {}
|
||||
|
||||
@@ -24,6 +24,8 @@ _EDITABLE = (
|
||||
"ccip_match_threshold",
|
||||
"ccip_auto_apply_enabled",
|
||||
"ccip_auto_apply_threshold",
|
||||
"embedder_model_name",
|
||||
"embedder_model_version",
|
||||
)
|
||||
|
||||
|
||||
@@ -54,6 +56,7 @@ async def get_settings():
|
||||
"ccip_match_threshold": s.ccip_match_threshold,
|
||||
"ccip_auto_apply_enabled": s.ccip_auto_apply_enabled,
|
||||
"ccip_auto_apply_threshold": s.ccip_auto_apply_threshold,
|
||||
"embedder_model_name": s.embedder_model_name,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -125,6 +128,11 @@ def _validate(p: dict) -> str | None:
|
||||
return "ccip_match_threshold must be between 0.5 and 0.999"
|
||||
if not (0.5 <= float(p["ccip_auto_apply_threshold"]) <= 0.999):
|
||||
return "ccip_auto_apply_threshold must be between 0.5 and 0.999"
|
||||
# Embedder model swap (#1190): both must be non-empty. Changing them means a
|
||||
# different embedding space — the operator must re-embed + retrain after.
|
||||
for key in ("embedder_model_name", "embedder_model_version"):
|
||||
if not str(p[key]).strip():
|
||||
return f"{key} must not be empty"
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@@ -131,6 +131,11 @@ def make_celery() -> Celery:
|
||||
"schedule": 86400.0, # drain the concept-crop back-catalogue +
|
||||
"args": ("siglip",), # retry failed embeds, no button needed
|
||||
},
|
||||
"enqueue-embed-backfill-daily": {
|
||||
"task": "backend.app.tasks.ml.enqueue_gpu_backfill",
|
||||
"schedule": 86400.0, # whole-image re-embed under the current
|
||||
"args": ("embed",), # model (an operator swap) drains via agent
|
||||
},
|
||||
"ccip-auto-apply-daily": {
|
||||
"task": "backend.app.tasks.ml.scheduled_ccip_auto_apply",
|
||||
"schedule": 86400.0, # no-op unless ccip_auto_apply_enabled
|
||||
|
||||
@@ -107,6 +107,12 @@ class MLSettings(Base):
|
||||
embedder_model_version: Mapped[str] = mapped_column(
|
||||
String(128), nullable=False, default="siglip-so400m-patch14-384"
|
||||
)
|
||||
# The HF model NAME the embedder loads (server CPU embed + announced to the
|
||||
# GPU agent in the lease). Operator-settable so the embedder is a choice, not
|
||||
# a hardcode (#1190): set name + version together, then re-embed + retrain.
|
||||
embedder_model_name: Mapped[str] = mapped_column(
|
||||
String(128), nullable=False, default="google/siglip-so400m-patch14-384"
|
||||
)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
|
||||
@@ -18,9 +18,11 @@ ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||||
# N_replicas × this within the cores allotted to ML to avoid oversubscription.
|
||||
_INTRA_OP_THREADS = 4
|
||||
|
||||
MODEL_NAME = os.environ.get(
|
||||
DEFAULT_MODEL_NAME = os.environ.get(
|
||||
"SIGLIP_MODEL_NAME", "google/siglip-so400m-patch14-384"
|
||||
)
|
||||
# Back-compat alias (api/gpu imported this name as the fallback embedder id).
|
||||
MODEL_NAME = DEFAULT_MODEL_NAME
|
||||
MODEL_VERSION = os.environ.get(
|
||||
"SIGLIP_MODEL_VERSION", "siglip-so400m-patch14-384"
|
||||
)
|
||||
@@ -29,35 +31,42 @@ _LOCAL_DIR = Path(os.environ.get("ML_MODEL_DIR", "/models")) / "siglip"
|
||||
|
||||
|
||||
class Embedder:
|
||||
def __init__(self, model_dir: Path | None = None):
|
||||
self._model_dir = model_dir or _LOCAL_DIR
|
||||
"""Loads whatever SigLIP-family model it's given by HF NAME. For the default
|
||||
model it prefers the pre-downloaded local dir (no re-download on existing
|
||||
deploys); any other name resolves as an HF repo id (downloaded + cached on
|
||||
first use), so an operator model swap (#1190) just works server-side."""
|
||||
|
||||
def __init__(self, model_name: str | None = None, model_dir: Path | None = None):
|
||||
self.model_name = model_name or DEFAULT_MODEL_NAME
|
||||
self._explicit_dir = model_dir
|
||||
self._model = None
|
||||
self._processor = None
|
||||
self._torch = None
|
||||
|
||||
def _source(self) -> str:
|
||||
if self._explicit_dir is not None:
|
||||
return str(self._explicit_dir)
|
||||
if self.model_name == DEFAULT_MODEL_NAME and _LOCAL_DIR.exists():
|
||||
return str(_LOCAL_DIR)
|
||||
return self.model_name
|
||||
|
||||
def load(self) -> None:
|
||||
if self._model is not None:
|
||||
return
|
||||
import torch
|
||||
from transformers import AutoModel, SiglipImageProcessor
|
||||
from transformers import AutoImageProcessor, AutoModel
|
||||
|
||||
self._torch = torch
|
||||
# Bound torch's CPU thread pool (see _INTRA_OP_THREADS) so each replica
|
||||
# stays a predictable core consumer on a shared node.
|
||||
torch.set_num_threads(_INTRA_OP_THREADS)
|
||||
# FC's embedder only does IMAGE inference — never text. AutoProcessor
|
||||
# loads the full processor including SiglipTokenizer, which requires
|
||||
# the sentencepiece library at import time even if we never call it.
|
||||
# SiglipImageProcessor loads ONLY preprocessor_config.json (image
|
||||
# side) and skips the tokenizer config entirely. Operator hit the
|
||||
# ImportError 2026-05-25 once the ml-worker started actually running
|
||||
# tag_and_embed; switching to the image-only loader avoids the
|
||||
# tokenizer dep without adding ~30 MB of unused C++ build to the
|
||||
# lean ml-worker image.
|
||||
self._processor = SiglipImageProcessor.from_pretrained(
|
||||
str(self._model_dir)
|
||||
)
|
||||
self._model = AutoModel.from_pretrained(str(self._model_dir))
|
||||
# IMAGE inference only — AutoImageProcessor loads just the image side
|
||||
# (preprocessor_config.json), skipping the SigLIP tokenizer + its
|
||||
# sentencepiece dep (operator hit that ImportError 2026-05-25). Works
|
||||
# for any SigLIP-family model, keeping the embedder model-agnostic.
|
||||
src = self._source()
|
||||
self._processor = AutoImageProcessor.from_pretrained(src)
|
||||
self._model = AutoModel.from_pretrained(src)
|
||||
self._model.eval()
|
||||
|
||||
def infer(self, image_path: Path) -> np.ndarray:
|
||||
@@ -74,8 +83,12 @@ class Embedder:
|
||||
_default_embedder: Embedder | None = None
|
||||
|
||||
|
||||
def get_embedder() -> Embedder:
|
||||
def get_embedder(model_name: str | None = None) -> Embedder:
|
||||
"""Cached embedder for `model_name` (default if None). Rebuilds the singleton
|
||||
when the requested name changes, so an operator model swap takes effect
|
||||
without restarting the worker."""
|
||||
global _default_embedder
|
||||
if _default_embedder is None:
|
||||
_default_embedder = Embedder()
|
||||
name = model_name or DEFAULT_MODEL_NAME
|
||||
if _default_embedder is None or _default_embedder.model_name != name:
|
||||
_default_embedder = Embedder(model_name=name)
|
||||
return _default_embedder
|
||||
|
||||
@@ -308,25 +308,36 @@ async def score_image(
|
||||
import numpy as np
|
||||
|
||||
img = await session.get(ImageRecord, image_id)
|
||||
if img is None or img.siglip_embedding is None:
|
||||
if img is None:
|
||||
return []
|
||||
settings = await _settings_async(session)
|
||||
heads = await _current_heads(session, settings.embedder_model_version)
|
||||
cur_version = settings.embedder_model_version
|
||||
heads = await _current_heads(session, cur_version)
|
||||
if heads["W"] is None:
|
||||
return []
|
||||
|
||||
bag = [np.asarray(img.siglip_embedding, dtype=np.float32)]
|
||||
# Only embeddings in the CURRENT model's space enter the bag. Mid model-swap
|
||||
# (#1190), an image still carrying the OLD-version whole-image vector is
|
||||
# skipped rather than scored by heads trained in a different space; a legacy
|
||||
# NULL version is treated as current (those predate per-row stamping).
|
||||
bag = []
|
||||
if img.siglip_embedding is not None and img.siglip_model_version in (
|
||||
cur_version, None,
|
||||
):
|
||||
bag.append(np.asarray(img.siglip_embedding, dtype=np.float32))
|
||||
region_vecs = (
|
||||
await session.execute(
|
||||
select(ImageRegion.siglip_embedding)
|
||||
.where(ImageRegion.image_record_id == image_id)
|
||||
.where(ImageRegion.siglip_embedding.is_not(None))
|
||||
.where(ImageRegion.embedding_version == settings.embedder_model_version)
|
||||
.where(ImageRegion.embedding_version == cur_version)
|
||||
)
|
||||
).all()
|
||||
for (vec,) in region_vecs:
|
||||
if vec is not None:
|
||||
bag.append(np.asarray(vec, dtype=np.float32))
|
||||
if not bag:
|
||||
return []
|
||||
|
||||
X = np.vstack(bag) # (B, D)
|
||||
norms = np.linalg.norm(X, axis=1, keepdims=True)
|
||||
|
||||
+33
-10
@@ -95,7 +95,7 @@ def tag_and_embed(self, image_id: int) -> dict:
|
||||
|
||||
phase = "load_models"
|
||||
tagger = get_tagger()
|
||||
embedder = get_embedder()
|
||||
embedder = get_embedder(settings.embedder_model_name)
|
||||
|
||||
if is_vid:
|
||||
# Layer-3 isolation: ffprobe (a separate process) validates
|
||||
@@ -330,10 +330,10 @@ def backfill(self) -> int:
|
||||
!= settings.tagger_model_version
|
||||
)
|
||||
| (ImageRecord.siglip_embedding.is_(None))
|
||||
| (
|
||||
ImageRecord.siglip_model_version
|
||||
!= settings.embedder_model_version
|
||||
)
|
||||
# NB: a siglip MODEL-VERSION mismatch (an operator model swap,
|
||||
# #1190) is intentionally NOT re-embedded here — the CPU
|
||||
# ml-worker can't churn the whole library at 384/512px. The
|
||||
# GPU agent owns version re-embeds via the 'embed' job.
|
||||
)
|
||||
.order_by(ImageRecord.id.asc())
|
||||
.limit(500)
|
||||
@@ -750,17 +750,40 @@ def enqueue_gpu_backfill(task_name: str) -> int:
|
||||
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
|
||||
from sqlalchemy import exists, insert, literal, or_
|
||||
from sqlalchemy import select as sa_select
|
||||
|
||||
from ..models import GpuJob, ImageRecord, ImageRegion
|
||||
from ..models import GpuJob, ImageRecord, ImageRegion, MLSettings
|
||||
|
||||
SessionLocal = _sync_session_factory()
|
||||
with SessionLocal() as session:
|
||||
if task_name == "siglip":
|
||||
has_concept = exists().where(
|
||||
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,
|
||||
@@ -769,7 +792,7 @@ def enqueue_gpu_backfill(task_name: str) -> int:
|
||||
)
|
||||
sel = sa_select(
|
||||
ImageRecord.id, literal("siglip"), literal("pending")
|
||||
).where(~has_concept).where(~queued)
|
||||
).where(~has_current_concept).where(~queued)
|
||||
else:
|
||||
already = exists().where(
|
||||
GpuJob.image_record_id == ImageRecord.id,
|
||||
|
||||
Reference in New Issue
Block a user