128 lines
3.7 KiB
Python
128 lines
3.7 KiB
Python
import glob
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
|
|
from app.sample import Sample
|
|
|
|
_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(s: Sample) -> bool:
|
|
devices = sorted(glob.glob("/sys/class/drm/card[0-9]*/device/gpu_busy_percent"))
|
|
if not devices:
|
|
return False
|
|
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 False
|
|
s.gpu = round(busy_sum / count, 1)
|
|
s.vram_used = vram_used
|
|
s.vram_total = vram_total
|
|
s.vram_pct = round(vram_used / vram_total * 100, 1) if vram_total else None
|
|
s.gpu_temp = max(temps) if temps else None
|
|
s.gpu_name = _gpu_name()
|
|
return True
|
|
|
|
|
|
def _nvidia(s: Sample) -> bool:
|
|
if not shutil.which("nvidia-smi"):
|
|
return False
|
|
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 False
|
|
lines = [l for l in out.splitlines() if l.strip()]
|
|
if not lines:
|
|
return False
|
|
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
|
|
s.gpu = round(busy / len(lines), 1)
|
|
s.vram_used = vram_used
|
|
s.vram_total = vram_total
|
|
s.vram_pct = round(vram_used / vram_total * 100, 1) if vram_total else None
|
|
s.gpu_temp = float(temp)
|
|
s.gpu_name = name
|
|
return True
|
|
|
|
|
|
def fill(s: Sample) -> None:
|
|
_ = _amd(s) or _nvidia(s)
|