From 5ce65e8193025c087d74640e171d8a0a17b11ef3 Mon Sep 17 00:00:00 2001
From: Joseph Magly <1159087+jmagly@users.noreply.github.com>
Date: Sun, 23 Aug 2026 20:06:08 -0400
Subject: [PATCH] fix(ui): show VRAM for every visible GPU
---
app.py | 84 ++++++++++++++++++-------------
ci/test-risk-map.json | 1 +
obliteratus/device.py | 8 +--
tests/test_app_vram.py | 87 +++++++++++++++++++++++++++++++++
tests/test_device_boundaries.py | 23 +++++++--
5 files changed, 161 insertions(+), 42 deletions(-)
create mode 100644 tests/test_app_vram.py
diff --git a/app.py b/app.py
index 1370b7a..a5f269d 100644
--- a/app.py
+++ b/app.py
@@ -935,45 +935,63 @@ def _cleanup_disk(
# ---------------------------------------------------------------------------
def _get_vram_html() -> str:
- """Return an HTML snippet showing GPU/accelerator memory usage as a styled bar."""
+ """Return per-device GPU/accelerator memory usage as styled bars."""
if not dev.is_gpu_available():
return (
'CPU ONLY — NO GPU DETECTED
'
)
try:
- mem = dev.get_memory_info()
- used = mem.used_gb
- total = mem.total_gb
- pct = (used / total * 100) if total > 0 else 0
- # Color shifts from green → yellow → red
- if pct < 50:
- bar_color = "#00ff41"
- elif pct < 80:
- bar_color = "#ffcc00"
- else:
- bar_color = "#ff003c"
- device_name = mem.device_name
- reserved_html = (
- f'reserved: {mem.reserved_gb:.1f} GB'
- if mem.reserved_gb > 0
- else 'unified memory'
- )
- return (
- f''
- f'
'
- f'{device_name}'
- f'{used:.1f} / {total:.1f} GB ({pct:.0f}%)
'
- f'
'
- f'
'
- f'{reserved_html}
'
- f'
'
+ f'
'
+ f'{device_name}'
+ f'{used:.1f} / {total:.1f} GB ({pct:.0f}%)
'
+ f'
'
+ f'
'
+ f'{reserved_html}
'
+ f'
Automatic sharding uses additional GPUs as model size '
+ 'requires; smaller models may remain on GPU 0.
'
+ )
+ return "".join(rows)
except Exception:
return 'Memory: unavailable
'
diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json
index f1be700..bd9f8fc 100644
--- a/ci/test-risk-map.json
+++ b/ci/test-risk-map.json
@@ -31,6 +31,7 @@
"obliteratus/strategies/__init__.py"
],
"required_tests": [
+ "tests/test_app_vram.py",
"tests/test_cli.py",
"tests/test_cli_boundaries.py",
"tests/test_module_imports.py",
diff --git a/obliteratus/device.py b/obliteratus/device.py
index cbab713..34fc572 100644
--- a/obliteratus/device.py
+++ b/obliteratus/device.py
@@ -83,10 +83,10 @@ def get_device(preference: str = "auto") -> str:
)
-def get_device_name() -> str:
- """Human-readable name of the current accelerator."""
+def get_device_name(device_index: int = 0) -> str:
+ """Human-readable name of the selected accelerator."""
if is_cuda():
- return torch.cuda.get_device_name(0)
+ return torch.cuda.get_device_name(device_index)
if is_mps():
# Apple doesn't expose a per-chip name via MPS; use platform info.
chip = platform.processor() or "Apple Silicon"
@@ -136,7 +136,7 @@ def _system_memory_gb() -> tuple[float, float]:
def get_memory_info(device_index: int = 0) -> MemoryInfo:
"""Query memory for the given accelerator (or system RAM for MPS/CPU)."""
- name = get_device_name()
+ name = get_device_name(device_index) if is_cuda() else get_device_name()
if is_cuda():
try:
diff --git a/tests/test_app_vram.py b/tests/test_app_vram.py
new file mode 100644
index 0000000..4c878db
--- /dev/null
+++ b/tests/test_app_vram.py
@@ -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
diff --git a/tests/test_device_boundaries.py b/tests/test_device_boundaries.py
index bc60004..e493050 100644
--- a/tests/test_device_boundaries.py
+++ b/tests/test_device_boundaries.py
@@ -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):