Coverage for cuda / pathfinder / _dynamic_libs / load_nvidia_dynamic_lib.py: 86%

1# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# SPDX-License-Identifier: Apache-2.0 

3 

4import functools 

5import struct 

6import sys 

7 

8from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import _FindNvidiaDynamicLib 

9from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL, load_dependencies 

10from cuda.pathfinder._utils.platform_aware import IS_WINDOWS 

11 

12if IS_WINDOWS: 

13 from cuda.pathfinder._dynamic_libs.load_dl_windows import ( 

14 check_if_already_loaded_from_elsewhere, 

15 load_with_abs_path, 

16 load_with_system_search, 

17 ) 

18else: 

19 from cuda.pathfinder._dynamic_libs.load_dl_linux import ( 

20 check_if_already_loaded_from_elsewhere, 

21 load_with_abs_path, 

22 load_with_system_search, 

23 ) 

24 

25 

26def _load_lib_no_cache(libname: str) -> LoadedDL: 

27 finder = _FindNvidiaDynamicLib(libname) 

28 abs_path = finder.try_site_packages() 

29 if abs_path is not None: 

30 found_via = "site-packages" 

31 else: 

32 abs_path = finder.try_with_conda_prefix() 

33 if abs_path is not None: 

34 found_via = "conda" 

35 

36 # If the library was already loaded by someone else, reproduce any OS-specific 

37 # side-effects we would have applied on a direct absolute-path load (e.g., 

38 # AddDllDirectory on Windows for libs that require it). 

39 loaded = check_if_already_loaded_from_elsewhere(libname, abs_path is not None) 

40 

41 # Load dependencies regardless of who loaded the primary lib first. 

42 # Doing this *after* the side-effect ensures dependencies resolve consistently 

43 # relative to the actually loaded location. 

44 load_dependencies(libname, load_nvidia_dynamic_lib) 

45 

46 if loaded is not None: 

47 return loaded 

48 

49 if abs_path is None: 

50 loaded = load_with_system_search(libname) 

51 if loaded is not None: 

52 return loaded 

53 abs_path = finder.try_with_cuda_home() 

54 if abs_path is None: 

55 finder.raise_not_found_error() 

56 else: 

57 found_via = "CUDA_HOME" 

58 

59 return load_with_abs_path(libname, abs_path, found_via) 

60 

61 

62@functools.cache 

63def load_nvidia_dynamic_lib(libname: str) -> LoadedDL: 

64 """Load an NVIDIA dynamic library by name. 

65 

66 Args: 

67 libname (str): The short name of the library to load (e.g., ``"cudart"``, 

68 ``"nvvm"``, etc.). 

69 

70 Returns: 

71 LoadedDL: Object containing the OS library handle and absolute path. 

72 

73 **Important:** 

74 

75 **Never close the returned handle.** Do **not** call ``dlclose`` (Linux) or 

76 ``FreeLibrary`` (Windows) on the ``LoadedDL._handle_uint``. 

77 

78 **Why:** the return value is cached (``functools.cache``) and shared across the 

79 process. Closing the handle can unload the module while other code still uses 

80 it, leading to crashes or subtle failures. 

81 

82 This applies to Linux and Windows. For context, see issue #1011: 

83 https://github.com/NVIDIA/cuda-python/issues/1011 

84 

85 Raises: 

86 DynamicLibNotFoundError: If the library cannot be found or loaded. 

87 RuntimeError: If Python is not 64-bit. 

88 

89 Search order: 

90 0. **Already loaded in the current process** 

91 

92 - If a matching library is already loaded by some other component, 

93 return its absolute path and handle and skip the rest of the search. 

94 

95 1. **NVIDIA Python wheels** 

96 

97 - Scan installed distributions (``site-packages``) to find libraries 

98 shipped in NVIDIA wheels. 

99 

100 2. **Conda environment** 

101 

102 - Conda installations are discovered via ``CONDA_PREFIX``, which is 

103 defined automatically in activated conda environments (see 

104 https://docs.conda.io/projects/conda-build/en/stable/user-guide/environment-variables.html). 

105 

106 3. **OS default mechanisms** 

107 

108 - Fall back to the native loader: 

109 

110 - Linux: ``dlopen()`` 

111 

112 - Windows: ``LoadLibraryW()`` 

113 

114 - CUDA Toolkit (CTK) system installs with system config updates are often 

115 discovered via: 

116 

117 - Linux: ``/etc/ld.so.conf.d/*cuda*.conf`` 

118 

119 - Windows: ``C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\vX.Y\\bin`` 

120 on the system ``PATH``. 

121 

122 4. **Environment variables** 

123 

124 - If set, use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order). 

125 

126 Notes: 

127 The search is performed **per library**. There is currently no mechanism to 

128 guarantee that multiple libraries are all resolved from the same location. 

129 

130 """ 

131 pointer_size_bits = struct.calcsize("P") * 8 

132 if pointer_size_bits != 64: 

133 raise RuntimeError( 

134 f"cuda.pathfinder.load_nvidia_dynamic_lib() requires 64-bit Python." 

135 f" Currently running: {pointer_size_bits}-bit Python" 

136 f" {sys.version_info.major}.{sys.version_info.minor}" 

137 ) 

138 return _load_lib_no_cache(libname)