From afef95a87de4bf1fa124e1f1ed7d13c2b1a0be1a Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Tue, 30 Jun 2026 23:34:12 -0400 Subject: [PATCH] feat(agent): download/GPU producer-consumer pipeline + fix detector fuse crash MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The agent workload is download-bound (download 400–5462ms vs GPU ~300–600ms), so the old N-slot serial chain (each slot: lease→download→decode→GPU→submit) left the fast GPU idle during every download. Rearchitect worker.py into a producer/consumer pipeline: downloader pool (autoscaled by BUFFER OCCUPANCY) → bounded queue → 1–2 GPU consumers (detect+embed→submit) - Downloaders are I/O-bound → many overlap; the autoscaler now tunes DOWNLOADER count by buffer fill (empty = GPU starving → add; full = outpacing GPU → add a 2nd consumer if it has util/VRAM headroom and lifts throughput, else trim). - Bounded buffer (12) = backpressure: a full buffer blocks downloaders, capping RAM + lease look-ahead. VRAM pressure sheds a consumer immediately. - Heartbeat thread keeps every held lease alive (buffered jobs wait on the GPU; curator's 180s TTL would otherwise reclaim them mid-buffer). - Preserves all resilience: lease exp-backoff, submit-path retry (#169), release-on-stop, region caps + video early-exit (#171). Stop drains BOTH pools and releases every held lease at once (single held-set as source of truth). - Consumers SHARE one embedder + proposers instance (a 2nd consumer adds concurrent inference, not N× VRAM — bounds the VRAM creep seen with N slots). - UI reworked for the pipeline: tiles show downloaders · buffer · on-GPU · processed · errors, a buffer-occupancy meter, and a consumers/waited-out line; the dial now tunes downloaders. Build marker 2026-07-01.1. Also fix the operator-flagged detector warning: yolo11n + the comic-panel model threw "'Conv' object has no attribute 'bn'" on every image (ultralytics' load- time Conv+BN fusion on a version-mismatched graph), silently disabling 2 of 3 crop proposers and spamming the log per image. Disable that fusion (unfused inference is correct, marginally slower) and permanently self-disable a proposer on the first inference failure instead of re-throwing forever. Refs milestone 122. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa --- agent/fc_agent/app.py | 39 +- agent/fc_agent/detectors.py | 20 +- agent/fc_agent/worker.py | 692 ++++++++++++++++++++++-------------- 3 files changed, 478 insertions(+), 273 deletions(-) diff --git a/agent/fc_agent/app.py b/agent/fc_agent/app.py index ee67b7b..ec1defe 100644 --- a/agent/fc_agent/app.py +++ b/agent/fc_agent/app.py @@ -1,8 +1,10 @@ """FastAPI control surface for the agent (served on localhost). -Start / stop the worker pool, tune the worker count live (trades desktop -responsiveness for throughput), and watch GPU load + progress + the server-side -queue. Config is env-seeded; the worker count is adjustable here on the fly. +Start / stop the download→GPU pipeline, tune the downloader count live (the +workload is download-bound, so downloaders are the dial that trades desktop +bandwidth for throughput), and watch GPU load + buffer occupancy + progress + +the server-side queue. Config is env-seeded; the downloader count is adjustable +here on the fly (GPU consumers autoscale between 1 and 2 on their own). """ from fastapi import FastAPI, Request from fastapi.responses import HTMLResponse, JSONResponse @@ -15,7 +17,7 @@ from .worker import Worker # Bump on every agent change. The page embeds this and /status reports it; the UI # warns to reload when they differ — so a stale browser-cached page can't be # mistaken for "the new image didn't deploy". (Belt-and-braces with no-store.) -VERSION = "2026-06-30.10 · video region early-exit" +VERSION = "2026-07-01.1 · download/GPU pipeline" logbuf.install() cfg = Config.from_env() @@ -156,9 +158,9 @@ _PAGE = """ #conc{width:3.4rem;height:32px;text-align:center;font:700 16px system-ui;background:#11151a; color:var(--fg);border:1px solid var(--bd);border-radius:8px} .hint{color:var(--mut);font-size:12px;margin-top:12px} - .tiles{display:grid;grid-template-columns:repeat(5,1fr);gap:10px;margin-bottom:16px} - .tile{background:#13171d;border:1px solid var(--bd);border-radius:10px;padding:12px 10px;text-align:center} - .tile .n{font:800 24px ui-monospace,monospace;line-height:1.1} + .tiles{display:grid;grid-template-columns:repeat(6,1fr);gap:8px;margin-bottom:16px} + .tile{background:#13171d;border:1px solid var(--bd);border-radius:10px;padding:12px 8px;text-align:center} + .tile .n{font:800 22px ui-monospace,monospace;line-height:1.1} .tile .n.warn{color:var(--red)} .tile .n.ok{color:var(--grn)} .tile .l{font-size:10px;text-transform:uppercase;letter-spacing:.06em;color:var(--mut);margin-top:4px} .meters{display:flex;flex-direction:column;gap:10px;margin-bottom:14px} @@ -167,6 +169,7 @@ _PAGE = """ .bar{height:9px;border-radius:5px;background:#11151a;border:1px solid var(--bd);overflow:hidden} .bar>i{display:block;height:100%;width:0;background:linear-gradient(90deg,#3a7d57,var(--grn));transition:width .4s} #utilbar{background:linear-gradient(90deg,#9a5a1f,var(--acc))} + #bufbar{background:linear-gradient(90deg,#2f5a9a,#4a86d8)} .queue{font:13px ui-monospace,monospace;color:var(--mut)} .banner{margin:0 0 14px;padding:.7rem .9rem;border-radius:10px;background:#3a2f12; border:1px solid #5a4a17;color:#ffd98a;font-size:13px} @@ -208,24 +211,28 @@ _PAGE = """ -
auto-tuning to fill the GPU · max 8
+
auto-tuning downloaders to keep the GPU fed · max 8
Status
—
state
-
0
active
+
0
downloaders
+
—
buffer
+
0
on GPU
0
processed
0
errors
-
0
waited out
GPU util—
VRAM—
+
buffer occupancy—
+
+
consumers — · waited out 0
queue —
@@ -281,13 +288,19 @@ _PAGE = """ const draining=!running && s.active>0 state.textContent=draining?'stopping':s.state state.className='n'+(draining?' busy':'') - active.textContent=s.active; done.textContent=s.processed + dln.textContent=(s.downloaders!=null?s.downloaders:'—') + bufn.textContent=(s.buffer!=null?(s.buffer+'/'+s.buffer_max):'—') + active.textContent=s.active; active.className='n'+(s.active>0?' busy':'') + done.textContent=s.processed err.textContent=s.errors; err.className='n'+(s.errors>0?' warn':'') - wait.textContent=s.transient||0 + pipe.textContent='consumers '+(s.consumers!=null?s.consumers:'—')+' · waited out '+(s.transient||0) + // Buffer occupancy bar (also driven here so it tracks the /status cadence). + if(s.buffer!=null && s.buffer_max){ const p=Math.round(100*s.buffer/s.buffer_max) + buflbl.textContent=s.buffer+' / '+s.buffer_max; bufbar.style.width=p+'%' } // Auto on → dial reflects the auto-chosen count (read-only); off → manual. if(document.activeElement!==autochk) autochk.checked=!!s.auto conc.disabled=!!s.auto; conc.style.opacity=s.auto?0.55:1 - conchint.textContent=s.auto?('auto-tuning to fill the GPU · max '+CAP):('manual · max '+CAP) + conchint.textContent=s.auto?('auto-tuning downloaders to keep the GPU fed · max '+CAP):('manual downloaders · max '+CAP) if(document.activeElement!==conc) conc.value=s.concurrency conc.max=CAP // Connection pill + queue come only from the /status poll (the Start/Stop POST diff --git a/agent/fc_agent/detectors.py b/agent/fc_agent/detectors.py index 3338850..0119490 100644 --- a/agent/fc_agent/detectors.py +++ b/agent/fc_agent/detectors.py @@ -17,6 +17,7 @@ cached under HF_HOME so the download happens once. import logging import os import threading +import types from pathlib import Path log = logging.getLogger("fc_agent.detectors") @@ -93,6 +94,18 @@ class YoloProposer: self._ok = False return self._model = YOLO(path) + # Disable ultralytics' load-time Conv+BN fusion. AutoBackend fuses + # the graph on the first predict; some checkpoints (yolo11n, the + # comic-panel model) crash that step with "'Conv' object has no + # attribute 'bn'" (a partially-fused / version-mismatched graph), + # which silently disabled those proposers (operator-flagged + # 2026-07-01). Unfused inference is correct — only marginally + # slower — and this is robust across ultralytics versions; if a + # future version ignores the override, the detect() guard below + # still self-disables the proposer instead of spamming per image. + inner = getattr(self._model, "model", None) + if inner is not None: + inner.fuse = types.MethodType(lambda self, *a, **k: self, inner) log.info("detector %s loaded (%s)", self.name, path) except Exception as exc: # noqa: BLE001 log.warning("detector %s disabled (load failed): %s", self.name, exc) @@ -105,7 +118,12 @@ class YoloProposer: try: res = self._model.predict(image, conf=self._conf, verbose=False)[0] except Exception as exc: # noqa: BLE001 - log.warning("detector %s inference failed: %s", self.name, exc) + # Permanently self-disable on the FIRST inference failure rather than + # re-throwing (and re-logging) on every image forever — an unfixable + # model fault degrades to "this proposer is off", logged once. + log.warning("detector %s disabled (inference failed): %s", self.name, exc) + self._ok = False + self._model = None return [] iw, ih = image.size names = getattr(res, "names", None) or {} diff --git a/agent/fc_agent/worker.py b/agent/fc_agent/worker.py index 2441491..359da89 100644 --- a/agent/fc_agent/worker.py +++ b/agent/fc_agent/worker.py @@ -1,15 +1,35 @@ -"""The lease → fetch → detect+embed → submit loop, run by a pool of worker -slots whose count is tunable live from the UI. +"""The lease → download → detect+embed → submit pipeline. -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. +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: -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. + 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 @@ -22,7 +42,7 @@ 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 +# 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 @@ -39,10 +59,13 @@ def _is_transient(exc: requests.RequestException) -> bool: 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 + +# 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 @@ -50,34 +73,36 @@ MAX_CONCURRENCY = 32 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) +# 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 grow must lift smoothed jobs/s by this to "help" -REPROBE_TICKS = 8 # decisions to hold after settling before re-probing +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 — the -# data that decides whether a download/compute split is worth building. +# 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 work loop's own lease/ +# 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 @@ -85,58 +110,88 @@ 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._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 - self._slots: list[_Slot] = [] + # 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 # slots currently mid-image + 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 - self._queue_thread = threading.Thread(target=self._queue_poll_loop, daemon=True) - self._queue_thread.start() + 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() - 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. + 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() - # Region proposers (extra YOLO detectors) — lazily built once, shared. 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 work loop is curator's only visitor. + 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: @@ -167,7 +222,9 @@ class Worker: 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.""" + 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: @@ -181,20 +238,15 @@ class Worker: 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) + 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(): @@ -204,58 +256,109 @@ class Worker: 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. + # 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, [] - self._active = 0 # no slots left → the meter reads 0 at once; any - # lagging decrement is clamped (see _bump) - for s in slots: - s.stop.set() # each slot releases its inflight on exit + 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): - # A manual set is an override → leave Auto. + # 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._target = max(1, min(MAX_CONCURRENCY, int(n))) + self._dl_target = max(1, min(DL_MAX, 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.""" + def _apply_downloaders(self, delta: int) -> bool: with self._lock: - new = max(1, min(MAX_CONCURRENCY, self._target + delta)) - if new == self._target: + new = max(1, min(DL_MAX, self._dl_target + delta)) + if new == self._dl_target: return False - self._target = new + 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): - 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() + """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 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. + # 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._target, - "max_concurrency": MAX_CONCURRENCY, + "concurrency": self._dl_target, # the UI dial = downloader count + "max_concurrency": DL_MAX, "auto": self._auto, - "workers": len(self._slots), + "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, @@ -267,14 +370,18 @@ class Worker: self.processed += processed self.errors += errors self.transient += transient - # Clamp at 0: a Stop resets _active to 0, so a slot that was mid-image - # decrements afterwards — that must not drive the meter negative. + # 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) - # --- per-slot loop ----------------------------------------------------- - def _loop(self, slot: _Slot): + # --- 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 not slot.stop.is_set() and self._running: + while self._running and not stop_evt.is_set(): try: _t = time.monotonic() jobs = self.client.lease(self.cfg.batch_size) @@ -283,132 +390,108 @@ class Worker: 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) + if stop_evt.wait(backoff): + break backoff = min(backoff * 2, MAX_BACKOFF_SECONDS) continue if not jobs: - self._interruptible_sleep(slot, self.cfg.poll_idle_seconds) + if stop_evt.wait(self.cfg.poll_idle_seconds): + break continue - slot.inflight = [j["job_id"] for j in jobs] + 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: - if slot.stop.is_set() or not self._running: + jid = job["job_id"] + if not self._running or stop_evt.is_set(): 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 = [] + 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 _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 + 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 - 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 + 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))] + else: + frames = [(None, media.load_image(data))] + self._record("decode", time.monotonic() - _t) + return frames - # 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 + # --- 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 - - tick += 1 - if tick < SAMPLES_PER_DECISION: + if item is None: # stop sentinel 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)) + 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: @@ -428,34 +511,18 @@ class Worker: 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) + 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: - _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 + 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 @@ -479,10 +546,14 @@ class Worker: 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(job["job_id"], vec, embed_version) + self.client.submit_embedding(jid, vec, embed_version) self._record("submit", time.monotonic() - _t) - self._bump(processed=1) + self._unhold(jid) return True # task picks what to produce per crop: @@ -506,6 +577,10 @@ class Worker: _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. @@ -566,6 +641,14 @@ class Worker: 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 @@ -575,36 +658,127 @@ class Worker: dropped = len(regions) - self.cfg.max_regions regions = regions[: self.cfg.max_regions] log.info("job %s: capped regions %d→%d (dropped %d)", - job.get("job_id"), len(regions) + dropped, - len(regions), dropped) + jid, len(regions) + dropped, len(regions), dropped) _t = time.monotonic() - self.client.submit(job["job_id"], regions, replace_kinds) + self.client.submit(jid, regions, replace_kinds) self._record("submit", time.monotonic() - _t) - self._bump(processed=1) + 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; no-ops if the server is still down, then the server's - # orphan-recovery reclaims it) and signal the loop to wait. + # effort; then the server's orphan-recovery reclaims it if down). self._bump(transient=1) - log.info("curator unreachable — released job %s, backing off", - job.get("job_id")) - self.client.release([job["job_id"]]) + log.info("curator unreachable — released job %s (post-GPU)", jid) + self.client.release([jid]) + self._unhold(jid) 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 + 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", - 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) + 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))