100 lines
2.8 KiB
Python
100 lines
2.8 KiB
Python
|
import importlib.util
|
||
|
|
from pathlib import Path
|
||
|
|
import sys
|
||
|
|
import types
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
|
||
|
|
ROOT = Path(__file__).resolve().parents[2]
|
||
|
|
EXPORTER_PATH = ROOT / "services" / "monitoring" / "scripts" / "nvidia_process_exporter.py"
|
||
|
|
|
||
|
|
|
||
|
|
def load_exporter(monkeypatch):
|
||
|
|
"""Load the exporter without requiring an NVIDIA driver on the test host."""
|
||
|
|
|
||
|
|
pynvml = types.ModuleType("pynvml")
|
||
|
|
|
||
|
|
class NVMLError(Exception):
|
||
|
|
pass
|
||
|
|
|
||
|
|
class NVMLErrorNotFound(NVMLError):
|
||
|
|
pass
|
||
|
|
|
||
|
|
class NVMLErrorNotSupported(NVMLError):
|
||
|
|
pass
|
||
|
|
|
||
|
|
pynvml.NVMLError = NVMLError
|
||
|
|
pynvml.NVMLError_NotFound = NVMLErrorNotFound
|
||
|
|
pynvml.NVMLError_NotSupported = NVMLErrorNotSupported
|
||
|
|
for name in (
|
||
|
|
"nvmlDeviceGetComputeRunningProcesses_v3",
|
||
|
|
"nvmlDeviceGetCount",
|
||
|
|
"nvmlDeviceGetGraphicsRunningProcesses_v3",
|
||
|
|
"nvmlDeviceGetHandleByIndex",
|
||
|
|
"nvmlDeviceGetName",
|
||
|
|
"nvmlDeviceGetProcessUtilization",
|
||
|
|
"nvmlDeviceGetUUID",
|
||
|
|
"nvmlDeviceGetUtilizationRates",
|
||
|
|
"nvmlInit",
|
||
|
|
):
|
||
|
|
setattr(pynvml, name, lambda *args, **kwargs: None)
|
||
|
|
|
||
|
|
monkeypatch.setitem(sys.modules, "pynvml", pynvml)
|
||
|
|
spec = importlib.util.spec_from_file_location("nvidia_process_exporter_test", EXPORTER_PATH)
|
||
|
|
module = importlib.util.module_from_spec(spec)
|
||
|
|
spec.loader.exec_module(module)
|
||
|
|
return module
|
||
|
|
|
||
|
|
|
||
|
|
def test_process_sample_window_uses_nvml_microseconds(monkeypatch):
|
||
|
|
exporter = load_exporter(monkeypatch)
|
||
|
|
observed = []
|
||
|
|
monkeypatch.setattr(exporter.time, "time", lambda: 1_700_000_000.0)
|
||
|
|
monkeypatch.setattr(
|
||
|
|
exporter,
|
||
|
|
"nvmlDeviceGetProcessUtilization",
|
||
|
|
lambda handle, since: observed.append(since) or [],
|
||
|
|
)
|
||
|
|
|
||
|
|
samples, supported = exporter.process_utilization_samples(object())
|
||
|
|
|
||
|
|
assert samples == {}
|
||
|
|
assert supported == 1
|
||
|
|
assert observed == [1_700_000_000_000_000 - 30_000_000]
|
||
|
|
|
||
|
|
|
||
|
|
def test_namespace_attribution_scales_to_current_device_total(monkeypatch):
|
||
|
|
exporter = load_exporter(monkeypatch)
|
||
|
|
|
||
|
|
result = exporter.reconcile_namespace_utilization(
|
||
|
|
{"game-stream": 20, "hermes": 10},
|
||
|
|
device_util=3,
|
||
|
|
)
|
||
|
|
|
||
|
|
assert sum(result.values()) == pytest.approx(3)
|
||
|
|
assert result["game-stream"] == pytest.approx(2)
|
||
|
|
assert result["hermes"] == pytest.approx(1)
|
||
|
|
|
||
|
|
|
||
|
|
def test_namespace_attribution_assigns_unexplained_compute_to_host(monkeypatch):
|
||
|
|
exporter = load_exporter(monkeypatch)
|
||
|
|
|
||
|
|
result = exporter.reconcile_namespace_utilization(
|
||
|
|
{"hermes": 1},
|
||
|
|
device_util=3,
|
||
|
|
)
|
||
|
|
|
||
|
|
assert result == {"hermes": 1, "host": 2}
|
||
|
|
|
||
|
|
|
||
|
|
def test_zero_device_utilization_clears_stale_process_samples(monkeypatch):
|
||
|
|
exporter = load_exporter(monkeypatch)
|
||
|
|
|
||
|
|
result = exporter.reconcile_namespace_utilization(
|
||
|
|
{"hermes": 40, "game-stream": 5},
|
||
|
|
device_util=0,
|
||
|
|
)
|
||
|
|
|
||
|
|
assert result == {"hermes": 0, "game-stream": 0}
|