feat(agent): GPU load readout + live worker-count tuning (#114)

Control UI gains what the operator asked for: - GPU load (nvidia-smi): util %, VRAM used/total + bar, temp — so you can see how hard the card is working while you're at the desktop. - Worker count is now a live − / + control (POST /concurrency), not just an env: the worker is a pool of independent slots (shared model, so slots add concurrent inference, not N× VRAM). Dial up for speed, down to free the card. Replaces pause/resume with Start/Stop + the worker dial. - Graceful release on stop / pool-shrink: a slot hands its still-leased jobs back via client.release() so they're re-picked immediately (pairs with the server recovery sweep). Not CI-tested (agent/ outside CI) — verified by running. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
2026-06-29 19:07:40 -04:00
parent 2cb0427868
commit 4a1a9ec5a7
5 changed files with 168 additions and 82 deletions
@@ -1,13 +1,14 @@
 """FastAPI control surface for the agent (served on localhost).
-Start / pause / resume / stop the worker, set nothing else here (config is env),
+Start / stop the worker pool, tune the worker count live (trades desktop
-and watch progress + the server-side queue. The container exposes this on a
+responsiveness for throughput), and watch GPU load + progress + the server-side
-localhost port; stopping the worker frees the GPU.
+queue. Config is env-seeded; the worker count is adjustable here on the fly.
 """
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from fastapi.responses import HTMLResponse, JSONResponse
 from .config import Config
 from .gpu import read_gpu
 from .worker import Worker
 cfg = Config.from_env()
@@ -26,29 +27,25 @@ def start():
    return JSONResponse(worker.status())
@app.post("/pause")
 def pause():
    worker.pause()
    return JSONResponse(worker.status())
@app.post("/resume")
 def resume():
    worker.resume()
    return JSONResponse(worker.status())
@app.post("/stop")
 def stop():
    worker.stop()
    return JSONResponse(worker.status())
@app.post("/concurrency")
 async def concurrency(request: Request):
    body = await request.json()
    worker.set_concurrency(int(body.get("value", 1)))
    return JSONResponse(worker.status())
@app.get("/status")
 def status():
    s = worker.status()
    s["fc_url"] = cfg.fc_url
    s["configured"] = bool(cfg.token)
    s["gpu"] = read_gpu()
    try:
        s["queue"] = worker.client.queue_status()
    except Exception:
@@ -59,35 +56,54 @@ def status():
 _PAGE = """<!doctype html><html><head><meta charset=utf-8>
 <title>FabledCurator GPU agent</title>
 <style>
- body{font:14px system-ui;margin:2rem;max-width:640px;background:#14171a;color:#e8e8e8}
+ body{font:14px system-ui;margin:2rem;max-width:680px;background:#14171a;color:#e8e8e8}
 h1{font-size:18px} button{font:14px system-ui;padding:.5rem 1rem;border:0;border-radius:6px;
- margin-right:.5rem;cursor:pointer;color:#fff} .start{background:#2e7d32}.pause{background:#b26a00}
+ margin-right:.5rem;cursor:pointer;color:#fff} .start{background:#2e7d32}.stop{background:#b3261e}
- .stop{background:#b3261e} .stat{display:inline-block;margin-right:1.5rem}
+ .step{background:#33373b;padding:.4rem .7rem;font-weight:700}
 .stat{display:inline-block;margin-right:1.5rem;vertical-align:top}
 .n{font-size:22px;font-weight:700} code{background:#222;padding:2px 6px;border-radius:4px}
- .q{margin-top:1rem;color:#9aa}
+ .q,.gpu{margin-top:1rem;color:#9aa} .bar{height:8px;border-radius:4px;background:#222;overflow:hidden;
 max-width:320px;margin-top:4px} .bar>i{display:block;height:100%;background:#3f7d3f}
 .row{margin:.8rem 0}
 </style></head><body>
 <h1>FabledCurator GPU agent</h1>
 <p>FC: <code id=fc>—</code> · token <code id=cfg>—</code></p>
-<p>
+<div class=row>
 <button class=start onclick=act('start')>Start</button>
 <button class=pause onclick=act('pause')>Pause</button>
 <button class=pause onclick=act('resume')>Resume</button>
 <button class=stop onclick=act('stop')>Stop</button>
-</p>
+</div>
-<p>
+<div class=row>
- <span class=stat><span class=n id=state>idle</span><br>state</span>
+ workers
 <button class=step onclick=setc(-1)>−</button>
 <b id=conc style=margin:0+.5rem>1</b>
 <button class=step onclick=setc(1)>+</button>
 <span class=cap style=color:#9aa>(more = faster + more GPU)</span>
 </div>
 <div class=row>
 <span class=stat><span class=n id=state>stopped</span><br>state</span>
 <span class=stat><span class=n id=active>0</span><br>active now</span>
 <span class=stat><span class=n id=done>0</span><br>processed</span>
 <span class=stat><span class=n id=err>0</span><br>errors</span>
- <span class=stat><span class=n id=cur>—</span><br>current image</span>
+</div>
-</p>
+<div class=gpu id=gpu>GPU — …</div>
 <div class=bar><i id=gpubar style=width:0%></i></div>
 <div class=q id=queue></div>
 <script>
 async function act(p){await fetch('/'+p,{method:'POST'});refresh()}
 async function setc(d){
  const v=Math.max(1,Math.min(8,parseInt(conc.textContent||'1')+d))
  await fetch('/concurrency',{method:'POST',headers:{'Content-Type':'application/json'},
   body:JSON.stringify({value:v})});refresh()
 }
 async function refresh(){
  const s=await (await fetch('/status')).json()
-  state.textContent=s.state; done.textContent=s.processed; err.textContent=s.errors
+  state.textContent=s.state; active.textContent=s.active; done.textContent=s.processed
-  cur.textContent=s.current??'—'; fc.textContent=s.fc_url
+  err.textContent=s.errors; conc.textContent=s.concurrency; fc.textContent=s.fc_url
  cfg.textContent=s.configured?'set':'MISSING'
  if(s.gpu){
   gpu.textContent=`GPU — ${s.gpu.util_pct}% util · VRAM ${s.gpu.mem_used_mb}/${s.gpu.mem_total_mb} MB · ${s.gpu.temp_c}°C`
   gpubar.style.width=Math.round(100*s.gpu.mem_used_mb/s.gpu.mem_total_mb)+'%'
  } else { gpu.textContent='GPU — n/a (CPU fallback?)'; gpubar.style.width='0%' }
  queue.textContent=s.queue?`queue — pending ${s.queue.pending} · in flight ${s.queue.leased} · done ${s.queue.done} · errored ${s.queue.error}`:'queue — unreachable'
 }
 refresh(); setInterval(refresh,3000)
@@ -54,6 +54,19 @@ class FcClient:
        except requests.RequestException:
            pass
    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
    def fetch_image(self, image_url: str) -> bytes:
        # image_url is a server-relative path ("/images/...").
        r = self.s.get(f"{self.base}{image_url}", timeout=180)
@@ -8,7 +8,8 @@ class Config:
    fc_url: str          # base URL of the FabledCurator web service
    token: str           # the bearer token from Settings → Tagging → GPU agent
    agent_id: str        # identifies this agent's leases
-    batch_size: int      # jobs leased per round (concurrency is still 1)
+    batch_size: int      # jobs a worker leases per round
    concurrency: int     # INITIAL parallel workers (tunable live from the UI)
    ccip_model: str      # imgutils CCIP model name ("" → imgutils default)
    detector_level: str  # imgutils person-detector level: n|s|m|x
    poll_idle_seconds: float   # wait between empty leases
@@ -20,6 +21,7 @@ class Config:
            token=os.environ.get("FC_TOKEN", ""),
            agent_id=os.environ.get("AGENT_ID", "desktop-agent"),
            batch_size=int(os.environ.get("BATCH_SIZE", "4")),
            concurrency=int(os.environ.get("CONCURRENCY", "1")),
            ccip_model=os.environ.get("CCIP_MODEL", ""),
            detector_level=os.environ.get("DETECTOR_LEVEL", "m"),
            poll_idle_seconds=float(os.environ.get("POLL_IDLE_SECONDS", "10")),
@@ -0,0 +1,30 @@
 """GPU load readout via nvidia-smi (present in the container thanks to the
 NVIDIA Container Toolkit's `utility` capability). Returns None if unavailable —
 the UI just shows n/a (e.g. CPU-fallback run)."""
 import subprocess
 def read_gpu() -> dict | None:
    try:
        out = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu",
                "--format=csv,noheader,nounits",
            ],
            capture_output=True, text=True, timeout=5, check=True,
        ).stdout.strip().splitlines()
    except (OSError, subprocess.SubprocessError):
        return None
    if not out:
        return None
    parts = [p.strip() for p in out[0].split(",")]
    try:
        return {
            "util_pct": int(float(parts[0])),
            "mem_used_mb": int(float(parts[1])),
            "mem_total_mb": int(float(parts[2])),
            "temp_c": int(float(parts[3])),
        }
    except (ValueError, IndexError):
        return None
@@ -1,8 +1,13 @@
-"""The lease → fetch → detect+embed → submit loop, with start/pause/stop control.
+"""The lease → fetch → detect+embed → submit loop, run by a pool of worker
 slots whose count is tunable live from the UI.
-Concurrency is 1 (one image at a time) so the GPU footprint stays small and a
+Each slot is an independent loop (its own leases; the server's SKIP-LOCKED lease
-stop frees the card promptly. Stop halts leasing + finishes the current item;
+keeps them from colliding). More slots = more GPU load + throughput; the model is
-unprocessed leases expire and the server re-queues them — nothing is lost.
+loaded once and shared, so slots add concurrent inference, not N× model VRAM.
 That's the dial the operator turns to trade desktop responsiveness for speed.
 Stop (or shrinking the pool) RELEASES a slot's still-leased jobs immediately so
 orphaned work is re-picked at once rather than waiting out the lease.
 """
 import threading
 import time
@@ -12,61 +17,78 @@ from .client import FcClient
 from .config import Config
 from .crops import crop_region
 MAX_CONCURRENCY = 8
 class _Slot:
    """One worker loop. `inflight` = jobs leased but not yet processed, so a
    graceful stop can hand them back."""
    __slots__ = ("stop", "inflight")
    def __init__(self):
        self.stop = threading.Event()
        self.inflight: list[int] = []
 class Worker:
    def __init__(self, cfg: Config):
        self.cfg = cfg
        self.client = FcClient(cfg.fc_url, cfg.token, cfg.agent_id)
        self._state = "idle"        # idle | running | paused | stopping
        self._lock = threading.Lock()
-        self._thread: threading.Thread | None = None
+        self._running = False
        self._target = max(1, min(MAX_CONCURRENCY, cfg.concurrency))
        self._slots: list[_Slot] = []
        self.processed = 0
        self.errors = 0
-        self.current = None
+        self._active = 0          # slots currently mid-image
    # --- control -----------------------------------------------------------
    def start(self):
        with self._lock:
-            if self._state in ("running", "paused"):
+            self._running = True
-                self._state = "running"
+            self._reconcile_locked()
                return
            self._state = "running"
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()
    def pause(self):
        with self._lock:
            if self._state == "running":
                self._state = "paused"
    def resume(self):
        with self._lock:
            if self._state == "paused":
                self._state = "running"
    def stop(self):
        with self._lock:
-            if self._state in ("running", "paused"):
+            self._running = False
-                self._state = "stopping"
+            slots, self._slots = self._slots, []
        for s in slots:
            s.stop.set()          # each slot releases its inflight on exit
    def set_concurrency(self, n: int):
        with self._lock:
            self._target = max(1, min(MAX_CONCURRENCY, int(n)))
            if self._running:
                self._reconcile_locked()
    def _reconcile_locked(self):
        while len(self._slots) < self._target:
            slot = _Slot()
            self._slots.append(slot)
            threading.Thread(target=self._loop, args=(slot,), daemon=True).start()
        while len(self._slots) > self._target:
            self._slots.pop().stop.set()
    def status(self) -> dict:
        with self._lock:
-            state = self._state
+            return {
-        return {
+                "state": "running" if self._running else "stopped",
-            "state": state, "processed": self.processed,
+                "concurrency": self._target,
-            "errors": self.errors, "current": self.current,
+                "workers": len(self._slots),
-        }
+                "active": self._active,
                "processed": self.processed,
                "errors": self.errors,
            }
-    # --- loop --------------------------------------------------------------
+    def _bump(self, *, processed=0, errors=0, active=0):
-    def _run(self):
+        with self._lock:
-        while True:
+            self.processed += processed
-            with self._lock:
+            self.errors += errors
-                st = self._state
+            self._active += active
-            if st == "stopping":
+
-                break
+    # --- per-slot loop -----------------------------------------------------
-            if st == "paused":
+    def _loop(self, slot: _Slot):
-                time.sleep(1)
+        while not slot.stop.is_set() and self._running:
                continue
            try:
                jobs = self.client.lease(self.cfg.batch_size)
            except Exception:
@@ -75,18 +97,21 @@ class Worker:
            if not jobs:
                time.sleep(self.cfg.poll_idle_seconds)
                continue
-            ids = [j["job_id"] for j in jobs]
+            slot.inflight = [j["job_id"] for j in jobs]
            for job in jobs:
-                with self._lock:
+                if slot.stop.is_set() or not self._running:
-                    if self._state == "stopping":
+                    break
                        break
                self._process(job)
-                self.client.heartbeat(ids)   # keep the rest of the batch alive
+                slot.inflight = [i for i in slot.inflight if i != job["job_id"]]
-        with self._lock:
+                if slot.inflight:
-            self._state = "idle"
+                    self.client.heartbeat(slot.inflight)
        # Graceful hand-back of anything leased but not processed.
        if slot.inflight:
            self.client.release(slot.inflight)
            slot.inflight = []
    def _process(self, job: dict):
-        self.current = job.get("image_id")
+        self._bump(active=1)
        try:
            data = self.client.fetch_image(job["image_url"])
            if media.is_video(job.get("mime", "")):
@@ -119,9 +144,9 @@ class Worker:
                        "detector_version": dv,
                    })
            self.client.submit(job["job_id"], regions, ["figure", "face"])
-            self.processed += 1
+            self._bump(processed=1)
        except Exception as exc:   # noqa: BLE001 — report + move on
-            self.errors += 1
+            self._bump(errors=1)
            self.client.fail(job["job_id"], str(exc)[:500])
        finally:
-            self.current = None
+            self._bump(active=-1)