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
This commit is contained in:
+46
-30
@@ -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),
|
||||
and watch progress + the server-side queue. The container exposes this on a
|
||||
localhost port; stopping the worker frees the GPU.
|
||||
Start / stop the worker pool, tune the worker count live (trades desktop
|
||||
responsiveness for throughput), and watch GPU load + progress + the server-side
|
||||
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}
|
||||
.stop{background:#b3261e} .stat{display:inline-block;margin-right:1.5rem}
|
||||
margin-right:.5rem;cursor:pointer;color:#fff} .start{background:#2e7d32}.stop{background:#b3261e}
|
||||
.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>
|
||||
<p>
|
||||
<span class=stat><span class=n id=state>idle</span><br>state</span>
|
||||
</div>
|
||||
<div class=row>
|
||||
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>
|
||||
</p>
|
||||
</div>
|
||||
<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
|
||||
cur.textContent=s.current??'—'; fc.textContent=s.fc_url
|
||||
state.textContent=s.state; active.textContent=s.active; done.textContent=s.processed
|
||||
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
|
||||
+76
-51
@@ -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
|
||||
stop frees the card promptly. Stop halts leasing + finishes the current item;
|
||||
unprocessed leases expire and the server re-queues them — nothing is lost.
|
||||
Each slot is an independent loop (its own leases; the server's SKIP-LOCKED lease
|
||||
keeps them from colliding). More slots = more GPU load + throughput; the model is
|
||||
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._state = "running"
|
||||
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"
|
||||
self._running = True
|
||||
self._reconcile_locked()
|
||||
|
||||
def stop(self):
|
||||
with self._lock:
|
||||
if self._state in ("running", "paused"):
|
||||
self._state = "stopping"
|
||||
self._running = False
|
||||
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 {
|
||||
"state": state, "processed": self.processed,
|
||||
"errors": self.errors, "current": self.current,
|
||||
}
|
||||
return {
|
||||
"state": "running" if self._running else "stopped",
|
||||
"concurrency": self._target,
|
||||
"workers": len(self._slots),
|
||||
"active": self._active,
|
||||
"processed": self.processed,
|
||||
"errors": self.errors,
|
||||
}
|
||||
|
||||
# --- loop --------------------------------------------------------------
|
||||
def _run(self):
|
||||
while True:
|
||||
with self._lock:
|
||||
st = self._state
|
||||
if st == "stopping":
|
||||
break
|
||||
if st == "paused":
|
||||
time.sleep(1)
|
||||
continue
|
||||
def _bump(self, *, processed=0, errors=0, active=0):
|
||||
with self._lock:
|
||||
self.processed += processed
|
||||
self.errors += errors
|
||||
self._active += active
|
||||
|
||||
# --- per-slot loop -----------------------------------------------------
|
||||
def _loop(self, slot: _Slot):
|
||||
while not slot.stop.is_set() and self._running:
|
||||
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 self._state == "stopping":
|
||||
break
|
||||
if slot.stop.is_set() or not self._running:
|
||||
break
|
||||
self._process(job)
|
||||
self.client.heartbeat(ids) # keep the rest of the batch alive
|
||||
with self._lock:
|
||||
self._state = "idle"
|
||||
slot.inflight = [i for i in slot.inflight if i != job["job_id"]]
|
||||
if slot.inflight:
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user