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),
+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)
@@ -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
+OR expired-leased jobs with FOR UPDATE SKIP LOCKED, so concurrent agents/workers
-retry after an agent died) never grab the same job and the queue self-heals
+never grab the same job. Orphan recovery is three-layered: a graceful agent stop
-without a separate recovery sweep. Result-writing (regions) is done by the API
+calls release() to hand its in-flight jobs back instantly; a hard crash is caught
-handler via RegionService; complete() just closes the job.
+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