feat(agent): idle-unload GPU models to free VRAM when the queue is idle
The SigLIP embedder + YOLO proposers load lazily then stay resident for the container's whole lifetime — a 24/7 agent with an empty queue squats on ~5GB of VRAM doing nothing (operator-observed: 4900MiB held at GPU-util 8% / P8). Sleep mode only sheds downloaders + poll cadence; even a UI Stop left the models loaded. Add a monitor thread that unloads the torch-owned models after cfg.idle_unload_seconds (env IDLE_UNLOAD_SECONDS, default 300; 0 disables) with the GPU genuinely idle (active==0, buffer drained, no job completed in the window), then torch.cuda.empty_cache() to hand the blocks back to the driver. They reload lazily on the next job via the existing _ensure_embedder / _proposers_for. Covers both sleep-mode idle and a full Stop. Surfaced in /status (models_loaded) and the agent UI pipe line; the VRAM meter drops too. Residual: imgutils CCIP/person ONNX sessions + the CUDA context stay resident (no clean unload API) — idle VRAM drops substantially, not to zero. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TbrA36zNczjVhrM6cWThQa
This commit is contained in:
@@ -57,6 +57,15 @@ MAX_BACKOFF_SECONDS = 60.0
|
||||
# up on their own.
|
||||
IDLE_POLL_MAX_SECONDS = 900.0
|
||||
|
||||
# Idle VRAM reclaim (operator 2026-07-17): the SigLIP embedder + YOLO proposers
|
||||
# load lazily and then stay warm for fast job bursts — but a 24/7 agent with an
|
||||
# empty queue would otherwise squat on that VRAM (~5GB on the operator's card)
|
||||
# indefinitely while doing nothing. So a monitor unloads them after
|
||||
# cfg.idle_unload_seconds with the GPU genuinely idle (nothing in flight, buffer
|
||||
# drained); they reload lazily on the next job. This is just how often the
|
||||
# monitor wakes to check — it bounds how soon past the threshold the unload fires.
|
||||
IDLE_UNLOAD_CHECK_INTERVAL = 30.0
|
||||
|
||||
# A job whose fetch dies transiently this many times IN ONE SESSION stops being
|
||||
# handed back and is failed instead. Transient handbacks (release) burn no
|
||||
# attempts on the server, so a poisoned transfer — an original that stalls the
|
||||
@@ -268,6 +277,11 @@ class Worker:
|
||||
self._proposers_sig = None # detector-config signature the current
|
||||
# proposers were built for (#134)
|
||||
self._proposers_lock = threading.Lock()
|
||||
# Monotonic time of the last GPU activity (a consumer finishing a job).
|
||||
# The idle monitor unloads the warm models once this goes stale by
|
||||
# cfg.idle_unload_seconds — see _idle_unload_loop.
|
||||
self._last_gpu_activity = time.monotonic()
|
||||
threading.Thread(target=self._idle_unload_loop, daemon=True).start()
|
||||
|
||||
# --- held-lease bookkeeping --------------------------------------------
|
||||
def _hold(self, job_ids) -> None:
|
||||
@@ -608,6 +622,9 @@ class Worker:
|
||||
"net_mb_s": round(self._net_mb_s, 1), # observed aggregate rate
|
||||
"bw_capped": self._bw_capped, # autoscaler holding at the cap (UI hint)
|
||||
"idle": self._idle, # queue empty → poll backed off (UI hint)
|
||||
# Whether the GPU models are currently resident (False after an idle
|
||||
# unload freed their VRAM) — a plain bool read, UI hint only.
|
||||
"models_loaded": self._embedder is not None or self._proposers is not None,
|
||||
}
|
||||
|
||||
def _bump(self, *, processed=0, downloaded=0, errors=0, active=0, transient=0):
|
||||
@@ -788,6 +805,9 @@ class Worker:
|
||||
self._bump(processed=1)
|
||||
finally:
|
||||
self._bump(active=-1)
|
||||
# Mark the GPU busy-until-now so the idle monitor starts its
|
||||
# unload countdown from when work actually stopped, not before.
|
||||
self._last_gpu_activity = time.monotonic()
|
||||
|
||||
def _ensure_embedder(self, model_name: str):
|
||||
if self._embedder is not None:
|
||||
@@ -845,6 +865,61 @@ class Worker:
|
||||
self._proposers_sig = sig
|
||||
return self._proposers
|
||||
|
||||
def _unload_models(self) -> bool:
|
||||
"""Release the GPU-resident models (SigLIP embedder + YOLO proposers) so an
|
||||
idle agent hands their VRAM back instead of squatting on the card. They
|
||||
reload lazily on the next job (_ensure_embedder / _proposers_for) — a
|
||||
few seconds' cost paid only when work actually resumes. Dropping the
|
||||
shared instances under their build locks means a concurrent job either
|
||||
sees the old instance (before) or rebuilds a fresh one (after); the idle
|
||||
monitor only calls this with nothing in flight, so no inference is using
|
||||
them. Returns True if anything was released."""
|
||||
released = False
|
||||
with self._embedder_lock:
|
||||
if self._embedder is not None:
|
||||
self._embedder.unload()
|
||||
self._embedder = None
|
||||
released = True
|
||||
with self._proposers_lock:
|
||||
if self._proposers is not None:
|
||||
self._proposers.unload()
|
||||
self._proposers = None
|
||||
self._proposers_sig = None
|
||||
released = True
|
||||
if released:
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
# torch's caching allocator holds freed blocks; hand them back
|
||||
# to the driver so nvidia-smi actually reflects the drop.
|
||||
torch.cuda.empty_cache()
|
||||
except Exception: # noqa: BLE001 — torch absent / CPU-only → nothing to free
|
||||
pass
|
||||
return released
|
||||
|
||||
def _idle_unload_loop(self) -> None:
|
||||
"""Unload the warm GPU models after a stretch of inactivity so a 24/7
|
||||
agent with an empty queue doesn't hold ~5GB of VRAM doing nothing. Fires
|
||||
only when nothing is in flight (active == 0 AND the buffer is drained) and
|
||||
no job has completed for cfg.idle_unload_seconds — a window long enough
|
||||
that a brief lull between bursts doesn't thrash reload/unload. Covers BOTH
|
||||
sleep mode (queue empty, pipeline still running) and a full Stop; the
|
||||
models reload lazily on the next job. idle_unload_seconds <= 0 disables it."""
|
||||
idle_after = self.cfg.idle_unload_seconds
|
||||
if idle_after <= 0:
|
||||
return
|
||||
while True:
|
||||
time.sleep(IDLE_UNLOAD_CHECK_INTERVAL)
|
||||
if self._embedder is None and self._proposers is None:
|
||||
continue # nothing loaded → nothing to free
|
||||
if self._active != 0 or not self._buffer.empty():
|
||||
continue # work in flight → keep them warm
|
||||
if time.monotonic() - self._last_gpu_activity < idle_after:
|
||||
continue # not idle long enough yet
|
||||
if self._unload_models():
|
||||
log.info("idle %.0fs — unloaded GPU models, freed VRAM "
|
||||
"(reload on next job)", idle_after)
|
||||
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user