diff --git a/README.md b/README.md index 5d6db44..212ea84 100644 --- a/README.md +++ b/README.md @@ -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` diff --git a/server.py b/server.py index 9b12adf..075d10f 100644 --- a/server.py +++ b/server.py @@ -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 diff --git a/tenants.json b/tenants.json new file mode 100644 index 0000000..8d68c6b --- /dev/null +++ b/tenants.json @@ -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." + } +] \ No newline at end of file diff --git a/tenants.py b/tenants.py new file mode 100644 index 0000000..1006b1f --- /dev/null +++ b/tenants.py @@ -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 diff --git a/tests/conftest.py b/tests/conftest.py index 9b6c399..1c8b3f4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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}) diff --git a/tests/test_tenants.py b/tests/test_tenants.py new file mode 100644 index 0000000..aa371b0 --- /dev/null +++ b/tests/test_tenants.py @@ -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) diff --git a/vram_arbitrator.py b/vram_arbitrator.py index 35a4fa0..2a7118b 100644 --- a/vram_arbitrator.py +++ b/vram_arbitrator.py @@ -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"],