refactor(agent): DRY pass on the GPU agent worker package
Consolidate genuine duplication in agent/fc_agent into single-source helpers (behavior-preserving; DRY Pass process #594): worker.py - _fail(jid, image_id, exc, verb) — 4 terminal "fail this job" blocks (downloader HTTP-fault + decode, consumer non-transient + generic). - _release(job_ids) (was _release_owned) — the one lease hand-back path; 6 inline release([jid])+unhold sites now route through it. - _stopped(stop_evt) + _abort_if_stopped(jid, stop_evt) — 4 stop-check -and-release blocks and every bare stop-check. - _timed(stage) contextmanager — ~8 monotonic()/_record() timing pairs; records only on clean exit, matching the old skip-on-raise behavior. - _ewma(prev, x, alpha) module fn — 3 EWMA updates in the autoscaler. client.py - _submit(path, payload) — submit / submit_embedding (retrying session). - _post_quiet(path, payload) — heartbeat / fail / release fire-and-forget. detectors.py - Proposers._top(detector, image, cap) — merges components() and panels(). config.py - _bool_env(name, default) — auto_start / auto_scale env parsing. Left alone (recorded): the xyxy→norm-xywh conversion duplicated across models.py/detectors.py (2 copies, independent wrapper modules — sharing would couple them), and the _ensure_embedder/_ensure_proposers pair (same lock shape, different concepts). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
+34
-44
@@ -44,6 +44,31 @@ class FcClient:
|
||||
s.mount("https://", adapter)
|
||||
return s
|
||||
|
||||
def _submit(self, path: str, payload: dict) -> dict:
|
||||
"""POST to a submit endpoint on the RETRYING session (by submit time the
|
||||
GPU work is done — a blip must not throw it away), raise on a hard error,
|
||||
and return the parsed JSON. `agent_id` is added to every body."""
|
||||
r = self._submit_s.post(
|
||||
f"{self.base}{path}",
|
||||
json={"agent_id": self.agent_id, **payload},
|
||||
timeout=120,
|
||||
)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
def _post_quiet(self, path: str, payload: dict) -> None:
|
||||
"""Fire-and-forget POST on the main session — heartbeat/fail/release are
|
||||
best-effort, so a transport error is swallowed (the worker's own retry and
|
||||
the server's orphan-recovery cover a lost call). `agent_id` is added."""
|
||||
try:
|
||||
self.s.post(
|
||||
f"{self.base}{path}",
|
||||
json={"agent_id": self.agent_id, **payload},
|
||||
timeout=30,
|
||||
)
|
||||
except requests.RequestException:
|
||||
pass
|
||||
|
||||
def lease(self, batch_size: int) -> list[dict]:
|
||||
r = self.s.post(
|
||||
f"{self.base}/api/gpu/jobs/lease",
|
||||
@@ -54,62 +79,27 @@ class FcClient:
|
||||
return r.json().get("jobs", [])
|
||||
|
||||
def submit(self, job_id: int, regions: list[dict], replace_kinds: list[str]) -> dict:
|
||||
r = self._submit_s.post(
|
||||
f"{self.base}/api/gpu/jobs/submit",
|
||||
json={
|
||||
"agent_id": self.agent_id, "job_id": job_id,
|
||||
"regions": regions, "replace_kinds": replace_kinds,
|
||||
},
|
||||
timeout=120,
|
||||
)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
return self._submit("/api/gpu/jobs/submit", {
|
||||
"job_id": job_id, "regions": regions, "replace_kinds": replace_kinds,
|
||||
})
|
||||
|
||||
def submit_embedding(self, job_id: int, embedding: list, version: str) -> dict:
|
||||
"""Post a whole-image SigLIP embedding (the 'embed' task) → image_record."""
|
||||
r = self._submit_s.post(
|
||||
f"{self.base}/api/gpu/jobs/submit_embedding",
|
||||
json={
|
||||
"agent_id": self.agent_id, "job_id": job_id,
|
||||
"embedding": embedding, "embedding_version": version,
|
||||
},
|
||||
timeout=120,
|
||||
)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
return self._submit("/api/gpu/jobs/submit_embedding", {
|
||||
"job_id": job_id, "embedding": embedding, "embedding_version": version,
|
||||
})
|
||||
|
||||
def heartbeat(self, job_ids: list[int]) -> None:
|
||||
try:
|
||||
self.s.post(
|
||||
f"{self.base}/api/gpu/jobs/heartbeat",
|
||||
json={"agent_id": self.agent_id, "job_ids": job_ids},
|
||||
timeout=30,
|
||||
)
|
||||
except requests.RequestException:
|
||||
pass
|
||||
self._post_quiet("/api/gpu/jobs/heartbeat", {"job_ids": job_ids})
|
||||
|
||||
def fail(self, job_id: int, error: str) -> None:
|
||||
try:
|
||||
self.s.post(
|
||||
f"{self.base}/api/gpu/jobs/fail",
|
||||
json={"agent_id": self.agent_id, "job_id": job_id, "error": error},
|
||||
timeout=30,
|
||||
)
|
||||
except requests.RequestException:
|
||||
pass
|
||||
self._post_quiet("/api/gpu/jobs/fail", {"job_id": job_id, "error": error})
|
||||
|
||||
def release(self, job_ids: list[int]) -> None:
|
||||
# Graceful hand-back on stop so orphaned work is re-leased at once.
|
||||
if not job_ids:
|
||||
return
|
||||
try:
|
||||
self.s.post(
|
||||
f"{self.base}/api/gpu/jobs/release",
|
||||
json={"agent_id": self.agent_id, "job_ids": job_ids},
|
||||
timeout=30,
|
||||
)
|
||||
except requests.RequestException:
|
||||
pass
|
||||
self._post_quiet("/api/gpu/jobs/release", {"job_ids": job_ids})
|
||||
|
||||
def fetch_image(self, image_url: str) -> bytes:
|
||||
# image_url is a server-relative path ("/images/...").
|
||||
|
||||
@@ -8,6 +8,11 @@ import os
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
def _bool_env(name: str, default: str = "") -> bool:
|
||||
"""A boolean env var — present + truthy ('1'/'true'/'yes') → True."""
|
||||
return os.environ.get(name, default).lower() in ("1", "true", "yes")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
fc_url: str # base URL of the FabledCurator web service
|
||||
@@ -57,8 +62,8 @@ class Config:
|
||||
poll_idle_seconds=float(os.environ.get("POLL_IDLE_SECONDS", "10")),
|
||||
embed_dtype=os.environ.get("SIGLIP_DTYPE", "float16"),
|
||||
embed_model_override=os.environ.get("EMBED_MODEL_NAME", ""),
|
||||
auto_start=os.environ.get("AUTO_START", "").lower() in ("1", "true", "yes"),
|
||||
auto_scale=os.environ.get("AUTO_SCALE", "true").lower() in ("1", "true", "yes"),
|
||||
auto_start=_bool_env("AUTO_START"),
|
||||
auto_scale=_bool_env("AUTO_SCALE", "true"),
|
||||
person_weights=os.environ.get("PERSON_WEIGHTS", "yolo11n.pt"),
|
||||
person_conf=float(os.environ.get("PERSON_CONF", "0.35")),
|
||||
anatomy_weights=os.environ.get("ANATOMY_WEIGHTS", ""),
|
||||
|
||||
@@ -202,14 +202,17 @@ class Proposers:
|
||||
boxes += self._person.detect(image)
|
||||
return nms_merge(boxes)[: self.cfg.max_figures] # nms_merge is score-desc
|
||||
|
||||
def components(self, image):
|
||||
if self._anatomy is None:
|
||||
@staticmethod
|
||||
def _top(detector, image, cap: int):
|
||||
"""Top-`cap` detections by score from an optional proposer (None → the
|
||||
proposer is off → []). Shared by the anatomy + panel proposers, which
|
||||
differ only in which detector and which cap."""
|
||||
if detector is None:
|
||||
return []
|
||||
items = sorted(self._anatomy.detect(image), key=lambda b: b[1], reverse=True)
|
||||
return items[: self.cfg.max_components]
|
||||
return sorted(detector.detect(image), key=lambda b: b[1], reverse=True)[:cap]
|
||||
|
||||
def components(self, image):
|
||||
return self._top(self._anatomy, image, self.cfg.max_components)
|
||||
|
||||
def panels(self, image):
|
||||
if self._panel is None:
|
||||
return []
|
||||
items = sorted(self._panel.detect(image), key=lambda b: b[1], reverse=True)
|
||||
return items[: self.cfg.max_panels]
|
||||
return self._top(self._panel, image, self.cfg.max_panels)
|
||||
|
||||
+156
-156
@@ -28,6 +28,7 @@ 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
|
||||
@@ -74,6 +75,13 @@ def _transient_reason(exc: requests.RequestException) -> str:
|
||||
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
|
||||
@@ -188,16 +196,38 @@ class Worker:
|
||||
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."""
|
||||
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)
|
||||
|
||||
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
|
||||
@@ -243,6 +273,17 @@ class Worker:
|
||||
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.
|
||||
@@ -404,11 +445,10 @@ class Worker:
|
||||
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():
|
||||
while not self._stopped(stop_evt):
|
||||
try:
|
||||
_t = time.monotonic()
|
||||
jobs = self.client.lease(self.cfg.batch_size)
|
||||
self._record("lease", time.monotonic() - _t)
|
||||
with self._timed("lease"):
|
||||
jobs = self.client.lease(self.cfg.batch_size)
|
||||
backoff = self.cfg.poll_idle_seconds # server answered → reset
|
||||
except Exception:
|
||||
# curator unreachable (redeploy, network drop): wait it out with
|
||||
@@ -425,7 +465,7 @@ class Worker:
|
||||
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():
|
||||
if self._stopped(stop_evt):
|
||||
break
|
||||
try:
|
||||
frames = self._download_decode(job)
|
||||
@@ -436,45 +476,36 @@ class Worker:
|
||||
# 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)
|
||||
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)
|
||||
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._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)
|
||||
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._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)
|
||||
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):
|
||||
owned.remove(jid) # ownership handed to the buffer/consumer
|
||||
else:
|
||||
break # stopped while waiting for buffer space
|
||||
self._release_owned(owned)
|
||||
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 self._running and not stop_evt.is_set():
|
||||
while not self._stopped(stop_evt):
|
||||
try:
|
||||
self._buffer.put(item, timeout=0.5)
|
||||
return True
|
||||
@@ -490,14 +521,13 @@ class Worker:
|
||||
# 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.
|
||||
_t = time.monotonic()
|
||||
url = f"{self.cfg.fc_url}{job['image_url']}"
|
||||
frames = 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,
|
||||
)
|
||||
self._record("decode", time.monotonic() - _t)
|
||||
with self._timed("decode"):
|
||||
frames = 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,
|
||||
)
|
||||
if not frames:
|
||||
# Couldn't sample. If curator is up, the file is unprocessable →
|
||||
# a job fault (fail it, don't re-lease forever). If curator is
|
||||
@@ -516,18 +546,16 @@ class Worker:
|
||||
frames = kept
|
||||
return frames
|
||||
# Stills: download the bytes and decode.
|
||||
_t = time.monotonic()
|
||||
data = self.client.fetch_image(job["image_url"])
|
||||
self._record("download", time.monotonic() - _t)
|
||||
_t = time.monotonic()
|
||||
frames = [(None, media.load_image(data))]
|
||||
self._record("decode", time.monotonic() - _t)
|
||||
with self._timed("download"):
|
||||
data = self.client.fetch_image(job["image_url"])
|
||||
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 self._running and not stop_evt.is_set():
|
||||
while not self._stopped(stop_evt):
|
||||
try:
|
||||
item = self._buffer.get(timeout=1.0)
|
||||
except queue.Empty:
|
||||
@@ -535,9 +563,7 @@ class Worker:
|
||||
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"])
|
||||
if self._abort_if_stopped(job["job_id"], stop_evt):
|
||||
continue
|
||||
self._bump(active=1)
|
||||
try:
|
||||
@@ -572,9 +598,7 @@ class Worker:
|
||||
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)
|
||||
if self._abort_if_stopped(jid, stop_evt):
|
||||
return False
|
||||
|
||||
task = job.get("task") or "ccip"
|
||||
@@ -590,22 +614,16 @@ class Worker:
|
||||
# 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)
|
||||
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
|
||||
_t = time.monotonic()
|
||||
self.client.submit_embedding(jid, vec, embed_version)
|
||||
self._record("submit", time.monotonic() - _t)
|
||||
with self._timed("submit"):
|
||||
self.client.submit_embedding(jid, vec, embed_version)
|
||||
self._unhold(jid)
|
||||
return True
|
||||
|
||||
@@ -628,84 +646,82 @@ class Worker:
|
||||
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
|
||||
# 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:
|
||||
# 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
|
||||
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
|
||||
self._record("gpu", time.monotonic() - _t_gpu)
|
||||
# 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 not self._running or stop_evt.is_set():
|
||||
self.client.release([jid])
|
||||
self._unhold(jid)
|
||||
if self._abort_if_stopped(jid, stop_evt):
|
||||
return False
|
||||
|
||||
# Backstop: never submit an unbounded pile of regions (a pathological
|
||||
@@ -718,9 +734,8 @@ class Worker:
|
||||
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)
|
||||
with self._timed("submit"):
|
||||
self.client.submit(jid, regions, replace_kinds)
|
||||
self._unhold(jid)
|
||||
return True
|
||||
except requests.RequestException as exc:
|
||||
@@ -731,21 +746,12 @@ class Worker:
|
||||
self._bump(transient=1)
|
||||
log.info("submit failed job %s (%s) — released, re-lease later",
|
||||
jid, _transient_reason(exc))
|
||||
self.client.release([jid])
|
||||
self._unhold(jid)
|
||||
self._release([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)
|
||||
self._fail(jid, job.get("image_id"), exc)
|
||||
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)
|
||||
self._fail(jid, job.get("image_id"), exc)
|
||||
return False
|
||||
|
||||
# --- autoscaler --------------------------------------------------------
|
||||
@@ -786,16 +792,12 @@ class Worker:
|
||||
continue
|
||||
|
||||
occ = self._buffer.qsize() / BUFFER_MAX
|
||||
occ_ewma = occ if occ_ewma is None else (
|
||||
OCC_ALPHA * occ + (1 - OCC_ALPHA) * occ_ewma
|
||||
)
|
||||
occ_ewma = _ewma(occ_ewma, occ, OCC_ALPHA)
|
||||
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
|
||||
)
|
||||
util_ewma = _ewma(util_ewma, util, UTIL_ALPHA)
|
||||
self._util_smooth = util_ewma
|
||||
|
||||
# Memory pressure overrides the cadence — react immediately.
|
||||
@@ -815,9 +817,7 @@ class Worker:
|
||||
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
|
||||
)
|
||||
tput_ewma = _ewma(tput_ewma, inst, TPUT_ALPHA)
|
||||
fail_delta = self.transient - prev_fail
|
||||
prev_fail = self.transient
|
||||
|
||||
|
||||
Reference in New Issue
Block a user