90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
import shutil
|
|
import subprocess
|
|
import time
|
|
from typing import Any
|
|
|
|
import psutil
|
|
|
|
_prev_io: dict[int, tuple[float, float, float]] = {}
|
|
_gpu_procs: dict[int, int] | None = None
|
|
_gpu_probe_t = 0.0
|
|
|
|
|
|
def _gpu_per_proc() -> dict[int, int]:
|
|
global _gpu_procs, _gpu_probe_t
|
|
if not shutil.which("nvidia-smi"):
|
|
return {}
|
|
if _gpu_procs is not None and time.monotonic() - _gpu_probe_t < 10:
|
|
return _gpu_procs
|
|
_gpu_probe_t = time.monotonic()
|
|
_gpu_procs = {}
|
|
try:
|
|
out = subprocess.run(
|
|
[
|
|
"nvidia-smi",
|
|
"--query-compute-apps=pid,used_memory",
|
|
"--format=csv,noheader,nounits",
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=5,
|
|
check=False
|
|
).stdout
|
|
for line in out.splitlines():
|
|
parts = [p.strip() for p in line.split(",")]
|
|
if len(parts) >= 2:
|
|
try:
|
|
_gpu_procs[int(parts[0])] = int(parts[1])
|
|
except ValueError:
|
|
continue
|
|
except (OSError, subprocess.SubprocessError):
|
|
pass
|
|
return _gpu_procs
|
|
|
|
|
|
def sample() -> list[dict[str, Any]]:
|
|
now = time.monotonic()
|
|
mem_total = psutil.virtual_memory().total
|
|
gpu = _gpu_per_proc()
|
|
out: list[dict[str, Any]] = []
|
|
alive: set[int] = set()
|
|
for p in psutil.process_iter():
|
|
try:
|
|
with p.oneshot():
|
|
if p.ppid() in (0, 2):
|
|
continue
|
|
cpu = p.cpu_percent(None)
|
|
mem = p.memory_info()
|
|
name = p.name()
|
|
user = p.username()
|
|
try:
|
|
io = p.io_counters()
|
|
except (psutil.AccessDenied, psutil.NoSuchProcess, OSError):
|
|
io = None
|
|
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
|
|
continue
|
|
pid = p.pid
|
|
alive.add(pid)
|
|
entry = {
|
|
"pid": pid,
|
|
"name": name,
|
|
"user": user,
|
|
"cpu": cpu,
|
|
"rss": mem.rss,
|
|
"mem_pct": (mem.rss / mem_total * 100) if mem_total else 0.0,
|
|
"io_read": 0.0,
|
|
"io_write": 0.0,
|
|
"gpu": gpu.get(pid),
|
|
}
|
|
if io is not None:
|
|
prev = _prev_io.get(pid)
|
|
if prev is not None and now > prev[0]:
|
|
dt = now - prev[0]
|
|
entry["io_read"] = max(0.0, (io.read_bytes - prev[1]) / dt)
|
|
entry["io_write"] = max(0.0, (io.write_bytes - prev[2]) / dt)
|
|
_prev_io[pid] = (now, io.read_bytes, io.write_bytes)
|
|
out.append(entry)
|
|
for pid in list(_prev_io):
|
|
if pid not in alive:
|
|
del _prev_io[pid]
|
|
return out
|