other

Functions:

• get_nvidia_version –

  Check if nvidia drivers are installed and if available return the version.

• is_gpu_available –

  Check if any GPU is available.

Attributes:

• IS_DOCS –

  Whether the script is currently running in sphinx docs.

IS_DOCS module-attribute

IS_DOCS = (
    not TYPE_CHECKING
    and "sphinx" in modules
    or basename(argv[0]) in ["sphinx-build", "sphinx-build.exe"]
    or hasattr(builtins, "__sphinx_build__")
    and __sphinx_build__
)

Whether the script is currently running in sphinx docs.

get_nvidia_version

get_nvidia_version() -> tuple[int, int] | None

Check if nvidia drivers are installed and if available return the version.

Source code

def get_nvidia_version() -> tuple[int, int] | None:
    """Check if nvidia drivers are installed and if available return the version."""

    try:
        smi = run(['nvidia-smi', '-q'], capture_output=True)
    except FileNotFoundError:
        pass
    else:
        if not smi.returncode:
            return _str_to_ver(smi.stdout.splitlines()[5].decode().split(':')[-1])

    return None

is_gpu_available

is_gpu_available() -> bool

Check if any GPU is available.

Source code

def is_gpu_available() -> bool:
    """Check if any GPU is available."""

    try:
        smi = run(['nvidia-smi'], capture_output=True, text=True)
    except FileNotFoundError:
        pass
    else:
        if smi.returncode == 0:
            return True

    try:
        rocm_smi = run(['rocm-smi'], capture_output=True, text=True)
    except FileNotFoundError:
        pass
    else:
        if rocm_smi.returncode == 0:
            return True

    return False