FabledCurator/agent/fc_agent/worker.py

"""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 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"


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 = 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)
                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.client.submit_embedding(job["job_id"], vec, embed_version)
                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"] if task == "siglip" else ["figure", "face", "concept"]
            )
            embedder = self._ensure_embedder(model_name) if want_siglip else None

            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)