bvandeusen
c6f38b0dac
Lift recall on small/local concepts (glasses, cum, stomach-bulge, xray, lactation) that the whole-image SigLIP vector washes out: the GPU agent now embeds figure crops with SigLIP too, stored as kind='concept' regions, and the suggestion rail scores each image as a BAG (whole-image + every concept crop), taking each head's MAX over the bag. The whole-image vector is always in the bag, so this can never score lower than before. Model-agnostic by construction: the server ANNOUNCES the embedding model (HF name + version) in the lease, so the agent loads whatever the heads were trained in and stays in lock-step — a model swap is a server setting + a re-embed migration, never an agent change. - agent: model-agnostic CropEmbedder (torch/transformers get_image_features, fp16 on CUDA, inference-locked); worker branches on job.task — 'ccip' emits figure(CCIP)+concept(SigLIP) in one pass, 'siglip' emits concept-only so the back-catalogue backfill never churns figure/CCIP regions; torch cu124 + transformers in the image. - server: lease announces embed_model_name/embed_version; score_image is max-over-bag (version-filtered region embeddings); enqueue_gpu_backfill 'siglip' gates on a missing concept region (drains the back-catalogue, retries failures, no double-enqueue); daily siglip-backfill beat; UI button; /api/ccip/overview reports images_with_concept_siglip. - v1 scope: suggestion rail only — auto-apply stays whole-image (conservative; heads' thresholds were calibrated on whole-image). Bulk-apply bag = follow-up. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
211 lines
8.4 KiB
Python
211 lines
8.4 KiB
Python
"""The lease → fetch → detect+embed → submit loop, run by a pool of worker
|
||
slots whose count is tunable live from the UI.
|
||
|
||
Each slot is an independent loop (its own leases; the server's SKIP-LOCKED lease
|
||
keeps them from colliding). More slots = more GPU load + throughput; the model is
|
||
loaded once and shared, so slots add concurrent inference, not N× model VRAM.
|
||
That's the dial the operator turns to trade desktop responsiveness for speed.
|
||
|
||
Stop (or shrinking the pool) RELEASES a slot's still-leased jobs immediately so
|
||
orphaned work is re-picked at once rather than waiting out the lease.
|
||
"""
|
||
import threading
|
||
import time
|
||
|
||
from . import media, models
|
||
from .client import FcClient
|
||
from .config import Config
|
||
from .crops import crop_region
|
||
|
||
# Generous cap: the pipeline is usually I/O-bound (downloading + decoding images
|
||
# over HTTP), so the GPU stays underused until many workers overlap that I/O.
|
||
# Push it up while watching the GPU util + VRAM in the UI.
|
||
MAX_CONCURRENCY = 32
|
||
|
||
# Fallbacks only — the server ANNOUNCES the embedding model (name + version) in
|
||
# the lease so the agent stays model-agnostic and in lock-step with the space
|
||
# the heads were trained in. These cover an older server that doesn't send them.
|
||
DEFAULT_EMBED_MODEL = "google/siglip-so400m-patch14-384"
|
||
DEFAULT_EMBED_VERSION = "siglip-so400m-patch14-384"
|
||
|
||
|
||
class _Slot:
|
||
"""One worker loop. `inflight` = jobs leased but not yet processed, so a
|
||
graceful stop can hand them back."""
|
||
__slots__ = ("stop", "inflight")
|
||
|
||
def __init__(self):
|
||
self.stop = threading.Event()
|
||
self.inflight: list[int] = []
|
||
|
||
|
||
class Worker:
|
||
def __init__(self, cfg: Config):
|
||
self.cfg = cfg
|
||
self.client = FcClient(cfg.fc_url, cfg.token, cfg.agent_id)
|
||
self._lock = threading.Lock()
|
||
self._running = False
|
||
self._target = max(1, min(MAX_CONCURRENCY, cfg.concurrency))
|
||
self._slots: list[_Slot] = []
|
||
self.processed = 0
|
||
self.errors = 0
|
||
self._active = 0 # slots currently mid-image
|
||
# The crop embedder (SigLIP-family) is built lazily on the first job that
|
||
# needs it, from the model the server announces — one shared instance.
|
||
self._embedder = None
|
||
self._embedder_lock = threading.Lock()
|
||
|
||
# --- control -----------------------------------------------------------
|
||
def start(self):
|
||
with self._lock:
|
||
self._running = True
|
||
self._reconcile_locked()
|
||
|
||
def stop(self):
|
||
with self._lock:
|
||
self._running = False
|
||
slots, self._slots = self._slots, []
|
||
for s in slots:
|
||
s.stop.set() # each slot releases its inflight on exit
|
||
|
||
def set_concurrency(self, n: int):
|
||
with self._lock:
|
||
self._target = max(1, min(MAX_CONCURRENCY, int(n)))
|
||
if self._running:
|
||
self._reconcile_locked()
|
||
|
||
def _reconcile_locked(self):
|
||
while len(self._slots) < self._target:
|
||
slot = _Slot()
|
||
self._slots.append(slot)
|
||
threading.Thread(target=self._loop, args=(slot,), daemon=True).start()
|
||
while len(self._slots) > self._target:
|
||
self._slots.pop().stop.set()
|
||
|
||
def status(self) -> dict:
|
||
with self._lock:
|
||
return {
|
||
"state": "running" if self._running else "stopped",
|
||
"concurrency": self._target,
|
||
"max_concurrency": MAX_CONCURRENCY,
|
||
"workers": len(self._slots),
|
||
"active": self._active,
|
||
"processed": self.processed,
|
||
"errors": self.errors,
|
||
}
|
||
|
||
def _bump(self, *, processed=0, errors=0, active=0):
|
||
with self._lock:
|
||
self.processed += processed
|
||
self.errors += errors
|
||
self._active += active
|
||
|
||
# --- per-slot loop -----------------------------------------------------
|
||
def _loop(self, slot: _Slot):
|
||
while not slot.stop.is_set() and self._running:
|
||
try:
|
||
jobs = self.client.lease(self.cfg.batch_size)
|
||
except Exception:
|
||
time.sleep(self.cfg.poll_idle_seconds)
|
||
continue
|
||
if not jobs:
|
||
time.sleep(self.cfg.poll_idle_seconds)
|
||
continue
|
||
slot.inflight = [j["job_id"] for j in jobs]
|
||
for job in jobs:
|
||
if slot.stop.is_set() or not self._running:
|
||
break
|
||
self._process(job)
|
||
slot.inflight = [i for i in slot.inflight if i != job["job_id"]]
|
||
if slot.inflight:
|
||
self.client.heartbeat(slot.inflight)
|
||
# Graceful hand-back of anything leased but not processed.
|
||
if slot.inflight:
|
||
self.client.release(slot.inflight)
|
||
slot.inflight = []
|
||
|
||
def _ensure_embedder(self, model_name: str):
|
||
if self._embedder is not None:
|
||
return self._embedder
|
||
with self._embedder_lock:
|
||
if self._embedder is None:
|
||
from .embedder import CropEmbedder
|
||
self._embedder = CropEmbedder(model_name, self.cfg.embed_dtype)
|
||
return self._embedder
|
||
|
||
def _process(self, job: dict):
|
||
self._bump(active=1)
|
||
try:
|
||
data = self.client.fetch_image(job["image_url"])
|
||
if media.is_video(job.get("mime", "")):
|
||
frames = media.sample_frames(
|
||
data, job.get("frame_interval_seconds", 4.0),
|
||
job.get("max_frames", 64),
|
||
) or [(None, media.load_image(data))]
|
||
else:
|
||
frames = [(None, media.load_image(data))]
|
||
|
||
# task picks what to produce per crop:
|
||
# 'siglip' (backfill existing images) → concept (SigLIP) regions
|
||
# ONLY, so it never churns their figure/CCIP regions or the
|
||
# character-reference cache.
|
||
# 'ccip' / 'both' (a new image's first pass) → figure (CCIP) AND
|
||
# concept (SigLIP) in one go, off the same crop.
|
||
task = job.get("task") or "ccip"
|
||
want_ccip = task in ("ccip", "both")
|
||
want_siglip = task in ("ccip", "siglip", "both")
|
||
replace_kinds = (
|
||
["concept"] if task == "siglip" else ["figure", "face", "concept"]
|
||
)
|
||
|
||
embed_version = job.get("embed_version") or DEFAULT_EMBED_VERSION
|
||
embedder = None
|
||
if want_siglip:
|
||
model_name = (
|
||
self.cfg.embed_model_override
|
||
or job.get("embed_model_name")
|
||
or DEFAULT_EMBED_MODEL
|
||
)
|
||
embedder = self._ensure_embedder(model_name)
|
||
|
||
regions = []
|
||
ccip_ev = self.cfg.ccip_model or "ccip-default"
|
||
dv = f"person-{self.cfg.detector_level}"
|
||
for t, frame in frames:
|
||
figs = models.detect_figures(frame, self.cfg.detector_level)
|
||
if not figs:
|
||
figs = [((0.0, 0.0, 1.0, 1.0), None)] # whole-frame fallback
|
||
for bbox, score in figs:
|
||
crop = crop_region(frame, bbox)
|
||
if crop is None:
|
||
continue
|
||
if want_ccip:
|
||
regions.append({
|
||
"kind": "figure",
|
||
"bbox": list(bbox),
|
||
"frame_time": t,
|
||
"score": score,
|
||
"ccip_embedding": models.ccip_vector(
|
||
crop, self.cfg.ccip_model or None
|
||
),
|
||
"embedding_version": ccip_ev,
|
||
"detector_version": dv,
|
||
})
|
||
if want_siglip:
|
||
regions.append({
|
||
"kind": "concept",
|
||
"bbox": list(bbox),
|
||
"frame_time": t,
|
||
"score": score,
|
||
"siglip_embedding": embedder.embed(crop),
|
||
"embedding_version": embed_version,
|
||
"detector_version": dv,
|
||
})
|
||
self.client.submit(job["job_id"], regions, replace_kinds)
|
||
self._bump(processed=1)
|
||
except Exception as exc: # noqa: BLE001 — report + move on
|
||
self._bump(errors=1)
|
||
self.client.fail(job["job_id"], str(exc)[:500])
|
||
finally:
|
||
self._bump(active=-1)
|