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"