111 lines
3.6 KiB
Python
111 lines
3.6 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]:
|
|
"""Map PID to GPU memory used (MiB) for NVIDIA compute processes.
|
|
|
|
Runs `nvidia-smi --query-compute-apps` at most once per 10 seconds
|
|
(the probe result is cached). Returns an empty mapping when nvidia-smi
|
|
is missing, which is the case on AMD machines.
|
|
|
|
Returns:
|
|
A pid to used-memory-in-MiB mapping.
|
|
"""
|
|
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]]:
|
|
"""One pass over all processes collecting cpu, memory, IO rate, GPU.
|
|
|
|
Processes whose parent is swapper/kthreadd (ppid 0/2) are skipped.
|
|
Per-process IO rates are byte deltas between successive calls divided
|
|
by elapsed time; previous readings are pruned when a process exits.
|
|
GPU memory comes from _gpu_per_proc(). Entries that die mid-iteration
|
|
are dropped, and per-process access errors are tolerated.
|
|
|
|
Returns:
|
|
A list of per-process dicts (pid, name, user, cpu, rss, mem_pct,
|
|
io_read, io_write, gpu), one per live process.
|
|
"""
|
|
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
|