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,