"""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 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 # 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 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) # 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 # 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) 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 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) self._lock = threading.Lock() self._running = False 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() 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 # 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 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() # 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 _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 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 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) # --- 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(): self._ctrl_stop.clear() self._ctrl_thread = threading.Thread(target=self._control_loop, daemon=True) self._ctrl_thread.start() def stop(self): # Flip the flag FIRST (atomic bool), before any lock, so /status and the # 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: 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): # 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": "running" if self._running else "stopped", "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, "errors": self.errors, "transient": self.transient, } def _bump(self, *, processed=0, errors=0, active=0, transient=0): with self._lock: self.processed += processed self.errors += errors self.transient += transient # 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 self._running and not stop_evt.is_set(): try: _t = time.monotonic() jobs = self.client.lease(self.cfg.batch_size) self._record("lease", time.monotonic() - _t) backoff = self.cfg.poll_idle_seconds # server answered → reset except Exception: # curator unreachable (redeploy, network drop): wait it out with # exponential backoff, capped — resume on our own when it returns. 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 not self._running or stop_evt.is_set(): break 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 _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 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 # --- 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 if item is None: # stop sentinel continue 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: 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 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 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 _t = time.monotonic() vecs = [embedder.embed(frame) for _, frame in frames] if len(vecs) > 1: vec = np.mean( np.asarray(vecs, dtype=np.float32), axis=0 ).tolist() else: vec = vecs[0] self._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(jid, vec, embed_version) self._record("submit", time.monotonic() - _t) self._unhold(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}" _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. 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: 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 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 # (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) _t = time.monotonic() self.client.submit(jid, regions, replace_kinds) self._record("submit", time.monotonic() - _t) 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; then the server's orphan-recovery reclaims it if down). self._bump(transient=1) log.info("curator unreachable — released job %s (post-GPU)", jid) self.client.release([jid]) self._unhold(jid) return False 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) 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", 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))