Describe GPU tenants as data so any application can be arbitrated
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:
@@ -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
174
tests/test_tenants.py
Normal 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)
|
||||
Reference in New Issue
Block a user