mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-30 06:30:37 +02:00
fix(ui): show VRAM for every visible GPU
This commit is contained in:
@@ -0,0 +1,87 @@
|
||||
"""Deterministic contracts for the Gradio accelerator-memory display."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
|
||||
def test_vram_html_covers_cpu_single_and_multi_accelerator_topologies():
|
||||
"""Exercise app state away from Gradio's import-time worker sockets."""
|
||||
script = r'''
|
||||
from types import SimpleNamespace
|
||||
|
||||
import app
|
||||
|
||||
|
||||
def memory(used, reserved, total, name):
|
||||
return SimpleNamespace(
|
||||
used_gb=used,
|
||||
reserved_gb=reserved,
|
||||
total_gb=total,
|
||||
device_name=name,
|
||||
)
|
||||
|
||||
|
||||
# Zero accelerators preserves the CPU-only message.
|
||||
app.dev.is_gpu_available = lambda: False
|
||||
html = app._get_vram_html()
|
||||
assert "CPU ONLY — NO GPU DETECTED" in html
|
||||
assert "data-device-index" not in html
|
||||
|
||||
# A single non-CUDA accelerator preserves one unified-memory row.
|
||||
app.dev.is_gpu_available = lambda: True
|
||||
app.dev.is_cuda = lambda: False
|
||||
app.dev.get_memory_info = lambda _index=0: memory(4, 0, 16, "Apple M3 (MPS)")
|
||||
html = app._get_vram_html()
|
||||
assert html.count("data-device-index=") == 1
|
||||
assert 'data-device-index="single"' in html
|
||||
assert "Apple M3 (MPS)" in html
|
||||
assert "4.0 / 16.0 GB (25%)" in html
|
||||
assert "Automatic sharding" not in html
|
||||
|
||||
# Homogeneous CUDA devices each receive an indexed row and query.
|
||||
memories = [
|
||||
memory(4, 5, 80, "NVIDIA A100"),
|
||||
memory(16, 20, 80, "NVIDIA A100"),
|
||||
memory(72, 74, 80, "NVIDIA A100"),
|
||||
]
|
||||
queried = []
|
||||
app.dev.is_cuda = lambda: True
|
||||
app.dev.device_count = lambda: len(memories)
|
||||
|
||||
def memory_info(index):
|
||||
queried.append(index)
|
||||
return memories[index]
|
||||
|
||||
app.dev.get_memory_info = memory_info
|
||||
html = app._get_vram_html()
|
||||
assert queried == [0, 1, 2]
|
||||
assert html.count("data-device-index=") == 3
|
||||
assert "GPU 0 · NVIDIA A100" in html
|
||||
assert "GPU 1 · NVIDIA A100" in html
|
||||
assert "GPU 2 · NVIDIA A100" in html
|
||||
assert "Automatic sharding uses additional GPUs" in html
|
||||
|
||||
# Heterogeneous device names stay aligned with their CUDA indices.
|
||||
memories = [
|
||||
memory(2, 3, 24, "NVIDIA RTX 4090"),
|
||||
memory(8, 9, 80, "NVIDIA A100"),
|
||||
]
|
||||
app.dev.device_count = lambda: len(memories)
|
||||
app.dev.get_memory_info = memories.__getitem__
|
||||
html = app._get_vram_html()
|
||||
assert "GPU 0 · NVIDIA RTX 4090" in html
|
||||
assert "GPU 1 · NVIDIA A100" in html
|
||||
assert html.index("GPU 0 · NVIDIA RTX 4090") < html.index("GPU 1 · NVIDIA A100")
|
||||
'''
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
check=False,
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stdout + result.stderr
|
||||
@@ -45,9 +45,11 @@ def test_explicit_device_validation(monkeypatch):
|
||||
|
||||
def test_names_and_device_counts(monkeypatch):
|
||||
monkeypatch.setattr(device, "is_cuda", lambda: True)
|
||||
monkeypatch.setattr(device.torch.cuda, "get_device_name", lambda _index: "Test GPU")
|
||||
names = ["Test GPU 0", "Test GPU 1", "Test GPU 2", "Test GPU 3"]
|
||||
monkeypatch.setattr(device.torch.cuda, "get_device_name", names.__getitem__)
|
||||
monkeypatch.setattr(device.torch.cuda, "device_count", lambda: 4)
|
||||
assert device.get_device_name() == "Test GPU"
|
||||
assert device.get_device_name() == "Test GPU 0"
|
||||
assert device.get_device_name(3) == "Test GPU 3"
|
||||
assert device.device_count() == 4
|
||||
|
||||
monkeypatch.setattr(device, "is_cuda", lambda: False)
|
||||
@@ -88,11 +90,18 @@ def test_system_memory_sources_and_fallback(monkeypatch):
|
||||
def test_memory_info_for_cuda_and_cuda_fallback(monkeypatch):
|
||||
gib = 1024**3
|
||||
monkeypatch.setattr(device, "is_cuda", lambda: True)
|
||||
monkeypatch.setattr(device, "get_device_name", lambda: "GPU")
|
||||
queried_names = []
|
||||
|
||||
def device_name(index=0):
|
||||
queried_names.append(index)
|
||||
return f"GPU {index}"
|
||||
|
||||
monkeypatch.setattr(device, "get_device_name", device_name)
|
||||
monkeypatch.setattr(device.torch.cuda, "mem_get_info", lambda _index: (6 * gib, 8 * gib))
|
||||
monkeypatch.setattr(device.torch.cuda, "memory_allocated", lambda _index: 1 * gib)
|
||||
monkeypatch.setattr(device.torch.cuda, "memory_reserved", lambda _index: 2 * gib)
|
||||
assert device.get_memory_info(2) == device.MemoryInfo(1, 2, 8, 6, "GPU")
|
||||
assert device.get_memory_info(2) == device.MemoryInfo(1, 2, 8, 6, "GPU 2")
|
||||
assert queried_names == [2]
|
||||
|
||||
monkeypatch.setattr(device.torch.cuda, "mem_get_info", Mock(side_effect=RuntimeError("unsupported")))
|
||||
monkeypatch.setattr(
|
||||
@@ -100,7 +109,11 @@ def test_memory_info_for_cuda_and_cuda_fallback(monkeypatch):
|
||||
"get_device_properties",
|
||||
lambda _index: SimpleNamespace(total_memory=10 * gib),
|
||||
)
|
||||
assert device.get_memory_info(2) == device.MemoryInfo(total_gb=10, free_gb=10, device_name="GPU")
|
||||
assert device.get_memory_info(2) == device.MemoryInfo(
|
||||
total_gb=10,
|
||||
free_gb=10,
|
||||
device_name="GPU 2",
|
||||
)
|
||||
|
||||
|
||||
def test_memory_info_for_mps_cpu_and_total_free(monkeypatch):
|
||||
|
||||
Reference in New Issue
Block a user