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 <noreply@anthropic.com>
This commit is contained in:
14
tenants.py
14
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),
|
||||
|
||||
Reference in New Issue
Block a user