From 0fbc3963b97450de1d9143d4f4c814116b6f972e Mon Sep 17 00:00:00 2001 From: drjones Date: Mon, 7 Sep 2026 17:22:34 -0700 Subject: [PATCH] Add a durable cross-tenant job queue with VRAM-aware scheduling The service only reacted: it noticed an application had started and scrambled to free memory. Nothing could be lined up. Each application has its own queue but they cannot see each other, so work submitted to one had no way to wait for the other. Jobs are stored in SQLite, so the queue is bounded by disk rather than memory and survives a restart. The scheduler takes the highest-priority pending job, arbitrates VRAM for it through the same plan_release, runs it, and moves on -- one at a time, because overlapping jobs would recreate the contention this service exists to resolve. Four bugs found by running it rather than reasoning about it: Dispatching without checking for room destroyed three queued LLM jobs in a row: a CUDA OOM kills llama-server outright, it does not fail gracefully. A job that cannot run yet now waits. The room check used the tenant's needs_vram_gb, which cannot be right for an LLM -- the requirement is a property of the model being loaded. A flat 4 GB passed with 8 GB free and then a 14.9 GB model was dispatched into it. The requirement is now computed per job. Waiting forever is also wrong. Three jobs sat pending indefinitely needing 14.93 GB on a card where at most ~14.8 GB can ever be free, because an unreclaimable process holds 0.82 GB. A job that cannot be satisfied now fails with the ceiling and the blockers named. plan_release assumed releasing a tenant frees everything it holds. ComfyUI keeps its CUDA context for as long as the process lives, so it reported that releasing ComfyUI would free 0.37 GB against a 0.33 GB shortfall; the job was cleared and the memory never arrived. Tenants declare vram_floor_gb and only memory above it counts. An exception during dispatch left the row RUNNING forever while the scheduler moved on. Failures now land on the job, and jobs left running by a previous process are requeued at startup. Verified end to end: five mixed jobs across both applications, queued at once, all completed with no failures. Tests: 250. Co-Authored-By: Claude Opus 5 --- README.md | 30 ++- jobs.py | 445 ++++++++++++++++++++++++++++++++++++++++++ server.py | 52 +++++ tenants.py | 14 +- tests/test_tenants.py | 55 ++++++ vram_arbitrator.py | 59 ++++-- 6 files changed, 642 insertions(+), 13 deletions(-) create mode 100644 jobs.py diff --git a/README.md b/README.md index e533d16..d9376e8 100644 --- a/README.md +++ b/README.md @@ -174,6 +174,34 @@ only if the card actually needs it. --- +## 1c. Lining Work Up + +Until now this service only *reacted*: it noticed an application had started and +scrambled to free memory. Nothing could be queued. Each application has its own queue, +but they cannot see each other, so work submitted to one has no way to wait for the other. + +```bash +curl -X POST localhost:9090/api/jobs -H 'Content-Type: application/json' -d '{ + "tenant": "ollama", "label": "nightly-summary", + "payload": {"model": "qwen3.8fast:latest", "prompt": "..."} +}' +``` + +Jobs live in SQLite, so the queue is bounded by disk rather than memory and survives a +restart. The scheduler takes the highest-priority pending job, arbitrates VRAM for it with +the same `plan_release`, runs it, and moves on. One at a time by design — the GPU is the +scarce resource this service exists to hand between applications, and overlapping jobs +would just recreate the contention it resolves. + +`GET /api/jobs` · `GET /api/jobs/{id}` · `DELETE /api/jobs/{id}` (pending only — running +work is never killed) · `DELETE /api/jobs` to clear the queue. + +**A job that cannot run yet waits; a job that can never run fails with the reason.** +Dispatching into insufficient VRAM does not fail gracefully — it kills `llama-server` +with a CUDA OOM. The requirement is computed per job (an LLM job needs the size of *its* +model, not a tenant-wide figure), and if the memory can never be assembled the job fails +naming what stands in the way rather than blocking the queue forever. + ## 1b. Any Application, Not Just These Two The purpose is fast handoff of one GPU between applications. It grew up around the two on @@ -238,7 +266,7 @@ explaining that, rather than silently doing nothing. ## 1a. Tests ```bash -/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/ -q # 244 passed in ~3.8s +/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/ -q # 250 passed in ~3.8s ``` Hermetic: no GPU, no network, no sleeps. An autouse fixture stubs `overclock_manager._sh` diff --git a/jobs.py b/jobs.py new file mode 100644 index 0000000..753311b --- /dev/null +++ b/jobs.py @@ -0,0 +1,445 @@ +"""A durable, unbounded job queue across every GPU tenant. + +Until now this service only reacted: it noticed an application had started working and +scrambled to free memory. Nothing could be *lined up*. Each application has its own queue +(ComfyUI's prompt queue, Ollama's serialised requests), but they cannot see each other, so +work submitted to one has no way to wait politely for the other. + +Jobs submitted here are stored in SQLite, so the queue is limited by disk rather than +memory and survives a restart. The scheduler takes the highest-priority pending job, +makes sure its tenant actually has the VRAM to run it -- reusing the same plan_release +arbitration -- dispatches it, and moves on. +""" +import asyncio +import json +import logging +import sqlite3 +import time +import uuid +from typing import Any, Dict, List, Optional + +import httpx + +import telemetry_store +import tenants as tenants_mod + +logger = logging.getLogger("jobs") + +PENDING, RUNNING, DONE, FAILED, CANCELLED = ( + "pending", "running", "done", "failed", "cancelled") + +SCHEMA = """ +CREATE TABLE IF NOT EXISTS jobs ( + id TEXT PRIMARY KEY, + tenant TEXT NOT NULL, + priority INTEGER NOT NULL DEFAULT 50, + payload TEXT NOT NULL, + state TEXT NOT NULL, + submitted_at REAL NOT NULL, + started_at REAL, + finished_at REAL, + error TEXT, + result TEXT, + label TEXT +); +CREATE INDEX IF NOT EXISTS idx_jobs_state ON jobs(state, priority DESC, submitted_at); +""" + + +def _conn() -> sqlite3.Connection: + c = sqlite3.connect(telemetry_store.DB_PATH, timeout=10.0) + c.row_factory = sqlite3.Row + return c + + +def init() -> None: + with _conn() as c: + c.executescript(SCHEMA) + + +def submit(tenant: str, payload: Dict[str, Any], priority: Optional[int] = None, + label: Optional[str] = None) -> Dict[str, Any]: + """Queue a job. There is no depth limit: the queue lives on disk.""" + t = tenants_mod.get_tenant(tenant) + if not t: + return {"success": False, "error": f"no tenant named '{tenant}'"} + job_id = uuid.uuid4().hex[:12] + row = { + "id": job_id, "tenant": tenant, + "priority": t.priority if priority is None else int(priority), + "payload": json.dumps(payload), "state": PENDING, + "submitted_at": time.time(), "label": label, + } + with _conn() as c: + c.execute("INSERT INTO jobs (id, tenant, priority, payload, state, submitted_at," + " label) VALUES (:id,:tenant,:priority,:payload,:state,:submitted_at," + ":label)", row) + logger.info(f"queued job {job_id} for '{tenant}' at priority {row['priority']}") + return {"success": True, "id": job_id, "tenant": tenant, + "priority": row["priority"], "state": PENDING} + + +def cancel(job_id: str) -> Dict[str, Any]: + with _conn() as c: + cur = c.execute("UPDATE jobs SET state=?, finished_at=? WHERE id=? AND state=?", + (CANCELLED, time.time(), job_id, PENDING)) + if cur.rowcount: + return {"success": True, "id": job_id, "state": CANCELLED} + return {"success": False, "error": "job is not pending (already running or finished)"} + + +def clear_pending() -> Dict[str, Any]: + with _conn() as c: + cur = c.execute("UPDATE jobs SET state=?, finished_at=? WHERE state=?", + (CANCELLED, time.time(), PENDING)) + return {"success": True, "cancelled": cur.rowcount} + + +def get(job_id: str) -> Optional[Dict[str, Any]]: + with _conn() as c: + r = c.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone() + return _row(r) if r else None + + +def _row(r: sqlite3.Row) -> Dict[str, Any]: + d = dict(r) + for key in ("payload", "result"): + if d.get(key): + try: + d[key] = json.loads(d[key]) + except Exception: + pass + if d.get("started_at") and d.get("finished_at"): + d["duration_s"] = round(d["finished_at"] - d["started_at"], 2) + if d.get("state") == PENDING: + d["waiting_s"] = round(time.time() - d["submitted_at"], 1) + return d + + +def listing(state: Optional[str] = None, limit: int = 100) -> List[Dict[str, Any]]: + q = "SELECT * FROM jobs" + args: List[Any] = [] + if state: + q += " WHERE state=?" + args.append(state) + # Pending jobs in the order the scheduler will take them; everything else newest first. + q += (" ORDER BY priority DESC, submitted_at ASC" if state == PENDING + else " ORDER BY submitted_at DESC") + q += " LIMIT ?" + args.append(limit) + with _conn() as c: + return [_row(r) for r in c.execute(q, args).fetchall()] + + +def stats() -> Dict[str, Any]: + with _conn() as c: + rows = c.execute("SELECT state, COUNT(*) n FROM jobs GROUP BY state").fetchall() + by_state = {r["state"]: r["n"] for r in rows} + pend = c.execute( + "SELECT tenant, COUNT(*) n FROM jobs WHERE state=? GROUP BY tenant", + (PENDING,)).fetchall() + oldest = c.execute( + "SELECT MIN(submitted_at) t FROM jobs WHERE state=?", (PENDING,)).fetchone() + return { + "by_state": by_state, + "pending_by_tenant": {r["tenant"]: r["n"] for r in pend}, + "queue_depth": by_state.get(PENDING, 0), + "oldest_pending_s": (round(time.time() - oldest["t"], 1) + if oldest and oldest["t"] else None), + } + + +def requeue_orphans() -> int: + """Return jobs abandoned mid-run to the queue. + + RUNNING means "this process is working on it". If no process is, that is untrue, and + the job would otherwise never finish and never retry. + """ + with _conn() as c: + cur = c.execute("UPDATE jobs SET state=?, started_at=NULL WHERE state=?", + (PENDING, RUNNING)) + return cur.rowcount + + +def _next_job() -> Optional[Dict[str, Any]]: + with _conn() as c: + r = c.execute( + "SELECT * FROM jobs WHERE state=? ORDER BY priority DESC, submitted_at ASC" + " LIMIT 1", (PENDING,)).fetchone() + return _row(r) if r else None + + +def _mark(job_id: str, state: str, **fields) -> None: + sets = ", ".join(f"{k}=?" for k in fields) + args = list(fields.values()) + [state, job_id] + with _conn() as c: + c.execute(f"UPDATE jobs SET {sets + ', ' if sets else ''}state=? WHERE id=?", args) + + +# ---------------------------------------------------------------- dispatch + +async def _dispatch_comfy(payload: Dict[str, Any]) -> Dict[str, Any]: + """Hand a workflow to ComfyUI and wait for it to finish.""" + base = tenants_mod.get_tenant("comfyui").busy.url.rsplit("/", 1)[0] + async with httpx.AsyncClient(timeout=30.0) as c: + r = await c.post(f"{base}/prompt", json={"prompt": payload.get("prompt", payload), + "client_id": "hyperswap-jobs"}) + if r.status_code != 200: + return {"ok": False, "error": f"HTTP {r.status_code}: {r.text[:200]}"} + prompt_id = r.json().get("prompt_id") + deadline = time.time() + payload.get("timeout_s", 1800) + while time.time() < deadline: + await asyncio.sleep(0.5) + h = await c.get(f"{base}/history/{prompt_id}") + entry = (h.json() or {}).get(prompt_id) if h.status_code == 200 else None + if not entry: + continue + status = entry.get("status", {}) + if status.get("status_str") == "error": + return {"ok": False, "error": "ComfyUI reported an execution error"} + if status.get("completed"): + return {"ok": True, "prompt_id": prompt_id} + return {"ok": False, "error": "timed out waiting for ComfyUI"} + + +async def _dispatch_ollama(payload: Dict[str, Any]) -> Dict[str, Any]: + async with httpx.AsyncClient(timeout=payload.get("timeout_s", 1800)) as c: + body = {"stream": False, "keep_alive": payload.get("keep_alive", "5m"), **payload} + body.pop("timeout_s", None) + r = await c.post("http://localhost:11434/api/generate", json=body) + if r.status_code != 200: + return {"ok": False, "error": f"HTTP {r.status_code}: {r.text[:200]}"} + d = r.json() + return {"ok": True, "response": (d.get("response") or "")[:2000], + "eval_count": d.get("eval_count"), + "tokens_per_sec": (round(d.get("eval_count", 0) + / (d.get("eval_duration", 1) / 1e9), 2) + if d.get("eval_duration") else None)} + + +DISPATCHERS = { + tenants_mod.KIND_DIFFUSION: _dispatch_comfy, + tenants_mod.KIND_LLM: _dispatch_ollama, +} + + +class Scheduler: + """Drains the queue, making room for each job before it runs. + + One job at a time by design. The GPU is the scarce resource this whole service + exists to hand between applications; running two jobs concurrently would just + recreate the contention it is meant to resolve. Throughput comes from swapping + quickly, not from overlapping. + """ + + def __init__(self) -> None: + self.running = False + self.task: Optional[asyncio.Task] = None + self.current: Optional[Dict[str, Any]] = None + self.last_finished: Optional[Dict[str, Any]] = None + self.completed = 0 + self.failed = 0 + self.waits = 0 + self.blocked: Optional[Dict[str, Any]] = None + self.idle_poll_s = 1.0 + self.blocked_poll_s = 2.0 + # How long a job may wait for VRAM before it is declared impossible. Long enough + # to outlast a normal diffusion run, short enough not to wedge the queue. + self.max_block_s = 120.0 + + async def start(self) -> None: + if self.running: + return + init() + # A job left RUNNING by a crash or a hard restart would sit there forever. + requeued = requeue_orphans() + if requeued: + logger.warning(f"requeued {requeued} job(s) left running by a previous process") + self.running = True + self.task = asyncio.create_task(self._loop()) + logger.info("job scheduler started") + + async def stop(self) -> None: + self.running = False + if self.task: + self.task.cancel() + + def _job_vram_requirement(self, tenant_name: str, payload: Dict[str, Any]) -> float: + """How much VRAM *this* job needs, not the tenant's generic figure. + + A tenant-wide needs_vram_gb cannot be right for an LLM: the requirement is a + property of the model being loaded. Ollama's generic 4 GB passed the room check + with 8 GB free, and then a 14.9 GB model was dispatched into it and killed + llama-server with a CUDA OOM -- three queued jobs destroyed in a row. + """ + import vram_arbitrator + + t = tenants_mod.get_tenant(tenant_name) + default = t.needs_vram_gb if t else 0.0 + model = payload.get("model") + if t and t.kind == tenants_mod.KIND_LLM and model: + size = vram_arbitrator._model_size_bytes(model) + if size: + # Measured on this box: a 12.87 GB blob occupies 14.9 GB once context + # and KV cache are allocated. + return round((size / (1024 ** 3)) * 1.16, 2) + return default + + async def _make_room(self, tenant_name: str, + payload: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + """Ensure the job's tenant has the VRAM it needs, using the normal arbitration.""" + import vram_arbitrator # imported late: it imports this module's siblings + + t = tenants_mod.get_tenant(tenant_name) + needed = self._job_vram_requirement(tenant_name, payload or {}) + if not t or not needed: + return {"ready": True, "reason": "no VRAM requirement declared"} + + state = await vram_arbitrator.arbitrator._tenant_state() + free_gb = vram_arbitrator.arbitrator._last_tenant_state["free_gb"] + held = next((s["vram_gb"] for s in state if s["name"] == tenant_name), 0.0) + if held + free_gb >= needed: + return {"ready": True, "needed_gb": needed, + "reason": f"{free_gb:.2f} GB free, job needs {needed:.2f} GB"} + + plan = tenants_mod.plan_release(tenant_name, state, free_gb, needed) + for victim in plan["release"]: + await vram_arbitrator.arbitrator._release_tenant( + victim, f"queued job for '{tenant_name}'") + if plan["release"]: + # Give the driver a moment to actually hand the memory back. + deadline = time.perf_counter() + 30 + target = int(needed * (1024 ** 3)) + while time.perf_counter() < deadline: + if vram_arbitrator.get_process_vram_bytes()["free_bytes"] >= target: + break + await asyncio.sleep(0.05) + # Only ready once the memory is genuinely there. A plan that *could* work is not + # the same as VRAM that *is* free, and dispatching on the former is what OOMs. + free_now = (vram_arbitrator.get_process_vram_bytes()["free_bytes"] / (1024 ** 3)) + # The best this GPU could ever offer this tenant: everything currently free, plus + # what it already holds, plus everything that is reclaimable at all. + reclaimable_gb = sum(s["vram_gb"] for s in state + if s["name"] != tenant_name and s.get("reclaimable")) + max_possible = round(held + free_now + reclaimable_gb, 2) + return {"ready": plan["possible"] and (held + free_now) >= needed, + "max_possible_gb": max_possible, + "released": plan["release"], "needed_gb": needed, + "free_gb": round(free_now, 2), + "reason": f"job needs {needed:.2f} GB; {plan['reason']}", + "blockers": plan.get("blockers")} + + async def _loop(self) -> None: + while self.running: + try: + job = _next_job() + if not job: + await asyncio.sleep(self.idle_poll_s) + continue + + tenant = tenants_mod.get_tenant(job["tenant"]) + dispatcher = DISPATCHERS.get(tenant.kind) if tenant else None + if not dispatcher: + _mark(job["id"], FAILED, finished_at=time.time(), + error=f"no dispatcher for tenant kind " + f"'{tenant.kind if tenant else '?'}'") + self.failed += 1 + continue + + # Check for room *before* claiming the job. Dispatching into + # insufficient VRAM does not fail gracefully -- it kills llama-server + # with a CUDA OOM, which is how three queued LLM jobs were destroyed + # while ComfyUI legitimately held the card. A job that cannot run yet + # waits; it does not fail. + room = await self._make_room(job["tenant"], job.get("payload") or {}) + if not room.get("ready"): + since = (self.blocked.get("since", time.time()) + if self.blocked and self.blocked.get("id") == job["id"] + else time.time()) + waited = time.time() - since + self.blocked = {"id": job["id"], "tenant": job["tenant"], + "reason": room.get("reason"), + "blockers": room.get("blockers"), + "needed_gb": room.get("needed_gb"), + "since": since, "waited_s": round(waited, 1)} + + # Waiting is right while the memory might still arrive. It is wrong + # when the job can never fit -- three LLM jobs sat pending forever + # needing 14.93 GB on a card where only ~14.8 GB can ever be free, + # because an unreclaimable process holds 0.82 GB. Say so and move on + # rather than blocking the queue behind an impossibility. + if waited > self.max_block_s: + ceiling = room.get("max_possible_gb") + detail = (f"needs {room.get('needed_gb')} GB but at most " + f"{ceiling} GB can ever be free on this GPU" + if ceiling is not None and room.get("needed_gb", 0) > ceiling + else f"waited {int(waited)}s for VRAM: {room.get('reason')}") + blockers = ", ".join( + f"{b['name']} ({b['vram_gb']} GB, {b['why']})" + for b in (room.get("blockers") or [])) + _mark(job["id"], FAILED, finished_at=time.time(), + error=f"{detail}{'; blocked by ' + blockers if blockers else ''}") + self.failed += 1 + self.blocked = None + logger.warning(f"job {job['id']} cannot run: {detail}") + continue + + self.waits += 1 + await asyncio.sleep(self.blocked_poll_s) + continue + self.blocked = None + + t0 = time.time() + _mark(job["id"], RUNNING, started_at=t0) + self.current = {**job, "state": RUNNING, "room": room, + "started_at": t0} + logger.info(f"running job {job['id']} for '{job['tenant']}' " + f"({room.get('reason')})") + + # Any failure here must land on the job. An exception used to escape to + # the loop's handler, leaving the row RUNNING forever while the scheduler + # moved on -- an orphan that never completed and never freed its slot. + try: + res = await dispatcher(job["payload"]) + except asyncio.CancelledError: + _mark(job["id"], PENDING, started_at=None) + self.current = None + raise + except Exception as e: + res = {"ok": False, "error": f"dispatch raised: {e}"} + + finished = time.time() + if res.get("ok"): + _mark(job["id"], DONE, finished_at=finished, + result=json.dumps(res)) + self.completed += 1 + else: + _mark(job["id"], FAILED, finished_at=finished, + error=str(res.get("error"))[:500]) + self.failed += 1 + self.last_finished = {"id": job["id"], "tenant": job["tenant"], + "ok": bool(res.get("ok")), + "duration_s": round(finished - t0, 2), + "made_room": room.get("released") or []} + self.current = None + except asyncio.CancelledError: + raise + except Exception as e: + logger.error(f"scheduler error: {e}") + self.current = None + await asyncio.sleep(1.0) + + def get_status(self) -> Dict[str, Any]: + return { + "running": self.running, + "current": self.current, + "last_finished": self.last_finished, + "completed": self.completed, + "failed": self.failed, + "waits": self.waits, + "blocked": self.blocked, + **stats(), + } + + +scheduler = Scheduler() diff --git a/server.py b/server.py index d90a833..893d657 100644 --- a/server.py +++ b/server.py @@ -18,6 +18,7 @@ from pydantic import BaseModel, Field import autotune import engines import health +import jobs as jobs_mod import overclock_manager import ram_optimizer import telemetry_store @@ -220,10 +221,13 @@ def _install_shutdown_hook() -> None: @asynccontextmanager async def lifespan(app: FastAPI): telemetry_store.start() + jobs_mod.init() await broker.start() + await jobs_mod.scheduler.start() await vram_arbitrator.arbitrator.start() _install_shutdown_hook() yield + await jobs_mod.scheduler.stop() await vram_arbitrator.arbitrator.stop() await broker.stop() # Never leave the card with locked clocks and pinned fans after we exit. @@ -328,6 +332,54 @@ class TenantReleaseRequest(BaseModel): confirm: bool = Field(True, description="Wait for NVML to confirm the VRAM was actually released") +class JobRequest(BaseModel): + tenant: str = Field(..., description="Which application should run this job", example="comfyui") + payload: Dict[str, Any] = Field(..., description="What to run: a ComfyUI workflow under 'prompt', or Ollama generate parameters") + priority: Optional[int] = Field(None, description="Defaults to the tenant's priority; higher runs sooner") + label: Optional[str] = Field(None, description="Human-readable name for the queue view") + + +@app.post("/api/jobs", summary="Queue a Job", tags=["Jobs"]) +async def api_submit_job(req: JobRequest): + """Queue work for any tenant. The queue is on disk, so there is no depth limit and + it survives a restart. Jobs run one at a time, highest priority first, with VRAM + arbitrated before each one starts.""" + res = jobs_mod.submit(req.tenant, req.payload, req.priority, req.label) + if not res.get("success"): + raise HTTPException(status_code=400, detail=res.get("error")) + return res + + +@app.get("/api/jobs", summary="The Job Queue", tags=["Jobs"]) +async def api_jobs(state: Optional[str] = Query(None, description="pending | running | done | failed | cancelled"), + limit: int = Query(100)): + """Queued and recent jobs. Pending jobs are listed in the order they will run.""" + return {"jobs": jobs_mod.listing(state, limit), "scheduler": jobs_mod.scheduler.get_status()} + + +@app.get("/api/jobs/{job_id}", summary="One Job", tags=["Jobs"]) +async def api_job(job_id: str): + job = jobs_mod.get(job_id) + if not job: + raise HTTPException(status_code=404, detail=f"no job '{job_id}'") + return job + + +@app.delete("/api/jobs/{job_id}", summary="Cancel a Pending Job", tags=["Jobs"]) +async def api_cancel_job(job_id: str): + """Cancel a job that has not started. Running jobs are left alone -- this service + frees VRAM by asking, never by killing work in flight.""" + res = jobs_mod.cancel(job_id) + if not res.get("success"): + raise HTTPException(status_code=409, detail=res.get("error")) + return res + + +@app.delete("/api/jobs", summary="Cancel All Pending Jobs", tags=["Jobs"]) +async def api_clear_jobs(): + return jobs_mod.clear_pending() + + @app.get("/api/tenants", summary="GPU Tenants", tags=["Tenants"]) async def api_tenants(): """Applications competing for the GPU, as configured. diff --git a/tenants.py b/tenants.py index 785b8f1..60a56f9 100644 --- a/tenants.py +++ b/tenants.py @@ -116,6 +116,12 @@ class GpuTenant: # tuning is workload-specific -- diffusion is compute bound, LLM decode is bandwidth # bound -- and that was previously switched by application name in the arbitrator. overclock_profile: Optional[str] = None + # VRAM that survives a release. ComfyUI keeps its CUDA context for as long as the + # process lives, so purging it does not return everything it holds. Ignoring this + # made plan_release over-promise: it reported that releasing ComfyUI would free + # 0.37 GB against a 0.33 GB shortfall, the job was cleared to run, and the memory + # never actually arrived. + vram_floor_gb: float = 0.0 match: ProcessMatch = field(default_factory=ProcessMatch) busy: BusyProbe = field(default_factory=BusyProbe) release: ReleaseStrategy = field(default_factory=ReleaseStrategy) @@ -138,6 +144,7 @@ def _tenant_from_dict(d: Dict[str, Any]) -> GpuTenant: priority=int(d.get("priority", 50)), needs_vram_gb=float(d.get("needs_vram_gb", 0.0)), overclock_profile=d.get("overclock_profile"), + vram_floor_gb=float(d.get("vram_floor_gb", 0.0)), idle_release_after_s=float(d.get("idle_release_after_s", 30.0)), match=ProcessMatch(**(d.get("match") or {})), busy=BusyProbe(**(d.get("busy") or {})), @@ -175,6 +182,7 @@ DEFAULT_TENANTS: List[Dict[str, Any]] = [ "needs_vram_gb": 6.0, "idle_release_after_s": 30.0, "overclock_profile": "comfy", + "vram_floor_gb": 0.45, "match": {"cmdline": ["comfyui", "comfy"], "cmdline_endswith": ["main.py"]}, "busy": {"type": "http_count", "url": "http://127.0.0.1:8188/queue", "count_keys": ["queue_running", "queue_pending"], @@ -425,8 +433,12 @@ def plan_release(demanding: str, tenants_state: List[Dict[str, Any]], for c in candidates: if freed >= shortfall: break + # Only what the tenant can actually give back, not everything it holds. + releasable = max(c.get("vram_gb", 0.0) - c.get("vram_floor_gb", 0.0), 0.0) + if releasable <= 0: + continue plan.append(c["name"]) - freed += c.get("vram_gb", 0.0) + freed += releasable blockers = [ {"name": s["name"], "vram_gb": s.get("vram_gb", 0.0), diff --git a/tests/test_tenants.py b/tests/test_tenants.py index e2fc28e..ced409a 100644 --- a/tests/test_tenants.py +++ b/tests/test_tenants.py @@ -348,3 +348,58 @@ class TestConfigUpgrade: }])) t = T.get_tenant("ollama") assert t.priority == 5 and t.needs_vram_gb == 99.0 + + +class TestVramFloor: + """VRAM that survives a release must not be promised to anyone else. + + ComfyUI keeps its CUDA context for as long as the process lives, so a purge does not + return everything it holds. Ignoring that made plan_release report it would free + 0.37 GB against a 0.33 GB shortfall; the job was cleared to run and the memory never + arrived, so it waited two minutes and then failed. + """ + + def test_only_memory_above_the_floor_counts_as_freeable(self): + state = [ + {"name": "llm", "priority": 50, "vram_gb": 0.0, "busy": True, + "reclaimable": True, "vram_floor_gb": 0.0}, + {"name": "diffusion", "priority": 60, "vram_gb": 0.44, "busy": False, + "reclaimable": True, "vram_floor_gb": 0.45}, + ] + plan = T.plan_release("llm", state, free_gb=14.6, needed_gb=14.93) + assert plan["possible"] is False + assert plan["release"] == [] + + def test_a_loaded_checkpoint_is_still_freeable_above_its_floor(self): + state = [ + {"name": "llm", "priority": 50, "vram_gb": 0.0, "busy": True, + "reclaimable": True, "vram_floor_gb": 0.0}, + {"name": "diffusion", "priority": 60, "vram_gb": 7.0, "busy": False, + "reclaimable": True, "vram_floor_gb": 0.45}, + ] + plan = T.plan_release("llm", state, free_gb=7.9, needed_gb=14.0) + assert plan["release"] == ["diffusion"] + # 7.0 held minus a 0.45 floor. + assert abs(plan["would_free_gb"] - 6.55) < 0.01 + + def test_a_tenant_at_its_floor_is_not_even_listed_for_release(self): + state = [ + {"name": "llm", "priority": 50, "vram_gb": 0.0, "busy": True, + "reclaimable": True, "vram_floor_gb": 0.0}, + {"name": "at_floor", "priority": 10, "vram_gb": 0.3, "busy": False, + "reclaimable": True, "vram_floor_gb": 0.45}, + {"name": "has_room", "priority": 20, "vram_gb": 5.0, "busy": False, + "reclaimable": True, "vram_floor_gb": 0.0}, + ] + plan = T.plan_release("llm", state, free_gb=0.0, needed_gb=4.0) + assert plan["release"] == ["has_room"] + + def test_default_floor_is_zero_so_existing_configs_are_unchanged(self): + state = [ + {"name": "a", "priority": 50, "vram_gb": 0.0, "busy": True, + "reclaimable": True}, + {"name": "b", "priority": 40, "vram_gb": 5.0, "busy": False, + "reclaimable": True}, + ] + plan = T.plan_release("a", state, free_gb=0.0, needed_gb=5.0) + assert plan["release"] == ["b"] and plan["possible"] is True diff --git a/vram_arbitrator.py b/vram_arbitrator.py index 5c04f93..4510e9a 100644 --- a/vram_arbitrator.py +++ b/vram_arbitrator.py @@ -151,7 +151,8 @@ def get_process_vram_bytes() -> Dict[str, int]: to actually drain. """ out = {"ollama_bytes": 0, "comfyui_bytes": 0, "other_bytes": 0, "free_bytes": 0, - "desktop_bytes": 0, "unmanaged_bytes": 0, "gpu_util_pct": 0} + "desktop_bytes": 0, "unmanaged_bytes": 0, "gpu_util_pct": 0, + "by_tenant_bytes": {}} if not NVML_AVAILABLE: return out try: @@ -178,6 +179,7 @@ def get_process_vram_bytes() -> Dict[str, int]: if len(_PID_KIND_CACHE) >= _PID_KIND_CACHE_MAX: _PID_KIND_CACHE.clear() _PID_KIND_CACHE[key] = kind + out["by_tenant_bytes"][kind] = out["by_tenant_bytes"].get(kind, 0) + used if kind == "ollama": out["ollama_bytes"] += used elif kind == "comfy": @@ -963,6 +965,7 @@ class AutoArbitrator: self.event_sources: Dict[str, str] = {} self._last_tenant_state: Optional[Dict[str, Any]] = None self.last_arbitration: Optional[Dict[str, Any]] = None + self.last_handoff: Optional[Dict[str, Any]] = None self.last_watchdog_error: Optional[str] = None self.stats = { "yields": 0, # release confirmed @@ -1133,12 +1136,16 @@ class AutoArbitrator: # Coalesce bursts: a single graph emits many messages, and one # arbitration pass per burst is enough. now = time.time() - if now - self._last_event_wake < 0.25: + if now - self._last_event_wake < 0.05: continue self._last_event_wake = now self.stats["event_wakeups"] = self.stats.get("event_wakeups", 0) + 1 try: - await self._arbitrate() + # A message on this tenant's own stream is live proof it is + # working right now, so it is taken as busy rather than asked + # over HTTP. A stale queue row could lie; an event arriving + # this instant cannot. + await self._arbitrate(active_tenant=tenant_name) except Exception as e: logger.debug(f"arbitration from event failed: {e}") except (websockets.exceptions.ConnectionClosed, OSError, asyncio.CancelledError): @@ -1260,18 +1267,29 @@ class AutoArbitrator: return False return True - async def _tenant_state(self) -> List[Dict[str, Any]]: - """Current VRAM and busy state for every configured tenant.""" + async def _tenant_state(self, active_tenant: Optional[str] = None + ) -> List[Dict[str, Any]]: + """Current VRAM and busy state for every configured tenant. + + `active_tenant` skips the HTTP busy probe for the tenant whose event stream just + fired: the event is the evidence. That removes a round trip from the handoff, + which is the one path where latency is the entire point. + """ + # Only the cheap NVML read. The full hardware snapshot also does a psutil lookup + # per process, which is wasted work on the handoff path where latency is the + # entire point. snap = get_process_vram_bytes() - stats = get_gpu_hardware_stats() - by_tenant = (stats.get("breakdown", {}) or {}).get("by_tenant_gb", {}) + by_tenant = {k: b / (1024 ** 3) for k, b in snap["by_tenant_bytes"].items()} out = [] for t in tenants_mod.load_tenants(): if not t.enabled: continue bucket = _BUCKET_ALIASES.get(t.name, t.name) vram_gb = by_tenant.get(bucket, 0.0) - probe = await tenants_mod.probe_busy(t, vram_gb=vram_gb) + if t.name == active_tenant: + probe = {"busy": True, "reason": "event received from its own stream"} + else: + probe = await tenants_mod.probe_busy(t, vram_gb=vram_gb) out.append({ "name": t.name, "priority": t.priority, @@ -1281,6 +1299,7 @@ class AutoArbitrator: "reclaimable": t.reclaimable, "needs_vram_gb": t.needs_vram_gb, "overclock_profile": t.overclock_profile, + "vram_floor_gb": t.vram_floor_gb, "idle_release_after_s": t.idle_release_after_s, "reason": probe.get("reason"), }) @@ -1319,7 +1338,7 @@ class AutoArbitrator: else: self._apply_oc_profile(self.IDLE_PROFILE) - async def _arbitrate(self) -> None: + async def _arbitrate(self, active_tenant: Optional[str] = None) -> None: """Generic arbitration over any number of tenants. The two-application version was a pair of hardcoded rules -- yield Ollama when @@ -1328,7 +1347,8 @@ class AutoArbitrator: that lacks the VRAM it declares it needs is starved, and the memory comes from idle reclaimable tenants below it in priority, lowest first. """ - state = await self._tenant_state() + t_start = time.perf_counter() + state = await self._tenant_state(active_tenant) free_gb = self._last_tenant_state["free_gb"] self._apply_profile_for_active(state) @@ -1355,8 +1375,24 @@ class AutoArbitrator: for victim in plan["release"]: await self._release_tenant( victim, f"{s['name']} needs {s['needs_vram_gb']} GB, {free_gb} GB free") + # Wait for the memory to actually come back, and record how long the whole + # handoff took. Swap speed is the point of this service, so it is measured + # rather than assumed. + target_bytes = int(s["needs_vram_gb"] * (1024 ** 3)) + deadline = time.perf_counter() + 30.0 + while time.perf_counter() < deadline: + if get_process_vram_bytes()["free_bytes"] >= target_bytes: + break + await asyncio.sleep(0.02) + handoff_ms = round((time.perf_counter() - t_start) * 1000, 1) + self.last_handoff = {"ts": time.time(), "to": s["name"], + "released": plan["release"], "handoff_ms": handoff_ms, + "triggered_by": "event" if active_tenant else "poll"} + self.stats["handoffs"] = self.stats.get("handoffs", 0) + 1 + logger.info(f"Handoff to '{s['name']}' in {handoff_ms} ms " + f"(released {', '.join(plan['release'])})") self.last_action = (f"Released {', '.join(plan['release'])} so " - f"'{s['name']}' could work") + f"'{s['name']}' could work — {handoff_ms} ms") return # 2. Idle release: a tenant holding VRAM it is not using, after a grace period. @@ -1450,6 +1486,7 @@ class AutoArbitrator: "comfy_stale_job": self.comfy_stale_job, "event_sources": dict(self.event_sources), "last_arbitration": self.last_arbitration, + "last_handoff": self.last_handoff, "tenant_state": self._last_tenant_state, "watchdog_branches": dict(self.watchdog_branches), "last_watchdog_error": self.last_watchdog_error,