Skip to content

other

Functions:

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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def get_nvidia_version() -> tuple[int, int] | None:
    """Check if nvidia drivers are installed and if available return the version."""

    try:
        nvcc = run(['nvcc', '--version'], capture_output=True)
    except FileNotFoundError:
        pass
    else:
        if not nvcc.returncode:
            return _str_to_ver(nvcc.stdout.splitlines()[3].decode().split(',')[-2].replace('release', ''))

    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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
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