Merge pull request 'feat(agent): idle-unload GPU models to free VRAM when the queue is idle' (#233) from dev into main

2026-07-17 13:07:15 -04:00
parent ce0dac3524 57e52433d0
commit b5b437ca80
5 changed files with 119 additions and 1 deletions
@@ -21,7 +21,7 @@ log = logging.getLogger("fc_agent.app")
 # Bump on every agent change. The page embeds this and /status reports it; the UI
 # warns to reload when they differ — so a stale browser-cached page can't be
 # mistaken for "the new image didn't deploy". (Belt-and-braces with no-store.)
-VERSION = "2026-07-02.6 · sleep mode: an empty queue sheds to one downloader and backs the lease poll off to 15 min"
+VERSION = "2026-07-17.1 · idle model-unload: after ~5 min idle the GPU models release their VRAM and reload on the next job (env IDLE_UNLOAD_SECONDS, 0=off) · sleep mode sheds to one downloader"
 logbuf.install()
 cfg = Config.from_env()
@@ -334,9 +334,12 @@ _PAGE = """<!doctype html><html><head><meta charset=utf-8>
  waited.textContent=s.transient||0
  // Instantaneous pool state → demoted to the sub-line, where its jumpiness reads
  // as live churn rather than a "broken" headline metric.
  // '=== false' (not falsy) so a stale page that doesn't send models_loaded shows
  // nothing; when the idle monitor unloads, the VRAM meter drops alongside this.
  pipe.textContent='downloaders '+(s.downloaders!=null?s.downloaders:'—')+' · consumers '+(s.consumers!=null?s.consumers:'—')+' · on GPU '+(s.active||0)
   +' · net '+(s.net_mb_s!=null?s.net_mb_s.toFixed(1):'—')+' MB/s'
   +(s.bandwidth_limit_mb_s>0?(' / cap '+s.bandwidth_limit_mb_s):'')
   +(s.models_loaded===false?' · GPU models unloaded (idle — reload on next job)':'')
  if(document.activeElement!==bw && s.bandwidth_limit_mb_s!=null) bw.value=s.bandwidth_limit_mb_s
  // Buffer occupancy bar (also driven here so it tracks the /status cadence).
  if(s.buffer!=null && s.buffer_max){ const p=Math.round(100*s.buffer/s.buffer_max)
@@ -51,6 +51,12 @@ class Config:
    bandwidth_limit_mb_s: float  # aggregate download cap in MEGABYTES/s across
    #                              all downloaders + video streams (0 = unlimited);
    #                              tunable live from the agent UI
    idle_unload_seconds: float   # after this long with the GPU idle (nothing in
    #                              flight, queue empty or Stopped), unload the
    #                              SigLIP embedder + YOLO proposers to free their
    #                              VRAM; they reload lazily on the next job. A
    #                              24/7 agent otherwise squats on ~5GB doing
    #                              nothing. 0 disables (keep models warm forever).
    @classmethod
    def from_env(cls) -> Config:
@@ -87,4 +93,8 @@ class Config:
            # link to ~1-1.5 MB/s per stream, browser included). Raise it (or 0)
            # from the agent UI on wired/faster networks.
            bandwidth_limit_mb_s=float(os.environ.get("BANDWIDTH_LIMIT_MB_S", "8")),
            # 5 min: long enough that a lull between job bursts doesn't thrash the
            # (few-second) reload, short enough that an agent left running with an
            # empty queue hands its VRAM back promptly.
            idle_unload_seconds=float(os.environ.get("IDLE_UNLOAD_SECONDS", "300")),
        )
@@ -170,6 +170,13 @@ class YoloProposer:
            ))
        return out
    def unload(self) -> None:
        """Drop the loaded YOLO so its VRAM can be reclaimed; detect() reloads it
        lazily on the next job. Leaves _ok untouched — a healthy proposer comes
        back, but one that self-disabled on a fault stays off."""
        with self._lock:
            self._model = None
 class Proposers:
    """The agent's proposer set, built from config. Each detector is optional —
@@ -216,3 +223,11 @@ class Proposers:
    def panels(self, image):
        return self._top(self._panel, image, self.cfg.max_panels)
    def unload(self) -> None:
        """Release every loaded proposer's YOLO (idle VRAM reclaim). The worker
        also drops its reference to this Proposers and rebuilds a fresh one via
        _proposers_for on the next job, so this is belt-and-braces."""
        for p in (self._person, self._anatomy, self._panel):
            if p is not None:
                p.unload()
@@ -75,3 +75,18 @@ class CropEmbedder:
        pooled = out.pooler_output if hasattr(out, "pooler_output") else out
        arr = pooled.float().cpu().numpy().astype(np.float32)
        return [row.reshape(-1).tolist() for row in arr]
    def unload(self) -> bool:
        """Drop the loaded model so its VRAM can be reclaimed — the idle monitor
        calls this after a spell with no work so an idle agent doesn't squat on
        the card; the next embed() reloads it lazily (a few seconds). Held under
        BOTH the load and inference locks so it can never race a concurrent load
        or an in-flight forward pass. Returns True if a model was actually
        released (the caller then runs one empty_cache() to hand the freed blocks
        back to the driver)."""
        with self._load_lock, self._infer_lock:
            if self._model is None:
                return False
            self._model = None
            self._processor = None
            return True
@@ -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