"""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 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" 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.transient = 0 # jobs handed back due to a server outage (NOT # failed) — the "waiting out curator" counter 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, "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: jobs = self.client.lease(self.cfg.batch_size) 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.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) 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) -> 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: 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) 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) 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) 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) self.client.fail(job["job_id"], str(exc)[:500]) return True finally: self._bump(active=-1)