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


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

    # --- 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 _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))]

            regions = []
            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
                    vec = models.ccip_vector(crop, self.cfg.ccip_model or None)
                    regions.append({
                        "kind": "figure",
                        "bbox": list(bbox),
                        "frame_time": t,
                        "score": score,
                        "ccip_embedding": vec,
                        "embedding_version": ev,
                        "detector_version": dv,
                    })
            self.client.submit(job["job_id"], regions, ["figure", "face"])
            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)