"""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)