bvandeusen
7c19ad91ed
Autoscaler (agent 2026-07-02.5): the buffer-occupancy signal alone would peg downloaders at DL_MAX while the bandwidth CAP — not concurrency — is the real constraint (8 streams sharing 8 MB/s move no more data than 4). Growth is now gated on the pipe having headroom (net < 85% of cap) and a pipe pinned at the cap (>= 95%) sheds streams down to 3; dead band prevents flapping. The UI hint says 'holding at the bandwidth cap' and /status reports bw_capped, so the behavior is legible without tests that need the ML stack. Reset content tagging: stays a FULL-instance reset (operator's call), but now lives in a fenced 'Danger zone' section on Cleanup and the apply is gated by a preview-derived confirm token (mirrors the Tier-C bulk-delete pattern — stale counts are rejected server-side). Copy no longer claims suggestions repopulate: it says plainly the heads' training examples are deleted and re-tagging starts fresh. Moved out of TagMaintenanceCard into DangerZoneCard. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CDgx8bQS5YrGRK76v8HUnM
1087 lines
56 KiB
Python
1087 lines
56 KiB
Python
"""The lease → download → detect+embed → submit pipeline.
|
||
|
||
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:
|
||
|
||
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 contextlib
|
||
import logging
|
||
import queue
|
||
import threading
|
||
import time
|
||
|
||
import numpy as np
|
||
import requests
|
||
|
||
from . import media, models
|
||
from .client import FcClient
|
||
from .config import Config
|
||
from .crops import crop_region
|
||
from .detectors import dedupe_crops
|
||
from .throttle import TokenBucket
|
||
|
||
# Cap on the lease-retry backoff: when curator is unreachable (e.g. you redeploy
|
||
# 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
|
||
|
||
# A job whose fetch dies transiently this many times IN ONE SESSION stops being
|
||
# handed back and is failed instead. Transient handbacks (release) burn no
|
||
# attempts on the server, so a poisoned transfer — an original that stalls the
|
||
# download every single time — would otherwise release→re-lease forever, churning
|
||
# bandwidth without ever landing in the error queue (operator-observed jobs
|
||
# 99044/125288/131594/143131, 2026-07-01). Failing it lets curator's attempt cap
|
||
# tombstone it WITH the real reason. A genuine curator outage is unaffected:
|
||
# every job takes at most one strike per outage, and strikes clear on success.
|
||
TRANSIENT_JOB_CAP = 3
|
||
|
||
|
||
def _is_transient(exc: requests.RequestException) -> bool:
|
||
"""A server/transport problem (wait it out) vs a job-specific fault (fail it).
|
||
No response → connection refused/timeout → curator is down → transient. With
|
||
a response: 5xx, auth (401/403, e.g. a token blip on redeploy), 408/409/429
|
||
(timeout / our lease reclaimed / rate-limited) are all 'not this job's fault'.
|
||
A specific 4xx like 404 (image gone) / 400 IS the job's fault → fail it."""
|
||
resp = getattr(exc, "response", None)
|
||
if resp is None:
|
||
return True
|
||
return resp.status_code >= 500 or resp.status_code in (401, 403, 408, 409, 429)
|
||
|
||
|
||
def _transient_reason(exc: requests.RequestException) -> str:
|
||
"""A SPECIFIC label for a transient failure — so the log distinguishes a
|
||
stalled/slow transfer from an actually-unreachable curator, which need
|
||
different fixes. `HTTP <code>` for a 5xx/auth/conflict, else the exception
|
||
class: ReadTimeout (transfer stalled >60s between bytes — curator up, this
|
||
file/stream is slow), ConnectTimeout (curator didn't accept in 10s → web
|
||
workers/pool exhausted or down), ConnectionError (reset mid-transfer)."""
|
||
resp = getattr(exc, "response", None)
|
||
if resp is not None:
|
||
return f"HTTP {resp.status_code}"
|
||
return type(exc).__name__
|
||
|
||
|
||
def _ewma(prev: float | None, x: float, alpha: float) -> float:
|
||
"""Exponentially-weighted moving average: seed with the first sample, then
|
||
fold each new sample in at weight `alpha`. Smooths the noisy control signals
|
||
(buffer occupancy, GPU util, throughput) the autoscaler decides on."""
|
||
return x if prev is None else alpha * x + (1 - alpha) * prev
|
||
|
||
|
||
# Pipeline sizing. Downloaders are I/O-bound, but every download streams a full
|
||
# original (large videos included) THROUGH curator's single Python file-serving
|
||
# path — so the ceiling is deliberately modest: too many concurrent large-file
|
||
# GETs saturate curator's web workers + NFS and slow everything (including the
|
||
# browser). 8 keeps a fast GPU fed without stampeding curator. 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 = 8 # 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
|
||
# the heads were trained in. These cover an older server that doesn't send them.
|
||
DEFAULT_EMBED_MODEL = "google/siglip-so400m-patch14-384"
|
||
DEFAULT_EMBED_VERSION = "siglip-so400m-patch14-384"
|
||
|
||
# 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)
|
||
# Bandwidth-cap awareness (operator 2026-07-02): with the aggregate governor in
|
||
# place, the occupancy signal alone would peg downloaders at DL_MAX while the
|
||
# CAP — not concurrency — is the real constraint: 8 streams sharing 8 MB/s move
|
||
# no more data than 4, they just hold more leases + RAM and stretch every job's
|
||
# latency. So growth is gated on the pipe having headroom, and a pipe pinned at
|
||
# the cap sheds streams down to BW_MIN_DL (enough overlap to keep the cap
|
||
# filled through TTFB + decode gaps). The dead band between the two thresholds
|
||
# prevents add/trim flapping.
|
||
BW_ADD_HEADROOM = 0.85 # add a downloader only while net < 85% of the cap
|
||
BW_TRIM_AT = 0.95 # net ≥ 95% of the cap → shed toward BW_MIN_DL
|
||
BW_MIN_DL = 3
|
||
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 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.
|
||
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 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
|
||
|
||
# Throughput rates (jobs/min, downloads/min) are computed HERE, on a fixed
|
||
# wall-clock cadence, and reported ready-to-show — NOT derived in the browser
|
||
# from poll deltas. A background/unfocused tab throttles its timers to ~1/min,
|
||
# which made a client-side delta-rate reject every sample and read blank forever;
|
||
# a server-side EWMA is immune to how often (or seldom) the UI actually polls.
|
||
RATE_INTERVAL = 3.0
|
||
RATE_ALPHA = 0.3
|
||
|
||
# Stop backstop: pressing Stop signals every worker thread to wind down and hand
|
||
# its lease back; the UI shows "stopping" until those threads have actually
|
||
# exited, then "stopped" — so the state the operator sees is always truthful.
|
||
# If a thread wedges (a stuck submit/release to an overloaded curator), NEVER
|
||
# hang in "stopping": declare "stopped" after this many seconds and let the lone
|
||
# straggler finish detached. This is the guarantee the old active>0 UI hack lacked.
|
||
STOPPING_TIMEOUT = 20.0
|
||
|
||
# The lifecycle the operator sees. `_state` is the source of truth the UI reads;
|
||
# `_running` is the low-level flag the worker THREADS gate on (True while
|
||
# starting/running; False from the Stop press onward, so they wind down at once).
|
||
# stopped → starting (spun up, not yet reached curator)
|
||
# → running (leased at least once — curator is answering)
|
||
# → stopping (Stop pressed; threads winding down + handing work back)
|
||
# → stopped (threads exited, or the timeout backstop fired)
|
||
STOPPED, STARTING, RUNNING, STOPPING = "stopped", "starting", "running", "stopping"
|
||
|
||
log = logging.getLogger("fc_agent.worker")
|
||
|
||
|
||
class Worker:
|
||
def __init__(self, cfg: Config):
|
||
self.cfg = cfg
|
||
self.client = FcClient(cfg.fc_url, cfg.token, cfg.agent_id)
|
||
# Sent to ffmpeg-from-URL so video sampling works whether or not a
|
||
# deployment gates /images behind the bearer token (FC's is public).
|
||
self._auth_header = (
|
||
f"Authorization: Bearer {cfg.token}\r\n" if cfg.token else ""
|
||
)
|
||
self._lock = threading.Lock()
|
||
# ONE bandwidth budget for everything the agent pulls (still downloads
|
||
# and ffmpeg video streams): the agent shares a desktop's network, so
|
||
# the polite bound is on the aggregate — see throttle.py for why.
|
||
self.throttle = TokenBucket(cfg.bandwidth_limit_mb_s * 1_048_576)
|
||
self._running = False
|
||
# The lifecycle state the UI shows (see the STOPPED/STARTING/... consts).
|
||
# It stays STOPPING — a truthful "winding down" — from the Stop press until
|
||
# the worker threads have actually exited (or the timeout backstop fires),
|
||
# then STOPPED. `_running` flips False the instant Stop is pressed.
|
||
self._state = STOPPED
|
||
self._shutdown_thread: threading.Thread | None = None
|
||
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
|
||
# 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()
|
||
# job_id → in-session transient-handback count (guarded by _held_lock).
|
||
# Cleared on success or terminal fail; KEPT across releases — persisting
|
||
# through the release→re-lease cycle is what lets TRANSIENT_JOB_CAP
|
||
# catch a poisoned transfer.
|
||
self._transient_seen: dict[int, int] = {}
|
||
self.processed = 0
|
||
self.downloaded = 0 # jobs fetched+decoded into the buffer (monotonic);
|
||
# feeds the server-side downloads/min rate below.
|
||
self.errors = 0
|
||
self.transient = 0 # jobs handed back due to a server outage (NOT
|
||
# failed) — the "waiting out curator" counter
|
||
self._active = 0 # jobs currently mid-GPU (consumers busy)
|
||
# Smoothed throughput rates the UI just displays (jobs/min, downloads/min),
|
||
# computed on a fixed cadence by _rate_loop so they're independent of how
|
||
# often the browser polls (see RATE_INTERVAL). Decay to 0 when work stops.
|
||
self._jpm = 0.0
|
||
self._dpm = 0.0
|
||
self._net_mb_s = 0.0 # smoothed aggregate download rate (UI readout)
|
||
self._bw_capped = False # autoscaler is holding/shedding at the cap (UI)
|
||
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
|
||
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()
|
||
threading.Thread(target=self._stats_loop, daemon=True).start()
|
||
threading.Thread(target=self._rate_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()
|
||
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 _strike(self, jid: int) -> int:
|
||
"""Record a transient bounce for this job; returns the in-session total
|
||
(compared against TRANSIENT_JOB_CAP by both the fetch and submit paths)."""
|
||
with self._held_lock:
|
||
n = self._transient_seen.get(jid, 0) + 1
|
||
self._transient_seen[jid] = n
|
||
return n
|
||
|
||
def _forget_strikes(self, jid: int) -> None:
|
||
"""Terminal outcome (submitted or failed) → this job's strike count no
|
||
longer matters. Deliberately NOT called on release: strikes surviving
|
||
the release→re-lease cycle is what makes TRANSIENT_JOB_CAP work."""
|
||
with self._held_lock:
|
||
self._transient_seen.pop(jid, None)
|
||
|
||
def _release(self, job_ids: list[int]) -> None:
|
||
"""Hand still-held leases back to curator and drop them from the held set —
|
||
the single hand-back path for both a downloader exiting (stop/shrink) with
|
||
unbuffered leases and a consumer releasing one job it can't finish. Empty
|
||
list is a no-op; a duplicate release is a harmless server no-op."""
|
||
if not job_ids:
|
||
return
|
||
self.client.release(job_ids)
|
||
for jid in job_ids:
|
||
self._unhold(jid)
|
||
|
||
def _fail(self, jid: int, image_id, exc: Exception, verb: str = "failed") -> None:
|
||
"""Terminal job-fault path: count the error, log it, tell curator the job
|
||
failed, and drop the lease. `verb` lets a caller say 'failed to decode'."""
|
||
self._bump(errors=1)
|
||
log.warning("job %s (image %s) %s: %s", jid, image_id, verb, str(exc)[:200])
|
||
self.client.fail(jid, str(exc)[:500])
|
||
self._unhold(jid)
|
||
self._forget_strikes(jid)
|
||
|
||
def _stopped(self, stop_evt: threading.Event) -> bool:
|
||
"""The shared 'should I bail now?' check — the worker is stopping (global
|
||
flag) or this thread's own stop event is set. Lock-free atomic reads."""
|
||
return not self._running or stop_evt.is_set()
|
||
|
||
def _abort_if_stopped(self, jid: int, stop_evt: threading.Event) -> bool:
|
||
"""If we're stopping, hand this job's lease back (so a Stop drains promptly
|
||
instead of finishing heavy work) and return True so the caller bails."""
|
||
if self._stopped(stop_evt):
|
||
self._release([jid])
|
||
return True
|
||
return False
|
||
|
||
# --- 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 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:
|
||
try:
|
||
self._queue = self.client.queue_status()
|
||
except Exception:
|
||
self._queue = None
|
||
time.sleep(QUEUE_POLL_INTERVAL)
|
||
|
||
def note_ui(self) -> None:
|
||
"""The UI polled /status — keep the queue snapshot warm for a while."""
|
||
self._ui_seen = time.monotonic()
|
||
|
||
def latest_queue(self) -> dict | None:
|
||
return self._queue
|
||
|
||
def util_smooth(self) -> float | None:
|
||
return self._util_smooth
|
||
|
||
def _record(self, stage: str, seconds: float) -> None:
|
||
with self._timing_lock:
|
||
s = self._timing.get(stage)
|
||
if s is None:
|
||
self._timing[stage] = [seconds, 1]
|
||
else:
|
||
s[0] += seconds
|
||
s[1] += 1
|
||
|
||
@contextlib.contextmanager
|
||
def _timed(self, stage: str):
|
||
"""Time a pipeline stage and fold it into the per-stage breakdown — but
|
||
ONLY on a clean exit: an exception (a failed download, a stopped GPU pass)
|
||
propagates through the `yield` and skips the record, so it never skews the
|
||
stage average with work that didn't complete. Matches the old inline
|
||
monotonic()/_record() pairs, where _record ran only if no exception fired."""
|
||
t = time.monotonic()
|
||
yield
|
||
self._record(stage, time.monotonic() - t)
|
||
|
||
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.
|
||
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:
|
||
snap = {k: (v[0], v[1]) for k, v in self._timing.items() if v[1]}
|
||
self._timing = {}
|
||
if not snap:
|
||
continue
|
||
order = ["lease", "download", "decode", "gpu", "submit"]
|
||
parts = [
|
||
f"{st} {1000 * snap[st][0] / snap[st][1]:.0f}ms"
|
||
for st in order if st in snap
|
||
]
|
||
jobs = (snap.get("gpu") or snap.get("download") or (0, 0))[1]
|
||
log.info("timing/%ds — %s (%d jobs)",
|
||
int(STATS_INTERVAL), " · ".join(parts), jobs)
|
||
|
||
def _rate_loop(self) -> None:
|
||
"""Compute the jobs/min + downloads/min the UI shows, on a fixed cadence
|
||
from the monotonic counters — EWMA-smoothed, clamped so a counter reset
|
||
(agent restart) can't spike them, decaying to 0 when work stops. Doing
|
||
this here (not in the browser) makes the numbers independent of the poll
|
||
rate, so a throttled/unfocused tab still shows a real rate."""
|
||
prev_p, prev_d, prev_t = self.processed, self.downloaded, time.monotonic()
|
||
prev_b = self.throttle.consumed
|
||
while True:
|
||
time.sleep(RATE_INTERVAL)
|
||
now = time.monotonic()
|
||
dt = now - prev_t
|
||
if dt > 0:
|
||
jp = max(0.0, 60.0 * (self.processed - prev_p) / dt)
|
||
dp = max(0.0, 60.0 * (self.downloaded - prev_d) / dt)
|
||
nb = max(0.0, (self.throttle.consumed - prev_b) / dt / 1_048_576)
|
||
self._jpm = RATE_ALPHA * jp + (1 - RATE_ALPHA) * self._jpm
|
||
self._dpm = RATE_ALPHA * dp + (1 - RATE_ALPHA) * self._dpm
|
||
self._net_mb_s = RATE_ALPHA * nb + (1 - RATE_ALPHA) * self._net_mb_s
|
||
prev_p, prev_d, prev_t = self.processed, self.downloaded, now
|
||
prev_b = self.throttle.consumed
|
||
|
||
# --- control -----------------------------------------------------------
|
||
def start(self):
|
||
# Only a fully-stopped worker can start — ignore a click while already
|
||
# up or still winding down (the UI also disables the button then). This
|
||
# keeps a Start during "stopping" from racing the shutdown monitor's
|
||
# lease-release against a fresh pool's new leases.
|
||
with self._lock:
|
||
if self._state != STOPPED:
|
||
return
|
||
self._running = True
|
||
self._state = STARTING # → RUNNING on the first successful lease
|
||
self._dl_target = max(1, self._dl_target)
|
||
self._consumer_target = max(1, self._consumer_target)
|
||
self._reconcile_locked()
|
||
log.info("▶ start → starting (spinning up, awaiting curator)")
|
||
# (Re)start the autoscaler control loop.
|
||
if self._ctrl_thread is None or not self._ctrl_thread.is_alive():
|
||
self._ctrl_stop.clear()
|
||
self._ctrl_thread = threading.Thread(target=self._control_loop, daemon=True)
|
||
self._ctrl_thread.start()
|
||
|
||
def _mark_running(self) -> None:
|
||
"""STARTING → RUNNING once a downloader has actually leased from curator —
|
||
so 'running' means 'curator is answering', not just 'Start was clicked'.
|
||
First caller wins; the rest are a cheap no-op."""
|
||
if self._state != STARTING:
|
||
return
|
||
with self._lock:
|
||
if self._state != STARTING:
|
||
return
|
||
self._state = RUNNING
|
||
log.info("state → running (curator reachable — leasing)")
|
||
|
||
def stop(self):
|
||
# Enter STOPPING and signal every worker thread to wind down. The handler
|
||
# returns at once — the actual wait-for-exit + lease-release runs in a
|
||
# background monitor, so a slow curator can NEVER block the Stop button or
|
||
# wedge the state (the bug: the old UI faked "stopping" from active>0 and
|
||
# sat there forever when a consumer hung mid-release).
|
||
with self._lock:
|
||
if self._state not in (STARTING, RUNNING):
|
||
return # already stopping/stopped
|
||
self._state = STOPPING
|
||
self._running = False # every thread's _stopped() now trips → wind down
|
||
dls, self._dls = self._dls, []
|
||
cons, self._consumers = self._consumers, []
|
||
self._active = 0
|
||
self._ctrl_stop.set()
|
||
for _, ev in dls:
|
||
ev.set()
|
||
for _, ev in cons:
|
||
ev.set()
|
||
# Wake any consumer blocked on an empty buffer so it sees the stop at once.
|
||
for _ in range(CONSUMER_MAX):
|
||
try:
|
||
self._buffer.put_nowait(None)
|
||
except queue.Full:
|
||
break
|
||
log.info("■ stop → stopping (winding down, handing work back)")
|
||
threads = [t for t, _ in dls] + [t for t, _ in cons]
|
||
self._shutdown_thread = threading.Thread(
|
||
target=self._finish_stop, args=(threads,), daemon=True)
|
||
self._shutdown_thread.start()
|
||
|
||
def _finish_stop(self, threads: list[threading.Thread]) -> None:
|
||
"""Wait for the signalled worker threads to actually exit, then land on
|
||
STOPPED — so 'stopped' is truthful. Bounded by STOPPING_TIMEOUT so a
|
||
wedged submit/release can never hold the UI in 'stopping' forever."""
|
||
deadline = time.monotonic() + STOPPING_TIMEOUT
|
||
for t in threads:
|
||
t.join(timeout=max(0.0, deadline - time.monotonic()))
|
||
# Hand back every lease still held in one shot (background — no handler
|
||
# block); a straggler releasing its own job later is a harmless dup.
|
||
self._drain_and_release()
|
||
with self._lock:
|
||
if self._state == STOPPING: # a Start may have re-entered; if so, leave it
|
||
self._state = STOPPED
|
||
self._active = 0 # any straggler's -1 is clamped (see _bump)
|
||
stragglers = sum(1 for t in threads if t.is_alive())
|
||
if stragglers:
|
||
log.info("state → stopped (%d thread(s) still winding down past %ds — detached)",
|
||
stragglers, int(STOPPING_TIMEOUT))
|
||
else:
|
||
log.info("state → stopped (all work handed back)")
|
||
|
||
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_bandwidth(self, mb_s: float):
|
||
# Live-retunes the shared bucket; a downloader blocked mid-wait re-checks
|
||
# immediately (set_rate notifies), so raising the cap takes effect now.
|
||
self.throttle.set_rate(max(0.0, float(mb_s)) * 1_048_576)
|
||
log.info("bandwidth cap set to %s",
|
||
"unlimited" if self.throttle.rate <= 0
|
||
else f"{self.throttle.rate / 1_048_576:g} MB/s")
|
||
|
||
def set_concurrency(self, n: int):
|
||
# 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._dl_target = max(1, min(DL_MAX, int(n)))
|
||
if self._running:
|
||
self._reconcile_locked()
|
||
|
||
def _apply_downloaders(self, delta: int) -> bool:
|
||
with self._lock:
|
||
new = max(1, min(DL_MAX, self._dl_target + delta))
|
||
if new == self._dl_target:
|
||
return False
|
||
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):
|
||
"""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 / 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": self._state, # stopped|starting|running|stopping (truthful)
|
||
"concurrency": self._dl_target, # the UI dial = downloader count
|
||
"max_concurrency": DL_MAX,
|
||
"auto": self._auto,
|
||
"downloaders": len(self._dls),
|
||
"consumers": len(self._consumers),
|
||
"buffer": self._buffer.qsize(),
|
||
"buffer_max": BUFFER_MAX,
|
||
"active": self._active,
|
||
"processed": self.processed,
|
||
"downloaded": self.downloaded,
|
||
"jobs_per_min": round(self._jpm, 1), # ready-to-show throughput
|
||
"downloads_per_min": round(self._dpm, 1),
|
||
"errors": self.errors,
|
||
"transient": self.transient,
|
||
"bandwidth_limit_mb_s": round(self.throttle.rate / 1_048_576, 1),
|
||
"net_mb_s": round(self._net_mb_s, 1), # observed aggregate rate
|
||
"bw_capped": self._bw_capped, # autoscaler holding at the cap (UI hint)
|
||
}
|
||
|
||
def _bump(self, *, processed=0, downloaded=0, errors=0, active=0, transient=0):
|
||
with self._lock:
|
||
self.processed += processed
|
||
self.downloaded += downloaded
|
||
self.errors += errors
|
||
self.transient += transient
|
||
# 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)
|
||
|
||
# --- 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 self._stopped(stop_evt):
|
||
try:
|
||
with self._timed("lease"):
|
||
jobs = self.client.lease(self.cfg.batch_size)
|
||
backoff = self.cfg.poll_idle_seconds # server answered → reset
|
||
self._mark_running() # curator answered → leave "starting"
|
||
except Exception:
|
||
# curator unreachable (redeploy, network drop): wait it out with
|
||
# exponential backoff, capped — resume on our own when it returns.
|
||
if stop_evt.wait(backoff):
|
||
break
|
||
backoff = min(backoff * 2, MAX_BACKOFF_SECONDS)
|
||
continue
|
||
if not jobs:
|
||
if stop_evt.wait(self.cfg.poll_idle_seconds):
|
||
break
|
||
continue
|
||
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:
|
||
jid = job["job_id"]
|
||
if self._stopped(stop_evt):
|
||
break
|
||
try:
|
||
frames = self._download_decode(job, stop_evt)
|
||
except requests.RequestException as exc:
|
||
owned.remove(jid)
|
||
if _is_transient(exc):
|
||
# A Stop mid-fetch lands here too (a killed ffmpeg is
|
||
# not the job's fault) — no strike for that; only count
|
||
# bounces taken while genuinely running.
|
||
if (not self._stopped(stop_evt)
|
||
and self._strike(jid) >= TRANSIENT_JOB_CAP):
|
||
# THIS job keeps dying while others move: a poisoned
|
||
# transfer, not a curator outage. Fail it so the
|
||
# server's attempt cap tombstones it with the real
|
||
# reason instead of cycling it forever.
|
||
self._fail(
|
||
jid, job.get("image_id"), exc,
|
||
verb="gave up after repeated transient failures",
|
||
)
|
||
continue
|
||
# 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._release([jid])
|
||
# Name the ACTUAL failure — "curator unreachable" was
|
||
# printed for every transient, hiding whether a single
|
||
# file's transfer stalled (ReadTimeout, curator fine) or
|
||
# curator itself is down (ConnectTimeout/ConnectionError).
|
||
log.info("fetch failed job %s (image %s, %s) — released, backing off",
|
||
jid, job.get("image_id"), _transient_reason(exc))
|
||
self._release(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._fail(jid, job.get("image_id"), exc)
|
||
continue
|
||
except Exception as exc: # noqa: BLE001 — bad media → the job's fault
|
||
owned.remove(jid)
|
||
self._fail(jid, job.get("image_id"), exc, verb="failed to decode")
|
||
continue
|
||
# Blocks on a full buffer (backpressure) but wakes promptly on stop.
|
||
if self._put((job, frames), stop_evt):
|
||
self._bump(downloaded=1) # fetched+decoded into the buffer
|
||
owned.remove(jid) # ownership handed to the buffer/consumer
|
||
else:
|
||
break # stopped while waiting for buffer space
|
||
self._release(owned)
|
||
|
||
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 not self._stopped(stop_evt):
|
||
try:
|
||
self._buffer.put(item, timeout=0.5)
|
||
return True
|
||
except queue.Full:
|
||
continue
|
||
return False
|
||
|
||
def _download_decode(self, job: dict, stop_evt: threading.Event):
|
||
"""Fetch the image bytes and decode → [(frame_time, PIL.Image)]. Videos
|
||
are sampled into frames (ffmpeg). Records the download + decode timings."""
|
||
if media.is_video(job.get("mime", "")):
|
||
# Stream the video: ffmpeg reads the media URL directly and Range-reads
|
||
# only the frames it needs, so we NEVER pull the whole file (VR/4K
|
||
# originals are 800MB+ — buffering that in RAM and getting cut off
|
||
# mid-download was the failure loop). Environment-agnostic + resilient.
|
||
url = f"{self.cfg.fc_url}{job['image_url']}"
|
||
with self._timed("decode"):
|
||
frames, ffmpeg_err = media.sample_frames_from_url(
|
||
url, job.get("frame_interval_seconds", 4.0),
|
||
job.get("max_frames", 64),
|
||
headers=self._auth_header, timeout=self.cfg.ffmpeg_timeout,
|
||
should_stop=lambda: self._stopped(stop_evt),
|
||
governor=self.throttle,
|
||
)
|
||
if not frames:
|
||
# Stop killed ffmpeg → NOT the job's fault; raise transient so the
|
||
# downloader releases (not fails) it as it winds down.
|
||
if self._stopped(stop_evt):
|
||
raise requests.ConnectionError("stopped during video sampling")
|
||
# Else couldn't sample. If curator is up, the file is unprocessable
|
||
# → a job fault: fail it WITH ffmpeg's reason, so the job's stored
|
||
# error says e.g. "moov atom not found" instead of a bare
|
||
# "unprocessable". If curator is unreachable, it's transient → let
|
||
# the loop back off + retry (ConnectionError is caught as such).
|
||
if self.client.is_reachable():
|
||
raise RuntimeError(
|
||
f"no frames sampled from video — {ffmpeg_err or 'unknown reason'}"
|
||
)
|
||
raise requests.ConnectionError("curator unreachable during video sampling")
|
||
# Temporal dedup: a near-static video re-runs the whole detect+embed
|
||
# chain on ~identical frames — drop near-dups HERE (CPU) pre-GPU.
|
||
if len(frames) > 1 and self.cfg.frame_dedupe_distance > 0:
|
||
kept = media.dedupe_frames(frames, self.cfg.frame_dedupe_distance)
|
||
if len(kept) < len(frames):
|
||
log.info("job %s: video frames %d→%d (near-dup dedup)",
|
||
job.get("job_id"), len(frames), len(kept))
|
||
frames = kept
|
||
return frames
|
||
# Stills: download the bytes and decode.
|
||
with self._timed("download"):
|
||
data = self.client.fetch_image(job["image_url"], throttle=self.throttle)
|
||
with self._timed("decode"):
|
||
frames = [(None, media.load_image(data))]
|
||
return frames
|
||
|
||
# --- GPU consumer pool -------------------------------------------------
|
||
def _consumer(self, stop_evt: threading.Event):
|
||
"""Pull decoded jobs off the buffer and run detect + embed + submit."""
|
||
while not self._stopped(stop_evt):
|
||
try:
|
||
item = self._buffer.get(timeout=1.0)
|
||
except queue.Empty:
|
||
continue
|
||
if item is None: # stop sentinel
|
||
continue
|
||
job, frames = item
|
||
if self._abort_if_stopped(job["job_id"], stop_evt):
|
||
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:
|
||
return self._embedder
|
||
with self._embedder_lock:
|
||
if self._embedder is None:
|
||
from .embedder import CropEmbedder
|
||
self._embedder = CropEmbedder(model_name, self.cfg.embed_dtype)
|
||
return self._embedder
|
||
|
||
def _ensure_proposers(self):
|
||
if self._proposers is not None:
|
||
return self._proposers
|
||
with self._proposers_lock:
|
||
if self._proposers is None:
|
||
from .detectors import Proposers
|
||
self._proposers = Proposers(self.cfg)
|
||
return self._proposers
|
||
|
||
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:
|
||
if self._abort_if_stopped(jid, stop_evt):
|
||
return False
|
||
|
||
task = job.get("task") or "ccip"
|
||
embed_version = job.get("embed_version") or DEFAULT_EMBED_VERSION
|
||
model_name = (
|
||
self.cfg.embed_model_override
|
||
or job.get("embed_model_name")
|
||
or DEFAULT_EMBED_MODEL
|
||
)
|
||
|
||
# 'embed' = WHOLE-IMAGE SigLIP embedding (re-embed the library under a
|
||
# new model, #1190) → image_record.siglip_embedding. Mean-pool video
|
||
# frames, matching the server's tag_and_embed. No regions.
|
||
if task == "embed":
|
||
embedder = self._ensure_embedder(model_name) # one-time model load
|
||
with self._timed("gpu"):
|
||
vecs = [embedder.embed(frame) for _, frame in frames]
|
||
vec = (
|
||
np.mean(np.asarray(vecs, dtype=np.float32), axis=0).tolist()
|
||
if len(vecs) > 1 else vecs[0]
|
||
)
|
||
if self._abort_if_stopped(jid, stop_evt):
|
||
return False
|
||
with self._timed("submit"):
|
||
self.client.submit_embedding(jid, vec, embed_version)
|
||
self._unhold(jid)
|
||
self._forget_strikes(jid)
|
||
return True
|
||
|
||
# task picks what to produce per crop:
|
||
# 'siglip' (backfill existing images) → concept (SigLIP) regions
|
||
# ONLY, so it never churns their figure/CCIP regions or the
|
||
# character-reference cache.
|
||
# 'ccip' / 'both' (a new image's first pass) → figure (CCIP) AND
|
||
# concept (SigLIP) in one go, off the same crop.
|
||
want_ccip = task in ("ccip", "both")
|
||
want_siglip = task in ("ccip", "siglip", "both")
|
||
replace_kinds = (
|
||
["concept", "panel"] if task == "siglip"
|
||
else ["figure", "face", "concept", "panel"]
|
||
)
|
||
embedder = self._ensure_embedder(model_name) if want_siglip else None
|
||
proposers = self._ensure_proposers()
|
||
|
||
regions = []
|
||
ccip_ev = self.cfg.ccip_model or "ccip-default"
|
||
dv = f"person-{self.cfg.detector_level}"
|
||
|
||
# detect + CCIP + batched embed across every (deduped) frame = "gpu".
|
||
with self._timed("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 self._stopped(stop_evt):
|
||
break
|
||
# FIGURE boxes: imgutils detect_person ∪ general COCO person,
|
||
# NMS-merged → CCIP identity (+ a concept crop). Covers anime +
|
||
# Western/realistic figures.
|
||
base = models.detect_figures(frame, self.cfg.detector_level)
|
||
figs = proposers.figures(frame, base)
|
||
if not figs:
|
||
figs = [((0.0, 0.0, 1.0, 1.0), 1.0, "whole")] # whole-frame fallback
|
||
|
||
# Collect every crop that needs a SigLIP embedding, then embed
|
||
# them in ONE batched forward pass (huge GPU-util + throughput
|
||
# win vs one forward per crop). CCIP runs per figure inline.
|
||
pending = [] # (crop, region-template-without-embedding)
|
||
for bbox, score, _label in figs:
|
||
crop = crop_region(frame, bbox)
|
||
if crop is None:
|
||
continue
|
||
if want_ccip:
|
||
regions.append({
|
||
"kind": "figure", "bbox": list(bbox), "frame_time": t,
|
||
"score": score,
|
||
"ccip_embedding": models.ccip_vector(
|
||
crop, self.cfg.ccip_model or None
|
||
),
|
||
"embedding_version": ccip_ev, "detector_version": dv,
|
||
})
|
||
if want_siglip:
|
||
pending.append((crop, {
|
||
"kind": "concept", "bbox": list(bbox), "frame_time": t,
|
||
"score": score, "detector_version": dv,
|
||
}))
|
||
if not want_siglip:
|
||
continue
|
||
# ANATOMY components (booru_yolo) + PANELS → concept/panel crops.
|
||
for bbox, score, label in proposers.components(frame):
|
||
crop = crop_region(frame, bbox)
|
||
if crop is not None:
|
||
pending.append((crop, {
|
||
"kind": "concept", "bbox": list(bbox), "frame_time": t,
|
||
"score": score, "detector_version": f"booru:{label}",
|
||
}))
|
||
for bbox, score, _label in proposers.panels(frame):
|
||
crop = crop_region(frame, bbox)
|
||
if crop is not None:
|
||
pending.append((crop, {
|
||
"kind": "panel", "bbox": list(bbox), "frame_time": t,
|
||
"score": score, "detector_version": "panel",
|
||
}))
|
||
if pending:
|
||
# Drop near-duplicate crops (a figure box ≈ an anatomy
|
||
# component, overlapping booru classes) before the embed so
|
||
# we never SigLIP the same region twice — saves GPU and a
|
||
# slot against max_regions. High-IoU + kind-aware, so
|
||
# intentional nested crops (figure ⊃ head) survive.
|
||
pending = dedupe_crops(pending, self.cfg.dedupe_iou)
|
||
vecs = embedder.embed_batch([c for c, _ in pending])
|
||
for (_c, tmpl), vec in zip(pending, vecs, strict=True):
|
||
tmpl["siglip_embedding"] = vec
|
||
tmpl["embedding_version"] = embed_version
|
||
regions.append(tmpl)
|
||
# Stop once we have enough: a long video (image 81602 = a 156 MB
|
||
# mp4, 64 sampled frames × ~32 regions) would otherwise burn
|
||
# ~38s of GPU across every frame before the submit is even
|
||
# truncated. Bounds the WORK, not just the POST body.
|
||
if len(regions) >= self.cfg.max_regions:
|
||
break
|
||
|
||
# A Stop mid-frame-loop leaves partial regions — don't submit those;
|
||
# hand the whole job back so another agent redoes it cleanly.
|
||
if self._abort_if_stopped(jid, stop_evt):
|
||
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
|
||
# (operator-flagged: image 81602 looped on 413).
|
||
if len(regions) > self.cfg.max_regions:
|
||
regions.sort(key=lambda r: r.get("score", 0.0) or 0.0, reverse=True)
|
||
dropped = len(regions) - self.cfg.max_regions
|
||
regions = regions[: self.cfg.max_regions]
|
||
log.info("job %s: capped regions %d→%d (dropped %d)",
|
||
jid, len(regions) + dropped, len(regions), dropped)
|
||
with self._timed("submit"):
|
||
self.client.submit(jid, regions, replace_kinds)
|
||
self._unhold(jid)
|
||
self._forget_strikes(jid)
|
||
return True
|
||
except requests.RequestException as exc:
|
||
if _is_transient(exc):
|
||
if (not self._stopped(stop_evt)
|
||
and self._strike(jid) >= TRANSIENT_JOB_CAP):
|
||
# Same poison rationale as the fetch path: a job whose
|
||
# submit keeps dying transiently would otherwise re-lease →
|
||
# re-download → re-GPU forever.
|
||
self._fail(jid, job.get("image_id"), exc,
|
||
verb="gave up after repeated transient failures")
|
||
return False
|
||
# curator down/redeploying, a 5xx, or our lease was reclaimed
|
||
# while we worked. NOT the job's fault — hand it back (best
|
||
# effort; then the server's orphan-recovery reclaims it if down).
|
||
self._bump(transient=1)
|
||
log.info("submit failed job %s (%s) — released, re-lease later",
|
||
jid, _transient_reason(exc))
|
||
self._release([jid])
|
||
return False
|
||
self._fail(jid, job.get("image_id"), exc)
|
||
return False
|
||
except Exception as exc: # noqa: BLE001 — a genuine job fault: report it
|
||
self._fail(jid, job.get("image_id"), exc)
|
||
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.
|
||
|
||
Failure guard (critical): an empty buffer can mean the GPU is starving OR
|
||
that downloads are FAILING (curator slow/unreachable). In the latter case
|
||
adding downloaders piles more concurrent large-file requests onto a
|
||
struggling curator — a congestion collapse that slows curator (and its
|
||
browser) further and never recovers. So if transient download failures
|
||
rose since the last decision, SHRINK toward the floor instead of growing,
|
||
and let the pool ramp back up only once downloads succeed again."""
|
||
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()
|
||
prev_fail = self.transient
|
||
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()
|
||
prev_fail = self.transient
|
||
tick = 0
|
||
con_grew = False
|
||
self._util_smooth = None
|
||
self._bw_capped = False
|
||
continue
|
||
|
||
occ = self._buffer.qsize() / BUFFER_MAX
|
||
occ_ewma = _ewma(occ_ewma, occ, OCC_ALPHA)
|
||
# Bandwidth-cap position: compare the observed aggregate (the same
|
||
# EWMA the UI shows) against the governor's cap. `soft` gates
|
||
# growth; `hard` sheds streams (see BW_* rationale above).
|
||
bw_rate = self.throttle.rate
|
||
net_bytes = self._net_mb_s * 1_048_576
|
||
bw_soft = bw_rate > 0 and net_bytes >= BW_ADD_HEADROOM * bw_rate
|
||
bw_hard = bw_rate > 0 and net_bytes >= BW_TRIM_AT * bw_rate
|
||
self._bw_capped = bw_soft
|
||
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 = _ewma(util_ewma, util, UTIL_ALPHA)
|
||
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 = _ewma(tput_ewma, inst, TPUT_ALPHA)
|
||
fail_delta = self.transient - prev_fail
|
||
prev_fail = self.transient
|
||
|
||
d0, c0 = self._dl_target, self._consumer_target
|
||
if fail_delta > 0:
|
||
# Downloads are FAILING (curator slow/unreachable), so the empty
|
||
# buffer is NOT the GPU starving — growing would stampede a
|
||
# struggling curator. Back off toward the floor and let it recover.
|
||
self._apply_downloaders(-1)
|
||
con_grew = False
|
||
elif occ_ewma < OCC_LOW:
|
||
# Buffer starving → downloads are the bottleneck. WHICH kind
|
||
# decides the move: a pipe pinned at the bandwidth cap gains
|
||
# nothing from more streams (they'd split the same budget and
|
||
# stretch per-job latency) — shed toward BW_MIN_DL; with cap
|
||
# headroom, concurrency is genuinely short — add a feeder.
|
||
if bw_hard and self._dl_target > BW_MIN_DL:
|
||
self._apply_downloaders(-1)
|
||
elif not bw_soft:
|
||
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%% · "
|
||
"net %.1f MB/s%s)",
|
||
d0, self._dl_target, c0, self._consumer_target,
|
||
round(occ_ewma * 100), round(util_ewma), tput_ewma,
|
||
round(vram * 100), self._net_mb_s,
|
||
" — at cap" if bw_soft else "")
|