149 lines
4.7 KiB
Python
149 lines
4.7 KiB
Python
import glob
|
|
import re
|
|
import shutil
|
|
from typing import Any, cast
|
|
|
|
from app.sample import Sample
|
|
from app.utils import sysfs
|
|
from app.utils.gpu import shorten
|
|
from app.utils.subprocess import run, run_json
|
|
|
|
_name_cache: str | None = None
|
|
|
|
_LSPCI_QUOTED = re.compile(r'"([^"]*)"')
|
|
|
|
|
|
def _gpu_name() -> str:
|
|
"""Resolve the display GPU name, cached for the process lifetime.
|
|
|
|
Runs `lspci -mm` once (stable machine-readable format, one line per
|
|
device with double-quoted fields: slot, class, vendor, device, ...)
|
|
and takes the vendor + device of the first VGA / 3D-controller line,
|
|
shortened with shorten(). Falls back to "GPU" if lspci is missing or
|
|
no matching device line is found.
|
|
|
|
Returns:
|
|
The display name to put on the overview card and Sample.
|
|
"""
|
|
global _name_cache
|
|
if _name_cache is None:
|
|
_name_cache = "GPU"
|
|
if shutil.which("lspci"):
|
|
rc, out, _err = run(["lspci", "-mm"], timeout=5)
|
|
if rc == 0:
|
|
for line in out.splitlines():
|
|
f = _LSPCI_QUOTED.findall(line)
|
|
if len(f) >= 3 and ("VGA" in f[0] or "3D controller" in f[0]):
|
|
_name_cache = shorten(f"{f[1]} {f[2]}")
|
|
break
|
|
return _name_cache
|
|
|
|
|
|
def _amd(s: Sample) -> bool:
|
|
"""Fill GPU fields from AMD sysfs (amdgpu driver).
|
|
|
|
Reads gpu_busy_percent, mem_info_vram_used/total, and hwmon
|
|
temp1_input (millidegrees) from each /sys/class/drm/card*/device.
|
|
Busy percent is averaged across cards, VRAM summed, temperature is the
|
|
hottest card. The display name comes from _gpu_name().
|
|
|
|
Args:
|
|
s: sample to fill.
|
|
|
|
Returns:
|
|
True if at least one card reported a busy percent, else False
|
|
(leaving s untouched).
|
|
"""
|
|
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]
|
|
busy_sum += sysfs.read_int(busy_path) or 0
|
|
count += 1
|
|
vram_used += sysfs.read_int(f"{dev}/mem_info_vram_used") or 0
|
|
vram_total += sysfs.read_int(f"{dev}/mem_info_vram_total") or 0
|
|
for hwmon in glob.glob(f"{dev}/hwmon/hwmon*"):
|
|
t = sysfs.read_int(f"{hwmon}/temp1_input")
|
|
if t is not None:
|
|
temps.append(t / 1000.0)
|
|
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:
|
|
"""Fill GPU fields by querying nvidia-smi.
|
|
|
|
Runs `nvidia-smi --query-gpu=... --format=json` (5 s timeout) and
|
|
parses the JSON array (keys are underscored, values strings): busy
|
|
percent averaged across GPUs, VRAM summed (MiB converted to bytes),
|
|
temperature the hottest GPU, name from the first GPU. JSON keeps
|
|
names with commas intact, which the old CSV format split on.
|
|
|
|
Args:
|
|
s: sample to fill.
|
|
|
|
Returns:
|
|
True if nvidia-smi exists and returned usable data, else False.
|
|
"""
|
|
if not shutil.which("nvidia-smi"):
|
|
return False
|
|
data, _err = run_json(
|
|
[
|
|
"nvidia-smi",
|
|
"--query-gpu=utilization.gpu,memory.used,memory.total,temperature.gpu,name",
|
|
"--format=json",
|
|
],
|
|
timeout=5,
|
|
)
|
|
if not isinstance(data, list) or not data:
|
|
return False
|
|
rows = cast("list[dict[str, Any]]", data)
|
|
busy = used = total = 0
|
|
temp = 0
|
|
count = 0
|
|
for e in rows:
|
|
try:
|
|
busy += int(e["utilization_gpu"])
|
|
used += int(e["memory_used"])
|
|
total += int(e["memory_total"])
|
|
temp = max(temp, int(e["temperature_gpu"]))
|
|
count += 1
|
|
except (ValueError, TypeError, KeyError):
|
|
continue
|
|
if count == 0:
|
|
return False
|
|
name = str(rows[0].get("name") or "").strip() or "GPU"
|
|
vram_used = used * 1024 * 1024
|
|
vram_total = total * 1024 * 1024
|
|
s.gpu = round(busy / 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 = float(temp)
|
|
s.gpu_name = name
|
|
return True
|
|
|
|
|
|
def fill(s: Sample) -> None:
|
|
"""Fill the gpu / vram / gpu_temp / gpu_name fields of a Sample.
|
|
|
|
Tries the AMD sysfs path first (no subprocess), then nvidia-smi.
|
|
If neither applies, the fields keep their Sample defaults.
|
|
|
|
Args:
|
|
s: sample to fill.
|
|
"""
|
|
_ = _amd(s) or _nvidia(s)
|