Skip to content

darts.utils.cuda

Utility functions for working with CUDA devices.

logger module-attribute 

logger = logging.getLogger(__name__)

debug_info 

debug_info()

Print debug information about the CUDA devices and library installations.

Raises:

  • ImportError –

    If pynvml cannot be imported properly.

Source code in darts/src/darts/utils/cuda.py 
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def debug_info():  # noqa: C901
    """Print debug information about the CUDA devices and library installations.

    Raises:
        ImportError: If pynvml cannot be imported properly.

    """
    logger.debug("===vvv CUDA DEBUG INFO vvv===")
    important_env_vars = [
        "CUDA_HOME",
        "CUDA_PATH",
        "CUDA_VISIBLE_DEVICES",
        "LD_LIBRARY_PATH",
        "NUMBA_CUDA_DRIVER",
        "NUMBA_CUDA_INCLUDE_PATH",
    ]
    for v in important_env_vars:
        value = os.environ.get(v, "UNSET")
        logger.debug(f"{v}: {value}")

    logger.debug("Quicknote: CUDA driver is something different than CUDA runtime, hence versions can mismatch")
    try:
        from pynvml import (  # type: ignore
            NVMLError,
            nvmlDeviceGetCount,
            nvmlDeviceGetHandleByIndex,
            nvmlDeviceGetMemoryInfo,
            nvmlDeviceGetName,
            nvmlInit,
            nvmlShutdown,
            nvmlSystemGetCudaDriverVersion_v2,
            nvmlSystemGetDriverVersion,
        )

        try:
            nvmlInit()
            cuda_driver_version_legacy = nvmlSystemGetDriverVersion().decode()
            cuda_driver_version = nvmlSystemGetCudaDriverVersion_v2()
            logger.debug(f"CUDA driver version: {cuda_driver_version} ({cuda_driver_version_legacy})")
            ndevices = nvmlDeviceGetCount()
            logger.debug(f"Number of CUDA devices: {ndevices}")

            for i in range(ndevices):
                handle = nvmlDeviceGetHandleByIndex(i)
                device_name = nvmlDeviceGetName(handle).decode()
                meminfo = nvmlDeviceGetMemoryInfo(handle)
                logger.debug(f"Device {i} ({device_name}): {meminfo.used / meminfo.total:.2%} memory usage.")
            nvmlShutdown()
        except NVMLError:
            raise ImportError

    except ImportError:
        logger.debug("Module 'pynvml' could not be imported. darts is probably installed without CUDA support.")
    except Exception as e:
        logger.error(
            "Error while trying to get CUDA driver version or device info."
            " Is is possible that this device has no CUDA device?"
        )
        logger.exception(e, exc_info=True)

    try:
        import torch

        logger.debug(f"PyTorch version: {torch.__version__}")
        logger.debug(f"PyTorch CUDA available: {torch.cuda.is_available()}")
        logger.debug(f"PyTorch CUDA runtime version: {torch.version.cuda}")
    except ImportError as e:
        logger.error("Module 'torch' could not be imported:")
        logger.exception(e, exc_info=True)

    try:
        import cupy  # type: ignore
        import cupyx

        logger.debug(f"Cupy version: {cupy.__version__}")
        cupy_driver_version = cupy.cuda.runtime.driverGetVersion()
        logger.debug(f"Cupy CUDA driver version: {cupy_driver_version}")
        # This is the version which is installed (dynamically linked via PATH or LD_LIBRARY_PATH) in the environment
        env_runtime_version = cupy.cuda.get_local_runtime_version()
        logger.debug(f"Cupy CUDA runtime version: {env_runtime_version}")
        if cupy_driver_version < env_runtime_version:
            logger.warning(
                "CUDA runtime version is newer than CUDA driver version!"
                " The CUDA environment is probably not setup correctly!"
                " Consider linking CUDA to an older version with CUDA_HOME and LD_LIBRARY_PATH environment variables,"
                " or in case of a setup done by pixi choose a different environment with the -e flag."
            )
        # This is the version which is was used when cupy was compiled (statically linked)
        cupy_runtime_version = cupy.cuda.runtime.runtimeGetVersion()
        if env_runtime_version != cupy_runtime_version:
            logger.debug(
                "Cupy CUDA runtime versions don't match!\n"
                f"Got {env_runtime_version} as local (dynamically linked) runtime version.\n"
                f"Got {cupy_runtime_version} as by cupy statically linked runtime version.\n"
                "This can happen if cupy was compiled using a different CUDA runtime version. "
                "Things should still work, note that Cupy will use the dynamically linked version."
            )

        logger.debug(f"Cupy config:\n{cupyx.get_runtime_info(full=True)}")
    except ImportError:
        logger.debug("Module 'cupy' not found, darts is probably installed without CUDA support.")

    try:
        import numba.cuda

        cuda_available = numba.cuda.is_available()
        logger.debug(f"Numba CUDA is available: {cuda_available}")
        if cuda_available:
            logger.debug(f"Numba CUDA runtime version: {numba.cuda.runtime.get_version()}")
            # logger.debug(f"Numba CUDA has supported devices: {numba.cuda.detect()}")
    except ImportError:
        logger.debug("Module 'numba.cuda' not found, darts is probably installed without CUDA support.")

    from xrspatial.utils import has_cuda_and_cupy

    logger.debug(f"Cupy+Numba CUDA available: {has_cuda_and_cupy()}")

    try:
        import cucim  # type: ignore

        logger.debug(f"Cucim version: {cucim.__version__}")
    except ImportError:
        logger.debug("Module 'cucim' not found, darts is probably installed without CUDA support.")

    logger.debug("===^^^ CUDA DEBUG INFO ^^^===")

decide_device 

decide_device(
    device: typing.Literal["cuda", "cpu", "auto"]
    | int
    | None,
) -> typing.Literal["cuda", "cpu"] | int

Decide the device based on the input.

Parameters:

  • device (typing.Literal['cuda', 'cpu', 'auto'] | int) –

    The device to run the model on.

Returns:

  • typing.Literal['cuda', 'cpu'] | int –

    Literal["cuda", "cpu"] | int: The device to run the model on.

Source code in darts/src/darts/utils/cuda.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
def decide_device(device: Literal["cuda", "cpu", "auto"] | int | None) -> Literal["cuda", "cpu"] | int:
    """Decide the device based on the input.

    Args:
        device (Literal["cuda", "cpu", "auto"] | int): The device to run the model on.

    Returns:
        Literal["cuda", "cpu"] | int: The device to run the model on.

    """
    import torch
    from xrspatial.utils import has_cuda_and_cupy

    # We can't provide a default value for device in the parameter list because then we would need to import torch at
    # top-level, which would make the CLI slow.
    if device is None:
        device = "cuda" if torch.cuda.is_available() and has_cuda_and_cupy() else "cpu"
        logger.info(f"Device not provided. Using {device}.")
        return device

    # Automatically select a free GPU (<50% memory usage)
    if device == "auto":
        logger.info(f"{device=}. Trying to automatically select a free GPU. (<50% memory usage)")

        # Check if torch and cupy are available
        if not has_cuda_and_cupy() or not torch.cuda.is_available():
            logger.info("CUDA not available. Using CPU.")
            return "cpu"

        try:
            from pynvml import (  # type: ignore
                nvmlDeviceGetCount,
                nvmlDeviceGetHandleByIndex,
                nvmlDeviceGetMemoryInfo,
                nvmlInit,
                nvmlShutdown,
            )
        except ImportError:
            logger.warning("Module 'pynvml' not found. Using CPU.")
            return "cpu"

        nvmlInit()
        ndevices = nvmlDeviceGetCount()

        # If there are multiple devices, we need to check which one is free
        for i in range(ndevices):
            handle = nvmlDeviceGetHandleByIndex(i)
            meminfo = nvmlDeviceGetMemoryInfo(handle)
            perc_used = meminfo.used / meminfo.total
            logger.debug(f"Device {i}: {perc_used:.2%} memory usage.")
            # If the device is less than 50% used, we skip it
            if perc_used > 0.5:
                continue
            else:
                nvmlShutdown()
                logger.info(f"Using free GPU {i} ({perc_used:.2%} memory usage).")
                return i
        else:
            nvmlShutdown()
            logger.warning(
                "No free GPU found (<50% memory usage). Using CPU. "
                "If you want to use a GPU, please select a device manually with the 'device' parameter."
            )
            return "cpu"

    # If device is int or "cuda" or "cpu", we just return it
    logger.info(f"Using {device=}.")
    return device

set_pixi_cuda_env 

set_pixi_cuda_env(
    version: typing.Literal[
        "cuda121", "cuda124", "cuda126", "cuda128"
    ],
    prefix: str = "../",
)

Set the CUDA environment variables.

This is useful when working with a notebook wich does not load the Pixi environment.

Warnign

This does not work currently as expected!

Parameters:

  • version (str) –

    The version string of the pixi CUDA envirnonent.

  • prefix (str, default: '../' ) –

    The prefix to the PIXI installation. E.g. when in the ./notebooks directory, the prefix is ../

Source code in darts/src/darts/utils/cuda.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
def set_pixi_cuda_env(version: Literal["cuda121", "cuda124", "cuda126", "cuda128"], prefix: str = "../"):
    """Set the CUDA environment variables.

    This is useful when working with a notebook wich does not load the Pixi environment.

    !!! warnign

        This does not work currently as expected!

    Args:
        version (str): The version string of the pixi CUDA envirnonent.
        prefix (str): The prefix to the PIXI installation. E.g. when in the `./notebooks` directory, the prefix is `../`

    """
    prefix = str(Path(prefix).resolve())
    os.environ["PATH"] = f"{prefix}/.pixi/envs/{version}/bin:{os.environ.get('PATH', '')}"
    os.environ["CUDA_HOME"] = f"{prefix}/.pixi/envs/{version}"
    os.environ["CUDA_PATH"] = f"{prefix}/.pixi/envs/{version}"
    os.environ["LD_LIBRARY_PATH"] = f"{prefix}/.pixi/envs/{version}/lib:{prefix}/.pixi/envs/{version}/lib/stubs"
    os.environ["NUMBA_CUDA_DRIVER"] = f"{prefix}/.pixi/envs/{version}/lib/stubs/libcuda.so"
    os.environ["NUMBA_CUDA_INCLUDE_PATH"] = f"{prefix}/.pixi/envs/{version}/include"