Spaces:
Runtime error
Runtime error
File size: 2,382 Bytes
2857cf3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | from __future__ import annotations
from types import SimpleNamespace
import subprocess
import pytest
import kneiff.utils as ut
def test_list_nvidia_gpus_reads_nvidia_smi(
monkeypatch: pytest.MonkeyPatch,
) -> None:
def fake_run(*args: object, **kwargs: object) -> SimpleNamespace:
return SimpleNamespace(
stdout="0, NVIDIA GeForce RTX 3090\n1, NVIDIA RTX 4070 Ti SUPER\n"
)
monkeypatch.setattr(subprocess, "run", fake_run)
gpus = ut.list_nvidia_gpus()
assert gpus == (
ut.NvidiaGpu(index="0", name="NVIDIA GeForce RTX 3090"),
ut.NvidiaGpu(index="1", name="NVIDIA RTX 4070 Ti SUPER"),
)
def test_resolve_cuda_device_name_matches_case_insensitive_fragment() -> None:
gpus = (
ut.NvidiaGpu(index="0", name="NVIDIA GeForce RTX 3090"),
ut.NvidiaGpu(index="1", name="NVIDIA RTX 4070 Ti SUPER"),
)
assert ut.resolve_cuda_device_name("rtx 4070 ti super", gpus=gpus) == "1"
def test_resolve_cuda_device_name_rejects_missing_and_ambiguous_names() -> None:
gpus = (
ut.NvidiaGpu(index="0", name="NVIDIA GeForce RTX 4090"),
ut.NvidiaGpu(index="1", name="NVIDIA GeForce RTX 4070"),
)
with pytest.raises(ValueError, match="did not match"):
ut.resolve_cuda_device_name("RTX 5090", gpus=gpus)
with pytest.raises(ValueError, match="multiple"):
ut.resolve_cuda_device_name("NVIDIA", gpus=gpus)
def test_cuda_env_for_selection_respects_explicit_and_inherited_values() -> None:
explicit = ut.cuda_env_for_selection(cuda_device="1", environ={})
assert explicit.cuda_device == "1"
assert explicit.env == {
ut.CUDA_DEVICE_ORDER_ENV_KEY: ut.CUDA_DEVICE_ORDER_VALUE,
ut.CUDA_VISIBLE_DEVICES_ENV_KEY: "1",
}
inherited = ut.cuda_env_for_selection(
cuda_device="1",
environ={ut.CUDA_VISIBLE_DEVICES_ENV_KEY: "0"},
)
assert inherited.cuda_device == "0"
assert inherited.env == {}
assert inherited.user_override is True
def test_cuda_env_for_selection_resolves_name() -> None:
gpus = (ut.NvidiaGpu(index="2", name="NVIDIA RTX 4070 Ti SUPER"),)
selection = ut.cuda_env_for_selection(
cuda_device_name="4070 ti",
environ={},
gpus=gpus,
)
assert selection.cuda_device == "2"
assert selection.env[ut.CUDA_VISIBLE_DEVICES_ENV_KEY] == "2"
|