drjones/gpu-program-swapper
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add a durable cross-tenant job queue with VRAM-aware scheduling

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
drjones
2026-09-07 17:22:34 -07:00
parent 48bd0096ff
commit 0fbc3963b9
6 changed files with 642 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

445
jobs.py Normal file
View File

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