bvandeusen
79269da802
Addresses operator reports: Stop never finishes, the agent polls curator constantly, and stale-cached pages get mistaken for a failed deploy. - Stop is prompt: flip _running BEFORE any lock so /status + worker loops see "stopped" immediately, and add a stop/shrink checkpoint in _process (after decode, before the expensive detect+embed) that releases the job and bails — so a Stop doesn't wait out heavy GPU work. - Lazy curator polling: the queue snapshot is fetched only while a browser is actually watching (a /status hit within UI_IDLE_GRACE) and on a 5s cadence, not a constant background loop. The work loop's own lease/submit is curator's only visitor otherwise — nothing polls just to poll. - Build marker: VERSION is embedded in the page and reported on /status; the UI shows a "reload" banner when they differ, so a browser-cached page can't be mistaken for "the new image didn't deploy" (complements the no-store header). CI: the lint lane now also `ruff check`s agent/ and compileall-parses it, so the GPU agent is linted + syntax-checked before its image builds (build.yml only `docker build`s it). Fixed the agent's pre-existing UP037/B905 so it passes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
590 lines
28 KiB
Python
590 lines
28 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 logging
|
||
import threading
|
||
import time
|
||
|
||
import numpy as np
|
||
import requests
|
||
|
||
from . import media, models
|
||
from .client import FcClient
|
||
from .config import Config
|
||
from .crops import crop_region
|
||
|
||
# Cap on the lease-retry backoff: when curator is unreachable (e.g. you redeploy
|
||
# it while away), each slot retries leasing with exponential backoff up to this
|
||
# many seconds, then resumes within this window once the server is back — no
|
||
# restart needed.
|
||
MAX_BACKOFF_SECONDS = 60.0
|
||
|
||
|
||
def _is_transient(exc: requests.RequestException) -> bool:
|
||
"""A server/transport problem (wait it out) vs a job-specific fault (fail it).
|
||
No response → connection refused/timeout → curator is down → transient. With
|
||
a response: 5xx, auth (401/403, e.g. a token blip on redeploy), 408/409/429
|
||
(timeout / our lease reclaimed / rate-limited) are all 'not this job's fault'.
|
||
A specific 4xx like 404 (image gone) / 400 IS the job's fault → fail it."""
|
||
resp = getattr(exc, "response", None)
|
||
if resp is None:
|
||
return True
|
||
return resp.status_code >= 500 or resp.status_code in (401, 403, 408, 409, 429)
|
||
|
||
# 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"
|
||
|
||
# Autoscaler (when Auto is on): a SMOOTHED, throughput-aware climb that SETTLES.
|
||
# Raw GPU util swings wildly (a batched embed pegs it ~99%, then image decode/IO
|
||
# drops it ~0%), so a single reading is meaningless — util is sampled often and
|
||
# EWMA-smoothed. Each decision (spaced ~24s) grows the pool by one only while
|
||
# doing so keeps lifting *throughput* (jobs/s, also smoothed); when a grow stops
|
||
# helping the pool is IO/CPU/curator-bound, so it backs off one and SETTLES,
|
||
# holding there before an occasional re-probe. This finds the worker count that
|
||
# maximises real work — instead of flopping every cycle, or growing forever
|
||
# because util never reaches a fixed threshold on an IO-bound load.
|
||
CONTROL_INTERVAL = 2.0 # util sampling cadence (seconds)
|
||
SAMPLES_PER_DECISION = 12 # decide ~every 24s (12 × 2s) on averaged signals
|
||
UTIL_HI = 92 # smoothed util above this = saturated → shrink
|
||
UTIL_START = 85 # only begin a climb when smoothed util is below this
|
||
VRAM_HI = 0.88 # shrink above this fraction of VRAM (memory pressure)
|
||
VRAM_GROW_MAX = 0.80 # don't grow past this VRAM
|
||
UTIL_ALPHA = 0.25 # util EWMA weight on the newest sample (smoother)
|
||
TPUT_ALPHA = 0.5 # throughput EWMA weight
|
||
TPUT_MARGIN = 0.08 # a grow must lift smoothed jobs/s by this to "help"
|
||
REPROBE_TICKS = 8 # decisions to hold after settling before re-probing
|
||
|
||
# How often to log the per-stage timing breakdown (lease/download/decode/gpu/
|
||
# submit) so the operator can see where a job's wall-clock actually goes — the
|
||
# data that decides whether a download/compute split is worth building.
|
||
STATS_INTERVAL = 30.0
|
||
|
||
# The queue snapshot exists only to populate the UI's counts, so it's polled
|
||
# lazily — only while a browser is actually watching (a /status hit in the last
|
||
# UI_IDLE_GRACE seconds), and not on a tight loop. The work loop's own lease/
|
||
# submit calls are the real "is curator up?" signal; nothing polls just to poll.
|
||
QUEUE_POLL_INTERVAL = 5.0
|
||
UI_IDLE_GRACE = 20.0
|
||
|
||
log = logging.getLogger("fc_agent.worker")
|
||
|
||
|
||
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._auto = bool(cfg.auto_scale) # autoscale worker count
|
||
self._ctrl_stop = threading.Event()
|
||
self._ctrl_thread: threading.Thread | None = None
|
||
self._slots: list[_Slot] = []
|
||
self.processed = 0
|
||
self.errors = 0
|
||
self.transient = 0 # jobs handed back due to a server outage (NOT
|
||
# failed) — the "waiting out curator" counter
|
||
self._active = 0 # slots currently mid-image
|
||
self._util_smooth: float | None = None # EWMA GPU util (set by control loop)
|
||
# Curator queue snapshot, refreshed by a background poller so the UI
|
||
# /status read is instant — never an inline curator HTTP call (which
|
||
# stalls the whole status view when curator is busy).
|
||
self._queue: dict | None = None
|
||
self._ui_seen = 0.0 # monotonic time of the last UI /status hit
|
||
self._queue_thread = threading.Thread(target=self._queue_poll_loop, daemon=True)
|
||
self._queue_thread.start()
|
||
# Per-stage timing: stage -> [sum_seconds, count], summarised to the log
|
||
# every STATS_INTERVAL so we can see where wall-clock goes per job.
|
||
self._timing: dict[str, list[float]] = {}
|
||
self._timing_lock = threading.Lock()
|
||
self._stats_thread = threading.Thread(target=self._stats_loop, daemon=True)
|
||
self._stats_thread.start()
|
||
# 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()
|
||
# Region proposers (extra YOLO detectors) — lazily built once, shared.
|
||
self._proposers = None
|
||
self._proposers_lock = threading.Lock()
|
||
|
||
def _queue_poll_loop(self):
|
||
"""Refresh the curator queue snapshot so /status is a pure in-memory read
|
||
— but ONLY while the UI is being watched (a recent /status hit). No
|
||
browser open → no polling; the work loop is curator's only visitor.
|
||
Errors just leave the last snapshot (or None) — never blocks the UI."""
|
||
while True:
|
||
if time.monotonic() - self._ui_seen <= UI_IDLE_GRACE:
|
||
try:
|
||
self._queue = self.client.queue_status()
|
||
except Exception:
|
||
self._queue = None
|
||
time.sleep(QUEUE_POLL_INTERVAL)
|
||
|
||
def note_ui(self) -> None:
|
||
"""The UI polled /status — keep the queue snapshot warm for a while."""
|
||
self._ui_seen = time.monotonic()
|
||
|
||
def latest_queue(self) -> dict | None:
|
||
return self._queue
|
||
|
||
def util_smooth(self) -> float | None:
|
||
return self._util_smooth
|
||
|
||
def _record(self, stage: str, seconds: float) -> None:
|
||
with self._timing_lock:
|
||
s = self._timing.get(stage)
|
||
if s is None:
|
||
self._timing[stage] = [seconds, 1]
|
||
else:
|
||
s[0] += seconds
|
||
s[1] += 1
|
||
|
||
def _stats_loop(self) -> None:
|
||
"""Log a per-stage timing breakdown every STATS_INTERVAL (only when there
|
||
was work), so the operator can see the download/decode/gpu/submit split."""
|
||
while True:
|
||
time.sleep(STATS_INTERVAL)
|
||
with self._timing_lock:
|
||
snap = {k: (v[0], v[1]) for k, v in self._timing.items() if v[1]}
|
||
self._timing = {}
|
||
if not snap:
|
||
continue
|
||
order = ["lease", "download", "decode", "gpu", "submit"]
|
||
parts = [
|
||
f"{st} {1000 * snap[st][0] / snap[st][1]:.0f}ms"
|
||
for st in order if st in snap
|
||
]
|
||
jobs = (snap.get("gpu") or snap.get("download") or (0, 0))[1]
|
||
# Per-job wall time across the compute path (lease is per-batch, so
|
||
# it's shown separately above, not folded into this figure).
|
||
per_job = sum(
|
||
snap[st][0] for st in ("download", "decode", "gpu", "submit")
|
||
if st in snap
|
||
)
|
||
pj_ms = 1000 * per_job / jobs if jobs else 0
|
||
log.info("timing/%ds — %s | wall/job %.0fms (%d jobs)",
|
||
int(STATS_INTERVAL), " · ".join(parts), pj_ms, jobs)
|
||
|
||
# --- control -----------------------------------------------------------
|
||
def start(self):
|
||
with self._lock:
|
||
self._running = True
|
||
self._reconcile_locked()
|
||
# (Re)start the autoscaler control loop.
|
||
if self._ctrl_thread is None or not self._ctrl_thread.is_alive():
|
||
self._ctrl_stop.clear()
|
||
self._ctrl_thread = threading.Thread(target=self._control_loop, daemon=True)
|
||
self._ctrl_thread.start()
|
||
|
||
def stop(self):
|
||
# Flip the flag FIRST (atomic bool), before any lock, so /status and the
|
||
# worker loops observe "stopped" immediately even if _lock is momentarily
|
||
# held — the state can never lag behind the click.
|
||
self._running = False
|
||
self._ctrl_stop.set()
|
||
with self._lock:
|
||
slots, self._slots = self._slots, []
|
||
for s in slots:
|
||
s.stop.set() # each slot releases its inflight on exit
|
||
|
||
def set_auto(self, on: bool):
|
||
with self._lock:
|
||
self._auto = bool(on)
|
||
|
||
def set_concurrency(self, n: int):
|
||
# A manual set is an override → leave Auto.
|
||
with self._lock:
|
||
self._auto = False
|
||
self._target = max(1, min(MAX_CONCURRENCY, int(n)))
|
||
if self._running:
|
||
self._reconcile_locked()
|
||
|
||
def _apply_step(self, delta: int) -> bool:
|
||
"""Nudge the target by delta (bounded). Returns True if it changed."""
|
||
with self._lock:
|
||
new = max(1, min(MAX_CONCURRENCY, self._target + delta))
|
||
if new == self._target:
|
||
return False
|
||
self._target = new
|
||
if self._running:
|
||
self._reconcile_locked()
|
||
return True
|
||
|
||
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:
|
||
# Lock-free on purpose: these are plain int / bool reads (atomic under the
|
||
# GIL) and this backs the UI poll — it must NEVER be able to block behind
|
||
# a worker holding _lock, or the whole status view freezes.
|
||
return {
|
||
"state": "running" if self._running else "stopped",
|
||
"concurrency": self._target,
|
||
"max_concurrency": MAX_CONCURRENCY,
|
||
"auto": self._auto,
|
||
"workers": len(self._slots),
|
||
"active": self._active,
|
||
"processed": self.processed,
|
||
"errors": self.errors,
|
||
"transient": self.transient,
|
||
}
|
||
|
||
def _bump(self, *, processed=0, errors=0, active=0, transient=0):
|
||
with self._lock:
|
||
self.processed += processed
|
||
self.errors += errors
|
||
self.transient += transient
|
||
self._active += active
|
||
|
||
# --- per-slot loop -----------------------------------------------------
|
||
def _loop(self, slot: _Slot):
|
||
backoff = self.cfg.poll_idle_seconds
|
||
while not slot.stop.is_set() and self._running:
|
||
try:
|
||
_t = time.monotonic()
|
||
jobs = self.client.lease(self.cfg.batch_size)
|
||
self._record("lease", time.monotonic() - _t)
|
||
backoff = self.cfg.poll_idle_seconds # server answered → reset
|
||
except Exception:
|
||
# curator unreachable (redeploy, network drop): wait it out with
|
||
# exponential backoff, capped — resume on our own when it returns.
|
||
self._interruptible_sleep(slot, backoff)
|
||
backoff = min(backoff * 2, MAX_BACKOFF_SECONDS)
|
||
continue
|
||
if not jobs:
|
||
self._interruptible_sleep(slot, 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
|
||
ok = self._process(job, slot)
|
||
slot.inflight = [i for i in slot.inflight if i != job["job_id"]]
|
||
if not ok:
|
||
# Server went away mid-batch: hand the rest back (best effort)
|
||
# and back off instead of hammering a recovering server or
|
||
# burning the jobs' attempt budgets on fail().
|
||
if slot.inflight:
|
||
self.client.release(slot.inflight)
|
||
slot.inflight = []
|
||
self._interruptible_sleep(slot, backoff)
|
||
backoff = min(backoff * 2, MAX_BACKOFF_SECONDS)
|
||
break
|
||
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 _interruptible_sleep(self, slot: _Slot, seconds: float):
|
||
"""Sleep, but wake immediately if the slot is told to stop — so a Stop or
|
||
a pool-shrink doesn't hang for a full backoff window."""
|
||
slot.stop.wait(timeout=seconds)
|
||
|
||
# --- autoscaler --------------------------------------------------------
|
||
def _control_loop(self):
|
||
"""Smoothed, throughput-aware climb that settles (Auto mode). Samples GPU
|
||
util often and EWMA-smooths it (raw util swings 0↔99 between a batched
|
||
embed and the IO/decode around it, so one reading is noise). Every
|
||
SAMPLES_PER_DECISION ticks it makes ONE move: grow by one while each grow
|
||
keeps lifting smoothed throughput; when a grow stops helping (IO/CPU/
|
||
curator-bound) back off one and SETTLE, holding before an occasional
|
||
re-probe. Memory pressure and saturation always shrink immediately."""
|
||
from . import gpu as gpumod
|
||
|
||
util_ewma: float | None = None
|
||
tput_ewma: float | None = None
|
||
prev_p, prev_t = self.processed, time.monotonic()
|
||
tick = 0
|
||
settled = False
|
||
grew_last = False # did the previous decision grow the pool?
|
||
tput_before = 0.0 # smoothed jobs/s at the count before that grow
|
||
hold = 0 # decisions left to hold while settled
|
||
while not self._ctrl_stop.wait(CONTROL_INTERVAL):
|
||
if not (self._running and self._auto):
|
||
util_ewma = tput_ewma = None
|
||
prev_p, prev_t = self.processed, time.monotonic()
|
||
tick = 0
|
||
settled = grew_last = False
|
||
hold = 0
|
||
self._util_smooth = None
|
||
continue
|
||
|
||
g = gpumod.read_gpu() or {}
|
||
mt = g.get("mem_total_mb") or 0
|
||
vram = (g.get("mem_used_mb", 0) / mt) if mt else 0.0
|
||
util = g.get("util_pct", 0) or 0
|
||
util_ewma = util if util_ewma is None else (
|
||
UTIL_ALPHA * util + (1 - UTIL_ALPHA) * util_ewma
|
||
)
|
||
self._util_smooth = util_ewma
|
||
|
||
# Memory pressure overrides the cadence — react immediately.
|
||
if vram >= VRAM_HI:
|
||
if self._apply_step(-1):
|
||
log.info("autoscale: -1 → %d workers (vram %d%% — memory pressure)",
|
||
self._target, round(vram * 100))
|
||
tick, settled, grew_last, hold = 0, True, False, REPROBE_TICKS
|
||
continue
|
||
|
||
tick += 1
|
||
if tick < SAMPLES_PER_DECISION:
|
||
continue
|
||
tick = 0
|
||
|
||
now = time.monotonic()
|
||
inst = (self.processed - prev_p) / max(1e-3, now - prev_t)
|
||
prev_p, prev_t = self.processed, now
|
||
tput_ewma = inst if tput_ewma is None else (
|
||
TPUT_ALPHA * inst + (1 - TPUT_ALPHA) * tput_ewma
|
||
)
|
||
|
||
t0 = self._target
|
||
if util_ewma > UTIL_HI: # saturated → ease off
|
||
self._apply_step(-1)
|
||
settled, grew_last, hold = True, False, REPROBE_TICKS
|
||
elif settled:
|
||
hold -= 1
|
||
if hold <= 0: # re-probe: try one grow
|
||
if util_ewma < UTIL_START and vram < VRAM_GROW_MAX:
|
||
tput_before = tput_ewma
|
||
grew_last = self._apply_step(+1)
|
||
settled = not grew_last
|
||
else:
|
||
hold = REPROBE_TICKS # still no room → keep holding
|
||
elif grew_last:
|
||
if tput_ewma > tput_before * (1 + TPUT_MARGIN): # the grow helped
|
||
tput_before = tput_ewma
|
||
if util_ewma < UTIL_START and vram < VRAM_GROW_MAX:
|
||
grew_last = self._apply_step(+1)
|
||
settled = not grew_last
|
||
else:
|
||
settled, grew_last, hold = True, False, REPROBE_TICKS
|
||
else: # overshot → back off + settle
|
||
self._apply_step(-1)
|
||
settled, grew_last, hold = True, False, REPROBE_TICKS
|
||
elif util_ewma < UTIL_START and vram < VRAM_GROW_MAX: # start a climb
|
||
tput_before = tput_ewma
|
||
grew_last = self._apply_step(+1)
|
||
settled = not grew_last
|
||
else:
|
||
settled, hold = True, REPROBE_TICKS # nothing to do → settle
|
||
|
||
if self._target != t0:
|
||
log.info("autoscale: %d→%d workers (util~%d%% · %.2f j/s · vram %d%%)",
|
||
t0, self._target, round(util_ewma), tput_ewma, round(vram * 100))
|
||
|
||
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 _ensure_proposers(self):
|
||
if self._proposers is not None:
|
||
return self._proposers
|
||
with self._proposers_lock:
|
||
if self._proposers is None:
|
||
from .detectors import Proposers
|
||
self._proposers = Proposers(self.cfg)
|
||
return self._proposers
|
||
|
||
def _process(self, job: dict, slot: _Slot) -> bool:
|
||
"""Process one job. Returns True when handled (completed, or hard-failed
|
||
because the job itself is bad) and False on a TRANSPORT error (curator
|
||
unreachable / 5xx / our lease was reclaimed mid-flight) — which is not
|
||
the job's fault, so the caller backs off and the job is left to be
|
||
re-leased rather than fail()ed into its attempt budget."""
|
||
self._bump(active=1)
|
||
try:
|
||
_t = time.monotonic()
|
||
data = self.client.fetch_image(job["image_url"])
|
||
self._record("download", time.monotonic() - _t)
|
||
|
||
_t = time.monotonic()
|
||
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))]
|
||
self._record("decode", time.monotonic() - _t)
|
||
|
||
# Stop/shrink checkpoint: bail BEFORE the expensive detect+embed so a
|
||
# Stop finishes promptly instead of waiting out heavy GPU work. Hand
|
||
# the job back to pending for another agent.
|
||
if not self._running or slot.stop.is_set():
|
||
self.client.release([job["job_id"]])
|
||
return True
|
||
|
||
task = job.get("task") or "ccip"
|
||
embed_version = job.get("embed_version") or DEFAULT_EMBED_VERSION
|
||
model_name = (
|
||
self.cfg.embed_model_override
|
||
or job.get("embed_model_name")
|
||
or DEFAULT_EMBED_MODEL
|
||
)
|
||
|
||
# 'embed' = WHOLE-IMAGE SigLIP embedding (re-embed the library under a
|
||
# new model, #1190) → image_record.siglip_embedding. Mean-pool video
|
||
# frames, matching the server's tag_and_embed. No regions.
|
||
if task == "embed":
|
||
embedder = self._ensure_embedder(model_name) # one-time model load
|
||
_t = time.monotonic()
|
||
vecs = [embedder.embed(frame) for _, frame in frames]
|
||
if len(vecs) > 1:
|
||
vec = np.mean(
|
||
np.asarray(vecs, dtype=np.float32), axis=0
|
||
).tolist()
|
||
else:
|
||
vec = vecs[0]
|
||
self._record("gpu", time.monotonic() - _t)
|
||
_t = time.monotonic()
|
||
self.client.submit_embedding(job["job_id"], vec, embed_version)
|
||
self._record("submit", time.monotonic() - _t)
|
||
self._bump(processed=1)
|
||
return True
|
||
|
||
# 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.
|
||
want_ccip = task in ("ccip", "both")
|
||
want_siglip = task in ("ccip", "siglip", "both")
|
||
replace_kinds = (
|
||
["concept", "panel"] if task == "siglip"
|
||
else ["figure", "face", "concept", "panel"]
|
||
)
|
||
embedder = self._ensure_embedder(model_name) if want_siglip else None
|
||
proposers = self._ensure_proposers()
|
||
|
||
regions = []
|
||
ccip_ev = self.cfg.ccip_model or "ccip-default"
|
||
dv = f"person-{self.cfg.detector_level}"
|
||
|
||
_t_gpu = time.monotonic() # detect + CCIP + batched embed = "gpu"
|
||
for t, frame in frames:
|
||
# FIGURE boxes: imgutils detect_person ∪ general COCO person,
|
||
# NMS-merged → CCIP identity (+ a concept crop). Covers anime +
|
||
# Western/realistic figures.
|
||
base = models.detect_figures(frame, self.cfg.detector_level)
|
||
figs = proposers.figures(frame, base)
|
||
if not figs:
|
||
figs = [((0.0, 0.0, 1.0, 1.0), 1.0, "whole")] # whole-frame fallback
|
||
|
||
# Collect every crop that needs a SigLIP embedding, then embed
|
||
# them in ONE batched forward pass (huge GPU-util + throughput
|
||
# win vs one forward per crop). CCIP runs per figure inline.
|
||
pending = [] # (crop, region-template-without-embedding)
|
||
for bbox, score, _label 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:
|
||
pending.append((crop, {
|
||
"kind": "concept", "bbox": list(bbox), "frame_time": t,
|
||
"score": score, "detector_version": dv,
|
||
}))
|
||
if not want_siglip:
|
||
continue
|
||
# ANATOMY components (booru_yolo) + PANELS → concept/panel crops.
|
||
for bbox, score, label in proposers.components(frame):
|
||
crop = crop_region(frame, bbox)
|
||
if crop is not None:
|
||
pending.append((crop, {
|
||
"kind": "concept", "bbox": list(bbox), "frame_time": t,
|
||
"score": score, "detector_version": f"booru:{label}",
|
||
}))
|
||
for bbox, score, _label in proposers.panels(frame):
|
||
crop = crop_region(frame, bbox)
|
||
if crop is not None:
|
||
pending.append((crop, {
|
||
"kind": "panel", "bbox": list(bbox), "frame_time": t,
|
||
"score": score, "detector_version": "panel",
|
||
}))
|
||
if pending:
|
||
vecs = embedder.embed_batch([c for c, _ in pending])
|
||
for (_c, tmpl), vec in zip(pending, vecs, strict=True):
|
||
tmpl["siglip_embedding"] = vec
|
||
tmpl["embedding_version"] = embed_version
|
||
regions.append(tmpl)
|
||
self._record("gpu", time.monotonic() - _t_gpu)
|
||
_t = time.monotonic()
|
||
self.client.submit(job["job_id"], regions, replace_kinds)
|
||
self._record("submit", time.monotonic() - _t)
|
||
self._bump(processed=1)
|
||
return True
|
||
except requests.RequestException as exc:
|
||
if _is_transient(exc):
|
||
# curator down/redeploying, a 5xx, or our lease was reclaimed
|
||
# while we worked. NOT the job's fault — hand it back (best
|
||
# effort; no-ops if the server is still down, then the server's
|
||
# orphan-recovery reclaims it) and signal the loop to wait.
|
||
self._bump(transient=1)
|
||
log.info("curator unreachable — released job %s, backing off",
|
||
job.get("job_id"))
|
||
self.client.release([job["job_id"]])
|
||
return False
|
||
# A job-specific HTTP fault (404 image gone, 400) → fail it so it
|
||
# doesn't re-lease forever.
|
||
self._bump(errors=1)
|
||
log.warning("job %s (image %s) failed: %s",
|
||
job.get("job_id"), job.get("image_id"), str(exc)[:200])
|
||
self.client.fail(job["job_id"], str(exc)[:500])
|
||
return True
|
||
except Exception as exc: # noqa: BLE001 — a genuine job fault: report it
|
||
self._bump(errors=1)
|
||
log.warning("job %s (image %s) failed: %s",
|
||
job.get("job_id"), job.get("image_id"), str(exc)[:200])
|
||
self.client.fail(job["job_id"], str(exc)[:500])
|
||
return True
|
||
finally:
|
||
self._bump(active=-1)
|