75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""GPU telemetry for the studio status page, via NVML (nvidia-ml-py)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import logging
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger("fastvideo.studio.gpu")
|
|
|
|
_nvml_initialized = False
|
|
|
|
|
|
def _ensure_nvml() -> Any:
|
|
"""Import and initialize NVML once; raises on machines without it."""
|
|
global _nvml_initialized # noqa: PLW0603
|
|
import pynvml
|
|
if not _nvml_initialized:
|
|
pynvml.nvmlInit()
|
|
_nvml_initialized = True
|
|
return pynvml
|
|
|
|
|
|
def _device_snapshot(pynvml: Any, index: int) -> dict[str, Any]:
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
|
|
name = pynvml.nvmlDeviceGetName(handle)
|
|
if isinstance(name, bytes):
|
|
name = name.decode()
|
|
mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
|
util = pynvml.nvmlDeviceGetUtilizationRates(handle)
|
|
|
|
# Optional sensors: not every GPU/driver exposes them.
|
|
temperature: int | None = None
|
|
power_watts: float | None = None
|
|
power_limit_watts: float | None = None
|
|
with contextlib.suppress(pynvml.NVMLError):
|
|
temperature = int(pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU))
|
|
with contextlib.suppress(pynvml.NVMLError):
|
|
power_watts = pynvml.nvmlDeviceGetPowerUsage(handle) / 1000.0
|
|
power_limit_watts = (pynvml.nvmlDeviceGetEnforcedPowerLimit(handle) / 1000.0)
|
|
|
|
return {
|
|
"index": index,
|
|
"name": name,
|
|
"utilization": int(util.gpu),
|
|
"memory_used_mib": int(mem.used / (1024 * 1024)),
|
|
"memory_total_mib": int(mem.total / (1024 * 1024)),
|
|
"temperature_c": temperature,
|
|
"power_watts": power_watts,
|
|
"power_limit_watts": power_limit_watts,
|
|
}
|
|
|
|
|
|
def get_gpu_snapshot() -> dict[str, Any]:
|
|
"""Return {"available", "gpus", "error"} — never raises.
|
|
|
|
``available: False`` covers both "no NVIDIA driver/library on this host"
|
|
and transient NVML failures; the frontend shows ``error`` as-is.
|
|
"""
|
|
try:
|
|
pynvml = _ensure_nvml()
|
|
count = pynvml.nvmlDeviceGetCount()
|
|
gpus = [_device_snapshot(pynvml, i) for i in range(count)]
|
|
return {"available": True, "gpus": gpus, "error": None}
|
|
except ImportError:
|
|
return {
|
|
"available": False,
|
|
"gpus": [],
|
|
"error": "nvidia-ml-py is not installed on the API server host.",
|
|
}
|
|
except Exception as exc: # NVMLError, driver issues, …
|
|
logger.warning("GPU snapshot failed: %s", exc)
|
|
return {"available": False, "gpus": [], "error": str(exc)}
|