"""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()