drjones/gpu-program-swapper
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
48bd0096ff28bf9c2d4b74f9c7a73b9e5fcba9fd
gpu-program-swapper/tenants.py
drjones 48bd0096ff Let a higher-priority job preempt lower-priority work in progress 
Measured failure: a diffusion job took 46 s instead of 3 s, squeezed into 1.6 GB.
The LLM yielded correctly, an inference reloaded it two seconds later, and
plan_release then refused to touch it because it was "busy" -- so ComfyUI crawled
while Ollama held 13 GB for the whole run.

"Never interrupt busy work" looks like the safe rule and is not. Preempting a
lower-priority tenant is safe precisely because releasing is asynchronous: an Ollama
unload queues behind its running request and applies when that finishes, so nothing
is killed mid-flight. That is what makes fast handoff possible at all, and refusing
to do it defeats the purpose of the service.

A tenant may now be asked for memory if it is idle, whatever its rank, or if it is
busy and ranks strictly below the demander. Equal or higher priority is never
interrupted, so peers cannot fight. Idle tenants are still preferred over preempting
busy ones.

Verified against the real contention: with Ollama at 12.38 GB and 98% utilisation, a
diffusion job released it within three seconds and completed in 18 s rather than 46 s.

Tests updated to the corrected rule, and the fixture's priorities aligned with what
actually ships -- it still had the LLM outranking diffusion from before that was
swapped.

Tests: 246.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-07 16:57:08 -07:00

450 lines
18 KiB
Python
Raw Blame History

"""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 EventSource:
"""A stream that tells us *when* to look, not what to think.
ComfyUI publishes a websocket, and the original listener parsed its message types to
decide what was happening -- which meant understanding one application's schema. Any
message is instead treated purely as a wake-up: re-run this tenant's busy probe now
rather than waiting for the next poll. That gives sub-second reaction to any
application with an event stream, with no knowledge of what it emits.
"""
type: str = "none" # none | websocket
url: Optional[str] = None
reconnect_backoff_s: float = 2.0
max_backoff_s: float = 15.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
# GPU profile to apply while this tenant is the active workload. Clock and power
# 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
match: ProcessMatch = field(default_factory=ProcessMatch)
busy: BusyProbe = field(default_factory=BusyProbe)
release: ReleaseStrategy = field(default_factory=ReleaseStrategy)
events: EventSource = field(default_factory=EventSource)
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)),
overclock_profile=d.get("overclock_profile"),
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 {})),
events=EventSource(**(d.get("events") 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,
# Lower than ComfyUI on purpose: an interactive diffusion job preempts the LLM,
# whose weights stay in the page cache and reload in seconds. Getting this the
# wrong way round silently disabled the service's central behaviour -- ComfyUI
# could never reclaim from Ollama.
"priority": 50,
"needs_vram_gb": 4.0,
"idle_release_after_s": 0.0,
"overclock_profile": "ollama",
"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": 60,
"needs_vram_gb": 6.0,
"idle_release_after_s": 30.0,
"overclock_profile": "comfy",
"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},
"events": {"type": "websocket", "url": "ws://127.0.0.1:8188/ws?clientId=hyperswap"},
"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", "events"):
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}
# Who may be asked for memory:
#
# * any idle reclaimable tenant, whatever its rank -- idle memory is not in use;
# * a *busy* tenant that ranks strictly below the demander.
#
# That second clause is the point of the whole service and was nearly lost. Refusing
# to touch anything busy looks safe and is not: a diffusion job measured here ran for
# 46 s instead of 3 s, squeezed into 1.6 GB, because the LLM reloaded straight after
# yielding and was then protected as "busy" while ComfyUI starved. Preempting a
# lower-priority tenant is safe precisely because releasing is asynchronous -- an
# Ollama unload queues behind its running request and applies when that finishes, so
# nothing is killed mid-flight.
#
# Equal or higher priority is never interrupted, so peers cannot fight.
demander_priority = demander.get("priority", 0)
candidates = [
s for s in tenants_state
if s["name"] != demanding
and s.get("reclaimable")
and s.get("vram_gb", 0) > 0
and (not s.get("busy") or s.get("priority", 0) < demander_priority)
]
# Idle tenants first, then lowest priority: never disturb working software while
# something idle still has memory to give.
candidates.sort(key=lambda s: (bool(s.get("busy")), 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 and ranks at or above the demander" if s.get("busy") else
"declares no release mechanism" if not s.get("reclaimable") else
"enough was freed without it")}
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"),
}
Reference in New Issue View Git Blame Copy Permalink