light-heart-labs--ods
wehub-resource-sync
404 行
16 KiB
Python
404 行
16 KiB
Python
"""Tests for gpu.py — tier classification, nvidia-smi parsing, Apple Silicon detection."""
|
||
|
||
import subprocess
|
||
from unittest.mock import MagicMock
|
||
|
||
import pytest
|
||
|
||
from gpu import (
|
||
get_gpu_tier, get_gpu_info_nvidia, get_gpu_info_apple,
|
||
get_gpu_info_amd, get_gpu_info, run_command,
|
||
)
|
||
|
||
|
||
# --- get_gpu_tier (pure function, no I/O) ---
|
||
|
||
|
||
class TestGetGpuTierDiscrete:
|
||
"""Discrete GPU tier boundaries."""
|
||
|
||
@pytest.mark.parametrize("vram_gb,expected", [
|
||
(4, "Minimal"),
|
||
(7.9, "Minimal"),
|
||
(8, "Entry"),
|
||
(15.9, "Entry"),
|
||
(16, "Standard"),
|
||
(23.9, "Standard"),
|
||
(24, "Prosumer"),
|
||
(79.9, "Prosumer"),
|
||
(80, "Professional"),
|
||
(128, "Professional"),
|
||
])
|
||
def test_tiers(self, vram_gb, expected):
|
||
assert get_gpu_tier(vram_gb) == expected
|
||
|
||
|
||
class TestGetGpuTierUnified:
|
||
"""Strix Halo (unified memory) tier boundaries."""
|
||
|
||
@pytest.mark.parametrize("vram_gb,expected", [
|
||
(64, "Strix Halo Compact"),
|
||
(89.9, "Strix Halo Compact"),
|
||
(90, "Strix Halo 90+"),
|
||
(96, "Strix Halo 90+"),
|
||
(128, "Strix Halo 90+"),
|
||
])
|
||
def test_tiers(self, vram_gb, expected):
|
||
assert get_gpu_tier(vram_gb, memory_type="unified") == expected
|
||
|
||
|
||
# --- get_gpu_info_nvidia (mock subprocess) ---
|
||
|
||
|
||
class TestGetGpuInfoNvidia:
|
||
|
||
def test_parses_valid_output(self, monkeypatch):
|
||
csv = "NVIDIA GeForce RTX 4090, 2048, 24564, 35, 62, 285.5"
|
||
monkeypatch.setattr("gpu.run_command", lambda cmd, **kw: (True, csv))
|
||
|
||
info = get_gpu_info_nvidia()
|
||
assert info is not None
|
||
assert info.name == "NVIDIA GeForce RTX 4090"
|
||
assert info.memory_used_mb == 2048
|
||
assert info.memory_total_mb == 24564
|
||
assert info.utilization_percent == 35
|
||
assert info.temperature_c == 62
|
||
assert info.power_w == 285.5
|
||
assert info.gpu_backend == "nvidia"
|
||
|
||
def test_handles_na_power(self, monkeypatch):
|
||
csv = "Tesla T4, 1024, 16384, 10, 45, [N/A]"
|
||
monkeypatch.setattr("gpu.run_command", lambda cmd, **kw: (True, csv))
|
||
|
||
info = get_gpu_info_nvidia()
|
||
assert info is not None
|
||
assert info.power_w is None
|
||
|
||
def test_returns_none_on_command_failure(self, monkeypatch):
|
||
monkeypatch.setattr("gpu.run_command", lambda cmd, **kw: (False, ""))
|
||
|
||
assert get_gpu_info_nvidia() is None
|
||
|
||
def test_returns_none_on_empty_output(self, monkeypatch):
|
||
monkeypatch.setattr("gpu.run_command", lambda cmd, **kw: (True, ""))
|
||
|
||
assert get_gpu_info_nvidia() is None
|
||
|
||
def test_returns_none_on_malformed_output(self, monkeypatch):
|
||
monkeypatch.setattr("gpu.run_command", lambda cmd, **kw: (True, "garbage"))
|
||
|
||
assert get_gpu_info_nvidia() is None
|
||
|
||
def test_multi_gpu_aggregation(self, monkeypatch):
|
||
csv = (
|
||
"NVIDIA GeForce RTX 4090, 2048, 24564, 35, 62, 285.5\n"
|
||
"NVIDIA GeForce RTX 4090, 4096, 24564, 50, 70, 300.0"
|
||
)
|
||
monkeypatch.setattr("gpu.run_command", lambda cmd, **kw: (True, csv))
|
||
|
||
info = get_gpu_info_nvidia()
|
||
assert info is not None
|
||
assert "× 2" in info.name
|
||
assert info.memory_used_mb == 2048 + 4096
|
||
assert info.memory_total_mb == 24564 * 2
|
||
|
||
def test_unified_memory_fallback_for_na_memory(self, monkeypatch):
|
||
# GB10 / GB200 unified-memory: nvidia-smi reports [N/A] for memory.
|
||
csv = "NVIDIA GB10, [N/A], [N/A], 6, 43, 8.2"
|
||
monkeypatch.setattr("gpu.run_command", lambda cmd, **kw: (True, csv))
|
||
monkeypatch.setattr("gpu._read_meminfo_mb", lambda: (12000, 124000))
|
||
|
||
info = get_gpu_info_nvidia()
|
||
assert info is not None
|
||
assert info.name == "NVIDIA GB10"
|
||
assert info.memory_used_mb == 12000
|
||
assert info.memory_total_mb == 124000
|
||
assert info.memory_type == "unified"
|
||
assert info.utilization_percent == 6
|
||
assert info.temperature_c == 43
|
||
assert info.power_w == 8.2
|
||
|
||
def test_unified_memory_fallback_skips_when_meminfo_unavailable(self, monkeypatch):
|
||
csv = "NVIDIA GB10, [N/A], [N/A], 6, 43, 8.2"
|
||
monkeypatch.setattr("gpu.run_command", lambda cmd, **kw: (True, csv))
|
||
monkeypatch.setattr("gpu._read_meminfo_mb", lambda: None)
|
||
|
||
assert get_gpu_info_nvidia() is None
|
||
|
||
|
||
# --- get_gpu_info_apple (mock subprocess) ---
|
||
|
||
|
||
class TestGetGpuInfoApple:
|
||
|
||
def test_returns_none_on_non_darwin(self, monkeypatch):
|
||
monkeypatch.setattr("gpu.platform.system", lambda: "Linux")
|
||
assert get_gpu_info_apple() is None
|
||
|
||
def test_parses_apple_silicon(self, monkeypatch):
|
||
monkeypatch.setattr("gpu.platform.system", lambda: "Darwin")
|
||
|
||
def mock_run_command(cmd, **kw):
|
||
if "machdep.cpu.brand_string" in cmd:
|
||
return True, "Apple M4 Max"
|
||
if "hw.memsize" in cmd:
|
||
return True, str(64 * 1024**3) # 64 GB
|
||
if cmd == ["vm_stat"]:
|
||
return True, (
|
||
"Mach Virtual Memory Statistics: (page size of 16384 bytes)\n"
|
||
"Pages active: 500000.\n"
|
||
"Pages wired down: 300000.\n"
|
||
"Pages occupied by compressor: 100000.\n"
|
||
)
|
||
return False, ""
|
||
|
||
monkeypatch.setattr("gpu.run_command", mock_run_command)
|
||
|
||
info = get_gpu_info_apple()
|
||
assert info is not None
|
||
assert info.name == "Apple M4 Max"
|
||
assert info.memory_total_mb == 64 * 1024
|
||
assert info.gpu_backend == "apple"
|
||
assert info.memory_type == "unified"
|
||
assert info.memory_used_mb > 0
|
||
|
||
def test_returns_none_when_sysctl_fails(self, monkeypatch):
|
||
monkeypatch.setattr("gpu.platform.system", lambda: "Darwin")
|
||
monkeypatch.setattr("gpu.run_command", lambda cmd, **kw: (False, ""))
|
||
assert get_gpu_info_apple() is None
|
||
|
||
|
||
# --- run_command ---
|
||
|
||
|
||
class TestRunCommand:
|
||
|
||
def test_returns_success_and_output(self, monkeypatch):
|
||
mock_result = MagicMock()
|
||
mock_result.returncode = 0
|
||
mock_result.stdout = " hello world "
|
||
monkeypatch.setattr("gpu.subprocess.run", lambda *a, **kw: mock_result)
|
||
|
||
ok, output = run_command(["echo", "hello"])
|
||
assert ok is True
|
||
assert output == "hello world"
|
||
|
||
def test_returns_false_on_timeout(self, monkeypatch):
|
||
def raise_timeout(*a, **kw):
|
||
raise subprocess.TimeoutExpired(cmd=["slow"], timeout=5)
|
||
monkeypatch.setattr("gpu.subprocess.run", raise_timeout)
|
||
|
||
ok, output = run_command(["slow"], timeout=5)
|
||
assert ok is False
|
||
assert output == "timeout"
|
||
|
||
def test_returns_false_on_file_not_found(self, monkeypatch):
|
||
def raise_fnf(*a, **kw):
|
||
raise FileNotFoundError("No such file: 'missing'")
|
||
monkeypatch.setattr("gpu.subprocess.run", raise_fnf)
|
||
|
||
ok, output = run_command(["missing"])
|
||
assert ok is False
|
||
assert "No such file" in output
|
||
|
||
|
||
# --- get_gpu_info_amd ---
|
||
|
||
|
||
class TestGetGpuInfoAmd:
|
||
|
||
def test_parses_discrete_gpu(self, monkeypatch):
|
||
"""Discrete GPU: 16 GB VRAM, small GTT."""
|
||
monkeypatch.setattr("gpu._find_amd_gpu_sysfs", lambda: "/sys/class/drm/card0/device")
|
||
monkeypatch.setattr("gpu._find_hwmon_dir", lambda base: "/sys/class/drm/card0/device/hwmon/hwmon0")
|
||
|
||
sysfs_values = {
|
||
"/sys/class/drm/card0/device/mem_info_vram_total": str(16 * 1024**3),
|
||
"/sys/class/drm/card0/device/mem_info_vram_used": str(4 * 1024**3),
|
||
"/sys/class/drm/card0/device/mem_info_gtt_total": str(8 * 1024**3),
|
||
"/sys/class/drm/card0/device/mem_info_gtt_used": str(1 * 1024**3),
|
||
"/sys/class/drm/card0/device/gpu_busy_percent": "45",
|
||
"/sys/class/drm/card0/device/product_name": "AMD Radeon RX 7900 XTX",
|
||
"/sys/class/drm/card0/device/hwmon/hwmon0/temp1_input": "62000",
|
||
"/sys/class/drm/card0/device/hwmon/hwmon0/power1_average": "200000000",
|
||
}
|
||
monkeypatch.setattr("gpu._read_sysfs", lambda path: sysfs_values.get(path))
|
||
|
||
info = get_gpu_info_amd()
|
||
assert info is not None
|
||
assert info.name == "AMD Radeon RX 7900 XTX"
|
||
assert info.memory_total_mb == 16 * 1024
|
||
assert info.memory_used_mb == 4 * 1024
|
||
assert info.memory_type == "discrete"
|
||
assert info.utilization_percent == 45
|
||
assert info.temperature_c == 62
|
||
assert info.power_w == 200.0
|
||
assert info.gpu_backend == "amd"
|
||
|
||
def test_parses_unified_apu(self, monkeypatch):
|
||
"""APU: small VRAM + large GTT signals unified memory."""
|
||
monkeypatch.setattr("gpu._find_amd_gpu_sysfs", lambda: "/sys/class/drm/card0/device")
|
||
monkeypatch.setattr("gpu._find_hwmon_dir", lambda base: None)
|
||
|
||
sysfs_values = {
|
||
"/sys/class/drm/card0/device/mem_info_vram_total": str(4 * 1024**3), # small VRAM partition
|
||
"/sys/class/drm/card0/device/mem_info_vram_used": str(1 * 1024**3),
|
||
"/sys/class/drm/card0/device/mem_info_gtt_total": str(96 * 1024**3), # large GTT => unified
|
||
"/sys/class/drm/card0/device/mem_info_gtt_used": str(20 * 1024**3),
|
||
"/sys/class/drm/card0/device/gpu_busy_percent": "10",
|
||
"/sys/class/drm/card0/device/product_name": None,
|
||
}
|
||
monkeypatch.setattr("gpu._read_sysfs", lambda path: sysfs_values.get(path))
|
||
|
||
info = get_gpu_info_amd()
|
||
assert info is not None
|
||
assert info.memory_type == "unified"
|
||
assert info.memory_total_mb == 96 * 1024
|
||
assert info.memory_used_mb == 20 * 1024
|
||
assert "Strix Halo" in info.name # default name
|
||
|
||
def test_reads_hwmon_temp_power(self, monkeypatch):
|
||
"""Ensure hwmon reads for temperature and power work."""
|
||
monkeypatch.setattr("gpu._find_amd_gpu_sysfs", lambda: "/sys/class/drm/card0/device")
|
||
monkeypatch.setattr("gpu._find_hwmon_dir", lambda base: "/hw")
|
||
|
||
sysfs_values = {
|
||
"/sys/class/drm/card0/device/mem_info_vram_total": str(16 * 1024**3),
|
||
"/sys/class/drm/card0/device/mem_info_vram_used": str(2 * 1024**3),
|
||
"/sys/class/drm/card0/device/mem_info_gtt_total": str(8 * 1024**3),
|
||
"/sys/class/drm/card0/device/mem_info_gtt_used": str(0),
|
||
"/sys/class/drm/card0/device/gpu_busy_percent": "20",
|
||
"/sys/class/drm/card0/device/product_name": "Test GPU",
|
||
"/hw/temp1_input": "75000",
|
||
"/hw/power1_average": "150000000",
|
||
}
|
||
monkeypatch.setattr("gpu._read_sysfs", lambda path: sysfs_values.get(path))
|
||
|
||
info = get_gpu_info_amd()
|
||
assert info is not None
|
||
assert info.temperature_c == 75
|
||
assert info.power_w == 150.0
|
||
|
||
def test_returns_none_when_no_device(self, monkeypatch):
|
||
monkeypatch.setattr("gpu._find_amd_gpu_sysfs", lambda: None)
|
||
assert get_gpu_info_amd() is None
|
||
|
||
def test_returns_none_when_vram_missing(self, monkeypatch):
|
||
monkeypatch.setattr("gpu._find_amd_gpu_sysfs", lambda: "/sys/class/drm/card0/device")
|
||
monkeypatch.setattr("gpu._find_hwmon_dir", lambda base: None)
|
||
monkeypatch.setattr("gpu._read_sysfs", lambda path: None)
|
||
|
||
assert get_gpu_info_amd() is None
|
||
|
||
|
||
# --- get_gpu_info_apple container path ---
|
||
|
||
|
||
class TestGetGpuInfoAppleContainer:
|
||
|
||
def test_host_ram_gb_returns_gpu_info(self, monkeypatch):
|
||
"""Linux container with GPU_BACKEND=apple and HOST_RAM_GB=64."""
|
||
monkeypatch.setattr("gpu.platform.system", lambda: "Linux")
|
||
monkeypatch.setattr("gpu.os.environ", {
|
||
"GPU_BACKEND": "apple",
|
||
"HOST_RAM_GB": "64",
|
||
})
|
||
|
||
info = get_gpu_info_apple()
|
||
assert info is not None
|
||
assert info.memory_total_mb == 64 * 1024
|
||
assert info.memory_type == "unified"
|
||
assert info.gpu_backend == "apple"
|
||
assert "64" in info.name
|
||
|
||
def test_no_host_ram_returns_none(self, monkeypatch):
|
||
"""Linux container with GPU_BACKEND=apple but no HOST_RAM_GB."""
|
||
monkeypatch.setattr("gpu.platform.system", lambda: "Linux")
|
||
monkeypatch.setattr("gpu.os.environ", {
|
||
"GPU_BACKEND": "apple",
|
||
})
|
||
|
||
assert get_gpu_info_apple() is None
|
||
|
||
def test_invalid_host_ram_returns_none(self, monkeypatch):
|
||
"""Linux container with GPU_BACKEND=apple and invalid HOST_RAM_GB."""
|
||
monkeypatch.setattr("gpu.platform.system", lambda: "Linux")
|
||
monkeypatch.setattr("gpu.os.environ", {
|
||
"GPU_BACKEND": "apple",
|
||
"HOST_RAM_GB": "not-a-number",
|
||
})
|
||
|
||
assert get_gpu_info_apple() is None
|
||
|
||
|
||
# --- get_gpu_info dispatcher ---
|
||
|
||
|
||
class TestGetGpuInfoDispatcher:
|
||
|
||
def _make_gpu_info(self, name, backend):
|
||
from models import GPUInfo
|
||
return GPUInfo(
|
||
name=name, memory_used_mb=1024, memory_total_mb=8192,
|
||
memory_percent=12.5, utilization_percent=50, temperature_c=60,
|
||
gpu_backend=backend,
|
||
)
|
||
|
||
def test_nvidia_backend_tries_nvidia(self, monkeypatch):
|
||
monkeypatch.setattr("gpu.os.environ", {"GPU_BACKEND": "nvidia"})
|
||
gpu = self._make_gpu_info("RTX 4090", "nvidia")
|
||
monkeypatch.setattr("gpu.get_gpu_info_nvidia", lambda: gpu)
|
||
monkeypatch.setattr("gpu.get_gpu_info_amd", lambda: None)
|
||
monkeypatch.setattr("gpu.get_gpu_info_apple", lambda: None)
|
||
|
||
result = get_gpu_info()
|
||
assert result is not None
|
||
assert result.name == "RTX 4090"
|
||
|
||
def test_amd_backend_tries_amd_first(self, monkeypatch):
|
||
monkeypatch.setattr("gpu.os.environ", {"GPU_BACKEND": "amd"})
|
||
gpu = self._make_gpu_info("Radeon RX 7900", "amd")
|
||
monkeypatch.setattr("gpu.get_gpu_info_amd", lambda: gpu)
|
||
monkeypatch.setattr("gpu.get_gpu_info_nvidia", lambda: None)
|
||
monkeypatch.setattr("gpu.get_gpu_info_apple", lambda: None)
|
||
|
||
result = get_gpu_info()
|
||
assert result is not None
|
||
assert result.name == "Radeon RX 7900"
|
||
|
||
def test_apple_on_darwin_autodetects(self, monkeypatch):
|
||
monkeypatch.setattr("gpu.os.environ", {"GPU_BACKEND": ""})
|
||
monkeypatch.setattr("gpu.platform.system", lambda: "Darwin")
|
||
gpu = self._make_gpu_info("Apple M4 Max", "apple")
|
||
monkeypatch.setattr("gpu.get_gpu_info_nvidia", lambda: None)
|
||
monkeypatch.setattr("gpu.get_gpu_info_amd", lambda: None)
|
||
monkeypatch.setattr("gpu.get_gpu_info_apple", lambda: gpu)
|
||
|
||
result = get_gpu_info()
|
||
assert result is not None
|
||
assert result.name == "Apple M4 Max"
|
||
|
||
def test_falls_back_through_backends(self, monkeypatch):
|
||
"""nvidia backend -> nvidia fails -> amd fails -> still tries nvidia first then amd."""
|
||
monkeypatch.setattr("gpu.os.environ", {"GPU_BACKEND": "nvidia"})
|
||
monkeypatch.setattr("gpu.platform.system", lambda: "Linux")
|
||
|
||
amd_gpu = self._make_gpu_info("AMD Fallback", "amd")
|
||
monkeypatch.setattr("gpu.get_gpu_info_nvidia", lambda: None)
|
||
monkeypatch.setattr("gpu.get_gpu_info_amd", lambda: amd_gpu)
|
||
monkeypatch.setattr("gpu.get_gpu_info_apple", lambda: None)
|
||
|
||
result = get_gpu_info()
|
||
assert result is not None
|
||
assert result.name == "AMD Fallback"
|
||
|
||
def test_returns_none_when_nothing_found(self, monkeypatch):
|
||
monkeypatch.setattr("gpu.os.environ", {"GPU_BACKEND": "nvidia"})
|
||
monkeypatch.setattr("gpu.platform.system", lambda: "Linux")
|
||
monkeypatch.setattr("gpu.get_gpu_info_nvidia", lambda: None)
|
||
monkeypatch.setattr("gpu.get_gpu_info_amd", lambda: None)
|
||
monkeypatch.setattr("gpu.get_gpu_info_apple", lambda: None)
|
||
|
||
result = get_gpu_info()
|
||
assert result is None
|