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

Describe GPU tenants as data so any application can be arbitrated

Browse Source 
The point of this service is fast handoff of one 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 -- about 385 references
across five modules. That made it a script for Ollama and ComfyUI rather than a GPU
arbitrator.

tenants.py describes an application as data: how to recognise its processes, how to
tell whether it is genuinely working, how to ask it for VRAM back, and how much it
matters when two want the card. Ollama, ComfyUI and the desktop compositor ship as
defaults in tenants.json, so behaviour is unchanged, but the arbitration logic no
longer knows any particular name. Endpoints are generic: GET /api/tenants,
GET /api/tenants/{name}, POST /api/tenants/{name}/release -- the last being the
general form of both the Ollama soft-yield and the ComfyUI purge.

Verified by registering a third application on this machine with no code change: the
speech relay that had been showing up only as anonymous "unmanaged VRAM" is now named,
attributed, and probed by the VRAM it holds rather than by an API it does not have.
Because it declares no release strategy, a release request returns 409 explaining that
its memory cannot be reclaimed, instead of reporting a success that did nothing.

Busy probes deliberately cannot use GPU utilisation. It is shared by every tenant, so
it cannot attribute work to one of them -- the mistake that made a stale ComfyUI queue
entry undetectable earlier in this branch. A tenant's own VRAM is the signal.

Writing the tests exposed that the suite had become non-hermetic: classification is now
configuration, so a test asserting "a third-party process is unmanaged" started failing
the moment the speech relay was registered on this machine. An autouse fixture now
isolates every test from the operator's live tenants.json.

Tests: 231 (was 206).

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
drjones
2026-09-07 11:20:44 -07:00
parent c3d9b36035
commit 1bcfbb2335
7 changed files with 764 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
README.md
View File

@@ -174,10 +174,46 @@ only if the card actually needs it.
---
## 1b. Any Application, Not Just These Two
The purpose is fast handoff of one 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 — about 385 references. That made it a script for Ollama
and ComfyUI rather than a GPU arbitrator.
A tenant is now **described as data** in `tenants.json`:
```json
{
"name": "trainer",
"kind": "other",
"priority": 80,
"match": { "cmdline": ["train.py"] },
"busy": { "type": "vram", "vram_busy_gb": 1.0 },
"release": { "type": "http_post", "url": "http://localhost:9999/release" }
}
```
| Field | What it answers |
| :--- | :--- |
| `match` | Which GPU processes belong to this application (name, cmdline substring, or suffix — ComfyUI is a bare `python main.py`) |
| `busy` | Whether it is *genuinely* working. `http_count` sums queue lists; `vram` needs no API at all. `vram_floor_gb` catches a queue that claims work while nothing is loaded |
| `release` | How to ask for VRAM back — `http_post` with a body, `per_model` for Ollama's per-model unload, or `none` |
| `priority` | Who wins contention |
Ollama, ComfyUI and the desktop compositor ship as defaults, so behaviour is unchanged —
but nothing in the arbitration logic knows their names. Endpoints are generic:
`GET /api/tenants`, `GET /api/tenants/{name}`, `POST /api/tenants/{name}/release`.
A tenant with `"release": {"type": "none"}` is still worth declaring. The 842 MB speech
relay on this box cannot be reclaimed, and naming it turns anonymous "unmanaged VRAM" into
"held by stt-relay, which exposes no release API" — and a release request returns **409**
explaining that, rather than silently doing nothing.
## 1a. Tests
```bash
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/ -q # 206 passed in ~3.7s
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/ -q # 231 passed in ~3.8s
```
Hermetic: no GPU, no network, no sleeps. An autouse fixture stubs `overclock_manager._sh`

81
server.py
View File

@@ -21,6 +21,7 @@ import health
import overclock_manager
import ram_optimizer
import telemetry_store
import tenants as tenants_mod
import thermal_governor
import vram_arbitrator
@@ -322,6 +323,86 @@ async def api_health():
return await health.run_health_checks()
class TenantReleaseRequest(BaseModel):
models: Optional[List[str]] = Field(None, description="For per-model tenants (Ollama), which to unload; defaults to everything resident")
confirm: bool = Field(True, description="Wait for NVML to confirm the VRAM was actually released")
@app.get("/api/tenants", summary="GPU Tenants", tags=["Tenants"])
async def api_tenants():
"""Applications competing for the GPU, as configured.
Each entry declares how its processes are recognised, how to tell whether it is
working, and how to ask it for VRAM back. Adding an application is a config change
in tenants.json, not a code change.
"""
gpu = vram_arbitrator.get_gpu_hardware_stats()
by_tenant = (gpu.get("breakdown", {}) or {}).get("by_tenant_gb", {})
out = []
for t in tenants_mod.describe():
name = t["name"]
# The two original tenants are reported under the bucket names the API has
# always used.
bucket = {"comfyui": "comfy"}.get(name, name)
t["vram_gb"] = by_tenant.get(bucket, 0.0)
out.append(t)
return {"tenants": out, "config_path": tenants_mod.CONFIG_PATH,
"unmanaged_gb": (gpu.get("breakdown", {}) or {}).get("unmanaged_gb", 0.0)}
@app.get("/api/tenants/{name}", summary="One GPU Tenant", tags=["Tenants"])
async def api_tenant(name: str):
"""A single tenant's definition, current VRAM, and whether it is genuinely busy."""
t = tenants_mod.get_tenant(name)
if not t:
raise HTTPException(status_code=404, detail=f"no tenant named '{name}'")
gpu = vram_arbitrator.get_gpu_hardware_stats()
bucket = {"comfyui": "comfy"}.get(name, name)
vram_gb = (gpu.get("breakdown", {}) or {}).get("by_tenant_gb", {}).get(bucket, 0.0)
busy = await tenants_mod.probe_busy(t, vram_gb=vram_gb)
d = t.to_dict()
d.update({"vram_gb": vram_gb, "reclaimable": t.reclaimable, "busy": busy})
return d
@app.post("/api/tenants/{name}/release", summary="Ask a Tenant for its VRAM", tags=["Tenants"])
async def api_tenant_release(name: str, req: Optional[TenantReleaseRequest] = None):
"""Release a tenant's VRAM using whatever mechanism that tenant declares.
This is the generic form of the Ollama soft-yield and the ComfyUI purge: the same
request works for any application in the registry, including ones added later.
"""
t = tenants_mod.get_tenant(name)
if not t:
raise HTTPException(status_code=404, detail=f"no tenant named '{name}'")
if not t.reclaimable:
raise HTTPException(status_code=409,
detail=f"'{name}' declares no way to release VRAM; its "
f"memory cannot be reclaimed by this service")
models = req.models if req else None
if t.release.per_model and not models:
state = await vram_arbitrator.get_ollama_live_state()
models = [m.get("name") for m in state.get("loaded_models", []) if m.get("name")]
before = vram_arbitrator.get_process_vram_bytes()
res = await tenants_mod.release_vram(t, models=models)
if (req is None or req.confirm) and res.get("released"):
bucket = {"comfyui": "comfy"}.get(name, name)
key = {"ollama": "ollama_bytes", "comfy": "comfyui_bytes"}.get(bucket)
if key:
baseline = before[key]
barrier = await vram_arbitrator._await_vram_release(baseline) \
if key == "ollama_bytes" else None
if barrier:
res.update({"outcome": barrier.get("outcome"),
"confirm_ms": barrier.get("confirm_ms")})
after = vram_arbitrator.get_process_vram_bytes()
res["free_vram_gb"] = round(after["free_bytes"] / (1024**3), 2)
return res
@app.get("/api/engines", summary="Live Engine Configuration", tags=["Telemetry"])
async def api_engines():
"""Real configuration of Ollama and ComfyUI, with what each setting implies for

108
tenants.json Normal file
View File

@@ -0,0 +1,108 @@
[
{
"name": "ollama",
"kind": "llm",
"priority": 60,
"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": "diffusion",
"priority": 50,
"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": "stt-relay",
"kind": "other",
"priority": 70,
"match": {
"cmdline": [
"stt_relay.py"
]
},
"busy": {
"type": "vram",
"vram_busy_gb": 1.0
},
"release": {
"type": "none"
},
"notes": "Long-running speech relay. Holds ~0.8 GB permanently and exposes no release API, so its VRAM is headroom this service can never offer. Declared so it is named rather than lumped into 'unmanaged'."
},
{
"name": "desktop",
"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."
}
]

313
tenants.py Normal file
View File

@@ -0,0 +1,313 @@
"""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
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)),
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,
"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,
"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}")
tenants = []
for d in raw:
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

21
tests/conftest.py
View File

@@ -67,3 +67,24 @@ def temp_db(tmp_path, monkeypatch):
telemetry_store.stop()
except Exception:
pass
@pytest.fixture(autouse=True)
def isolated_tenant_registry(tmp_path, monkeypatch):
"""Never let tests read the operator's live tenants.json.
Classification is now configuration, which means a test that reads the real config
changes result when someone adds an application to their own machine -- exactly what
happened when stt-relay was registered and a "third party is unmanaged" test started
seeing it as a named tenant. Every test gets the shipped defaults unless it opts out
by pointing CONFIG_PATH somewhere itself.
"""
import json as _json
import tenants as _tenants
path = tmp_path / "tenants-default.json"
path.write_text(_json.dumps(_tenants.DEFAULT_TENANTS))
monkeypatch.setattr(_tenants, "CONFIG_PATH", str(path))
_tenants._cache.update({"ts": 0.0, "tenants": None, "mtime": None})
yield
_tenants._cache.update({"ts": 0.0, "tenants": None, "mtime": None})

174
tests/test_tenants.py Normal file
View File

@@ -0,0 +1,174 @@
"""Tests for the GPU tenant registry.
The point of this service is fast handoff of one GPU between applications, and it should
work for any application -- not only the two it grew up around. Their names had ended up
compiled into process matching, VRAM attribution, busy detection and release calls alike.
These tests pin the properties that make the registry generic: adding an application is
configuration, and nothing in the arbitration logic knows a particular name.
"""
import asyncio
import json
import pytest
import tenants as T
@pytest.fixture
def cfg(tmp_path, monkeypatch):
path = tmp_path / "tenants.json"
monkeypatch.setattr(T, "CONFIG_PATH", str(path))
T._cache.update({"ts": 0.0, "tenants": None, "mtime": None})
return path
class TestProcessMatching:
def test_matches_by_process_name(self):
m = T.ProcessMatch(names=["ollama"])
assert m.matches("ollama", "/usr/bin/ollama serve")
assert not m.matches("python", "main.py")
def test_matches_by_cmdline_substring(self):
m = T.ProcessMatch(cmdline=["llama-server"])
assert m.matches("python", "/usr/local/lib/ollama/llama-server --model x")
def test_matches_by_cmdline_suffix(self):
# ComfyUI is a bare `python main.py`, with nothing else distinguishing it.
m = T.ProcessMatch(cmdline_endswith=["main.py"])
assert m.matches("python", "/opt/ComfyUI/venv/bin/python main.py")
assert not m.matches("python", "/opt/other/main.py --serve")
def test_matching_is_case_insensitive(self):
assert T.ProcessMatch(names=["Xorg"]).matches("XORG", "")
class TestDefaultsPreserveExistingBehaviour:
"""The shipped defaults must classify exactly as the hardcoded version did."""
@pytest.mark.parametrize("pname,cmdline,expected", [
("llama-server", "/usr/local/lib/ollama/llama-server --model x", "ollama"),
("ollama", "/usr/bin/ollama serve", "ollama"),
("python", "/home/u/ComfyUI/venv/bin/python main.py --listen", "comfyui"),
("gnome-shell", "/usr/bin/gnome-shell --mode=ubuntu", "desktop"),
("Xorg", "/usr/lib/xorg/Xorg :8", "desktop"),
("python", "/home/u/robopest-venv/bin/python /home/u/stt_relay.py", "unmanaged"),
("trainer", "/opt/ml/bin/trainer --epochs 3", "unmanaged"),
])
def test_classification(self, cfg, pname, cmdline, expected):
assert T.classify_process(pname, cmdline) == expected
def test_unknown_process_is_unmanaged_not_silently_owned(self, cfg):
# Misattributing a third party's VRAM to a tenant would make this service
# promise headroom it cannot deliver.
assert T.classify_process("weird", "/opt/x/weird --run") == "unmanaged"
class TestAddingAnApplicationIsConfiguration:
def test_a_new_tenant_is_recognised_without_code_changes(self, cfg):
cfg.write_text(json.dumps(T.DEFAULT_TENANTS + [{
"name": "trainer",
"kind": "other",
"priority": 80,
"match": {"cmdline": ["train.py"]},
"release": {"type": "http_post", "url": "http://localhost:9999/release"},
}]))
assert T.classify_process("python", "/opt/ml/train.py --epochs 3") == "trainer"
t = T.get_tenant("trainer")
assert t.priority == 80 and t.reclaimable
def test_first_run_writes_the_defaults(self, cfg):
assert not cfg.exists()
T.load_tenants(force=True)
assert cfg.exists()
assert {t["name"] for t in json.loads(cfg.read_text())} == {
"ollama", "comfyui", "desktop"}
def test_a_malformed_entry_is_skipped_not_fatal(self, cfg):
cfg.write_text(json.dumps([{"name": "ok", "match": {"names": ["a"]}},
{"no_name": True}]))
names = [t.name for t in T.load_tenants(force=True)]
assert names == ["ok"]
def test_corrupt_config_falls_back_to_defaults(self, cfg):
cfg.write_text("{ not json")
assert {t.name for t in T.load_tenants(force=True)} >= {"ollama", "comfyui"}
class TestReclaimability:
def test_a_tenant_with_no_release_strategy_is_not_reclaimable(self, cfg):
t = T.GpuTenant(name="x", release=T.ReleaseStrategy(type="none"))
assert t.reclaimable is False
def test_release_refuses_rather_than_reporting_success(self, cfg):
t = T.GpuTenant(name="x", release=T.ReleaseStrategy(type="none"))
res = asyncio.run(T.release_vram(t))
assert res["success"] is False and res["released"] is False
assert "no way to release" in res["reason"]
def test_per_model_release_with_nothing_loaded_is_a_no_op(self, cfg):
t = T.GpuTenant(name="ollama", release=T.ReleaseStrategy(
type="http_post", url="http://x/api", per_model=True))
res = asyncio.run(T.release_vram(t, models=[]))
assert res["success"] is True and res["released"] is False
class TestBusyProbe:
def _probe(self, monkeypatch, payload, status=200):
class _R:
status_code = status
def json(self_inner): return payload
class _C:
async def __aenter__(self): return self
async def __aexit__(self, *a): return False
async def get(self, url): return _R()
monkeypatch.setattr(T.httpx, "AsyncClient", lambda **k: _C())
def test_empty_queue_is_not_busy(self, monkeypatch):
self._probe(monkeypatch, {"queue_running": [], "queue_pending": []})
t = T.GpuTenant(name="c", busy=T.BusyProbe(
type="http_count", url="http://x/queue",
count_keys=["queue_running", "queue_pending"]))
assert asyncio.run(T.probe_busy(t))["busy"] is False
def test_queued_work_while_holding_no_vram_is_flagged_below_floor(self, monkeypatch):
# ComfyUI leaves dead jobs in queue_running; only its VRAM reveals that nothing
# is loaded.
self._probe(monkeypatch, {"queue_running": [[1, "abc"]], "queue_pending": []})
t = T.GpuTenant(name="c", busy=T.BusyProbe(
type="http_count", url="http://x/queue", count_keys=["queue_running"],
vram_floor_gb=1.5))
res = asyncio.run(T.probe_busy(t, vram_gb=0.56))
assert res["busy"] is True and res.get("below_floor") is True
def test_queued_work_with_a_checkpoint_loaded_is_plainly_busy(self, monkeypatch):
self._probe(monkeypatch, {"queue_running": [[1, "abc"]], "queue_pending": []})
t = T.GpuTenant(name="c", busy=T.BusyProbe(
type="http_count", url="http://x/queue", count_keys=["queue_running"],
vram_floor_gb=1.5))
res = asyncio.run(T.probe_busy(t, vram_gb=6.8))
assert res["busy"] is True and not res.get("below_floor")
def test_vram_probe_needs_no_http_endpoint(self):
# An application with no API can still be observed by what it holds.
t = T.GpuTenant(name="x", busy=T.BusyProbe(type="vram", vram_busy_gb=1.0))
assert asyncio.run(T.probe_busy(t, vram_gb=2.0))["busy"] is True
assert asyncio.run(T.probe_busy(t, vram_gb=0.5))["busy"] is False
def test_an_unreachable_probe_reports_not_busy_rather_than_raising(self, monkeypatch):
class _C:
async def __aenter__(self): return self
async def __aexit__(self, *a): return False
async def get(self, url): raise ConnectionError("refused")
monkeypatch.setattr(T.httpx, "AsyncClient", lambda **k: _C())
t = T.GpuTenant(name="c", busy=T.BusyProbe(type="http_count", url="http://x",
count_keys=["q"]))
res = asyncio.run(T.probe_busy(t))
assert res["busy"] is False and "failed" in res["reason"]
class TestPriority:
def test_describe_orders_by_priority(self, cfg):
rows = T.describe()
prios = [r["priority"] for r in rows]
assert prios == sorted(prios, reverse=True)
assert all("reclaimable" in r for r in rows)

49
vram_arbitrator.py
View File

@@ -12,6 +12,7 @@ import websockets
import overclock_manager
import ram_optimizer
import tenants as tenants_mod
import telemetry_store
try:
@@ -198,6 +199,19 @@ _PID_KIND_CACHE: Dict[tuple, str] = {}
_PID_KIND_CACHE_MAX = 512
def _pid_key(pid: int) -> Optional[tuple]:
try:
return (pid, psutil.Process(pid).create_time())
except Exception:
return None
# Tenant names as used by this module's buckets. The tenant registry is the source of
# truth for *which* application a process belongs to; these two names are kept because
# the REST payloads and the dashboard have used them since the beginning.
_BUCKET_ALIASES = {"comfyui": "comfy"}
def _pid_key(pid: int) -> Optional[tuple]:
try:
return (pid, psutil.Process(pid).create_time())
@@ -214,26 +228,16 @@ DESKTOP_PROCESS_HINTS = (
def _classify_pid(pid: int) -> str:
"""Bucket a GPU process into ollama | comfy | desktop | unmanaged.
"""Which tenant owns this GPU process.
The old version had one catch-all "other" bucket, which put a 3.9 MB compositor and
an 842 MB long-running inference script in the same number. That matters: this
service can reclaim VRAM from ComfyUI, but it cannot touch a third-party workload,
and pretending otherwise makes it promise headroom it cannot deliver.
The matching rules used to be substrings compiled into this function, which made the
two applications on this box part of the arbitrator rather than input to it. They now
come from the tenant registry, so a third application is a config entry.
"unmanaged" still means something specific and useful: VRAM held by something with no
declared way to release it, and therefore headroom this service can never offer.
"""
try:
proc = psutil.Process(pid)
pname = proc.name().lower()
cmdline = " ".join(proc.cmdline()).lower()
except Exception:
return "unmanaged"
if "ollama" in pname or "llama-server" in cmdline:
return "ollama"
if "comfyui" in cmdline or "comfy" in cmdline or cmdline.rstrip().endswith("main.py"):
return "comfy"
if any(hint in pname or hint in cmdline for hint in DESKTOP_PROCESS_HINTS):
return "desktop"
return "unmanaged"
return _BUCKET_ALIASES.get(tenants_mod.classify_pid(pid), tenants_mod.classify_pid(pid))
def get_gpu_hardware_stats() -> Dict[str, Any]:
@@ -336,7 +340,10 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
"desktop_bytes": 0,
"unmanaged_bytes": 0,
"unmanaged": [],
"processes": []
"processes": [],
# Generic attribution: one entry per tenant, so an application added to the
# registry is reported without any change here.
"by_tenant": {},
}
try:
@@ -376,6 +383,8 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
"vram_mb": round(used_mem / (1024**2), 1),
})
proc_breakdown["by_tenant"][kind] = (
proc_breakdown["by_tenant"].get(kind, 0) + used_mem)
proc_breakdown["processes"].append({
"pid": pid,
"name": pname,
@@ -431,6 +440,8 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
# reclaimed, so it is permanently unavailable headroom.
"unmanaged_gb": round(proc_breakdown["unmanaged_bytes"] / (1024**3), 2),
"unmanaged": proc_breakdown["unmanaged"],
"by_tenant_gb": {k: round(b / (1024**3), 2)
for k, b in proc_breakdown["by_tenant"].items()},
"free_mb": round(free_vram / (1024**2), 1),
"free_gb": round(free_vram / (1024**3), 2),
"processes": proc_breakdown["processes"],