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)
|
s.mount("https://", adapter)
|
||||||
return s
|
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]:
|
def lease(self, batch_size: int) -> list[dict]:
|
||||||
r = self.s.post(
|
r = self.s.post(
|
||||||
f"{self.base}/api/gpu/jobs/lease",
|
f"{self.base}/api/gpu/jobs/lease",
|
||||||
@@ -54,62 +79,27 @@ class FcClient:
|
|||||||
return r.json().get("jobs", [])
|
return r.json().get("jobs", [])
|
||||||
|
|
||||||
def submit(self, job_id: int, regions: list[dict], replace_kinds: list[str]) -> dict:
|
def submit(self, job_id: int, regions: list[dict], replace_kinds: list[str]) -> dict:
|
||||||
r = self._submit_s.post(
|
return self._submit("/api/gpu/jobs/submit", {
|
||||||
f"{self.base}/api/gpu/jobs/submit",
|
"job_id": job_id, "regions": regions, "replace_kinds": replace_kinds,
|
||||||
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()
|
|
||||||
|
|
||||||
def submit_embedding(self, job_id: int, embedding: list, version: str) -> dict:
|
def submit_embedding(self, job_id: int, embedding: list, version: str) -> dict:
|
||||||
"""Post a whole-image SigLIP embedding (the 'embed' task) → image_record."""
|
"""Post a whole-image SigLIP embedding (the 'embed' task) → image_record."""
|
||||||
r = self._submit_s.post(
|
return self._submit("/api/gpu/jobs/submit_embedding", {
|
||||||
f"{self.base}/api/gpu/jobs/submit_embedding",
|
"job_id": job_id, "embedding": embedding, "embedding_version": version,
|
||||||
json={
|
})
|
||||||
"agent_id": self.agent_id, "job_id": job_id,
|
|
||||||
"embedding": embedding, "embedding_version": version,
|
|
||||||
},
|
|
||||||
timeout=120,
|
|
||||||
)
|
|
||||||
r.raise_for_status()
|
|
||||||
return r.json()
|
|
||||||
|
|
||||||
def heartbeat(self, job_ids: list[int]) -> None:
|
def heartbeat(self, job_ids: list[int]) -> None:
|
||||||
try:
|
self._post_quiet("/api/gpu/jobs/heartbeat", {"job_ids": job_ids})
|
||||||
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
|
|
||||||
|
|
||||||
def fail(self, job_id: int, error: str) -> None:
|
def fail(self, job_id: int, error: str) -> None:
|
||||||
try:
|
self._post_quiet("/api/gpu/jobs/fail", {"job_id": job_id, "error": error})
|
||||||
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
|
|
||||||
|
|
||||||
def release(self, job_ids: list[int]) -> None:
|
def release(self, job_ids: list[int]) -> None:
|
||||||
# Graceful hand-back on stop so orphaned work is re-leased at once.
|
# Graceful hand-back on stop so orphaned work is re-leased at once.
|
||||||
if not job_ids:
|
if not job_ids:
|
||||||
return
|
return
|
||||||
try:
|
self._post_quiet("/api/gpu/jobs/release", {"job_ids": job_ids})
|
||||||
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
|
|
||||||
|
|
||||||
def fetch_image(self, image_url: str) -> bytes:
|
def fetch_image(self, image_url: str) -> bytes:
|
||||||
# image_url is a server-relative path ("/images/...").
|
# image_url is a server-relative path ("/images/...").
|
||||||
|
|||||||
@@ -8,6 +8,11 @@ import os
|
|||||||
from dataclasses import dataclass
|
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
|
@dataclass
|
||||||
class Config:
|
class Config:
|
||||||
fc_url: str # base URL of the FabledCurator web service
|
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")),
|
poll_idle_seconds=float(os.environ.get("POLL_IDLE_SECONDS", "10")),
|
||||||
embed_dtype=os.environ.get("SIGLIP_DTYPE", "float16"),
|
embed_dtype=os.environ.get("SIGLIP_DTYPE", "float16"),
|
||||||
embed_model_override=os.environ.get("EMBED_MODEL_NAME", ""),
|
embed_model_override=os.environ.get("EMBED_MODEL_NAME", ""),
|
||||||
auto_start=os.environ.get("AUTO_START", "").lower() in ("1", "true", "yes"),
|
auto_start=_bool_env("AUTO_START"),
|
||||||
auto_scale=os.environ.get("AUTO_SCALE", "true").lower() in ("1", "true", "yes"),
|
auto_scale=_bool_env("AUTO_SCALE", "true"),
|
||||||
person_weights=os.environ.get("PERSON_WEIGHTS", "yolo11n.pt"),
|
person_weights=os.environ.get("PERSON_WEIGHTS", "yolo11n.pt"),
|
||||||
person_conf=float(os.environ.get("PERSON_CONF", "0.35")),
|
person_conf=float(os.environ.get("PERSON_CONF", "0.35")),
|
||||||
anatomy_weights=os.environ.get("ANATOMY_WEIGHTS", ""),
|
anatomy_weights=os.environ.get("ANATOMY_WEIGHTS", ""),
|
||||||
|
|||||||
@@ -202,14 +202,17 @@ class Proposers:
|
|||||||
boxes += self._person.detect(image)
|
boxes += self._person.detect(image)
|
||||||
return nms_merge(boxes)[: self.cfg.max_figures] # nms_merge is score-desc
|
return nms_merge(boxes)[: self.cfg.max_figures] # nms_merge is score-desc
|
||||||
|
|
||||||
def components(self, image):
|
@staticmethod
|
||||||
if self._anatomy is None:
|
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 []
|
return []
|
||||||
items = sorted(self._anatomy.detect(image), key=lambda b: b[1], reverse=True)
|
return sorted(detector.detect(image), key=lambda b: b[1], reverse=True)[:cap]
|
||||||
return items[: self.cfg.max_components]
|
|
||||||
|
def components(self, image):
|
||||||
|
return self._top(self._anatomy, image, self.cfg.max_components)
|
||||||
|
|
||||||
def panels(self, image):
|
def panels(self, image):
|
||||||
if self._panel is None:
|
return self._top(self._panel, image, self.cfg.max_panels)
|
||||||
return []
|
|
||||||
items = sorted(self._panel.detect(image), key=lambda b: b[1], reverse=True)
|
|
||||||
return items[: 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
|
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.
|
held lease immediately so orphaned work is re-picked without waiting out the TTL.
|
||||||
"""
|
"""
|
||||||
|
import contextlib
|
||||||
import logging
|
import logging
|
||||||
import queue
|
import queue
|
||||||
import threading
|
import threading
|
||||||
@@ -74,6 +75,13 @@ def _transient_reason(exc: requests.RequestException) -> str:
|
|||||||
return type(exc).__name__
|
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
|
# Pipeline sizing. Downloaders are I/O-bound, but every download streams a full
|
||||||
# original (large videos included) THROUGH curator's single Python file-serving
|
# original (large videos included) THROUGH curator's single Python file-serving
|
||||||
# path — so the ceiling is deliberately modest: too many concurrent large-file
|
# path — so the ceiling is deliberately modest: too many concurrent large-file
|
||||||
@@ -188,16 +196,38 @@ class Worker:
|
|||||||
with self._held_lock:
|
with self._held_lock:
|
||||||
self._held.discard(job_id)
|
self._held.discard(job_id)
|
||||||
|
|
||||||
def _release_owned(self, job_ids: list[int]) -> None:
|
def _release(self, job_ids: list[int]) -> None:
|
||||||
"""Hand a set of still-held leases back to curator and drop them from the
|
"""Hand still-held leases back to curator and drop them from the held set —
|
||||||
held set — used when a downloader exits (stop/shrink) still owning leases
|
the single hand-back path for both a downloader exiting (stop/shrink) with
|
||||||
it hadn't yet buffered."""
|
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:
|
if not job_ids:
|
||||||
return
|
return
|
||||||
self.client.release(job_ids)
|
self.client.release(job_ids)
|
||||||
for jid in job_ids:
|
for jid in job_ids:
|
||||||
self._unhold(jid)
|
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 ---------------------------------------------------
|
# --- background loops ---------------------------------------------------
|
||||||
def _heartbeat_loop(self) -> None:
|
def _heartbeat_loop(self) -> None:
|
||||||
"""Keep every held lease alive so buffered jobs waiting on the GPU aren't
|
"""Keep every held lease alive so buffered jobs waiting on the GPU aren't
|
||||||
@@ -243,6 +273,17 @@ class Worker:
|
|||||||
s[0] += seconds
|
s[0] += seconds
|
||||||
s[1] += 1
|
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:
|
def _stats_loop(self) -> None:
|
||||||
"""Log a per-stage timing breakdown every STATS_INTERVAL (only when there
|
"""Log a per-stage timing breakdown every STATS_INTERVAL (only when there
|
||||||
was work), so the operator can see the download/decode/gpu/submit split.
|
was work), so the operator can see the download/decode/gpu/submit split.
|
||||||
@@ -404,11 +445,10 @@ class Worker:
|
|||||||
on any exit path it releases whatever it still owns so nothing is stranded
|
on any exit path it releases whatever it still owns so nothing is stranded
|
||||||
holding a lease."""
|
holding a lease."""
|
||||||
backoff = self.cfg.poll_idle_seconds
|
backoff = self.cfg.poll_idle_seconds
|
||||||
while self._running and not stop_evt.is_set():
|
while not self._stopped(stop_evt):
|
||||||
try:
|
try:
|
||||||
_t = time.monotonic()
|
with self._timed("lease"):
|
||||||
jobs = self.client.lease(self.cfg.batch_size)
|
jobs = self.client.lease(self.cfg.batch_size)
|
||||||
self._record("lease", time.monotonic() - _t)
|
|
||||||
backoff = self.cfg.poll_idle_seconds # server answered → reset
|
backoff = self.cfg.poll_idle_seconds # server answered → reset
|
||||||
except Exception:
|
except Exception:
|
||||||
# curator unreachable (redeploy, network drop): wait it out with
|
# 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
|
owned = [j["job_id"] for j in jobs] # released on any early exit
|
||||||
for job in jobs:
|
for job in jobs:
|
||||||
jid = job["job_id"]
|
jid = job["job_id"]
|
||||||
if not self._running or stop_evt.is_set():
|
if self._stopped(stop_evt):
|
||||||
break
|
break
|
||||||
try:
|
try:
|
||||||
frames = self._download_decode(job)
|
frames = self._download_decode(job)
|
||||||
@@ -436,45 +476,36 @@ class Worker:
|
|||||||
# NOT the job's fault. Hand back this job + the rest of the
|
# NOT the job's fault. Hand back this job + the rest of the
|
||||||
# batch and back the whole loop off.
|
# batch and back the whole loop off.
|
||||||
self._bump(transient=1)
|
self._bump(transient=1)
|
||||||
self.client.release([jid])
|
self._release([jid])
|
||||||
self._unhold(jid)
|
|
||||||
# Name the ACTUAL failure — "curator unreachable" was
|
# Name the ACTUAL failure — "curator unreachable" was
|
||||||
# printed for every transient, hiding whether a single
|
# printed for every transient, hiding whether a single
|
||||||
# file's transfer stalled (ReadTimeout, curator fine) or
|
# file's transfer stalled (ReadTimeout, curator fine) or
|
||||||
# curator itself is down (ConnectTimeout/ConnectionError).
|
# curator itself is down (ConnectTimeout/ConnectionError).
|
||||||
log.info("fetch failed job %s (image %s, %s) — released, backing off",
|
log.info("fetch failed job %s (image %s, %s) — released, backing off",
|
||||||
jid, job.get("image_id"), _transient_reason(exc))
|
jid, job.get("image_id"), _transient_reason(exc))
|
||||||
self._release_owned(owned)
|
self._release(owned)
|
||||||
owned = []
|
owned = []
|
||||||
if not stop_evt.wait(backoff):
|
if not stop_evt.wait(backoff):
|
||||||
backoff = min(backoff * 2, MAX_BACKOFF_SECONDS)
|
backoff = min(backoff * 2, MAX_BACKOFF_SECONDS)
|
||||||
break
|
break
|
||||||
# a job-specific HTTP fault (404 image gone, 400) → fail it
|
# a job-specific HTTP fault (404 image gone, 400) → fail it
|
||||||
self._bump(errors=1)
|
self._fail(jid, job.get("image_id"), exc)
|
||||||
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
|
continue
|
||||||
except Exception as exc: # noqa: BLE001 — bad media → the job's fault
|
except Exception as exc: # noqa: BLE001 — bad media → the job's fault
|
||||||
owned.remove(jid)
|
owned.remove(jid)
|
||||||
self._bump(errors=1)
|
self._fail(jid, job.get("image_id"), exc, verb="failed to decode")
|
||||||
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
|
continue
|
||||||
# Blocks on a full buffer (backpressure) but wakes promptly on stop.
|
# Blocks on a full buffer (backpressure) but wakes promptly on stop.
|
||||||
if self._put((job, frames), stop_evt):
|
if self._put((job, frames), stop_evt):
|
||||||
owned.remove(jid) # ownership handed to the buffer/consumer
|
owned.remove(jid) # ownership handed to the buffer/consumer
|
||||||
else:
|
else:
|
||||||
break # stopped while waiting for buffer space
|
break # stopped while waiting for buffer space
|
||||||
self._release_owned(owned)
|
self._release(owned)
|
||||||
|
|
||||||
def _put(self, item, stop_evt: threading.Event) -> bool:
|
def _put(self, item, stop_evt: threading.Event) -> bool:
|
||||||
"""Push onto the bounded buffer, blocking while it's full but rechecking
|
"""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."""
|
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:
|
try:
|
||||||
self._buffer.put(item, timeout=0.5)
|
self._buffer.put(item, timeout=0.5)
|
||||||
return True
|
return True
|
||||||
@@ -490,14 +521,13 @@ class Worker:
|
|||||||
# only the frames it needs, so we NEVER pull the whole file (VR/4K
|
# 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
|
# originals are 800MB+ — buffering that in RAM and getting cut off
|
||||||
# mid-download was the failure loop). Environment-agnostic + resilient.
|
# mid-download was the failure loop). Environment-agnostic + resilient.
|
||||||
_t = time.monotonic()
|
|
||||||
url = f"{self.cfg.fc_url}{job['image_url']}"
|
url = f"{self.cfg.fc_url}{job['image_url']}"
|
||||||
frames = media.sample_frames_from_url(
|
with self._timed("decode"):
|
||||||
url, job.get("frame_interval_seconds", 4.0),
|
frames = media.sample_frames_from_url(
|
||||||
job.get("max_frames", 64),
|
url, job.get("frame_interval_seconds", 4.0),
|
||||||
headers=self._auth_header, timeout=self.cfg.ffmpeg_timeout,
|
job.get("max_frames", 64),
|
||||||
)
|
headers=self._auth_header, timeout=self.cfg.ffmpeg_timeout,
|
||||||
self._record("decode", time.monotonic() - _t)
|
)
|
||||||
if not frames:
|
if not frames:
|
||||||
# Couldn't sample. If curator is up, the file is unprocessable →
|
# Couldn't sample. If curator is up, the file is unprocessable →
|
||||||
# a job fault (fail it, don't re-lease forever). If curator is
|
# a job fault (fail it, don't re-lease forever). If curator is
|
||||||
@@ -516,18 +546,16 @@ class Worker:
|
|||||||
frames = kept
|
frames = kept
|
||||||
return frames
|
return frames
|
||||||
# Stills: download the bytes and decode.
|
# Stills: download the bytes and decode.
|
||||||
_t = time.monotonic()
|
with self._timed("download"):
|
||||||
data = self.client.fetch_image(job["image_url"])
|
data = self.client.fetch_image(job["image_url"])
|
||||||
self._record("download", time.monotonic() - _t)
|
with self._timed("decode"):
|
||||||
_t = time.monotonic()
|
frames = [(None, media.load_image(data))]
|
||||||
frames = [(None, media.load_image(data))]
|
|
||||||
self._record("decode", time.monotonic() - _t)
|
|
||||||
return frames
|
return frames
|
||||||
|
|
||||||
# --- GPU consumer pool -------------------------------------------------
|
# --- GPU consumer pool -------------------------------------------------
|
||||||
def _consumer(self, stop_evt: threading.Event):
|
def _consumer(self, stop_evt: threading.Event):
|
||||||
"""Pull decoded jobs off the buffer and run detect + embed + submit."""
|
"""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:
|
try:
|
||||||
item = self._buffer.get(timeout=1.0)
|
item = self._buffer.get(timeout=1.0)
|
||||||
except queue.Empty:
|
except queue.Empty:
|
||||||
@@ -535,9 +563,7 @@ class Worker:
|
|||||||
if item is None: # stop sentinel
|
if item is None: # stop sentinel
|
||||||
continue
|
continue
|
||||||
job, frames = item
|
job, frames = item
|
||||||
if not self._running or stop_evt.is_set():
|
if self._abort_if_stopped(job["job_id"], stop_evt):
|
||||||
self.client.release([job["job_id"]])
|
|
||||||
self._unhold(job["job_id"])
|
|
||||||
continue
|
continue
|
||||||
self._bump(active=1)
|
self._bump(active=1)
|
||||||
try:
|
try:
|
||||||
@@ -572,9 +598,7 @@ class Worker:
|
|||||||
it so a Stop drains promptly instead of finishing heavy GPU work."""
|
it so a Stop drains promptly instead of finishing heavy GPU work."""
|
||||||
jid = job["job_id"]
|
jid = job["job_id"]
|
||||||
try:
|
try:
|
||||||
if not self._running or stop_evt.is_set():
|
if self._abort_if_stopped(jid, stop_evt):
|
||||||
self.client.release([jid])
|
|
||||||
self._unhold(jid)
|
|
||||||
return False
|
return False
|
||||||
|
|
||||||
task = job.get("task") or "ccip"
|
task = job.get("task") or "ccip"
|
||||||
@@ -590,22 +614,16 @@ class Worker:
|
|||||||
# frames, matching the server's tag_and_embed. No regions.
|
# frames, matching the server's tag_and_embed. No regions.
|
||||||
if task == "embed":
|
if task == "embed":
|
||||||
embedder = self._ensure_embedder(model_name) # one-time model load
|
embedder = self._ensure_embedder(model_name) # one-time model load
|
||||||
_t = time.monotonic()
|
with self._timed("gpu"):
|
||||||
vecs = [embedder.embed(frame) for _, frame in frames]
|
vecs = [embedder.embed(frame) for _, frame in frames]
|
||||||
if len(vecs) > 1:
|
vec = (
|
||||||
vec = np.mean(
|
np.mean(np.asarray(vecs, dtype=np.float32), axis=0).tolist()
|
||||||
np.asarray(vecs, dtype=np.float32), axis=0
|
if len(vecs) > 1 else vecs[0]
|
||||||
).tolist()
|
)
|
||||||
else:
|
if self._abort_if_stopped(jid, stop_evt):
|
||||||
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
|
return False
|
||||||
_t = time.monotonic()
|
with self._timed("submit"):
|
||||||
self.client.submit_embedding(jid, vec, embed_version)
|
self.client.submit_embedding(jid, vec, embed_version)
|
||||||
self._record("submit", time.monotonic() - _t)
|
|
||||||
self._unhold(jid)
|
self._unhold(jid)
|
||||||
return True
|
return True
|
||||||
|
|
||||||
@@ -628,84 +646,82 @@ class Worker:
|
|||||||
ccip_ev = self.cfg.ccip_model or "ccip-default"
|
ccip_ev = self.cfg.ccip_model or "ccip-default"
|
||||||
dv = f"person-{self.cfg.detector_level}"
|
dv = f"person-{self.cfg.detector_level}"
|
||||||
|
|
||||||
_t_gpu = time.monotonic() # detect + CCIP + batched embed = "gpu"
|
# detect + CCIP + batched embed across every (deduped) frame = "gpu".
|
||||||
for t, frame in frames:
|
with self._timed("gpu"):
|
||||||
# Bail promptly on Stop instead of grinding through every frame of
|
for t, frame in frames:
|
||||||
# a long video before the caller can hand the job back.
|
# Bail promptly on Stop instead of grinding through every frame
|
||||||
if not self._running or stop_evt.is_set():
|
# of a long video before the caller can hand the job back.
|
||||||
break
|
if self._stopped(stop_evt):
|
||||||
# FIGURE boxes: imgutils detect_person ∪ general COCO person,
|
break
|
||||||
# NMS-merged → CCIP identity (+ a concept crop). Covers anime +
|
# FIGURE boxes: imgutils detect_person ∪ general COCO person,
|
||||||
# Western/realistic figures.
|
# NMS-merged → CCIP identity (+ a concept crop). Covers anime +
|
||||||
base = models.detect_figures(frame, self.cfg.detector_level)
|
# Western/realistic figures.
|
||||||
figs = proposers.figures(frame, base)
|
base = models.detect_figures(frame, self.cfg.detector_level)
|
||||||
if not figs:
|
figs = proposers.figures(frame, base)
|
||||||
figs = [((0.0, 0.0, 1.0, 1.0), 1.0, "whole")] # whole-frame fallback
|
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
|
# Collect every crop that needs a SigLIP embedding, then embed
|
||||||
# them in ONE batched forward pass (huge GPU-util + throughput
|
# them in ONE batched forward pass (huge GPU-util + throughput
|
||||||
# win vs one forward per crop). CCIP runs per figure inline.
|
# win vs one forward per crop). CCIP runs per figure inline.
|
||||||
pending = [] # (crop, region-template-without-embedding)
|
pending = [] # (crop, region-template-without-embedding)
|
||||||
for bbox, score, _label in figs:
|
for bbox, score, _label in figs:
|
||||||
crop = crop_region(frame, bbox)
|
crop = crop_region(frame, bbox)
|
||||||
if crop is None:
|
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
|
continue
|
||||||
if want_ccip:
|
# ANATOMY components (booru_yolo) + PANELS → concept/panel crops.
|
||||||
regions.append({
|
for bbox, score, label in proposers.components(frame):
|
||||||
"kind": "figure", "bbox": list(bbox), "frame_time": t,
|
crop = crop_region(frame, bbox)
|
||||||
"score": score,
|
if crop is not None:
|
||||||
"ccip_embedding": models.ccip_vector(
|
pending.append((crop, {
|
||||||
crop, self.cfg.ccip_model or None
|
"kind": "concept", "bbox": list(bbox), "frame_time": t,
|
||||||
),
|
"score": score, "detector_version": f"booru:{label}",
|
||||||
"embedding_version": ccip_ev, "detector_version": dv,
|
}))
|
||||||
})
|
for bbox, score, _label in proposers.panels(frame):
|
||||||
if want_siglip:
|
crop = crop_region(frame, bbox)
|
||||||
pending.append((crop, {
|
if crop is not None:
|
||||||
"kind": "concept", "bbox": list(bbox), "frame_time": t,
|
pending.append((crop, {
|
||||||
"score": score, "detector_version": dv,
|
"kind": "panel", "bbox": list(bbox), "frame_time": t,
|
||||||
}))
|
"score": score, "detector_version": "panel",
|
||||||
if not want_siglip:
|
}))
|
||||||
continue
|
if pending:
|
||||||
# ANATOMY components (booru_yolo) + PANELS → concept/panel crops.
|
# Drop near-duplicate crops (a figure box ≈ an anatomy
|
||||||
for bbox, score, label in proposers.components(frame):
|
# component, overlapping booru classes) before the embed so
|
||||||
crop = crop_region(frame, bbox)
|
# we never SigLIP the same region twice — saves GPU and a
|
||||||
if crop is not None:
|
# slot against max_regions. High-IoU + kind-aware, so
|
||||||
pending.append((crop, {
|
# intentional nested crops (figure ⊃ head) survive.
|
||||||
"kind": "concept", "bbox": list(bbox), "frame_time": t,
|
pending = dedupe_crops(pending, self.cfg.dedupe_iou)
|
||||||
"score": score, "detector_version": f"booru:{label}",
|
vecs = embedder.embed_batch([c for c, _ in pending])
|
||||||
}))
|
for (_c, tmpl), vec in zip(pending, vecs, strict=True):
|
||||||
for bbox, score, _label in proposers.panels(frame):
|
tmpl["siglip_embedding"] = vec
|
||||||
crop = crop_region(frame, bbox)
|
tmpl["embedding_version"] = embed_version
|
||||||
if crop is not None:
|
regions.append(tmpl)
|
||||||
pending.append((crop, {
|
# Stop once we have enough: a long video (image 81602 = a 156 MB
|
||||||
"kind": "panel", "bbox": list(bbox), "frame_time": t,
|
# mp4, 64 sampled frames × ~32 regions) would otherwise burn
|
||||||
"score": score, "detector_version": "panel",
|
# ~38s of GPU across every frame before the submit is even
|
||||||
}))
|
# truncated. Bounds the WORK, not just the POST body.
|
||||||
if pending:
|
if len(regions) >= self.cfg.max_regions:
|
||||||
# Drop near-duplicate crops (a figure box ≈ an anatomy
|
break
|
||||||
# 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)
|
|
||||||
|
|
||||||
# A Stop mid-frame-loop leaves partial regions — don't submit those;
|
# A Stop mid-frame-loop leaves partial regions — don't submit those;
|
||||||
# hand the whole job back so another agent redoes it cleanly.
|
# hand the whole job back so another agent redoes it cleanly.
|
||||||
if not self._running or stop_evt.is_set():
|
if self._abort_if_stopped(jid, stop_evt):
|
||||||
self.client.release([jid])
|
|
||||||
self._unhold(jid)
|
|
||||||
return False
|
return False
|
||||||
|
|
||||||
# Backstop: never submit an unbounded pile of regions (a pathological
|
# Backstop: never submit an unbounded pile of regions (a pathological
|
||||||
@@ -718,9 +734,8 @@ class Worker:
|
|||||||
regions = regions[: self.cfg.max_regions]
|
regions = regions[: self.cfg.max_regions]
|
||||||
log.info("job %s: capped regions %d→%d (dropped %d)",
|
log.info("job %s: capped regions %d→%d (dropped %d)",
|
||||||
jid, len(regions) + dropped, len(regions), dropped)
|
jid, len(regions) + dropped, len(regions), dropped)
|
||||||
_t = time.monotonic()
|
with self._timed("submit"):
|
||||||
self.client.submit(jid, regions, replace_kinds)
|
self.client.submit(jid, regions, replace_kinds)
|
||||||
self._record("submit", time.monotonic() - _t)
|
|
||||||
self._unhold(jid)
|
self._unhold(jid)
|
||||||
return True
|
return True
|
||||||
except requests.RequestException as exc:
|
except requests.RequestException as exc:
|
||||||
@@ -731,21 +746,12 @@ class Worker:
|
|||||||
self._bump(transient=1)
|
self._bump(transient=1)
|
||||||
log.info("submit failed job %s (%s) — released, re-lease later",
|
log.info("submit failed job %s (%s) — released, re-lease later",
|
||||||
jid, _transient_reason(exc))
|
jid, _transient_reason(exc))
|
||||||
self.client.release([jid])
|
self._release([jid])
|
||||||
self._unhold(jid)
|
|
||||||
return False
|
return False
|
||||||
self._bump(errors=1)
|
self._fail(jid, job.get("image_id"), exc)
|
||||||
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
|
return False
|
||||||
except Exception as exc: # noqa: BLE001 — a genuine job fault: report it
|
except Exception as exc: # noqa: BLE001 — a genuine job fault: report it
|
||||||
self._bump(errors=1)
|
self._fail(jid, job.get("image_id"), exc)
|
||||||
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
|
return False
|
||||||
|
|
||||||
# --- autoscaler --------------------------------------------------------
|
# --- autoscaler --------------------------------------------------------
|
||||||
@@ -786,16 +792,12 @@ class Worker:
|
|||||||
continue
|
continue
|
||||||
|
|
||||||
occ = self._buffer.qsize() / BUFFER_MAX
|
occ = self._buffer.qsize() / BUFFER_MAX
|
||||||
occ_ewma = occ if occ_ewma is None else (
|
occ_ewma = _ewma(occ_ewma, occ, OCC_ALPHA)
|
||||||
OCC_ALPHA * occ + (1 - OCC_ALPHA) * occ_ewma
|
|
||||||
)
|
|
||||||
g = gpumod.read_gpu() or {}
|
g = gpumod.read_gpu() or {}
|
||||||
mt = g.get("mem_total_mb") or 0
|
mt = g.get("mem_total_mb") or 0
|
||||||
vram = (g.get("mem_used_mb", 0) / mt) if mt else 0.0
|
vram = (g.get("mem_used_mb", 0) / mt) if mt else 0.0
|
||||||
util = g.get("util_pct", 0) or 0
|
util = g.get("util_pct", 0) or 0
|
||||||
util_ewma = util if util_ewma is None else (
|
util_ewma = _ewma(util_ewma, util, UTIL_ALPHA)
|
||||||
UTIL_ALPHA * util + (1 - UTIL_ALPHA) * util_ewma
|
|
||||||
)
|
|
||||||
self._util_smooth = util_ewma
|
self._util_smooth = util_ewma
|
||||||
|
|
||||||
# Memory pressure overrides the cadence — react immediately.
|
# Memory pressure overrides the cadence — react immediately.
|
||||||
@@ -815,9 +817,7 @@ class Worker:
|
|||||||
now = time.monotonic()
|
now = time.monotonic()
|
||||||
inst = (self.processed - prev_p) / max(1e-3, now - prev_t)
|
inst = (self.processed - prev_p) / max(1e-3, now - prev_t)
|
||||||
prev_p, prev_t = self.processed, now
|
prev_p, prev_t = self.processed, now
|
||||||
tput_ewma = inst if tput_ewma is None else (
|
tput_ewma = _ewma(tput_ewma, inst, TPUT_ALPHA)
|
||||||
TPUT_ALPHA * inst + (1 - TPUT_ALPHA) * tput_ewma
|
|
||||||
)
|
|
||||||
fail_delta = self.transient - prev_fail
|
fail_delta = self.transient - prev_fail
|
||||||
prev_fail = self.transient
|
prev_fail = self.transient
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user