test_system_system.py

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0


import os

import pytest

try:
    from cuda.bindings import driver, runtime
except ImportError:
    from cuda import cuda as driver
    from cuda import cudart as runtime

from cuda.core import Device, system
from cuda.core._utils.cuda_utils import handle_return

from .conftest import skip_if_nvml_unsupported


@skip_if_nvml_unsupported
def test_driver_version():
    umd, kmd = system.get_driver_version()

    # UMD: 2-tuple (major, minor), cross-check against cuDriverGetVersion
    assert isinstance(umd, tuple)
    assert len(umd) == 2
    version = handle_return(driver.cuDriverGetVersion())
    expected_umd = (version // 1000, (version % 1000) // 10)
    assert umd == expected_umd, "UMD driver version does not match expected value"

    # KMD: 3-tuple (major, minor, patch), or 2-tuple on WSL
    assert isinstance(kmd, tuple)
    assert len(kmd) in (2, 3)
    ver_maj, ver_min, *ver_patch = kmd
    assert 400 <= ver_maj < 1000
    assert ver_min >= 0
    if ver_patch:
        assert 0 <= ver_patch[0] <= 99


def test_num_devices():
    num_devices = system.get_num_devices()
    expected_num_devices = handle_return(runtime.cudaGetDeviceCount())
    assert num_devices == expected_num_devices, "Number of devices does not match expected value"


def test_devices():
    devices = Device.get_all_devices()
    expected_num_devices = handle_return(runtime.cudaGetDeviceCount())
    expected_devices = tuple(Device(device_id) for device_id in range(expected_num_devices))
    assert len(devices) == len(expected_devices), "Number of devices does not match expected value"
    for device, expected_device in zip(devices, expected_devices):
        assert device.device_id == expected_device.device_id, "Device ID does not match expected value"


def test_driver_version_requires_nvml():
    if system.CUDA_BINDINGS_NVML_IS_COMPATIBLE:
        pytest.skip("NVML is available, cannot test the error path")
    with pytest.raises(RuntimeError, match="requires NVML support"):
        system.get_driver_version()


@skip_if_nvml_unsupported
def test_nvml_version():
    nvml_version = system.get_nvml_version()
    assert isinstance(nvml_version, tuple)
    assert len(nvml_version) in (3, 4)

    (cuda_ver_maj, ver_maj, ver_min, *ver_patch) = nvml_version
    assert cuda_ver_maj >= 10
    assert 400 <= ver_maj < 1000
    assert ver_min >= 0
    if ver_patch:
        assert 0 <= ver_patch[0] <= 99


@skip_if_nvml_unsupported
def test_get_process_name():
    try:
        process_name = system.get_process_name(os.getpid())
    except system.NotFoundError:
        pytest.skip("Process not found")

    assert isinstance(process_name, str)
    assert "python" in process_name


def test_device_count():
    device_count = system.get_num_devices()
    assert isinstance(device_count, int)
    assert device_count >= 0


@skip_if_nvml_unsupported
def test_get_driver_branch():
    driver_branch = system.get_driver_branch()
    assert isinstance(driver_branch, str)
    assert len(driver_branch) > 0