Files
stemify/desktop/stemify/gpu_monitor.py
T

78 lines
2.8 KiB
Python

"""How busy the GPU is with work that isn't ours.
Two signals, both read from the kernel (no GPU calls):
- DRM fdinfo: per-client gfx/compute engine time. Games and other apps show
up here; our ROCm worker runs through KFD and our own UI is excluded by pid.
- amdgpu gpu_busy_percent: catches other ROCm/compute users that fdinfo
misses, but also counts our worker, so it only means something while our
worker is idle.
"""
import os
import re
_KV = re.compile(r"^(drm-[\w-]+):\s+(\S+)", re.M)
_ENGINES = ("drm-engine-gfx", "drm-engine-compute")
class GpuMonitor:
def __init__(self):
self.pdev = None # "0000:03:00.0"; None counts every GPU
self.exclude: set[int] = {os.getpid()}
self._last: dict[str, int] = {} # drm-client-id -> engine ns
self._last_t = None
def set_device(self, pci):
if pci != self.pdev:
self.pdev = pci
self._last, self._last_t = {}, None
def _busy_path(self):
if self.pdev:
p = f"/sys/bus/pci/devices/{self.pdev}/gpu_busy_percent"
return p if os.path.exists(p) else None
return None
def _clients(self):
out = {}
for pid in os.listdir("/proc"):
if not pid.isdigit() or int(pid) in self.exclude:
continue
try:
fds = os.listdir(f"/proc/{pid}/fd")
except OSError:
continue
for fd in fds:
try:
if not os.readlink(f"/proc/{pid}/fd/{fd}").startswith("/dev/dri/"):
continue
with open(f"/proc/{pid}/fdinfo/{fd}") as f:
kv = dict(_KV.findall(f.read()))
except OSError:
continue
cid = kv.get("drm-client-id")
if not cid or (self.pdev and kv.get("drm-pdev") != self.pdev):
continue
# one client can sit behind several dup'd fds; count it once
out[cid] = sum(int(kv[k]) for k in _ENGINES if kv.get(k, "").isdigit())
return out
def sample(self, now):
"""(other_fraction, busy_percent). other_fraction is engine time used by
other processes over the last interval, as a fraction of wall time."""
clients = self._clients()
other = 0.0
if self._last_t is not None and now > self._last_t:
delta = sum(ns - self._last[c] for c, ns in clients.items() if c in self._last and ns >= self._last[c])
other = delta / ((now - self._last_t) * 1e9)
self._last, self._last_t = clients, now
pct = -1
path = self._busy_path()
if path:
try:
with open(path) as f:
pct = int(f.read().strip())
except (OSError, ValueError):
pass
return other, pct