Merge pull request 'agent: download/GPU producer-consumer pipeline + detector fuse fix' (#173) from dev into main
Build images / sign-extension (push) Successful in 3s
CI / lint (push) Successful in 2s
Build images / build-ml (push) Successful in 7s
Build images / build-agent (push) Successful in 8s
Build images / build-web (push) Successful in 6s
CI / frontend-build (push) Successful in 19s
CI / backend-lint-and-test (push) Successful in 26s
CI / integration (push) Successful in 3m25s

This commit was merged in pull request #173.
This commit is contained in:
2026-06-30 23:39:25 -04:00
3 changed files with 478 additions and 273 deletions
+26 -13
View File
@@ -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 = """<!doctype html><html><head><meta charset=utf-8>
#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 = """<!doctype html><html><head><meta charset=utf-8>
.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 = """<!doctype html><html><head><meta charset=utf-8>
<button class=step onclick=setc(1)>+</button>
</div>
</div>
<div class=hint id=conchint>auto-tuning to fill the GPU · max <b id=capn>8</b></div>
<div class=hint id=conchint>auto-tuning downloaders to keep the GPU fed · max <b id=capn>8</b></div>
</section>
<section class=card>
<div class=card-h>Status</div>
<div class=tiles>
<div class=tile><div class=n id=state>—</div><div class=l>state</div></div>
<div class=tile><div class=n id=active>0</div><div class=l>active</div></div>
<div class=tile><div class=n id=dln>0</div><div class=l>downloaders</div></div>
<div class=tile><div class=n id=bufn>—</div><div class=l>buffer</div></div>
<div class=tile><div class=n id=active>0</div><div class=l>on GPU</div></div>
<div class=tile><div class="n ok" id=done>0</div><div class=l>processed</div></div>
<div class=tile><div class=n id=err>0</div><div class=l>errors</div></div>
<div class=tile><div class=n id=wait>0</div><div class=l>waited out</div></div>
</div>
<div class=meters>
<div class=meter><div class=meter-h><span>GPU util</span><b id=utillbl>—</b></div>
<div class=bar><i id=utilbar></i></div></div>
<div class=meter><div class=meter-h><span>VRAM</span><b id=vramlbl>—</b></div>
<div class=bar><i id=gpubar></i></div></div>
<div class=meter><div class=meter-h><span>buffer occupancy</span><b id=buflbl>—</b></div>
<div class=bar><i id=bufbar></i></div></div>
</div>
<div class=queue id=pipe>consumers — · waited out 0</div>
<div class=queue id=queue>queue —</div>
</section>
@@ -281,13 +288,19 @@ _PAGE = """<!doctype html><html><head><meta charset=utf-8>
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
+19 -1
View File
@@ -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 {}
+433 -259
View File
@@ -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 4005462ms,
GPU ~300600ms), 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) (12) ── 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))