bvandeusen
7cdce0c474
Near-static videos are the dominant GPU load: sampled into up to 64 frames, each re-runs the whole detect→CCIP→SigLIP chain on ~identical content. Add a CPU perceptual-hash frame dedup upstream of the GPU so the redundant frames are never processed at all (not just their embeds). - media.dedupe_frames() + _dhash(): 8×8 difference-hash (64-bit) per frame; greedy keep — a frame survives only if its hash differs from every kept frame by >= min_distance bits (Hamming). A static run collapses to one frame; genuinely distinct scenes all survive. Order + frame_time preserved. - Called in worker._download_decode right after sample_frames, so it runs in the decode stage on the downloader thread (CPU) — the GPU consumers only ever see deduped frames, and buffered video items shrink (less RAM too). - Env-tunable FRAME_DEDUPE_DISTANCE (default 8; higher keeps more frames for brief localized changes an 8×8 hash can miss; 0 disables). Logs `video frames N→M` when it drops any, so video load reduction is visible. Complements the spatial per-frame crop dedup (2026-07-01.2); this is the temporal axis. Build marker 2026-07-01.3. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
801 lines
38 KiB
Python
801 lines
38 KiB
Python
"""The lease → download → detect+embed → submit pipeline.
|
||
|
||
The workload is DOWNLOAD-BOUND (operator timing 2026-07-01: download 400–5462ms,
|
||
GPU ~300–600ms), so a design where each worker runs the whole serial chain leaves
|
||
the fast GPU idle during every download. This splits the chain into a producer/
|
||
consumer pipeline instead:
|
||
|
||
downloader pool (N threads) ── lease→download→decode ──▶ [bounded buffer]
|
||
│
|
||
┌────────────────────┘
|
||
▼
|
||
GPU consumer(s) (1–2) ── detect+embed(batched)→submit
|
||
|
||
* DOWNLOADERS are I/O-bound so many overlap well; the autoscaler tunes their
|
||
count by BUFFER OCCUPANCY — a near-empty buffer means the GPU is starving
|
||
(add downloaders); a near-full buffer means they outpace the GPU (hold/trim,
|
||
or add a 2nd consumer if the GPU has headroom).
|
||
* The BOUNDED BUFFER is backpressure: decoded frames are big, so a full buffer
|
||
blocks downloaders — capping RAM and how far leases run ahead of the GPU.
|
||
* CONSUMERS are GPU-bound and fast, so one usually keeps up; a 2nd is added
|
||
only when the buffer stays full and the GPU has spare util/VRAM.
|
||
* A HEARTBEAT thread keeps every still-held lease alive (buffered jobs wait for
|
||
the GPU and would otherwise hit curator's 180s lease TTL and be reclaimed).
|
||
|
||
Resilience carried over from the slot model: lease exponential backoff (ride out
|
||
a curator redeploy), submit-path retry (client.py — never discard finished GPU
|
||
work on a blip), release-on-stop (hand leases back at once), region caps + video
|
||
early-exit (bound pathological jobs). Stop drains BOTH pools and releases every
|
||
held lease immediately so orphaned work is re-picked without waiting out the TTL.
|
||
"""
|
||
import logging
|
||
import queue
|
||
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
|
||
from .detectors import dedupe_crops
|
||
|
||
# Cap on the lease-retry backoff: when curator is unreachable (e.g. you redeploy
|
||
# it while away), a downloader 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)
|
||
|
||
|
||
# Pipeline sizing. Downloaders are I/O-bound so the ceiling is generous; consumers
|
||
# are GPU-bound so a couple saturate the card. The buffer is small on purpose —
|
||
# each slot can hold many decoded video frames, so it bounds RAM, not just depth.
|
||
DL_MAX = 24 # max downloader threads
|
||
CONSUMER_MAX = 2 # max GPU consumer threads
|
||
BUFFER_MAX = 12 # bounded decoded-frame buffer (backpressure + RAM cap)
|
||
|
||
# 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 (Auto mode): scale DOWNLOADERS by buffer occupancy — the elegant
|
||
# control signal, since the buffer sits exactly between the two stages. Buffer
|
||
# mostly EMPTY → GPU starving → add downloaders. Buffer mostly FULL → downloaders
|
||
# outpace the GPU → the GPU is the bottleneck: add a 2nd consumer if it has
|
||
# util/VRAM headroom and doing so lifts throughput, else trim a downloader (it's
|
||
# only adding lease pressure). Occupancy + util are EWMA-smoothed (both are noisy
|
||
# tick-to-tick), and decisions are spaced so a move is judged on averaged signals.
|
||
CONTROL_INTERVAL = 2.0 # sampling cadence (seconds)
|
||
SAMPLES_PER_DECISION = 6 # decide ~every 12s on averaged signals
|
||
OCC_ALPHA = 0.3 # buffer-occupancy EWMA weight on the newest sample
|
||
OCC_LOW = 0.25 # below this = buffer starving → add a downloader
|
||
OCC_HIGH = 0.80 # above this = downloaders outpace the GPU
|
||
UTIL_ALPHA = 0.25 # GPU-util EWMA weight
|
||
UTIL_START = 85 # GPU has headroom below this (gate a 2nd consumer)
|
||
VRAM_HI = 0.90 # memory pressure → shed a consumer
|
||
VRAM_GROW_MAX = 0.82 # don't add a consumer past this VRAM
|
||
TPUT_ALPHA = 0.5 # throughput EWMA weight
|
||
TPUT_MARGIN = 0.08 # a consumer add must lift smoothed j/s by this to keep
|
||
|
||
# Keep buffered-but-unprocessed leases alive: they hold curator leases while they
|
||
# wait for the GPU, so heartbeat them well inside curator's 180s lease TTL.
|
||
HEARTBEAT_INTERVAL = 45.0
|
||
|
||
# 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.
|
||
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 pipeline'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 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._auto = bool(cfg.auto_scale) # autoscale the downloader count
|
||
self._dl_target = max(1, min(DL_MAX, cfg.concurrency))
|
||
self._consumer_target = 1 # GPU is fast — start with one
|
||
self._dls: list[tuple[threading.Thread, threading.Event]] = []
|
||
self._consumers: list[tuple[threading.Thread, threading.Event]] = []
|
||
self._ctrl_stop = threading.Event()
|
||
self._ctrl_thread: threading.Thread | None = None
|
||
# Decoded jobs waiting for the GPU: (job, frames). Bounded = backpressure.
|
||
self._buffer: queue.Queue = queue.Queue(maxsize=BUFFER_MAX)
|
||
# Every job leased and not yet terminal (submitted / failed / released) is
|
||
# "held" — the heartbeat thread keeps these alive, and stop() releases them
|
||
# all at once. Add on lease, discard on every terminal client call.
|
||
self._held: set[int] = set()
|
||
self._held_lock = threading.Lock()
|
||
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 # jobs currently mid-GPU (consumers busy)
|
||
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
|
||
threading.Thread(target=self._queue_poll_loop, daemon=True).start()
|
||
threading.Thread(target=self._heartbeat_loop, daemon=True).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()
|
||
threading.Thread(target=self._stats_loop, daemon=True).start()
|
||
# The crop embedder (SigLIP-family) and region proposers are built lazily
|
||
# on the first job that needs them and SHARED across all consumers — one
|
||
# instance, so a 2nd consumer adds concurrent inference, not N× VRAM.
|
||
self._embedder = None
|
||
self._embedder_lock = threading.Lock()
|
||
self._proposers = None
|
||
self._proposers_lock = threading.Lock()
|
||
|
||
# --- held-lease bookkeeping --------------------------------------------
|
||
def _hold(self, job_ids) -> None:
|
||
with self._held_lock:
|
||
self._held.update(job_ids)
|
||
|
||
def _unhold(self, job_id: int) -> None:
|
||
with self._held_lock:
|
||
self._held.discard(job_id)
|
||
|
||
def _release_owned(self, job_ids: list[int]) -> None:
|
||
"""Hand a set of still-held leases back to curator and drop them from the
|
||
held set — used when a downloader exits (stop/shrink) still owning leases
|
||
it hadn't yet buffered."""
|
||
if not job_ids:
|
||
return
|
||
self.client.release(job_ids)
|
||
for jid in job_ids:
|
||
self._unhold(jid)
|
||
|
||
# --- background loops ---------------------------------------------------
|
||
def _heartbeat_loop(self) -> None:
|
||
"""Keep every held lease alive so buffered jobs waiting on the GPU aren't
|
||
reclaimed by curator's 180s TTL. Errors are swallowed by client.heartbeat;
|
||
a reclaimed lease just re-leases elsewhere — never fatal."""
|
||
while True:
|
||
if self._running:
|
||
with self._held_lock:
|
||
ids = list(self._held)
|
||
if ids:
|
||
self.client.heartbeat(ids)
|
||
time.sleep(HEARTBEAT_INTERVAL)
|
||
|
||
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 pipeline 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.
|
||
In the pipeline these stages run on DIFFERENT threads concurrently, so the
|
||
figures are per-stage averages, not a single job's serial wall-clock."""
|
||
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]
|
||
log.info("timing/%ds — %s (%d jobs)",
|
||
int(STATS_INTERVAL), " · ".join(parts), jobs)
|
||
|
||
# --- control -----------------------------------------------------------
|
||
def start(self):
|
||
with self._lock:
|
||
self._running = True
|
||
self._dl_target = max(1, self._dl_target)
|
||
self._consumer_target = max(1, self._consumer_target)
|
||
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
|
||
# 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:
|
||
dls, self._dls = self._dls, []
|
||
cons, self._consumers = self._consumers, []
|
||
self._active = 0 # no consumers left → the meter reads 0 at once;
|
||
# any lagging decrement is clamped (see _bump)
|
||
for _, ev in dls:
|
||
ev.set()
|
||
for _, ev in cons:
|
||
ev.set()
|
||
# Wake any consumer blocked on an empty buffer.
|
||
for _ in range(CONSUMER_MAX):
|
||
try:
|
||
self._buffer.put_nowait(None)
|
||
except queue.Full:
|
||
break
|
||
# Drain the buffer + release every still-held lease in one shot so orphaned
|
||
# work is re-leased at once. A downloader/consumer mid-flight may also
|
||
# release its own job — a duplicate release is a harmless no-op.
|
||
self._drain_and_release()
|
||
|
||
def _drain_and_release(self) -> None:
|
||
while True:
|
||
try:
|
||
self._buffer.get_nowait()
|
||
except queue.Empty:
|
||
break
|
||
with self._held_lock:
|
||
ids = list(self._held)
|
||
self._held.clear()
|
||
if ids:
|
||
self.client.release(ids)
|
||
|
||
def set_auto(self, on: bool):
|
||
with self._lock:
|
||
self._auto = bool(on)
|
||
|
||
def set_concurrency(self, n: int):
|
||
# The UI dial tunes the DOWNLOADER count. A manual set is an override →
|
||
# leave Auto so the autoscaler stops fighting the operator.
|
||
with self._lock:
|
||
self._auto = False
|
||
self._dl_target = max(1, min(DL_MAX, int(n)))
|
||
if self._running:
|
||
self._reconcile_locked()
|
||
|
||
def _apply_downloaders(self, delta: int) -> bool:
|
||
with self._lock:
|
||
new = max(1, min(DL_MAX, self._dl_target + delta))
|
||
if new == self._dl_target:
|
||
return False
|
||
self._dl_target = new
|
||
if self._running:
|
||
self._reconcile_locked()
|
||
return True
|
||
|
||
def _apply_consumers(self, delta: int) -> bool:
|
||
with self._lock:
|
||
new = max(1, min(CONSUMER_MAX, self._consumer_target + delta))
|
||
if new == self._consumer_target:
|
||
return False
|
||
self._consumer_target = new
|
||
if self._running:
|
||
self._reconcile_locked()
|
||
return True
|
||
|
||
def _reconcile_locked(self):
|
||
"""Bring both thread pools to their target counts. New threads start; a
|
||
shrink sets a thread's stop event (it exits after its current iteration,
|
||
releasing any lease it still owns)."""
|
||
while len(self._dls) < self._dl_target:
|
||
ev = threading.Event()
|
||
th = threading.Thread(target=self._downloader, args=(ev,), daemon=True)
|
||
self._dls.append((th, ev))
|
||
th.start()
|
||
while len(self._dls) > self._dl_target:
|
||
_, ev = self._dls.pop()
|
||
ev.set()
|
||
while len(self._consumers) < self._consumer_target:
|
||
ev = threading.Event()
|
||
th = threading.Thread(target=self._consumer, args=(ev,), daemon=True)
|
||
self._consumers.append((th, ev))
|
||
th.start()
|
||
while len(self._consumers) > self._consumer_target:
|
||
_, ev = self._consumers.pop()
|
||
ev.set()
|
||
|
||
def status(self) -> dict:
|
||
# Lock-free on purpose: these are plain int / bool / len reads (atomic
|
||
# under the GIL) and this backs the UI poll — it must NEVER be able to
|
||
# block behind a thread holding _lock, or the whole status view freezes.
|
||
return {
|
||
"state": "running" if self._running else "stopped",
|
||
"concurrency": self._dl_target, # the UI dial = downloader count
|
||
"max_concurrency": DL_MAX,
|
||
"auto": self._auto,
|
||
"downloaders": len(self._dls),
|
||
"consumers": len(self._consumers),
|
||
"buffer": self._buffer.qsize(),
|
||
"buffer_max": BUFFER_MAX,
|
||
"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
|
||
# Clamp at 0: a Stop resets _active to 0, so a consumer that was
|
||
# mid-image decrements afterwards — that must not go negative.
|
||
self._active = max(0, self._active + active)
|
||
|
||
# --- downloader pool ---------------------------------------------------
|
||
def _downloader(self, stop_evt: threading.Event):
|
||
"""Lease a batch, download + decode each job, and hand it to the GPU
|
||
consumers via the bounded buffer. Owns its leases until they're buffered;
|
||
on any exit path it releases whatever it still owns so nothing is stranded
|
||
holding a lease."""
|
||
backoff = self.cfg.poll_idle_seconds
|
||
while self._running and not stop_evt.is_set():
|
||
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.
|
||
if stop_evt.wait(backoff):
|
||
break
|
||
backoff = min(backoff * 2, MAX_BACKOFF_SECONDS)
|
||
continue
|
||
if not jobs:
|
||
if stop_evt.wait(self.cfg.poll_idle_seconds):
|
||
break
|
||
continue
|
||
self._hold(j["job_id"] for j in jobs)
|
||
owned = [j["job_id"] for j in jobs] # released on any early exit
|
||
for job in jobs:
|
||
jid = job["job_id"]
|
||
if not self._running or stop_evt.is_set():
|
||
break
|
||
try:
|
||
frames = self._download_decode(job)
|
||
except requests.RequestException as exc:
|
||
owned.remove(jid)
|
||
if _is_transient(exc):
|
||
# curator down/redeploying or our lease was reclaimed —
|
||
# NOT the job's fault. Hand back this job + the rest of the
|
||
# batch and back the whole loop off.
|
||
self._bump(transient=1)
|
||
self.client.release([jid])
|
||
self._unhold(jid)
|
||
log.info("curator unreachable — released job %s, backing off", jid)
|
||
self._release_owned(owned)
|
||
owned = []
|
||
if not stop_evt.wait(backoff):
|
||
backoff = min(backoff * 2, MAX_BACKOFF_SECONDS)
|
||
break
|
||
# a job-specific HTTP fault (404 image gone, 400) → fail it
|
||
self._bump(errors=1)
|
||
log.warning("job %s (image %s) failed: %s",
|
||
jid, job.get("image_id"), str(exc)[:200])
|
||
self.client.fail(jid, str(exc)[:500])
|
||
self._unhold(jid)
|
||
continue
|
||
except Exception as exc: # noqa: BLE001 — bad media → the job's fault
|
||
owned.remove(jid)
|
||
self._bump(errors=1)
|
||
log.warning("job %s (image %s) failed to decode: %s",
|
||
jid, job.get("image_id"), str(exc)[:200])
|
||
self.client.fail(jid, str(exc)[:500])
|
||
self._unhold(jid)
|
||
continue
|
||
# Blocks on a full buffer (backpressure) but wakes promptly on stop.
|
||
if self._put((job, frames), stop_evt):
|
||
owned.remove(jid) # ownership handed to the buffer/consumer
|
||
else:
|
||
break # stopped while waiting for buffer space
|
||
self._release_owned(owned)
|
||
|
||
def _put(self, item, stop_evt: threading.Event) -> bool:
|
||
"""Push onto the bounded buffer, blocking while it's full but rechecking
|
||
stop so a shrink/Stop can't hang a full-buffer window. False = stopped."""
|
||
while self._running and not stop_evt.is_set():
|
||
try:
|
||
self._buffer.put(item, timeout=0.5)
|
||
return True
|
||
except queue.Full:
|
||
continue
|
||
return False
|
||
|
||
def _download_decode(self, job: dict):
|
||
"""Fetch the image bytes and decode → [(frame_time, PIL.Image)]. Videos
|
||
are sampled into frames (ffmpeg). Records the download + decode timings."""
|
||
_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))]
|
||
# Temporal dedup: a near-static video sampled into many frames re-runs
|
||
# the whole detect+embed chain on ~identical frames. Drop near-dup
|
||
# frames HERE (decode stage, CPU) so the GPU never sees them.
|
||
if len(frames) > 1 and self.cfg.frame_dedupe_distance > 0:
|
||
kept = media.dedupe_frames(frames, self.cfg.frame_dedupe_distance)
|
||
if len(kept) < len(frames):
|
||
log.info("job %s: video frames %d→%d (near-dup dedup)",
|
||
job.get("job_id"), len(frames), len(kept))
|
||
frames = kept
|
||
else:
|
||
frames = [(None, media.load_image(data))]
|
||
self._record("decode", time.monotonic() - _t)
|
||
return frames
|
||
|
||
# --- GPU consumer pool -------------------------------------------------
|
||
def _consumer(self, stop_evt: threading.Event):
|
||
"""Pull decoded jobs off the buffer and run detect + embed + submit."""
|
||
while self._running and not stop_evt.is_set():
|
||
try:
|
||
item = self._buffer.get(timeout=1.0)
|
||
except queue.Empty:
|
||
continue
|
||
if item is None: # stop sentinel
|
||
continue
|
||
job, frames = item
|
||
if not self._running or stop_evt.is_set():
|
||
self.client.release([job["job_id"]])
|
||
self._unhold(job["job_id"])
|
||
continue
|
||
self._bump(active=1)
|
||
try:
|
||
if self._consume(job, frames, stop_evt):
|
||
self._bump(processed=1)
|
||
finally:
|
||
self._bump(active=-1)
|
||
|
||
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 _consume(self, job: dict, frames: list, stop_evt: threading.Event) -> bool:
|
||
"""Detect + embed the decoded frames and submit the result. Returns True
|
||
when the job was completed (→ count it processed), False otherwise: a
|
||
transient transport fault releases the job (counted 'waited out'); a
|
||
job-specific fault fails it (counted an error); a stop mid-flight releases
|
||
it so a Stop drains promptly instead of finishing heavy GPU work."""
|
||
jid = job["job_id"]
|
||
try:
|
||
if not self._running or stop_evt.is_set():
|
||
self.client.release([jid])
|
||
self._unhold(jid)
|
||
return False
|
||
|
||
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)
|
||
if not self._running or stop_evt.is_set():
|
||
self.client.release([jid])
|
||
self._unhold(jid)
|
||
return False
|
||
_t = time.monotonic()
|
||
self.client.submit_embedding(jid, vec, embed_version)
|
||
self._record("submit", time.monotonic() - _t)
|
||
self._unhold(jid)
|
||
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:
|
||
# Bail promptly on Stop instead of grinding through every frame of
|
||
# a long video before the caller can hand the job back.
|
||
if not self._running or stop_evt.is_set():
|
||
break
|
||
# 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:
|
||
# Drop near-duplicate crops (a figure box ≈ an anatomy
|
||
# component, overlapping booru classes) before the embed so we
|
||
# never SigLIP the same region twice — saves GPU and a slot
|
||
# against max_regions. High-IoU + kind-aware, so intentional
|
||
# nested crops (figure ⊃ head) survive.
|
||
pending = dedupe_crops(pending, self.cfg.dedupe_iou)
|
||
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)
|
||
# Stop once we have enough: a long video (image 81602 = a 156 MB
|
||
# mp4, 64 sampled frames × ~32 regions) would otherwise burn ~38s
|
||
# of GPU across every frame before the submit is even truncated.
|
||
# Bounds the WORK, not just the POST body.
|
||
if len(regions) >= self.cfg.max_regions:
|
||
break
|
||
self._record("gpu", time.monotonic() - _t_gpu)
|
||
|
||
# A Stop mid-frame-loop leaves partial regions — don't submit those;
|
||
# hand the whole job back so another agent redoes it cleanly.
|
||
if not self._running or stop_evt.is_set():
|
||
self.client.release([jid])
|
||
self._unhold(jid)
|
||
return False
|
||
|
||
# Backstop: never submit an unbounded pile of regions (a pathological
|
||
# image / long video). Keep the highest-scoring max_regions so the
|
||
# POST body stays sane — curator rejects an oversized one with 413
|
||
# (operator-flagged: image 81602 looped on 413).
|
||
if len(regions) > self.cfg.max_regions:
|
||
regions.sort(key=lambda r: r.get("score", 0.0) or 0.0, reverse=True)
|
||
dropped = len(regions) - self.cfg.max_regions
|
||
regions = regions[: self.cfg.max_regions]
|
||
log.info("job %s: capped regions %d→%d (dropped %d)",
|
||
jid, len(regions) + dropped, len(regions), dropped)
|
||
_t = time.monotonic()
|
||
self.client.submit(jid, regions, replace_kinds)
|
||
self._record("submit", time.monotonic() - _t)
|
||
self._unhold(jid)
|
||
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; then the server's orphan-recovery reclaims it if down).
|
||
self._bump(transient=1)
|
||
log.info("curator unreachable — released job %s (post-GPU)", jid)
|
||
self.client.release([jid])
|
||
self._unhold(jid)
|
||
return False
|
||
self._bump(errors=1)
|
||
log.warning("job %s (image %s) failed: %s",
|
||
jid, job.get("image_id"), str(exc)[:200])
|
||
self.client.fail(jid, str(exc)[:500])
|
||
self._unhold(jid)
|
||
return False
|
||
except Exception as exc: # noqa: BLE001 — a genuine job fault: report it
|
||
self._bump(errors=1)
|
||
log.warning("job %s (image %s) failed: %s",
|
||
jid, job.get("image_id"), str(exc)[:200])
|
||
self.client.fail(jid, str(exc)[:500])
|
||
self._unhold(jid)
|
||
return False
|
||
|
||
# --- autoscaler --------------------------------------------------------
|
||
def _control_loop(self):
|
||
"""Scale DOWNLOADERS by buffer occupancy (Auto mode). The buffer sits
|
||
between the two stages, so its fill level is the direct signal: empty =
|
||
the GPU is starving (add downloaders); full = downloaders outpace the GPU
|
||
(the GPU is the bottleneck → add a 2nd consumer if it has headroom and the
|
||
add lifts throughput, else trim a downloader). Occupancy, util and
|
||
throughput are EWMA-smoothed and decisions spaced so moves ride averaged
|
||
signals, not tick-to-tick noise. VRAM pressure sheds a consumer at once."""
|
||
from . import gpu as gpumod
|
||
|
||
occ_ewma: float | None = None
|
||
util_ewma: float | None = None
|
||
tput_ewma: float | None = None
|
||
prev_p, prev_t = self.processed, time.monotonic()
|
||
tick = 0
|
||
con_grew = False # did the previous decision add a consumer?
|
||
tput_before = 0.0 # smoothed jobs/s before that consumer add
|
||
while not self._ctrl_stop.wait(CONTROL_INTERVAL):
|
||
if not (self._running and self._auto):
|
||
occ_ewma = util_ewma = tput_ewma = None
|
||
prev_p, prev_t = self.processed, time.monotonic()
|
||
tick = 0
|
||
con_grew = False
|
||
self._util_smooth = None
|
||
continue
|
||
|
||
occ = self._buffer.qsize() / BUFFER_MAX
|
||
occ_ewma = occ if occ_ewma is None else (
|
||
OCC_ALPHA * occ + (1 - OCC_ALPHA) * occ_ewma
|
||
)
|
||
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 and self._consumer_target > 1:
|
||
if self._apply_consumers(-1):
|
||
log.info("autoscale: consumers→%d (vram %d%% — memory pressure)",
|
||
self._consumer_target, round(vram * 100))
|
||
tick = 0
|
||
con_grew = False
|
||
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
|
||
)
|
||
|
||
d0, c0 = self._dl_target, self._consumer_target
|
||
if occ_ewma < OCC_LOW:
|
||
# Buffer starving → GPU idle waiting on downloads → add a feeder.
|
||
self._apply_downloaders(+1)
|
||
con_grew = False
|
||
elif occ_ewma > OCC_HIGH:
|
||
# Downloaders outpace the GPU. Prefer helping the GPU (add a 2nd
|
||
# consumer) when it has util + VRAM headroom and the last add
|
||
# actually paid off; otherwise trim a downloader.
|
||
if con_grew:
|
||
if tput_ewma > tput_before * (1 + TPUT_MARGIN):
|
||
con_grew = False # it helped → keep it, stop probing
|
||
else:
|
||
self._apply_consumers(-1) # no gain → revert
|
||
con_grew = False
|
||
elif (self._consumer_target < CONSUMER_MAX
|
||
and util_ewma < UTIL_START and vram < VRAM_GROW_MAX):
|
||
tput_before = tput_ewma
|
||
con_grew = self._apply_consumers(+1)
|
||
if not con_grew: # already maxed → trim a feeder
|
||
self._apply_downloaders(-1)
|
||
else:
|
||
self._apply_downloaders(-1)
|
||
else:
|
||
con_grew = False # balanced → settle
|
||
|
||
if self._dl_target != d0 or self._consumer_target != c0:
|
||
log.info(
|
||
"autoscale: dl %d→%d · consumers %d→%d "
|
||
"(buf %d%% · util~%d%% · %.2f j/s · vram %d%%)",
|
||
d0, self._dl_target, c0, self._consumer_target,
|
||
round(occ_ewma * 100), round(util_ewma), tput_ewma,
|
||
round(vram * 100))
|