tests/test_gpu_utils.py

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"