diff --git a/app.py b/app.py
index a5f269d..1ea2b66 100644
--- a/app.py
+++ b/app.py
@@ -59,6 +59,7 @@ import gradio as gr
import torch
from obliteratus import device as dev
from obliteratus.credential_sources import resolve_first, resolve_secret, secret_available
+from obliteratus.ui_vram import render_vram_html
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
# ── ZeroGPU support ─────────────────────────────────────────────────
@@ -936,64 +937,7 @@ def _cleanup_disk(
def _get_vram_html() -> str:
"""Return per-device GPU/accelerator memory usage as styled bars."""
- if not dev.is_gpu_available():
- return (
- 'CPU ONLY — NO GPU DETECTED
'
- )
- try:
- cuda_count = dev.device_count() if dev.is_cuda() else 0
- device_indices = range(cuda_count) if cuda_count else (None,)
- rows = []
- for device_index in device_indices:
- query_index = 0 if device_index is None else device_index
- device_key = "single" if device_index is None else device_index
- mem = dev.get_memory_info(query_index)
- 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 = (
- f"GPU {device_index} · {mem.device_name}"
- if device_index is not None
- else mem.device_name
- )
- reserved_html = (
- f'reserved: {mem.reserved_gb:.1f} GB'
- if mem.reserved_gb > 0
- else 'unified memory'
- )
- rows.append(
- 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
'
+ return render_vram_html(dev)
# ---------------------------------------------------------------------------
diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json
index bd9f8fc..a914e95 100644
--- a/ci/test-risk-map.json
+++ b/ci/test-risk-map.json
@@ -28,6 +28,7 @@
"obliteratus/evaluation/__init__.py",
"obliteratus/models/__init__.py",
"obliteratus/reporting/__init__.py",
+ "obliteratus/ui_vram.py",
"obliteratus/strategies/__init__.py"
],
"required_tests": [
diff --git a/obliteratus/ui_vram.py b/obliteratus/ui_vram.py
new file mode 100644
index 0000000..ad67c83
--- /dev/null
+++ b/obliteratus/ui_vram.py
@@ -0,0 +1,66 @@
+"""HTML rendering for accelerator-memory status."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def render_vram_html(device: Any) -> str:
+ """Return per-device GPU/accelerator memory usage as styled bars."""
+ if not device.is_gpu_available():
+ return (
+ 'CPU ONLY — NO GPU DETECTED
'
+ )
+ try:
+ cuda_count = device.device_count() if device.is_cuda() else 0
+ device_indices = range(cuda_count) if cuda_count else (None,)
+ rows = []
+ for device_index in device_indices:
+ query_index = 0 if device_index is None else device_index
+ device_key = "single" if device_index is None else device_index
+ mem = device.get_memory_info(query_index)
+ used = mem.used_gb
+ total = mem.total_gb
+ pct = (used / total * 100) if total > 0 else 0
+ if pct < 50:
+ bar_color = "#00ff41"
+ elif pct < 80:
+ bar_color = "#ffcc00"
+ else:
+ bar_color = "#ff003c"
+ device_name = (
+ f"GPU {device_index} · {mem.device_name}"
+ if device_index is not None
+ else mem.device_name
+ )
+ reserved_html = (
+ f'reserved: {mem.reserved_gb:.1f} GB'
+ if mem.reserved_gb > 0
+ else 'unified memory'
+ )
+ rows.append(
+ 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/tests/test_app_vram.py b/tests/test_app_vram.py
index 4c878db..8bfa9f4 100644
--- a/tests/test_app_vram.py
+++ b/tests/test_app_vram.py
@@ -1,20 +1,13 @@
-"""Deterministic contracts for the Gradio accelerator-memory display."""
+"""Deterministic contracts for the 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
+from obliteratus.ui_vram import render_vram_html
-def memory(used, reserved, total, name):
+def _memory(used, reserved, total, name):
return SimpleNamespace(
used_gb=used,
reserved_gb=reserved,
@@ -23,65 +16,66 @@ def memory(used, reserved, total, 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
+def test_vram_html_covers_cpu_single_and_multi_accelerator_topologies():
+ device = SimpleNamespace(is_gpu_available=lambda: False)
+ html = render_vram_html(device)
+ 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,
+ device = SimpleNamespace(
+ is_gpu_available=lambda: True,
+ is_cuda=lambda: False,
+ get_memory_info=lambda _index=0: _memory(4, 0, 16, "Apple M3 (MPS)"),
)
+ html = render_vram_html(device)
+ 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
- assert result.returncode == 0, result.stdout + result.stderr
+ memories = [
+ _memory(4, 5, 80, "NVIDIA A100"),
+ _memory(16, 20, 80, "NVIDIA A100"),
+ _memory(72, 74, 80, "NVIDIA A100"),
+ ]
+ queried = []
+
+ def memory_info(index):
+ queried.append(index)
+ return memories[index]
+
+ device = SimpleNamespace(
+ is_gpu_available=lambda: True,
+ is_cuda=lambda: True,
+ device_count=lambda: len(memories),
+ get_memory_info=memory_info,
+ )
+ html = render_vram_html(device)
+ 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
+
+ memories = [
+ _memory(2, 3, 24, "NVIDIA RTX 4090"),
+ _memory(8, 9, 80, "NVIDIA A100"),
+ ]
+ device.device_count = lambda: len(memories)
+ device.get_memory_info = memories.__getitem__
+ html = render_vram_html(device)
+ 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")
+
+
+def test_vram_html_handles_memory_errors():
+ device = SimpleNamespace(
+ is_gpu_available=lambda: True,
+ is_cuda=lambda: True,
+ device_count=lambda: 1,
+ get_memory_info=lambda _index: (_ for _ in ()).throw(RuntimeError("unavailable")),
+ )
+ assert "Memory: unavailable" in render_vram_html(device)