4a38cd68b3
Completes the generalisation. Classification and release were already data; the decision loop was still two hardcoded rules -- yield Ollama when ComfyUI is busy, purge ComfyUI when Ollama is starved -- which could not express a third participant. plan_release() works from the registry instead. A busy tenant that cannot reach its declared needs_vram_gb is starved, and the memory comes from idle reclaimable tenants below it in priority, lowest first, stopping once enough is freed. Tenants that cannot be released are named as blockers rather than passed over, so an impossible plan says which process is in the way. The plan is returned before being acted on, so the decision is testable and is logged before anything is released. Idle release is now per-tenant too, replacing the ComfyUI-specific purge timer. Two bugs found by running it against the live machine rather than only in tests: Starvation was measured against free VRAM alone, so a tenant working perfectly well on 13 GB was flagged as demanding simply because little was left over -- which is the normal state of a busy GPU, and would have caused pointless releases from everything else. A tenant is starved only if it cannot reach what it needs counting what it already holds. Fields added to the tenant schema were silently absent from the config already written to disk, so needs_vram_gb defaulted to 0 and starvation could never trigger for the two tenants that mattered. Shipped defaults are now merged into an existing config on load, with explicit user values still winning. Tests: 242 (was 231), including a three-application contention case -- the property the hardcoded pair of rules could not express. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
403 lines
16 KiB
Python
403 lines
16 KiB
Python
"""GPU tenants: the applications competing for the card, described as data.
|
|
|
|
The point of this service is fast handoff of a single GPU between applications. It grew
|
|
up around the two on this box, and their names ended up compiled into process matching,
|
|
VRAM attribution, busy detection and release calls alike -- roughly 385 references. That
|
|
makes it a script for Ollama and ComfyUI rather than a GPU arbitrator.
|
|
|
|
A tenant is described here instead:
|
|
|
|
* how to recognise its processes (match)
|
|
* how to tell whether it is actually working (busy probe)
|
|
* how to ask it to give VRAM back (release strategy)
|
|
* how much it matters when two want the card (priority)
|
|
|
|
Ollama and ComfyUI ship as defaults so behaviour is unchanged, but nothing about the
|
|
arbitration logic knows their names. A third application -- a training run, a speech
|
|
model, another inference server -- is a config entry, not a code change. A tenant that
|
|
cannot be released (no API to ask) is still worth declaring, because naming it turns
|
|
"unmanaged VRAM" into "held by X, which cannot be reclaimed".
|
|
"""
|
|
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
from dataclasses import dataclass, field, asdict
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import httpx
|
|
import psutil
|
|
|
|
logger = logging.getLogger("tenants")
|
|
|
|
_BASE = os.path.dirname(os.path.abspath(__file__))
|
|
CONFIG_PATH = os.environ.get("HYPERSWAP_TENANTS", os.path.join(_BASE, "tenants.json"))
|
|
|
|
# Kinds are advisory: they drive presentation and sensible defaults, never control flow.
|
|
KIND_LLM, KIND_DIFFUSION, KIND_DESKTOP, KIND_OTHER = "llm", "diffusion", "desktop", "other"
|
|
|
|
|
|
@dataclass
|
|
class ProcessMatch:
|
|
"""How to recognise a tenant's processes among those NVML reports."""
|
|
names: List[str] = field(default_factory=list) # matched against process name
|
|
cmdline: List[str] = field(default_factory=list) # substrings of the full cmdline
|
|
cmdline_endswith: List[str] = field(default_factory=list)
|
|
|
|
def matches(self, pname: str, cmdline: str) -> bool:
|
|
pname, cmdline = pname.lower(), cmdline.lower()
|
|
if any(n.lower() in pname for n in self.names):
|
|
return True
|
|
if any(c.lower() in cmdline for c in self.cmdline):
|
|
return True
|
|
return any(cmdline.rstrip().endswith(c.lower()) for c in self.cmdline_endswith)
|
|
|
|
|
|
@dataclass
|
|
class BusyProbe:
|
|
"""How to tell whether a tenant is genuinely working.
|
|
|
|
`vram_floor_gb` exists because a queue flag can lie: ComfyUI leaves dead jobs in
|
|
queue_running, and only its VRAM reveals that nothing is loaded. GPU utilisation is
|
|
deliberately unavailable as a signal -- it is shared by every tenant, so it cannot
|
|
attribute work to one of them.
|
|
"""
|
|
type: str = "none" # none | http_count | vram
|
|
url: Optional[str] = None
|
|
count_keys: List[str] = field(default_factory=list) # keys whose lists are summed
|
|
vram_busy_gb: float = 0.0 # busy when its VRAM exceeds this
|
|
vram_floor_gb: float = 0.0 # below this it holds no real work
|
|
stale_after_s: float = 90.0
|
|
|
|
|
|
@dataclass
|
|
class ReleaseStrategy:
|
|
"""How to ask a tenant to give VRAM back."""
|
|
type: str = "none" # none | http_post
|
|
url: Optional[str] = None
|
|
body: Dict[str, Any] = field(default_factory=dict)
|
|
# Set when the call must name the loaded model (Ollama unloads per model).
|
|
per_model: bool = False
|
|
timeout_s: float = 120.0
|
|
confirm: bool = True # wait for NVML to show the memory released
|
|
|
|
|
|
@dataclass
|
|
class GpuTenant:
|
|
name: str
|
|
kind: str = KIND_OTHER
|
|
enabled: bool = True
|
|
# Higher wins contention; a tenant yields to anything above it.
|
|
priority: int = 50
|
|
# How much free VRAM this application needs before it can work. Used to decide
|
|
# whether a busy tenant is actually being starved, rather than merely busy.
|
|
needs_vram_gb: float = 0.0
|
|
# How long a reclaimable tenant may sit idle holding VRAM before it is asked for it
|
|
# back. Iterating on a ComfyUI workflow should not pay a reload between every run,
|
|
# so this is deliberately not immediate.
|
|
idle_release_after_s: float = 30.0
|
|
match: ProcessMatch = field(default_factory=ProcessMatch)
|
|
busy: BusyProbe = field(default_factory=BusyProbe)
|
|
release: ReleaseStrategy = field(default_factory=ReleaseStrategy)
|
|
notes: str = ""
|
|
|
|
@property
|
|
def reclaimable(self) -> bool:
|
|
return self.release.type != "none"
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return asdict(self)
|
|
|
|
|
|
def _tenant_from_dict(d: Dict[str, Any]) -> GpuTenant:
|
|
return GpuTenant(
|
|
name=d["name"],
|
|
kind=d.get("kind", KIND_OTHER),
|
|
enabled=d.get("enabled", True),
|
|
priority=int(d.get("priority", 50)),
|
|
needs_vram_gb=float(d.get("needs_vram_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 {})),
|
|
release=ReleaseStrategy(**(d.get("release") or {})),
|
|
notes=d.get("notes", ""),
|
|
)
|
|
|
|
|
|
# Defaults reproduce today's behaviour exactly; they are data, not special cases.
|
|
DEFAULT_TENANTS: List[Dict[str, Any]] = [
|
|
{
|
|
"name": "ollama",
|
|
"kind": KIND_LLM,
|
|
"priority": 60,
|
|
"needs_vram_gb": 4.0,
|
|
"idle_release_after_s": 0.0,
|
|
"match": {"names": ["ollama"], "cmdline": ["llama-server", "ollama"]},
|
|
"busy": {"type": "http_count", "url": "http://localhost:11434/api/ps",
|
|
"count_keys": ["models"]},
|
|
"release": {"type": "http_post", "url": "http://localhost:11434/api/generate",
|
|
"body": {"keep_alive": 0}, "per_model": True, "timeout_s": 120.0},
|
|
"notes": "Unloads per model. With OLLAMA_NUM_PARALLEL=1 the request queues "
|
|
"behind any running generation and applies when it finishes.",
|
|
},
|
|
{
|
|
"name": "comfyui",
|
|
"kind": KIND_DIFFUSION,
|
|
"priority": 50,
|
|
"needs_vram_gb": 6.0,
|
|
"idle_release_after_s": 30.0,
|
|
"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"],
|
|
"vram_floor_gb": 1.5, "stale_after_s": 90.0},
|
|
"release": {"type": "http_post", "url": "http://127.0.0.1:8188/free",
|
|
"body": {"unload_models": True, "free_memory": True},
|
|
"timeout_s": 30.0},
|
|
"notes": "Leaves dead jobs in queue_running; the queue flag is corroborated "
|
|
"against its own VRAM before being believed.",
|
|
},
|
|
{
|
|
"name": "desktop",
|
|
"kind": KIND_DESKTOP,
|
|
"priority": 90,
|
|
"match": {"names": ["gnome-shell", "xorg", "mutter", "kwin", "plasmashell",
|
|
"gnome-remote-desktop", "sddm", "gdm", "picom", "weston"]},
|
|
"release": {"type": "none"},
|
|
"notes": "Compositor and display server. Small, permanent, never reclaimable.",
|
|
},
|
|
]
|
|
|
|
|
|
_cache: Dict[str, Any] = {"ts": 0.0, "tenants": None, "mtime": None}
|
|
CACHE_TTL_S = 10.0
|
|
|
|
|
|
def load_tenants(force: bool = False) -> List[GpuTenant]:
|
|
"""Load tenant definitions, writing the defaults out on first run."""
|
|
now = time.time()
|
|
try:
|
|
mtime = os.path.getmtime(CONFIG_PATH) if os.path.exists(CONFIG_PATH) else None
|
|
except OSError:
|
|
mtime = None
|
|
if (not force and _cache["tenants"] is not None
|
|
and mtime == _cache["mtime"] and (now - _cache["ts"]) < CACHE_TTL_S):
|
|
return _cache["tenants"]
|
|
|
|
raw: List[Dict[str, Any]]
|
|
if os.path.exists(CONFIG_PATH):
|
|
try:
|
|
with open(CONFIG_PATH) as f:
|
|
raw = json.load(f)
|
|
except Exception as e:
|
|
logger.error(f"could not read {CONFIG_PATH}, using defaults: {e}")
|
|
raw = DEFAULT_TENANTS
|
|
else:
|
|
raw = DEFAULT_TENANTS
|
|
try:
|
|
with open(CONFIG_PATH, "w") as f:
|
|
json.dump(DEFAULT_TENANTS, f, indent=2)
|
|
logger.info(f"wrote default tenant definitions to {CONFIG_PATH}")
|
|
except Exception as e:
|
|
logger.warning(f"could not write {CONFIG_PATH}: {e}")
|
|
|
|
# Merge in any fields a shipped default has gained since the config was written.
|
|
# Without this, adding a field silently disables the behaviour it controls for every
|
|
# existing install -- needs_vram_gb defaulted to 0, which made starvation
|
|
# undetectable for the two tenants that had been written out before it existed.
|
|
defaults_by_name = {d["name"]: d for d in DEFAULT_TENANTS}
|
|
tenants = []
|
|
for d in raw:
|
|
base = defaults_by_name.get(d.get("name"))
|
|
if base:
|
|
merged = {**base, **d}
|
|
for key in ("match", "busy", "release"):
|
|
if isinstance(base.get(key), dict):
|
|
merged[key] = {**base[key], **(d.get(key) or {})}
|
|
d = merged
|
|
try:
|
|
tenants.append(_tenant_from_dict(d))
|
|
except Exception as e:
|
|
logger.error(f"skipping malformed tenant {d!r}: {e}")
|
|
_cache.update({"ts": now, "tenants": tenants, "mtime": mtime})
|
|
return tenants
|
|
|
|
|
|
def get_tenant(name: str) -> Optional[GpuTenant]:
|
|
return next((t for t in load_tenants() if t.name == name), None)
|
|
|
|
|
|
def save_tenants(tenants: List[Dict[str, Any]]) -> bool:
|
|
try:
|
|
with open(CONFIG_PATH, "w") as f:
|
|
json.dump(tenants, f, indent=2)
|
|
_cache["tenants"] = None
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"save_tenants failed: {e}")
|
|
return False
|
|
|
|
|
|
def classify_process(pname: str, cmdline: str) -> str:
|
|
"""Return the owning tenant's name, or 'unmanaged'.
|
|
|
|
'unmanaged' is meaningful rather than a dumping ground: it is VRAM this service has
|
|
no way to reclaim, and it is reported as such.
|
|
"""
|
|
for t in load_tenants():
|
|
if t.enabled and t.match.matches(pname, cmdline):
|
|
return t.name
|
|
return "unmanaged"
|
|
|
|
|
|
def classify_pid(pid: int) -> str:
|
|
try:
|
|
proc = psutil.Process(pid)
|
|
return classify_process(proc.name(), " ".join(proc.cmdline()))
|
|
except Exception:
|
|
return "unmanaged"
|
|
|
|
|
|
async def probe_busy(tenant: GpuTenant, vram_gb: float = 0.0,
|
|
state: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
|
"""Is this tenant actually working? Returns {busy, reason, stale}."""
|
|
probe = tenant.busy
|
|
if probe.type == "vram":
|
|
busy = vram_gb > probe.vram_busy_gb
|
|
return {"busy": busy, "reason": f"{vram_gb:.2f} GB held", "stale": False}
|
|
if probe.type != "http_count" or not probe.url:
|
|
return {"busy": False, "reason": "no busy probe configured", "stale": False}
|
|
|
|
try:
|
|
async with httpx.AsyncClient(timeout=3.0) as c:
|
|
r = await c.get(probe.url)
|
|
if r.status_code != 200:
|
|
return {"busy": False, "reason": f"probe HTTP {r.status_code}", "stale": False}
|
|
data = r.json()
|
|
count = sum(len(data.get(k) or []) for k in probe.count_keys)
|
|
except Exception as e:
|
|
return {"busy": False, "reason": f"probe failed: {str(e)[:60]}", "stale": False}
|
|
|
|
if count == 0:
|
|
return {"busy": False, "reason": "queue empty", "stale": False}
|
|
# A queue that claims work while the tenant holds no VRAM is not doing work.
|
|
if probe.vram_floor_gb and vram_gb < probe.vram_floor_gb:
|
|
return {"busy": True, "reason": f"{count} queued, holding {vram_gb:.2f} GB",
|
|
"stale": None, "below_floor": True}
|
|
return {"busy": True, "reason": f"{count} queued/running", "stale": False}
|
|
|
|
|
|
async def release_vram(tenant: GpuTenant, models: Optional[List[str]] = None
|
|
) -> Dict[str, Any]:
|
|
"""Ask a tenant to give its VRAM back, however that tenant expects to be asked."""
|
|
strategy = tenant.release
|
|
if strategy.type == "none" or not strategy.url:
|
|
return {"success": False, "tenant": tenant.name, "released": False,
|
|
"reason": "this tenant exposes no way to release VRAM"}
|
|
|
|
t0 = time.perf_counter()
|
|
payloads: List[Dict[str, Any]] = []
|
|
if strategy.per_model:
|
|
for m in (models or []):
|
|
payloads.append({**strategy.body, "model": m})
|
|
if not payloads:
|
|
return {"success": True, "tenant": tenant.name, "released": False,
|
|
"reason": "nothing loaded to release"}
|
|
else:
|
|
payloads.append(dict(strategy.body))
|
|
|
|
errors = []
|
|
try:
|
|
async with httpx.AsyncClient(timeout=strategy.timeout_s) as c:
|
|
for body in payloads:
|
|
try:
|
|
await c.post(strategy.url, json=body)
|
|
except Exception as e:
|
|
errors.append(str(e)[:80])
|
|
except Exception as e:
|
|
errors.append(str(e)[:80])
|
|
|
|
return {
|
|
"success": not errors,
|
|
"tenant": tenant.name,
|
|
"released": True,
|
|
"requests": len(payloads),
|
|
"duration_ms": round((time.perf_counter() - t0) * 1000, 2),
|
|
"errors": errors or None,
|
|
}
|
|
|
|
|
|
def describe() -> List[Dict[str, Any]]:
|
|
"""Tenant definitions for the API, with what each can and cannot do."""
|
|
out = []
|
|
for t in sorted(load_tenants(), key=lambda x: -x.priority):
|
|
d = t.to_dict()
|
|
d["reclaimable"] = t.reclaimable
|
|
d["busy_probe"] = t.busy.type
|
|
d["release_via"] = t.release.type
|
|
out.append(d)
|
|
return out
|
|
|
|
|
|
# ---------------------------------------------------------------- arbitration
|
|
|
|
def plan_release(demanding: str, tenants_state: List[Dict[str, Any]],
|
|
free_gb: float, needed_gb: float) -> Dict[str, Any]:
|
|
"""Decide who should give up VRAM so a starved tenant can work.
|
|
|
|
Generic over any number of applications: candidates are every *reclaimable* tenant
|
|
that is not itself busy and ranks below the demanding one, taken lowest priority
|
|
first, until enough would be freed. The two-application version of this was a pair
|
|
of hardcoded rules -- yield Ollama for ComfyUI, purge ComfyUI for Ollama -- which
|
|
could not express a third participant at all.
|
|
|
|
Returns the plan rather than performing it, so the decision is testable and can be
|
|
logged before anything is actually released.
|
|
"""
|
|
by_name = {s["name"]: s for s in tenants_state}
|
|
demander = by_name.get(demanding)
|
|
if not demander:
|
|
return {"possible": False, "reason": f"unknown tenant '{demanding}'", "release": []}
|
|
|
|
# The demander keeps what it already holds; only the remainder must be found.
|
|
shortfall = needed_gb - free_gb - demander.get("vram_gb", 0.0)
|
|
if shortfall <= 0:
|
|
return {"possible": True, "reason": "enough VRAM is already free",
|
|
"release": [], "shortfall_gb": 0.0}
|
|
|
|
candidates = [
|
|
s for s in tenants_state
|
|
if s["name"] != demanding
|
|
and s.get("reclaimable")
|
|
and not s.get("busy")
|
|
and s.get("priority", 0) <= demander.get("priority", 0)
|
|
and s.get("vram_gb", 0) > 0
|
|
]
|
|
candidates.sort(key=lambda s: (s.get("priority", 0), -s.get("vram_gb", 0)))
|
|
|
|
plan, freed = [], 0.0
|
|
for c in candidates:
|
|
if freed >= shortfall:
|
|
break
|
|
plan.append(c["name"])
|
|
freed += c.get("vram_gb", 0.0)
|
|
|
|
blockers = [
|
|
{"name": s["name"], "vram_gb": s.get("vram_gb", 0.0),
|
|
"why": ("busy" if s.get("busy") else
|
|
"declares no release mechanism" if not s.get("reclaimable") else
|
|
"higher priority")}
|
|
for s in tenants_state
|
|
if s["name"] != demanding and s.get("vram_gb", 0) > 0 and s["name"] not in plan
|
|
]
|
|
|
|
return {
|
|
"possible": freed >= shortfall,
|
|
"shortfall_gb": round(shortfall, 2),
|
|
"would_free_gb": round(freed, 2),
|
|
"release": plan,
|
|
"blockers": blockers,
|
|
"reason": (f"releasing {', '.join(plan)} frees {freed:.2f} GB of the "
|
|
f"{shortfall:.2f} GB shortfall" if plan else
|
|
"no reclaimable idle tenant holds enough VRAM"),
|
|
}
|