Merge pull request 'Agent: GPU load + live worker tuning + fast orphan recovery' (#147) from dev into main

2026-06-29 19:11:22 -04:00
parent d52ac0a0e2 4a1a9ec5a7
commit 18300e1f8a
12 changed files with 328 additions and 88 deletions
@@ -7,6 +7,17 @@ Redis stay private; the agent never touches them.

 You run it when you want a burst and stop it to reclaim the card.

+## 0. Host prerequisite — NVIDIA Container Toolkit
+Docker needs the toolkit to hand the GPU to a container (else: *"could not select
+device driver nvidia with capabilities [[gpu]]"*). On Arch/CachyOS:
+```sh
+sudo pacman -S nvidia-container-toolkit
+sudo nvidia-ctk runtime configure --runtime=docker
+sudo systemctl restart docker
+# verify:
+docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi
+```
+
 ## 1. Get a token
 In FC: **Settings → Tagging → GPU agent → Generate token** (or Rotate). Copy it.

@@ -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
@@ -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)
@@ -196,3 +196,18 @@ async def fail():
        )
        await session.commit()
    return jsonify({"ok": ok})
+
+
+@gpu_bp.route("/jobs/release", methods=["POST"])
+async def release():
+    """Graceful stop: the agent hands its still-leased jobs back to pending so
+    they're picked up immediately instead of waiting out the lease."""
+    body = await request.get_json(silent=True) or {}
+    agent_id = str(body.get("agent_id") or "agent")
+    job_ids = [int(x) for x in (body.get("job_ids") or [])]
+    async with get_session() as session:
+        if not await _agent_authed(session):
+            return jsonify({"error": "unauthorized"}), 401
+        n = await GpuJobService(session).release(agent_id, job_ids)
+        await session.commit()
+    return jsonify({"released": n})
@@ -117,6 +117,10 @@ def make_celery() -> Celery:
                "task": "backend.app.tasks.ml.scheduled_apply_head_tags",
                "schedule": 86400.0,  # no-op unless head_auto_apply_enabled
            },
+            "recover-orphaned-gpu-jobs": {
+                "task": "backend.app.tasks.ml.recover_orphaned_gpu_jobs",
+                "schedule": 60.0,  # quick pickup of work a dead agent orphaned
+            },
            "snapshot-head-metrics-daily": {
                "task": "backend.app.tasks.maintenance.snapshot_head_metrics",
                "schedule": 86400.0,
@@ -1,10 +1,13 @@
-"""GPU-job queue engine (#114): enqueue / lease / heartbeat / complete / fail.
+"""GPU-job queue engine (#114): enqueue / lease / heartbeat / complete / fail
+/ release / recover_orphaned.

 Backs the HTTP API the desktop agent pulls work from. The lease claims pending
-OR expired-leased jobs with FOR UPDATE SKIP LOCKED, so concurrent agents (or a
-retry after an agent died) never grab the same job and the queue self-heals
-without a separate recovery sweep. Result-writing (regions) is done by the API
-handler via RegionService; complete() just closes the job.
+OR expired-leased jobs with FOR UPDATE SKIP LOCKED, so concurrent agents/workers
+never grab the same job. Orphan recovery is three-layered: a graceful agent stop
+calls release() to hand its in-flight jobs back instantly; a hard crash is caught
+by recover_orphaned() (a 60s beat sweep) which resets expired leases to pending;
+and the lease itself reclaims expired leases as a final backstop. Result-writing
+(regions) is done by the API handler via RegionService; complete() just closes.
 """

 from datetime import UTC, datetime, timedelta
@@ -14,7 +17,10 @@ from sqlalchemy.ext.asyncio import AsyncSession

 from ...models import GpuJob

-DEFAULT_LEASE_TTL = 300     # seconds an agent holds a job before it can be re-leased
+# Lease window. Kept comfortably above any single job (a capped-frame video embed
+# is tens of seconds) so a live, heartbeating worker is never falsely expired,
+# but short enough that a hard crash recovers fast once the sweep fires.
+DEFAULT_LEASE_TTL = 180     # seconds an agent holds a job before it can be re-leased
 DEFAULT_BATCH = 8
 MAX_ATTEMPTS = 3

@@ -132,3 +138,40 @@ class GpuJobService:
        job.error = (error or "")[:1000]
        job.updated_at = datetime.now(UTC)
        return True
+
+    async def release(self, token: str, job_ids: list[int]) -> int:
+        """Hand the agent's still-leased jobs back to pending NOW (graceful stop),
+        so another worker picks them up immediately instead of waiting out the
+        lease. Scoped to the token's own leases. Returns rows released."""
+        if not job_ids:
+            return 0
+        now = datetime.now(UTC)
+        res = await self.session.execute(
+            update(GpuJob)
+            .where(
+                GpuJob.id.in_(job_ids),
+                GpuJob.lease_token == token,
+                GpuJob.status == "leased",
+            )
+            .values(
+                status="pending", lease_token=None, leased_at=None,
+                lease_expires_at=None, updated_at=now,
+            )
+        )
+        return res.rowcount or 0
+
+    async def recover_orphaned(self) -> int:
+        """Reset every expired lease back to pending — catches agents that died
+        mid-job (no graceful release). Run on a short beat so the queue recovers
+        + reads honestly even when no worker is actively leasing. Returns rows
+        recovered."""
+        now = datetime.now(UTC)
+        res = await self.session.execute(
+            update(GpuJob)
+            .where(GpuJob.status == "leased", GpuJob.lease_expires_at < now)
+            .values(
+                status="pending", lease_token=None, leased_at=None,
+                lease_expires_at=None, updated_at=now,
+            )
+        )
+        return res.rowcount or 0
@@ -768,3 +768,30 @@ def enqueue_gpu_backfill(task_name: str) -> int:
        ).fetchall()
        session.commit()
        return len(rows)
+
+
+@celery.task(name="backend.app.tasks.ml.recover_orphaned_gpu_jobs")
+def recover_orphaned_gpu_jobs() -> int:
+    """Reset expired GPU-job leases back to pending — recovers work orphaned by an
+    agent that died mid-job (no graceful release). Short beat cadence so orphans
+    get picked back up quickly + the queue counts read honestly. Returns the
+    number recovered."""
+    from datetime import UTC, datetime
+
+    from sqlalchemy import update
+
+    from ..models import GpuJob
+
+    SessionLocal = _sync_session_factory()
+    with SessionLocal() as session:
+        now = datetime.now(UTC)
+        res = session.execute(
+            update(GpuJob)
+            .where(GpuJob.status == "leased", GpuJob.lease_expires_at < now)
+            .values(
+                status="pending", lease_token=None, leased_at=None,
+                lease_expires_at=None, updated_at=now,
+            )
+        )
+        session.commit()
+        return res.rowcount or 0
@@ -96,3 +96,22 @@ async def test_backfill_enqueues_then_is_idempotent(db):
    n = enqueue_gpu_backfill("ccip")     # sync task, own session
    assert n >= 2
    assert enqueue_gpu_backfill("ccip") == 0   # all already pending
+
+
+@pytest.mark.asyncio
+async def test_release_hands_job_back_to_pending(client, db):
+    img = await _img(db, "e" * 64)
+    await GpuJobService(db).enqueue(img.id, "ccip")
+    await db.commit()
+    token = (await (await client.post("/api/gpu/token/rotate")).get_json())["token"]
+    hdr = {"Authorization": f"Bearer {token}"}
+    j = (await (await client.post(
+        "/api/gpu/jobs/lease", json={"agent_id": "a1"}, headers=hdr,
+    )).get_json())["jobs"][0]
+
+    resp = await client.post("/api/gpu/jobs/release", json={
+        "agent_id": "a1", "job_ids": [j["job_id"]],
+    }, headers=hdr)
+    assert resp.status_code == 200 and (await resp.get_json())["released"] == 1
+    st = await (await client.get("/api/gpu/status")).get_json()
+    assert st["pending"] == 1 and st["leased"] == 0
@@ -123,3 +123,38 @@ async def test_fail_requeues_until_cap(db):
    assert await svc.fail("agent-1", job.id, "boom again") is True
    await db.commit()
    assert (await db.get(GpuJob, job.id)).status == "error"
+
+
+@pytest.mark.asyncio
+async def test_release_returns_to_pending(db):
+    img = await _img(db, "01" + "a" * 62)
+    svc = GpuJobService(db)
+    await svc.enqueue(img.id, "ccip")
+    await db.commit()
+    job = (await svc.lease("agent-1"))[0]
+    await db.commit()
+
+    assert await svc.release("other", [job.id]) == 0      # not this token's lease
+    assert await svc.release("agent-1", [job.id]) == 1     # graceful hand-back
+    await db.commit()
+    fresh = await db.get(GpuJob, job.id)
+    assert fresh.status == "pending" and fresh.lease_token is None
+
+
+@pytest.mark.asyncio
+async def test_recover_orphaned_resets_only_expired(db):
+    img1 = await _img(db, "02" + "a" * 62)
+    img2 = await _img(db, "03" + "a" * 62)
+    svc = GpuJobService(db)
+    await svc.enqueue(img1.id, "ccip")
+    await svc.enqueue(img2.id, "ccip")
+    await db.commit()
+    expired, fresh = await svc.lease("dead", batch_size=2)
+    # One lease is in the past (orphaned), the other still valid.
+    expired.lease_expires_at = datetime.now(UTC) - timedelta(minutes=10)
+    await db.commit()
+
+    assert await svc.recover_orphaned() == 1
+    await db.commit()
+    assert (await db.get(GpuJob, expired.id)).status == "pending"
+    assert (await db.get(GpuJob, fresh.id)).status == "leased"   # untouched