140 lines
4 KiB
Python
140 lines
4 KiB
Python
import glob
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
from typing import Any
|
|
|
|
_name_cache: str | None = None
|
|
|
|
|
|
def _read(path: str) -> str | None:
|
|
try:
|
|
with open(path) as f:
|
|
return f.read().strip()
|
|
except OSError:
|
|
return None
|
|
|
|
|
|
def _shorten(name: str) -> str:
|
|
name = re.sub(r"\s*\(rev.*\)$", "", name).strip()
|
|
groups = re.findall(r"\[([^\]]+)\]", name)
|
|
if len(groups) >= 2:
|
|
brand = groups[0]
|
|
series = groups[-1].split(" / ")[0]
|
|
model = name.split("]", 1)[1].split("[", 1)[0].strip()
|
|
return f"{brand} {model} ({series})".strip()
|
|
return name[:50]
|
|
|
|
|
|
def _gpu_name() -> str:
|
|
global _name_cache
|
|
if _name_cache is None:
|
|
_name_cache = "GPU"
|
|
if shutil.which("lspci"):
|
|
try:
|
|
out = subprocess.run(
|
|
["lspci"], capture_output=True, text=True, timeout=5, check=False
|
|
).stdout
|
|
for line in out.splitlines():
|
|
if "VGA" in line or "3D controller" in line:
|
|
_name_cache = _shorten(line.split(":", 2)[-1].strip())
|
|
break
|
|
except (OSError, subprocess.SubprocessError):
|
|
pass
|
|
return _name_cache
|
|
|
|
|
|
def _amd_sample() -> dict[str, Any] | None:
|
|
devices = sorted(glob.glob("/sys/class/drm/card[0-9]*/device/gpu_busy_percent"))
|
|
if not devices:
|
|
return None
|
|
busy_sum = 0
|
|
count = 0
|
|
vram_used = 0
|
|
vram_total = 0
|
|
temps: list[float] = []
|
|
for busy_path in devices:
|
|
dev = busy_path.rsplit("/", 1)[0]
|
|
try:
|
|
busy_sum += int(_read(busy_path) or 0)
|
|
count += 1
|
|
except ValueError:
|
|
continue
|
|
vram_used += int(_read(f"{dev}/mem_info_vram_used") or 0)
|
|
vram_total += int(_read(f"{dev}/mem_info_vram_total") or 0)
|
|
for hwmon in glob.glob(f"{dev}/hwmon/hwmon*"):
|
|
t = _read(f"{hwmon}/temp1_input")
|
|
if t:
|
|
try:
|
|
temps.append(int(t) / 1000.0)
|
|
except ValueError:
|
|
pass
|
|
if count == 0:
|
|
return None
|
|
return {
|
|
"gpu": round(busy_sum / count, 1),
|
|
"vram_used": vram_used,
|
|
"vram_total": vram_total,
|
|
"vram_pct": round(vram_used / vram_total * 100, 1) if vram_total else None,
|
|
"gpu_temp": max(temps) if temps else None,
|
|
"gpu_name": _gpu_name(),
|
|
}
|
|
|
|
|
|
def _nvidia_sample() -> dict[str, Any] | None:
|
|
if not shutil.which("nvidia-smi"):
|
|
return None
|
|
try:
|
|
out = subprocess.run(
|
|
[
|
|
"nvidia-smi",
|
|
"--query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu,name",
|
|
"--format=csv,noheader,nounits",
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=5,
|
|
check=True,
|
|
).stdout
|
|
except (OSError, subprocess.SubprocessError):
|
|
return None
|
|
lines = [l for l in out.splitlines() if l.strip()]
|
|
if not lines:
|
|
return None
|
|
busy = used = total = 0
|
|
temp = 0
|
|
for line in lines:
|
|
parts = [p.strip() for p in line.split(",")]
|
|
try:
|
|
busy += int(parts[0])
|
|
used += int(parts[1])
|
|
total += int(parts[2])
|
|
temp = max(temp, int(parts[3]))
|
|
except ValueError:
|
|
continue
|
|
name = lines[0].split(",")[-1].strip()
|
|
vram_used = used * 1024 * 1024
|
|
vram_total = total * 1024 * 1024
|
|
return {
|
|
"gpu": round(busy / len(lines), 1),
|
|
"vram_used": vram_used,
|
|
"vram_total": vram_total,
|
|
"vram_pct": round(vram_used / vram_total * 100, 1) if vram_total else None,
|
|
"gpu_temp": float(temp),
|
|
"gpu_name": name,
|
|
}
|
|
|
|
|
|
def sample() -> dict[str, Any]:
|
|
return (
|
|
_amd_sample()
|
|
or _nvidia_sample()
|
|
or {
|
|
"gpu": None,
|
|
"vram_used": None,
|
|
"vram_total": None,
|
|
"vram_pct": None,
|
|
"gpu_temp": None,
|
|
"gpu_name": "no GPU detected",
|
|
}
|
|
)
|