diff --git a/app.py b/app.py index 1370b7a..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 ───────────────────────────────────────────────── @@ -935,47 +936,8 @@ def _cleanup_disk( # --------------------------------------------------------------------------- def _get_vram_html() -> str: - """Return an HTML snippet showing GPU/accelerator memory usage as a styled bar.""" - 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'
' - f'{reserved_html}
' - f'
' - ) - except Exception: - return '
Memory: unavailable
' + """Return per-device GPU/accelerator memory usage as styled bars.""" + return render_vram_html(dev) # --------------------------------------------------------------------------- diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index f1be700..a914e95 100644 --- a/ci/test-risk-map.json +++ b/ci/test-risk-map.json @@ -28,9 +28,11 @@ "obliteratus/evaluation/__init__.py", "obliteratus/models/__init__.py", "obliteratus/reporting/__init__.py", + "obliteratus/ui_vram.py", "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/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'
' + f'{reserved_html}
' + f'
' + ) + if cuda_count > 1: + rows.append( + '
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 new file mode 100644 index 0000000..8bfa9f4 --- /dev/null +++ b/tests/test_app_vram.py @@ -0,0 +1,81 @@ +"""Deterministic contracts for the accelerator-memory display.""" + +from __future__ import annotations + +from types import SimpleNamespace + +from obliteratus.ui_vram import render_vram_html + + +def _memory(used, reserved, total, name): + return SimpleNamespace( + used_gb=used, + reserved_gb=reserved, + total_gb=total, + device_name=name, + ) + + +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 + + 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 + + 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) 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):