mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
238 lines
9.8 KiB
Python
238 lines
9.8 KiB
Python
"""Deterministic contracts for accelerator detection and fallback behavior."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from obliteratus import device
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("cuda", "mps", "expected"),
|
|
[(True, False, "cuda"), (False, True, "mps"), (False, False, "cpu")],
|
|
)
|
|
def test_auto_device_preference(monkeypatch, cuda, mps, expected):
|
|
monkeypatch.setattr(device, "is_cuda", lambda: cuda)
|
|
monkeypatch.setattr(device, "is_mps", lambda: mps)
|
|
assert device.get_device() == expected
|
|
assert device.is_gpu_available() is (cuda or mps)
|
|
|
|
|
|
def test_explicit_device_validation(monkeypatch):
|
|
monkeypatch.setattr(device, "is_cuda", lambda: False)
|
|
monkeypatch.setattr(device, "is_mps", lambda: False)
|
|
|
|
assert device.get_device("cpu") == "cpu"
|
|
with pytest.raises(RuntimeError, match="CUDA was requested.*device='auto'"):
|
|
device.get_device("cuda:0")
|
|
with pytest.raises(RuntimeError, match="MPS was requested.*device='auto'"):
|
|
device.get_device("mps")
|
|
with pytest.raises(ValueError, match="Unknown device 'tpu'"):
|
|
device.get_device("tpu")
|
|
with pytest.raises(ValueError, match="Unknown device 'cuda:gpu'"):
|
|
device.get_device("cuda:gpu")
|
|
|
|
monkeypatch.setattr(device, "is_cuda", lambda: True)
|
|
assert device.get_device("cuda") == "cuda"
|
|
assert device.get_device("cuda:3") == "cuda:3"
|
|
monkeypatch.setattr(device, "is_mps", lambda: True)
|
|
assert device.get_device("mps") == "mps"
|
|
|
|
|
|
def test_names_and_device_counts(monkeypatch):
|
|
monkeypatch.setattr(device, "is_cuda", lambda: True)
|
|
monkeypatch.setattr(device.torch.cuda, "get_device_name", lambda _index: "Test GPU")
|
|
monkeypatch.setattr(device.torch.cuda, "device_count", lambda: 4)
|
|
assert device.get_device_name() == "Test GPU"
|
|
assert device.device_count() == 4
|
|
|
|
monkeypatch.setattr(device, "is_cuda", lambda: False)
|
|
monkeypatch.setattr(device, "is_mps", lambda: True)
|
|
monkeypatch.setattr(device.platform, "processor", lambda: "M3")
|
|
assert device.get_device_name() == "Apple M3 (MPS)"
|
|
assert device.device_count() == 1
|
|
|
|
monkeypatch.setattr(device, "is_mps", lambda: False)
|
|
assert device.get_device_name() == "CPU"
|
|
assert device.device_count() == 0
|
|
|
|
|
|
def test_system_memory_sources_and_fallback(monkeypatch):
|
|
gib = 1024**3
|
|
fake_psutil = SimpleNamespace(
|
|
virtual_memory=lambda: SimpleNamespace(total=32 * gib, available=12 * gib),
|
|
)
|
|
monkeypatch.setitem(__import__("sys").modules, "psutil", fake_psutil)
|
|
assert device._system_memory_gb() == (32.0, 12.0)
|
|
|
|
monkeypatch.delitem(__import__("sys").modules, "psutil", raising=False)
|
|
real_import = __import__("builtins").__import__
|
|
|
|
def reject_psutil(name, *args, **kwargs):
|
|
if name == "psutil":
|
|
raise ImportError
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr("builtins.__import__", reject_psutil)
|
|
monkeypatch.setattr(device.os, "sysconf", lambda name: {"SC_PHYS_PAGES": 4, "SC_PAGE_SIZE": gib}[name])
|
|
assert device._system_memory_gb() == (4.0, 2.4)
|
|
|
|
monkeypatch.setattr(device.os, "sysconf", Mock(side_effect=ValueError))
|
|
assert device._system_memory_gb() == (16.0, 8.0)
|
|
|
|
|
|
def test_memory_info_for_cuda_and_cuda_fallback(monkeypatch):
|
|
gib = 1024**3
|
|
monkeypatch.setattr(device, "is_cuda", lambda: True)
|
|
monkeypatch.setattr(device, "get_device_name", lambda: "GPU")
|
|
monkeypatch.setattr(device.torch.cuda, "mem_get_info", lambda _index: (6 * gib, 8 * gib))
|
|
monkeypatch.setattr(device.torch.cuda, "memory_allocated", lambda _index: 1 * gib)
|
|
monkeypatch.setattr(device.torch.cuda, "memory_reserved", lambda _index: 2 * gib)
|
|
assert device.get_memory_info(2) == device.MemoryInfo(1, 2, 8, 6, "GPU")
|
|
|
|
monkeypatch.setattr(device.torch.cuda, "mem_get_info", Mock(side_effect=RuntimeError("unsupported")))
|
|
monkeypatch.setattr(
|
|
device.torch.cuda,
|
|
"get_device_properties",
|
|
lambda _index: SimpleNamespace(total_memory=10 * gib),
|
|
)
|
|
assert device.get_memory_info(2) == device.MemoryInfo(total_gb=10, free_gb=10, device_name="GPU")
|
|
|
|
|
|
def test_memory_info_for_mps_cpu_and_total_free(monkeypatch):
|
|
monkeypatch.setattr(device, "is_cuda", lambda: False)
|
|
monkeypatch.setattr(device, "is_mps", lambda: True)
|
|
monkeypatch.setattr(device, "get_device_name", lambda: "MPS")
|
|
monkeypatch.setattr(device, "_system_memory_gb", lambda: (20.0, 8.0))
|
|
assert device.get_memory_info() == device.MemoryInfo(6, 0, 14, 8, "MPS")
|
|
assert device.get_total_free_gb() == pytest.approx(5.6)
|
|
|
|
monkeypatch.setattr(device, "is_mps", lambda: False)
|
|
monkeypatch.setattr(device, "get_device_name", lambda: "CPU")
|
|
assert device.get_memory_info() == device.MemoryInfo(total_gb=20, free_gb=8, device_name="CPU")
|
|
assert device.get_total_free_gb() == 0
|
|
|
|
|
|
def test_total_cuda_memory_sums_query_and_fallback(monkeypatch):
|
|
gib = 1024**3
|
|
monkeypatch.setattr(device, "is_cuda", lambda: True)
|
|
monkeypatch.setattr(device.torch.cuda, "device_count", lambda: 2)
|
|
monkeypatch.setattr(
|
|
device.torch.cuda,
|
|
"mem_get_info",
|
|
Mock(side_effect=[(3 * gib, 4 * gib), RuntimeError("no query")]),
|
|
)
|
|
monkeypatch.setattr(
|
|
device.torch.cuda,
|
|
"get_device_properties",
|
|
lambda _index: SimpleNamespace(total_memory=5 * gib),
|
|
)
|
|
assert device.get_total_free_gb() == 8.0
|
|
|
|
|
|
def test_cache_cleanup_paths_are_best_effort(monkeypatch):
|
|
cuda_empty = Mock()
|
|
monkeypatch.setattr(device, "is_cuda", lambda: True)
|
|
monkeypatch.setattr(device.torch.cuda, "empty_cache", cuda_empty)
|
|
device.empty_cache()
|
|
cuda_empty.assert_called_once_with()
|
|
|
|
sync = Mock()
|
|
reset = Mock()
|
|
monkeypatch.setattr(device.torch.cuda, "empty_cache", Mock(side_effect=RuntimeError("busy")))
|
|
monkeypatch.setattr(device.torch.cuda, "synchronize", sync)
|
|
monkeypatch.setattr(device.torch.cuda, "reset_peak_memory_stats", reset)
|
|
monkeypatch.setattr(device.gc, "collect", Mock())
|
|
device.free_gpu_memory()
|
|
sync.assert_called_once_with()
|
|
reset.assert_called_once_with()
|
|
|
|
mps_empty = Mock()
|
|
mps_sync = Mock(side_effect=RuntimeError("busy"))
|
|
monkeypatch.setattr(device, "is_cuda", lambda: False)
|
|
monkeypatch.setattr(device, "is_mps", lambda: True)
|
|
monkeypatch.setattr(device.torch, "mps", SimpleNamespace(empty_cache=mps_empty, synchronize=mps_sync))
|
|
device.empty_cache()
|
|
mps_empty.side_effect = RuntimeError("busy")
|
|
device.free_gpu_memory()
|
|
assert mps_empty.call_count == 2
|
|
mps_sync.assert_called_once_with()
|
|
|
|
|
|
def test_cuda_cleanup_tolerates_every_recovery_failure(monkeypatch):
|
|
monkeypatch.setattr(device, "is_cuda", lambda: True)
|
|
monkeypatch.setattr(device.torch.cuda, "empty_cache", Mock(side_effect=RuntimeError("busy")))
|
|
monkeypatch.setattr(device.torch.cuda, "synchronize", Mock(side_effect=RuntimeError("lost")))
|
|
monkeypatch.setattr(
|
|
device.torch.cuda,
|
|
"reset_peak_memory_stats",
|
|
Mock(side_effect=RuntimeError("lost")),
|
|
)
|
|
monkeypatch.setattr(device.gc, "collect", Mock())
|
|
device.free_gpu_memory()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("version", "expected"),
|
|
[("2.2.9", False), ("2.3.0", True), ("2.10.1", True)],
|
|
)
|
|
def test_mps_bfloat16_version_boundary(monkeypatch, version, expected):
|
|
monkeypatch.setattr(device, "is_cuda", lambda: False)
|
|
monkeypatch.setattr(device.torch, "__version__", version)
|
|
assert device.supports_bfloat16("mps") is expected
|
|
|
|
|
|
def test_seed_dtype_and_capability_contracts(monkeypatch):
|
|
manual_seed = Mock()
|
|
cuda_seed = Mock()
|
|
monkeypatch.setattr(device.torch, "manual_seed", manual_seed)
|
|
monkeypatch.setattr(device.torch.cuda, "manual_seed_all", cuda_seed)
|
|
monkeypatch.setattr(device, "is_cuda", lambda: True)
|
|
device.set_seed_all(7)
|
|
manual_seed.assert_called_once_with(7)
|
|
cuda_seed.assert_called_once_with(7)
|
|
|
|
assert device.default_dtype("cpu") is torch.float32
|
|
assert device.default_dtype("cuda") is torch.float16
|
|
monkeypatch.setattr(device.torch.cuda, "get_device_capability", lambda _index: (8, 0))
|
|
assert device.supports_bfloat16("cuda") is True
|
|
monkeypatch.setattr(device.torch.cuda, "get_device_capability", lambda _index: (7, 5))
|
|
assert device.supports_bfloat16("cuda") is False
|
|
monkeypatch.setattr(device, "is_cuda", lambda: False)
|
|
assert device.supports_bfloat16("cuda") is False
|
|
assert device.supports_bfloat16("cpu") is True
|
|
assert device.supports_float64("mps") is False
|
|
assert device.supports_float64("cpu") is True
|
|
assert device.supports_bitsandbytes("cuda:1") is True
|
|
assert device.supports_bitsandbytes("mps") is False
|
|
assert device.supports_device_map_auto("cuda") is True
|
|
assert device.supports_device_map_auto("cpu") is False
|
|
|
|
|
|
def test_svd_dtype_and_oom_matching():
|
|
assert device.safe_svd_dtype(torch.ones(1, dtype=torch.float64)) is torch.float64
|
|
assert device.safe_svd_dtype(torch.ones(1, dtype=torch.float16)) is torch.float32
|
|
assert device.is_oom_error(torch.cuda.OutOfMemoryError("oom")) is True
|
|
assert device.is_oom_error(RuntimeError("MPS backend out of memory")) is True
|
|
assert device.is_oom_error(RuntimeError("other")) is False
|
|
|
|
|
|
def test_mps_svd_dtype_uses_float32_without_requiring_mps_hardware():
|
|
tensor = SimpleNamespace(device=SimpleNamespace(type="mps"), dtype=torch.float64)
|
|
assert device.safe_svd_dtype(tensor) is torch.float32
|
|
|
|
|
|
def test_configure_cuda_allocator(monkeypatch):
|
|
monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False)
|
|
monkeypatch.setattr(device, "is_cuda", lambda: True)
|
|
device.configure_cuda_alloc()
|
|
assert device.os.environ["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True"
|
|
|
|
monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "existing")
|
|
device.configure_cuda_alloc()
|
|
assert device.os.environ["PYTORCH_CUDA_ALLOC_CONF"] == "existing"
|