"""Per-job CPU/GPU usage sampling for the monitor: best-effort, never fatal.
The dashboard's usage columns are cosmetic telemetry, and this module treats
them that way everywhere: every function returns ``None`` (or an empty
mapping) instead of raising, a dead PID or unreachable host is a dash on
screen, and nothing here is ever allowed to take the monitor down.
The seam mirrors backend health probing: backends that can measure their
jobs expose the duck-typed hooks
``sample_cpu_percents(handles) -> list[float | None]``
``sample_gpu_percents(handles) -> list[float | None]``
taking the ledger's serialized handle dicts (``asdict(Handle)``) in order
and answering one reading (or ``None``) per handle, never raising. For CPU
the ``local`` backend answers from ``/proc`` deltas (true instantaneous
usage of the job's whole process subtree, via ``LocalCpuTracker`` below) and
``ssh_exec`` from one batched remote ``ps`` (lifetime-average ``pcpu``,
``ps`` semantics -- honest enough for long-running science jobs). For GPU
both answer from ``nvidia-smi pmon`` (per-process SM utilisation, summed
over the job's subtree); a machine without ``nvidia-smi`` -- or a backend
without the hook -- simply contributes no readings, so a CPU-only rig shows
a dash, never an error. ``UsageSampler`` groups the in-flight ledger by
backend, applies a per-backend TTL (local sampling is free and can run
every tick; a remote sample is an SSH round-trip and is cached, the
``_HealthCache`` discipline), and folds everything into the
``{job_key: percent}`` mappings ``build_snapshot`` takes as ``cpu_by_job``
and ``gpu_by_job``.
The ``/proc``/``pmon`` parsing and subtree arithmetic live here as pure,
injectable functions so both the backends and the tests use one
implementation.
"""
from __future__ import annotations
import os
import time
from pathlib import Path
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from collections.abc import Callable, Mapping, Sequence
from matmelt.core.state import InFlightJob
#: Fallback clock ticks per second when ``sysconf`` cannot answer.
_DEFAULT_HZ = 100.0
#: How long one remote CPU sample stays fresh (seconds). Remote sampling is
#: an SSH round-trip; local sampling passes ``0.0`` (every tick).
REMOTE_CPU_TTL_S = 10.0
[docs]
def clock_ticks_per_s() -> float:
"""The kernel's ``USER_HZ``, the unit ``/proc`` CPU times are counted in.
Returns
-------
float
Clock ticks per second, or ``_DEFAULT_HZ`` (100) where ``sysconf``
cannot answer -- a wrong-but-sane divisor beats no reading at all for
cosmetic telemetry.
"""
try:
hz = float(os.sysconf("SC_CLK_TCK"))
except (ValueError, OSError, AttributeError):
return _DEFAULT_HZ
return hz if hz > 0 else _DEFAULT_HZ
[docs]
def parse_proc_stat(text: str) -> tuple[int, int, float] | None:
"""Parse one ``/proc/<pid>/stat`` line.
The comm field is parenthesised and may itself contain spaces or
parentheses (``(Web Content)``), so fields are taken after the LAST ``)``
-- the only parse that is correct for every process name.
Parameters
----------
text
The full contents of a ``/proc/<pid>/stat`` file.
Returns
-------
tuple or None
``(pid, ppid, cpu_ticks)``, where ``cpu_ticks`` is ``utime + stime``
(user + system jiffies of this process alone, excluding children).
Anything malformed is ``None``, never an exception.
Examples
--------
>>> parse_proc_stat("42 (my proc) S 7 42 42 0 -1 0 0 0 0 0 30 12 0 0")
(42, 7, 42.0)
>>> parse_proc_stat("not a stat line") is None
True
"""
head, sep, tail = text.rpartition(")")
if not sep:
return None
try:
pid = int(head.split("(", 1)[0])
fields = tail.split()
# After the comm field: state ppid pgrp ... utime(idx 11) stime(idx 12).
return pid, int(fields[1]), float(int(fields[11]) + int(fields[12]))
except (ValueError, IndexError):
return None
[docs]
def scan_proc(proc_root: Path = Path("/proc")) -> dict[int, tuple[int, float]]:
"""Read every process's parent and CPU time in one pass.
Parameters
----------
proc_root
Where to scan. Injectable so tests can build a fake ``/proc`` in
``tmp_path``; the default is the real one.
Returns
-------
dict
``{pid: (ppid, cpu_ticks)}``. Unreadable or vanished entries
(processes exit mid-scan constantly) are skipped silently, and an
absent ``proc_root`` yields an empty table.
"""
table: dict[int, tuple[int, float]] = {}
try:
entries = list(proc_root.iterdir())
except OSError:
return table
for entry in entries:
if not entry.name.isdigit():
continue
try:
parsed = parse_proc_stat((entry / "stat").read_text(encoding="ascii", errors="replace"))
except OSError:
continue
if parsed is not None:
table[parsed[0]] = (parsed[1], parsed[2])
return table
[docs]
def subtree_pids(table: Mapping[int, tuple[int, float]], root: int) -> set[int] | None:
"""Collect ``root`` and every process descended from it.
A job is a process *tree* -- the wrapper shell plus the runner plus
whatever ``multiprocessing`` spawned -- so per-job accounting must cover
the whole subtree, not just the pid the backend recorded.
Parameters
----------
table
``{pid: (ppid, value)}``, as returned by ``scan_proc``.
root
The pid to start from.
Returns
-------
set of int or None
``root`` plus all descendants, or ``None`` when ``root`` is not in the
table (the process is gone). A visited set guards traversal against a
corrupt table; real ppid chains cannot cycle.
Examples
--------
>>> table = {1: (0, 0.0), 2: (1, 0.0), 3: (2, 0.0), 9: (0, 0.0)}
>>> sorted(subtree_pids(table, 1))
[1, 2, 3]
>>> subtree_pids(table, 404) is None
True
"""
if root not in table:
return None
children: dict[int, list[int]] = {}
for pid, (ppid, _) in table.items():
children.setdefault(ppid, []).append(pid)
stack, seen = [root], {root}
while stack:
pid = stack.pop()
for child in children.get(pid, ()):
if child not in seen:
seen.add(child)
stack.append(child)
return seen
[docs]
def subtree_total(table: Mapping[int, tuple[int, float]], root: int) -> float | None:
"""Sum the per-process values over ``root`` and all its descendants.
Parameters
----------
table
``{pid: (ppid, value)}`` -- CPU ticks locally, ``ps`` ``pcpu``
remotely.
root
The pid whose subtree is summed.
Returns
-------
float or None
The subtree total, or ``None`` when ``root`` is not in the table (the
process is gone). See ``subtree_pids`` for the traversal contract.
"""
pids = subtree_pids(table, root)
if pids is None:
return None
return sum(table[pid][1] for pid in pids)
[docs]
def parse_pmon(text: str) -> dict[int, float]:
"""Parse ``nvidia-smi pmon -c 1 -s u`` output into per-process SM usage.
Parameters
----------
text
The command's stdout. Its table is ``gpu pid type sm mem enc dec
command``, with ``#`` header lines and ``-`` for idle or unavailable
cells.
Returns
-------
dict
``{pid: sm_percent}``. A ``-`` cell reads as ``0.0`` -- a listed
process doing no GPU work *is* at 0% -- and a pid appearing on several
GPUs is summed. Garbage lines are skipped; a machine without the tool
never reaches this parser, because the caller answers ``None`` instead.
Examples
--------
>>> parse_pmon("# gpu pid type sm mem enc dec command\\n0 1234 C 45 12 - - python3")
{1234: 45.0}
"""
usage: dict[int, float] = {}
for line in text.splitlines():
parts = line.split()
if not parts or parts[0].startswith("#"):
continue
try:
pid = int(parts[1])
sm = 0.0 if parts[3] == "-" else float(parts[3])
except (ValueError, IndexError):
continue
usage[pid] = usage.get(pid, 0.0) + sm
return usage
[docs]
def gpu_percent_for_roots(
process_table: Mapping[int, tuple[int, float]],
pmon: Mapping[int, float],
roots: Sequence[int | None],
) -> list[float | None]:
"""Per-root GPU percentages: pmon SM readings summed over each subtree.
Parameters
----------
process_table
``{pid: (ppid, value)}`` for the machine the jobs run on, used only for
its parent links.
pmon
``{pid: sm_percent}`` from ``parse_pmon``.
roots
One root pid per job, in the caller's order. ``None`` marks a job whose
pid is not known.
Returns
-------
list
One reading per root, in the same order. ``None`` roots and roots
absent from the process table (gone) stay ``None``; a live subtree with
no pmon entries is an honest ``0.0`` -- running, but not on the GPU.
"""
readings: list[float | None] = []
for root in roots:
if root is None:
readings.append(None)
continue
pids = subtree_pids(process_table, root)
if pids is None:
readings.append(None)
continue
readings.append(sum(pmon.get(pid, 0.0) for pid in pids))
return readings
[docs]
class LocalCpuTracker:
"""Delta-based CPU usage for local process subtrees, one ``/proc`` scan per call.
Readings are a percentage of *one core*, so a 4-worker job legitimately
reads ~400%. State for vanished pids is dropped, so a recycled PID can
never inherit a stale baseline.
Parameters
----------
clock
Monotonic time source, injectable for tests.
proc_root
Where to scan, injectable for tests.
ticks_per_s
Override for the kernel's ``USER_HZ``. ``None`` asks
``clock_ticks_per_s``.
"""
def __init__(
self,
*,
clock: Callable[[], float] = time.monotonic,
proc_root: Path = Path("/proc"),
ticks_per_s: float | None = None,
) -> None:
self._clock = clock
self._proc_root = proc_root
self._hz = clock_ticks_per_s() if ticks_per_s is None else ticks_per_s
self._prev: dict[int, tuple[float, float]] = {} # pid -> (ticks, at)
[docs]
def sample(self, root_pids: Sequence[int]) -> dict[int, float | None]:
"""Take one reading per requested pid, against the previous call.
Parameters
----------
root_pids
The job root pids to measure.
Returns
-------
dict
``{pid: percent or None}``. The percentage is the subtree's tick
delta since this tracker last saw that pid, over the wall-clock
delta. The first sighting of a pid has no delta and answers
``None`` -- an honest "measuring", never a fabricated ``0.0`` --
as does a pid that has exited.
"""
now = self._clock()
table = scan_proc(self._proc_root)
readings: dict[int, float | None] = {}
for pid in root_pids:
ticks = subtree_total(table, pid)
if ticks is None:
readings[pid] = None
self._prev.pop(pid, None)
continue
previous = self._prev.get(pid)
self._prev[pid] = (ticks, now)
if previous is None or now <= previous[1]:
readings[pid] = None # first sighting (or a stuck clock): no delta yet
continue
elapsed = now - previous[1]
readings[pid] = max(0.0, (ticks - previous[0]) / self._hz / elapsed * 100.0)
return readings
[docs]
class UsageSampler:
"""The monitor's per-tick usage sweep over the in-flight ledger.
Groups jobs by backend, asks each backend exposing the sampling hook for
readings, and folds the answers into ``{job_key: percent}``. Jobs without a
reading are simply absent, so the renderer's ``.get`` shows a dash. A hook
that raises despite its contract is treated as "no readings" for that sweep
-- telemetry must never take the dashboard down.
Parameters
----------
backends
``{backend_name: backend}``. A backend without the hook contributes
nothing, which is how a CPU-only rig shows dashes rather than errors.
hook_name
Which duck-typed hook to call: ``sample_cpu_percents`` by default.
Construct a second sampler with ``sample_gpu_percents`` for the GPU
column.
ttl_by_backend
Per-backend cache lifetime in seconds. The caller passes ``0.0`` for
backends whose sampling is free (local, every tick).
default_ttl_s
Lifetime for backends not named in ``ttl_by_backend``, so an SSH
round-trip happens at most every ``REMOTE_CPU_TTL_S`` regardless of the
dashboard's refresh interval.
"""
def __init__(
self,
backends: Mapping[str, object],
*,
hook_name: str = "sample_cpu_percents",
ttl_by_backend: Mapping[str, float] | None = None,
default_ttl_s: float = REMOTE_CPU_TTL_S,
) -> None:
self._backends = dict(backends)
self._hook_name = hook_name
self._ttls = dict(ttl_by_backend or {})
self._default_ttl_s = default_ttl_s
self._cached: dict[str, tuple[float, dict[str, float]]] = {}
[docs]
def sample(self, jobs: Sequence[InFlightJob], now: float) -> dict[str, float]:
"""Read usage for the given in-flight jobs, throttled per backend.
Parameters
----------
jobs
The in-flight ledger rows to measure, in any order.
now
The current monotonic time, against which each backend's cache TTL
is judged. The caller owns the clock so a tick measures once.
Returns
-------
dict
``{job_key: percent}``, merged across backends. A backend still
inside its TTL contributes its cached readings unchanged.
"""
by_backend: dict[str, list[InFlightJob]] = {}
for job in jobs:
by_backend.setdefault(job.backend_name, []).append(job)
merged: dict[str, float] = {}
for backend_name, group in by_backend.items():
cached = self._cached.get(backend_name)
ttl = self._ttls.get(backend_name, self._default_ttl_s)
if cached is not None and (now - cached[0]) <= ttl:
merged.update(cached[1])
continue
merged.update(self._sweep(backend_name, group, now))
return merged
def _sweep(self, backend_name: str, group: list[InFlightJob], now: float) -> dict[str, float]:
"""Probe one backend's jobs for real, then cache the answer.
Parameters
----------
backend_name
Which backend to probe.
group
That backend's in-flight jobs, whose serialized handles are passed
to the hook in order.
now
Timestamp stored with the cache entry.
Returns
-------
dict
``{job_key: percent}`` for this backend, empty when it has no hook,
the hook raised, or it returned fewer readings than handles.
"""
readings: dict[str, float] = {}
hook = getattr(self._backends.get(backend_name), self._hook_name, None)
if callable(hook):
handles: list[dict[str, Any]] = [job.handle for job in group]
try:
percents = list(hook(handles))
except Exception: # cosmetic telemetry: a broken hook must never be fatal
percents = []
for job, percent in zip(group, percents, strict=False):
if percent is not None:
readings[job.job_key] = float(percent)
self._cached[backend_name] = (now, readings)
return readings