fix(ui): refresh GPU memory gauge periodically

This commit is contained in:
Joseph Magly
2026-08-24 12:51:02 -04:00
parent f7dbcd4962
commit f354effc72
4 changed files with 171 additions and 9 deletions
+65 -3
View File
@@ -59,7 +59,12 @@ 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 obliteratus.ui_vram import (
DEFAULT_VRAM_REFRESH_INTERVAL,
VRAM_REFRESH_CHOICES,
render_vram_html,
resolve_vram_refresh_interval,
)
from obliteratus.model_load_settings import (
DTYPE_CHOICES,
QUANTIZATION_CHOICES,
@@ -3957,6 +3962,28 @@ html:not(.dark) body::after { display: none; }
margin-bottom: 8px;
}
/* ---- GPU MEMORY MONITOR ---- */
.vram-monitor-row { align-items: flex-end !important; }
.vram-display { flex-grow: 1 !important; }
.vram-refresh-control {
max-width: 108px !important;
margin-left: auto !important;
opacity: 0.72;
}
.vram-refresh-control:hover,
.vram-refresh-control:focus-within { opacity: 1; }
.vram-refresh-control label span {
font-size: 0.62rem !important;
letter-spacing: 0.06em !important;
}
.vram-refresh-control .wrap { min-height: 30px !important; }
@media (max-width: 640px) {
.vram-refresh-control {
max-width: 100px !important;
margin-left: auto !important;
}
}
/* ---- TAB STYLING ---- */
.tabs { border-bottom: 1px solid #1a1f2e !important; }
button.tab-nav {
@@ -4207,8 +4234,24 @@ with gr.Blocks(title="OBLITERATUS", fill_height=True) as demo:
</div>
""")
# GPU VRAM monitor — refreshed on page load and after key operations
vram_display = gr.HTML(value=_get_vram_html())
# GPU VRAM monitor — periodically refreshed and updated after key operations
with gr.Row(elem_classes=["vram-monitor-row"]):
vram_display = gr.HTML(
value=_get_vram_html(),
scale=12,
elem_classes=["vram-display"],
)
vram_refresh_interval = gr.Dropdown(
choices=VRAM_REFRESH_CHOICES,
value=DEFAULT_VRAM_REFRESH_INTERVAL,
label="GPU refresh",
filterable=False,
allow_custom_value=False,
scale=1,
min_width=96,
elem_classes=["vram-refresh-control"],
)
vram_refresh_timer = gr.Timer(value=DEFAULT_VRAM_REFRESH_INTERVAL, active=True)
# ZeroGPU info — only shown when running on HF Spaces with ZeroGPU
if _ZEROGPU_AVAILABLE:
@@ -5497,6 +5540,25 @@ Built on the shoulders of:
# Refresh VRAM on page load
demo.load(fn=_get_vram_html, outputs=[vram_display])
# Poll independently of long-running queued model work. Updating the one
# timer's interval avoids accumulating polling loops when the selector changes.
vram_refresh_timer.tick(
fn=_get_vram_html,
outputs=[vram_display],
queue=False,
concurrency_limit=1,
concurrency_id="vram-refresh",
api_visibility="private",
)
vram_refresh_interval.change(
fn=resolve_vram_refresh_interval,
inputs=[vram_refresh_interval],
outputs=[vram_refresh_timer],
queue=False,
show_progress="hidden",
api_visibility="private",
)
# ---------------------------------------------------------------------------
# Launch
+15 -5
View File
@@ -5,14 +5,24 @@ from __future__ import annotations
from typing import Any
VRAM_REFRESH_CHOICES = (("10 sec", 10.0), ("30 sec", 30.0), ("1 min", 60.0))
DEFAULT_VRAM_REFRESH_INTERVAL = 30.0
def resolve_vram_refresh_interval(value: float) -> float:
"""Return a supported VRAM polling interval, falling back to the default."""
allowed = {interval for _label, interval in VRAM_REFRESH_CHOICES}
return value if value in allowed else DEFAULT_VRAM_REFRESH_INTERVAL
def render_vram_html(device: Any) -> str:
"""Return per-device GPU/accelerator memory usage as styled bars."""
if not device.is_gpu_available():
return (
'<div style="text-align:center;color:#4a5568;font-size:0.72rem;'
'letter-spacing:1px;margin-top:6px;">CPU ONLY — NO GPU DETECTED</div>'
)
try:
if not device.is_gpu_available():
return (
'<div style="text-align:center;color:#4a5568;font-size:0.72rem;'
'letter-spacing:1px;margin-top:6px;">CPU ONLY — NO GPU DETECTED</div>'
)
cuda_count = device.device_count() if device.is_cuda() else 0
device_indices = range(cuda_count) if cuda_count else (None,)
rows = []
+69
View File
@@ -31,3 +31,72 @@ def test_gradio_application_constructs_without_public_listener():
assert probe.returncode == 0, probe.stderr
assert probe.stdout.strip() == "operator UI constructed without a listener"
assert probe.stderr == ""
def test_gpu_memory_refresh_controls_have_one_configurable_polling_loop():
"""Exercise the rendered Gradio graph rather than relying on source shape."""
pytest.importorskip("gradio", reason="install the locked spaces extra")
probe = subprocess.run(
[
sys.executable,
"-c",
r'''
import asyncio
import app
config = app.demo.get_config_file()
components = config["components"]
timers = [item for item in components if item["type"] == "timer"]
assert len(timers) == 1
timer = timers[0]
assert timer["props"]["value"] == 30.0
assert timer["props"]["active"] is True
selectors = [
item for item in components
if "vram-refresh-control" in item["props"].get("elem_classes", [])
]
assert len(selectors) == 1
selector = selectors[0]
assert selector["type"] == "dropdown"
assert selector["props"]["label"] == "GPU refresh"
assert selector["props"]["choices"] == [
("10 sec", 10.0), ("30 sec", 30.0), ("1 min", 60.0)
]
assert selector["props"]["value"] == 30.0
assert selector["props"]["filterable"] is False
assert selector["props"]["allow_custom_value"] is False
assert selector["props"]["min_width"] == 96
timer_ticks = [
dependency for dependency in config["dependencies"]
if (timer["id"], "tick") in dependency["targets"]
]
assert len(timer_ticks) == 1
assert timer_ticks[0]["queue"] is False
assert timer_ticks[0]["outputs"] == [app.vram_display._id]
interval_changes = [
dependency for dependency in config["dependencies"]
if (selector["id"], "change") in dependency["targets"]
and dependency["outputs"] == [timer["id"]]
]
assert len(interval_changes) == 1
assert interval_changes[0]["queue"] is False
async def verify_interval_changes():
for seconds in (10.0, 30.0, 60.0):
result = await app.demo.process_api(
interval_changes[0]["id"], [seconds], state=None
)
assert result["data"] == [seconds]
asyncio.run(verify_interval_changes())
''',
],
capture_output=True,
text=True,
timeout=120,
check=False,
)
assert probe.returncode == 0, probe.stdout + probe.stderr
+22 -1
View File
@@ -4,7 +4,12 @@ from __future__ import annotations
from types import SimpleNamespace
from obliteratus.ui_vram import render_vram_html
from obliteratus.ui_vram import (
DEFAULT_VRAM_REFRESH_INTERVAL,
VRAM_REFRESH_CHOICES,
render_vram_html,
resolve_vram_refresh_interval,
)
def _memory(used, reserved, total, name):
@@ -79,3 +84,19 @@ def test_vram_html_handles_memory_errors():
get_memory_info=lambda _index: (_ for _ in ()).throw(RuntimeError("unavailable")),
)
assert "Memory: unavailable" in render_vram_html(device)
device.is_gpu_available = lambda: (_ for _ in ()).throw(RuntimeError("unavailable"))
assert "Memory: unavailable" in render_vram_html(device)
def test_vram_refresh_interval_contract():
assert VRAM_REFRESH_CHOICES == (
("10 sec", 10.0),
("30 sec", 30.0),
("1 min", 60.0),
)
assert DEFAULT_VRAM_REFRESH_INTERVAL == 30.0
assert resolve_vram_refresh_interval(10.0) == 10.0
assert resolve_vram_refresh_interval(30.0) == 30.0
assert resolve_vram_refresh_interval(60.0) == 60.0
assert resolve_vram_refresh_interval(15.0) == 30.0