"""Tests for GPU monitor -- mock pynvml (no real GPU required).""" from __future__ import annotations import sys import time import types from dataclasses import dataclass from unittest.mock import MagicMock, patch import pytest # --------------------------------------------------------------------------- # Helpers: build a fake pynvml module # --------------------------------------------------------------------------- @dataclass class _FakeUtilization: gpu: int = 80 memory: int = 50 @dataclass class _FakeMemInfo: total: int = 24 * 1024**3 used: int = 12 * 1024**3 free: int = 12 * 1024**3 def _make_fake_pynvml(device_count: int = 1, power_mw: int = 300_000): """Return a fake pynvml module object.""" mod = types.ModuleType("pynvml") mod.nvmlInit = MagicMock() mod.nvmlShutdown = MagicMock() mod.nvmlDeviceGetCount = MagicMock(return_value=device_count) mod.nvmlDeviceGetHandleByIndex = MagicMock(side_effect=lambda i: f"handle-{i}") mod.nvmlDeviceGetPowerUsage = MagicMock(return_value=power_mw) mod.nvmlDeviceGetUtilizationRates = MagicMock(return_value=_FakeUtilization()) mod.nvmlDeviceGetMemoryInfo = MagicMock(return_value=_FakeMemInfo()) mod.nvmlDeviceGetTemperature = MagicMock(return_value=65) mod.NVML_TEMPERATURE_GPU = 0 return mod def _snap(power=300, util=80, mem=12, temp=65, dev=0): """Shorthand to build a GpuSnapshot.""" from openjarvis.telemetry.gpu_monitor import GpuSnapshot return GpuSnapshot( power_watts=power, utilization_pct=util, memory_used_gb=mem, temperature_c=temp, device_id=dev, ) # --------------------------------------------------------------------------- # Tests: GpuHardwareSpec lookup # --------------------------------------------------------------------------- class TestGpuHardwareSpec: def test_lookup_exact_key(self): from openjarvis.telemetry.gpu_monitor import lookup_gpu_spec spec = lookup_gpu_spec("A100-SXM") assert spec is not None assert spec.tflops_fp16 == 312 assert spec.bandwidth_gb_s == 2039 assert spec.tdp_watts == 400 def test_lookup_substring_match(self): from openjarvis.telemetry.gpu_monitor import lookup_gpu_spec spec = lookup_gpu_spec("NVIDIA H100-SXM 80GB") assert spec is not None assert spec.tflops_fp16 == 990 def test_lookup_case_insensitive(self): from openjarvis.telemetry.gpu_monitor import lookup_gpu_spec spec = lookup_gpu_spec("nvidia rtx 4090") assert spec is not None assert spec.tflops_fp16 == 165 def test_lookup_unknown_returns_none(self): from openjarvis.telemetry.gpu_monitor import lookup_gpu_spec assert lookup_gpu_spec("SOME-UNKNOWN-GPU-9999") is None def test_all_specs_present(self): from openjarvis.telemetry.gpu_monitor import GPU_SPECS expected = { "B200-SXM", "A100-SXM", "A100-PCIE", "H100-SXM", "H100-PCIE", "L40S", "A10", "RTX 4090", "RTX 3090", "MI300X", "MI250X", "M4 Max", "M2 Ultra", "Arc B580", "Arc B570", "Jetson Orin NX 16GB", "Jetson Orin NX 8GB", "Jetson AGX Orin", "Snapdragon X Elite", "Snapdragon X Plus", } assert set(GPU_SPECS.keys()) == expected def test_spec_is_frozen(self): from openjarvis.telemetry.gpu_monitor import GPU_SPECS spec = GPU_SPECS["A100-SXM"] with pytest.raises(AttributeError): spec.tflops_fp16 = 999 # --------------------------------------------------------------------------- # Tests: energy integration math (trapezoidal rule) # --------------------------------------------------------------------------- class TestEnergyIntegration: def test_constant_power(self): """Constant 300W for 10 seconds = 3000 J.""" from openjarvis.telemetry.gpu_monitor import GpuMonitor snapshots = [[_snap(power=300)] for _ in range(11)] timestamps = [float(i) for i in range(11)] sample = GpuMonitor._aggregate(snapshots, timestamps, wall_duration=10.0) assert sample.energy_joules == pytest.approx(3000.0, rel=1e-6) assert sample.mean_power_watts == pytest.approx(300.0, rel=1e-6) assert sample.peak_power_watts == pytest.approx(300.0, rel=1e-6) assert sample.duration_seconds == 10.0 assert sample.num_snapshots == 11 def test_linear_ramp(self): """Linear ramp 0W..400W over 4s => 800 J (trapezoidal).""" from openjarvis.telemetry.gpu_monitor import GpuMonitor snapshots = [ [_snap(power=100.0 * i, util=50, mem=8, temp=60)] for i in range(5) ] timestamps = [float(i) for i in range(5)] sample = GpuMonitor._aggregate(snapshots, timestamps, wall_duration=4.0) assert sample.energy_joules == pytest.approx(800.0, rel=1e-6) assert sample.mean_power_watts == pytest.approx(200.0, rel=1e-6) assert sample.peak_power_watts == pytest.approx(400.0) def test_empty_snapshots(self): """No snapshots yields zeroed sample.""" from openjarvis.telemetry.gpu_monitor import GpuMonitor sample = GpuMonitor._aggregate([], [], wall_duration=5.0) assert sample.energy_joules == 0.0 assert sample.num_snapshots == 0 assert sample.duration_seconds == 5.0 def test_single_snapshot(self): """One snapshot: energy is zero (no interval).""" from openjarvis.telemetry.gpu_monitor import GpuMonitor snapshots = [[_snap(power=250, util=90, mem=16, temp=70)]] timestamps = [0.0] sample = GpuMonitor._aggregate(snapshots, timestamps, wall_duration=0.05) assert sample.energy_joules == 0.0 assert sample.num_snapshots == 1 assert sample.mean_power_watts == pytest.approx(250.0) # --------------------------------------------------------------------------- # Tests: GpuSample aggregation # --------------------------------------------------------------------------- class TestGpuSampleAggregation: def test_peak_values(self): from openjarvis.telemetry.gpu_monitor import GpuMonitor snapshots = [ [_snap(power=200, util=60, mem=10, temp=55)], [_snap(power=400, util=95, mem=20, temp=80)], [_snap(power=300, util=70, mem=15, temp=65)], ] timestamps = [0.0, 1.0, 2.0] sample = GpuMonitor._aggregate(snapshots, timestamps, wall_duration=2.0) assert sample.peak_power_watts == pytest.approx(400.0) assert sample.peak_utilization_pct == pytest.approx(95.0) assert sample.peak_memory_used_gb == pytest.approx(20.0) assert sample.peak_temperature_c == pytest.approx(80.0) def test_mean_values(self): from openjarvis.telemetry.gpu_monitor import GpuMonitor snapshots = [ [_snap(power=100, util=40, mem=8, temp=50)], [_snap(power=200, util=60, mem=12, temp=60)], [_snap(power=300, util=80, mem=16, temp=70)], ] timestamps = [0.0, 1.0, 2.0] sample = GpuMonitor._aggregate(snapshots, timestamps, wall_duration=2.0) assert sample.mean_power_watts == pytest.approx(200.0) assert sample.mean_utilization_pct == pytest.approx(60.0) assert sample.mean_memory_used_gb == pytest.approx(12.0) assert sample.mean_temperature_c == pytest.approx(60.0) # --------------------------------------------------------------------------- # Tests: Multi-GPU aggregation # --------------------------------------------------------------------------- class TestMultiGpu: def test_multi_device_power_sum(self): """Power summed across devices; util/temp averaged.""" from openjarvis.telemetry.gpu_monitor import GpuMonitor tick = [ _snap(power=200, util=80, mem=10, temp=60, dev=0), _snap(power=300, util=90, mem=12, temp=70, dev=1), ] snapshots = [tick, tick] timestamps = [0.0, 1.0] sample = GpuMonitor._aggregate(snapshots, timestamps, wall_duration=1.0) assert sample.energy_joules == pytest.approx(500.0) assert sample.mean_power_watts == pytest.approx(500.0) assert sample.mean_utilization_pct == pytest.approx(85.0) assert sample.mean_memory_used_gb == pytest.approx(22.0) assert sample.mean_temperature_c == pytest.approx(65.0) def test_multi_device_varying_power(self): """Multi-GPU with varying power over time.""" from openjarvis.telemetry.gpu_monitor import GpuMonitor snapshots = [ [ _snap(power=100, util=50, mem=8, temp=55, dev=0), _snap(power=100, util=50, mem=8, temp=55, dev=1), ], [ _snap(power=300, util=90, mem=16, temp=75, dev=0), _snap(power=300, util=90, mem=16, temp=75, dev=1), ], ] timestamps = [0.0, 2.0] sample = GpuMonitor._aggregate(snapshots, timestamps, wall_duration=2.0) # Tick powers: 200, 600 => 0.5*(200+600)*2 = 800 assert sample.energy_joules == pytest.approx(800.0) assert sample.peak_power_watts == pytest.approx(600.0) # --------------------------------------------------------------------------- # Tests: context manager flow (with mocked pynvml) # --------------------------------------------------------------------------- class TestContextManager: def test_sample_context_manager(self): """sample() starts/stops polling and populates result.""" fake_pynvml = _make_fake_pynvml(device_count=1, power_mw=300_000) import openjarvis.telemetry.gpu_monitor as mod orig_avail = mod._PYNVML_AVAILABLE orig_pynvml = getattr(mod, "pynvml", None) mod._PYNVML_AVAILABLE = True mod.pynvml = fake_pynvml try: monitor = mod.GpuMonitor(poll_interval_ms=10) monitor._initialized = True monitor._device_count = 1 monitor._handles = ["handle-0"] with monitor.sample() as result: time.sleep(0.1) assert result.duration_seconds > 0 assert result.num_snapshots > 0 assert result.mean_power_watts > 0 finally: mod._PYNVML_AVAILABLE = orig_avail if orig_pynvml is not None: mod.pynvml = orig_pynvml def test_sample_no_gpu(self): """sample() yields empty GpuSample when no GPU.""" from openjarvis.telemetry.gpu_monitor import GpuMonitor monitor = GpuMonitor.__new__(GpuMonitor) monitor._poll_interval_s = 0.05 monitor._handles = [] monitor._device_count = 0 monitor._initialized = False with monitor.sample() as result: pass assert result.num_snapshots == 0 assert result.energy_joules == 0.0 assert result.duration_seconds >= 0 # --------------------------------------------------------------------------- # Tests: available() # --------------------------------------------------------------------------- class TestAvailable: def test_available_false_when_pynvml_missing(self): """available() returns False when pynvml not importable.""" import openjarvis.telemetry.gpu_monitor as mod orig = mod._PYNVML_AVAILABLE mod._PYNVML_AVAILABLE = False try: assert mod.GpuMonitor.available() is False finally: mod._PYNVML_AVAILABLE = orig def test_available_true_with_fake_pynvml(self): """available() returns True when pynvml can init.""" fake_pynvml = _make_fake_pynvml() with patch.dict(sys.modules, {"pynvml": fake_pynvml}): import openjarvis.telemetry.gpu_monitor as mod orig = mod._PYNVML_AVAILABLE mod._PYNVML_AVAILABLE = True mod.pynvml = fake_pynvml try: assert mod.GpuMonitor.available() is True fake_pynvml.nvmlInit.assert_called() fake_pynvml.nvmlShutdown.assert_called() finally: mod._PYNVML_AVAILABLE = orig def test_available_false_when_init_fails(self): """available() returns False when nvmlInit raises.""" fake_pynvml = _make_fake_pynvml() fake_pynvml.nvmlInit.side_effect = RuntimeError("no driver") with patch.dict(sys.modules, {"pynvml": fake_pynvml}): import openjarvis.telemetry.gpu_monitor as mod orig = mod._PYNVML_AVAILABLE mod._PYNVML_AVAILABLE = True mod.pynvml = fake_pynvml try: assert mod.GpuMonitor.available() is False finally: mod._PYNVML_AVAILABLE = orig # --------------------------------------------------------------------------- # Tests: dataclass defaults # --------------------------------------------------------------------------- class TestDataclasses: def test_gpu_snapshot_defaults(self): from openjarvis.telemetry.gpu_monitor import GpuSnapshot s = GpuSnapshot( power_watts=300, utilization_pct=80, memory_used_gb=12, temperature_c=65, ) assert s.device_id == 0 def test_gpu_sample_defaults(self): from openjarvis.telemetry.gpu_monitor import GpuSample s = GpuSample() assert s.energy_joules == 0.0 assert s.num_snapshots == 0 assert s.duration_seconds == 0.0