| from __future__ import annotations |
|
|
| import pytest |
| from fitcheck.gpu_db import get_gpu |
|
|
|
|
| @pytest.mark.parametrize( |
| ("gpu_name", "display_name", "vram_mib", "usable_mib"), |
| [ |
| ("t4", "Tesla T4", 14_912, 14_000), |
| ("4090", "RTX 4090", 24_576, 23_500), |
| ("a100-80", "A100 80GB", 81_920, 79_000), |
| ("3090", "RTX 3090", 24_576, 23_500), |
| ("a100-40", "A100 40GB", 40_960, 39_500), |
| ("h100", "H100 80GB", 81_920, 79_000), |
| ("l4", "L4 24GB", 24_576, 23_000), |
| ("3060-12", "RTX 3060 12GB", 12_288, 11_500), |
| ("4070ti", "RTX 4070 Ti", 12_288, 11_500), |
| ], |
| ) |
| def test_get_gpu_returns_known_gpu_spec( |
| gpu_name: str, |
| display_name: str, |
| vram_mib: int, |
| usable_mib: int, |
| ) -> None: |
| gpu = get_gpu(gpu_name) |
|
|
| assert gpu.name == display_name |
| assert gpu.vram_mib == vram_mib |
| assert gpu.usable_mib == usable_mib |
|
|
|
|
| @pytest.mark.parametrize("gpu_name", ["T4", " t4 ", "T4 ", " T4"]) |
| def test_get_gpu_name_lookup_is_case_and_whitespace_insensitive(gpu_name: str) -> None: |
| gpu = get_gpu(gpu_name) |
|
|
| assert gpu.name == "Tesla T4" |
|
|
|
|
| def test_get_gpu_unknown_name_has_actionable_error() -> None: |
| with pytest.raises(ValueError, match="Unknown GPU 'not-a-gpu'") as error: |
| get_gpu("not-a-gpu") |
|
|
| assert "Available GPUs:" in str(error.value) |
|
|
|
|
| def test_get_gpu_vram_override_returns_custom_spec() -> None: |
| gpu = get_gpu(vram_mib=24_000) |
|
|
| assert gpu.name == "Custom GPU" |
| assert gpu.vram_mib == 24_000 |
| assert gpu.usable_mib == 22_800 |
|
|
|
|
| def test_get_gpu_vram_override_with_custom_name() -> None: |
| gpu = get_gpu(name="My Rig", vram_mib=10_000) |
|
|
| assert gpu.name == "My Rig" |
| assert gpu.vram_mib == 10_000 |
| assert gpu.usable_mib == 9_500 |
|
|
|
|
| def test_get_gpu_vram_override_ignores_blank_name() -> None: |
| gpu = get_gpu(name=" ", vram_mib=10_000) |
| assert gpu.name == "Custom GPU" |
|
|
|
|
| @pytest.mark.parametrize("bad_vram", [0, -1, True, 12.5, "24000"]) |
| def test_get_gpu_rejects_invalid_vram_override(bad_vram: object) -> None: |
| with pytest.raises(ValueError, match="vram_mib must be a positive integer"): |
| get_gpu(vram_mib=bad_vram) |
|
|
|
|
| def test_get_gpu_requires_name_or_vram_override() -> None: |
| with pytest.raises(ValueError, match="Provide a GPU name or a positive vram_mib override"): |
| get_gpu() |
|
|
