codekingpro/portable-devtools
114k
1# -------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License.
4# --------------------------------------------------------------------------
5"""
6ONNX Runtime is a performance-focused scoring engine for Open Neural Network Exchange (ONNX) models.
7For more information on ONNX Runtime, please see `aka.ms/onnxruntime <https://aka.ms/onnxruntime/>`_
8or the `Github project <https://github.com/microsoft/onnxruntime/>`_.
9"""
10
11import contextlib
12
13__version__ = "1.26.0"
14__author__ = "Microsoft"
15
16# we need to do device version validation (for example to check Cuda version for an onnxruntime-training package).
17# in order to know whether the onnxruntime package is for training it needs
18# to do import onnxruntime.training.ortmodule first.
19# onnxruntime.capi._pybind_state is required before import onnxruntime.training.ortmodule.
20# however, import onnxruntime.capi._pybind_state will already raise an exception if a required Cuda version
21# is not found.
22# here we need to save the exception and continue with Cuda version validation in order to post
23# meaningful messages to the user.
24# the saved exception is raised after device version validation.
25try:
26 from onnxruntime.capi._pybind_state import (
27 ExecutionMode, # noqa: F401
28 ExecutionOrder, # noqa: F401
29 GraphOptimizationLevel, # noqa: F401
30 LoraAdapter, # noqa: F401
31 ModelMetadata, # noqa: F401
32 NodeArg, # noqa: F401
33 OrtAllocatorType, # noqa: F401
34 OrtArenaCfg, # noqa: F401
35 OrtCompileApiFlags, # noqa: F401
36 OrtDeviceMemoryType, # noqa: F401
37 OrtEpAssignedNode, # noqa: F401
38 OrtEpAssignedSubgraph, # noqa: F401
39 OrtEpDevice, # noqa: F401
40 OrtExecutionProviderDevicePolicy, # noqa: F401
41 OrtExternalInitializerInfo, # noqa: F401
42 OrtHardwareDevice, # noqa: F401
43 OrtHardwareDeviceType, # noqa: F401
44 OrtMemoryInfo, # noqa: F401
45 OrtMemoryInfoDeviceType, # noqa: F401
46 OrtMemType, # noqa: F401
47 OrtSparseFormat, # noqa: F401
48 OrtSyncStream, # noqa: F401
49 RunOptions, # noqa: F401
50 SessionIOBinding, # noqa: F401
51 SessionOptions, # noqa: F401
52 create_and_register_allocator, # noqa: F401
53 create_and_register_allocator_v2, # noqa: F401
54 disable_telemetry_events, # noqa: F401
55 enable_telemetry_events, # noqa: F401
56 get_all_providers, # noqa: F401
57 get_available_providers, # noqa: F401
58 get_build_info, # noqa: F401
59 get_device, # noqa: F401
60 get_ep_devices, # noqa: F401
61 get_version_string, # noqa: F401
62 has_collective_ops, # noqa: F401
63 register_execution_provider_library, # noqa: F401
64 set_default_logger_severity, # noqa: F401
65 set_default_logger_verbosity, # noqa: F401
66 set_global_thread_pool_sizes, # noqa: F401
67 set_seed, # noqa: F401
68 unregister_execution_provider_library, # noqa: F401
69 )
70
71 import_capi_exception = None
72except Exception as e:
73 import_capi_exception = e
74
75from onnxruntime.capi import onnxruntime_validation
76
77if import_capi_exception:
78 raise import_capi_exception
79
80from onnxruntime.capi.onnxruntime_inference_collection import (
81 AdapterFormat, # noqa: F401
82 InferenceSession, # noqa: F401
83 IOBinding, # noqa: F401
84 ModelCompiler, # noqa: F401
85 OrtDevice, # noqa: F401
86 OrtDeviceVendorId, # noqa: F401
87 OrtValue, # noqa: F401
88 SparseTensor, # noqa: F401
89 copy_tensors, # noqa: F401
90)
91
92# TODO: thiagofc: Temporary experimental namespace for new PyTorch front-end
93try: # noqa: SIM105
94 from . import experimental # noqa: F401
95except ImportError:
96 pass
97
98
99package_name, version, cuda_version = onnxruntime_validation.get_package_name_and_version_info()
100
101if version:
102 __version__ = version
103
104onnxruntime_validation.check_distro_info()
105
106
107def _get_package_version(package_name: str):
108 from importlib.metadata import PackageNotFoundError, version # noqa: PLC0415
109
110 try:
111 package_version = version(package_name)
112 except PackageNotFoundError:
113 package_version = None
114 return package_version
115
116
117def _get_package_root(package_name: str, directory_name: str | None = None):
118 from importlib.metadata import PackageNotFoundError, distribution # noqa: PLC0415
119
120 root_directory_name = directory_name or package_name
121 try:
122 dist = distribution(package_name)
123 files = dist.files or []
124
125 for file in files:
126 if file.name.endswith("__init__.py") and root_directory_name in file.parts:
127 return file.locate().parent
128
129 # Fallback to the first __init__.py
130 if not directory_name:
131 for file in files:
132 if file.name.endswith("__init__.py"):
133 return file.locate().parent
134 except PackageNotFoundError:
135 # package not found, do nothing
136 pass
137
138 return None
139
140
141def _extract_cuda_major_version(version_str: str) -> str:
142 """Extract CUDA major version from version string (e.g., '12.1' -> '12').
143
144 Args:
145 version_str: CUDA version string to parse
146
147 Returns:
148 Major version as string, or "12" if parsing fails
149 """
150 return version_str.split(".")[0] if version_str else "12"
151
152
153def _get_cufft_version(cuda_major: str) -> str:
154 """Get cufft library version based on CUDA major version.
155
156 Args:
157 cuda_major: CUDA major version as string (e.g., "12", "13")
158
159 Returns:
160 cufft version as string
161 """
162 # cufft versions: CUDA 12.x -> 11, CUDA 13.x -> 12
163 return "12" if cuda_major == "13" else "11"
164
165
166def _get_nvidia_dll_paths(is_windows: bool, cuda: bool = True, cudnn: bool = True):
167 # Dynamically determine CUDA major version from build info
168 cuda_major_version = _extract_cuda_major_version(cuda_version)
169 cufft_version = _get_cufft_version(cuda_major_version)
170
171 if is_windows:
172 # Path is relative to site-packages directory.
173 cuda_dll_paths = [
174 ("nvidia", "cublas", "bin", f"cublasLt64_{cuda_major_version}.dll"),
175 ("nvidia", "cublas", "bin", f"cublas64_{cuda_major_version}.dll"),
176 ("nvidia", "cufft", "bin", f"cufft64_{cufft_version}.dll"),
177 ("nvidia", "cuda_runtime", "bin", f"cudart64_{cuda_major_version}.dll"),
178 ]
179 cudnn_dll_paths = [
180 ("nvidia", "cudnn", "bin", "cudnn_engines_runtime_compiled64_9.dll"),
181 ("nvidia", "cudnn", "bin", "cudnn_engines_precompiled64_9.dll"),
182 ("nvidia", "cudnn", "bin", "cudnn_heuristic64_9.dll"),
183 ("nvidia", "cudnn", "bin", "cudnn_ops64_9.dll"),
184 ("nvidia", "cudnn", "bin", "cudnn_adv64_9.dll"),
185 ("nvidia", "cudnn", "bin", "cudnn_graph64_9.dll"),
186 ("nvidia", "cudnn", "bin", "cudnn64_9.dll"),
187 ]
188 else: # Linux
189 # cublas64 depends on cublasLt64, so cublasLt64 should be loaded first.
190 cuda_dll_paths = [
191 ("nvidia", "cublas", "lib", f"libcublasLt.so.{cuda_major_version}"),
192 ("nvidia", "cublas", "lib", f"libcublas.so.{cuda_major_version}"),
193 ("nvidia", "cuda_nvrtc", "lib", f"libnvrtc.so.{cuda_major_version}"),
194 ("nvidia", "curand", "lib", "libcurand.so.10"),
195 ("nvidia", "cufft", "lib", f"libcufft.so.{cufft_version}"),
196 ("nvidia", "cuda_runtime", "lib", f"libcudart.so.{cuda_major_version}"),
197 ]
198
199 # Do not load cudnn sub DLLs (they will be dynamically loaded later) to be consistent with PyTorch in Linux.
200 cudnn_dll_paths = [
201 ("nvidia", "cudnn", "lib", "libcudnn.so.9"),
202 ]
203
204 return (cuda_dll_paths if cuda else []) + (cudnn_dll_paths if cudnn else [])
205
206
207def print_debug_info():
208 """Print information to help debugging."""
209 import importlib.util # noqa: PLC0415
210 import os # noqa: PLC0415
211 import platform # noqa: PLC0415
212 from importlib.metadata import distributions # noqa: PLC0415
213
214 print(f"{package_name} version: {__version__}")
215 if cuda_version:
216 print(f"CUDA version used in build: {cuda_version}")
217 print("platform:", platform.platform())
218
219 print("\nPython package, version and location:")
220 ort_packages = []
221 for dist in distributions():
222 package = dist.metadata["Name"]
223 if package == "onnxruntime" or package.startswith(("onnxruntime-", "ort-")):
224 # Exclude packages whose root directory name is not onnxruntime.
225 location = _get_package_root(package, "onnxruntime")
226 if location and (package not in ort_packages):
227 ort_packages.append(package)
228 print(f"{package}=={dist.version} at {location}")
229
230 if len(ort_packages) > 1:
231 print(
232 "\033[33mWARNING: multiple onnxruntime packages are installed to the same location. "
233 "Please 'pip uninstall` all above packages, then `pip install` only one of them.\033[0m"
234 )
235
236 if cuda_version:
237 # Print version of installed packages that is related to CUDA or cuDNN DLLs.
238 cuda_major = _extract_cuda_major_version(cuda_version)
239
240 packages = [
241 "torch",
242 f"nvidia-cuda-runtime-cu{cuda_major}",
243 f"nvidia-cudnn-cu{cuda_major}",
244 f"nvidia-cublas-cu{cuda_major}",
245 f"nvidia-cufft-cu{cuda_major}",
246 f"nvidia-curand-cu{cuda_major}",
247 f"nvidia-cuda-nvrtc-cu{cuda_major}",
248 f"nvidia-nvjitlink-cu{cuda_major}",
249 ]
250 for package in packages:
251 directory_name = "nvidia" if package.startswith("nvidia-") else None
252 version = _get_package_version(package)
253 if version:
254 print(f"{package}=={version} at {_get_package_root(package, directory_name)}")
255 else:
256 print(f"{package} not installed")
257
258 if platform.system() == "Windows":
259 print(f"\nEnvironment variable:\nPATH={os.environ.get('PATH', '(unset)')}")
260 elif platform.system() == "Linux":
261 print(f"\nEnvironment variable:\nLD_LIBRARY_PATH={os.environ.get('LD_LIBRARY_PATH', '(unset)')}")
262
263 if importlib.util.find_spec("psutil"):
264
265 def is_target_dll(path: str):
266 target_keywords = ["vcruntime140", "msvcp140"]
267 if cuda_version:
268 target_keywords = ["cufft", "cublas", "cudart", "nvrtc", "curand", "cudnn", *target_keywords]
269 return any(keyword in path for keyword in target_keywords)
270
271 import psutil # noqa: PLC0415
272
273 p = psutil.Process(os.getpid())
274
275 print("\nList of loaded DLLs:")
276 for lib in p.memory_maps():
277 if is_target_dll(lib.path.lower()):
278 print(lib.path)
279
280 if cuda_version:
281 if importlib.util.find_spec("cpuinfo") and importlib.util.find_spec("py3nvml"):
282 from .transformers.machine_info import get_device_info # noqa: PLC0415
283
284 print("\nDevice information:")
285 print(get_device_info())
286 else:
287 print("please `pip install py-cpuinfo py3nvml` to show device information.")
288 else:
289 print("please `pip install psutil` to show loaded DLLs.")
290
291
292def preload_dlls(cuda: bool = True, cudnn: bool = True, msvc: bool = True, directory=None):
293 """Preload CUDA 12.x+ and cuDNN 9.x DLLs in Windows or Linux, and MSVC runtime DLLs in Windows.
294
295 When the installed PyTorch is compatible (using same major version of CUDA and cuDNN),
296 there is no need to call this function if `import torch` is done before `import onnxruntime`.
297
298 Args:
299 cuda (bool, optional): enable loading CUDA DLLs. Defaults to True.
300 cudnn (bool, optional): enable loading cuDNN DLLs. Defaults to True.
301 msvc (bool, optional): enable loading MSVC DLLs in Windows. Defaults to True.
302 directory(str, optional): a directory contains CUDA or cuDNN DLLs. It can be an absolute path,
303 or a path relative to the directory of this file.
304 If directory is None (default value), the search order: the lib directory of compatible PyTorch in Windows,
305 nvidia site packages, default DLL loading paths.
306 If directory is empty string (""), the search order: nvidia site packages, default DLL loading paths.
307 If directory is a path, the search order: the directory, default DLL loading paths.
308 """
309 import ctypes # noqa: PLC0415
310 import os # noqa: PLC0415
311 import platform # noqa: PLC0415
312 import sys # noqa: PLC0415
313
314 if platform.system() not in ["Windows", "Linux"]:
315 return
316
317 is_windows = platform.system() == "Windows"
318 if is_windows and msvc:
319 try:
320 ctypes.CDLL("vcruntime140.dll")
321 ctypes.CDLL("msvcp140.dll")
322 if platform.machine() != "ARM64":
323 ctypes.CDLL("vcruntime140_1.dll")
324 except OSError:
325 print("Microsoft Visual C++ Redistributable is not installed, this may lead to the DLL load failure.")
326 print("It can be downloaded at https://aka.ms/vs/17/release/vc_redist.x64.exe.")
327
328 # Check if CUDA version is supported (12.x or 13.x+)
329 ort_cuda_major = None
330 if cuda_version:
331 try:
332 ort_cuda_major = int(cuda_version.split(".")[0])
333 if ort_cuda_major < 12 and (cuda or cudnn):
334 print(
335 f"\033[33mWARNING: {package_name} is built with CUDA {cuda_version}, which is not supported for preloading. "
336 f"CUDA 12.x or newer is required. Call preload_dlls with cuda=False and cudnn=False.\033[0m"
337 )
338 return
339 except ValueError:
340 print(
341 f"\033[33mWARNING: Unable to parse CUDA version '{cuda_version}'. "
342 "Skipping DLL preloading. Call preload_dlls with cuda=False and cudnn=False.\033[0m"
343 )
344 return
345 elif cuda or cudnn:
346 # No CUDA version info available but CUDA/cuDNN preloading requested
347 return
348
349 is_cuda_cudnn_imported_by_torch = False
350
351 if is_windows:
352 torch_version = _get_package_version("torch")
353 # Check if torch CUDA version matches onnxruntime CUDA version
354 torch_cuda_major = None
355 if torch_version and "+cu" in torch_version:
356 with contextlib.suppress(ValueError):
357 # Extract CUDA version from torch (e.g., "2.0.0+cu121" -> 12)
358 cu_part = torch_version.split("+cu")[1]
359 torch_cuda_major = int(cu_part[:2]) # First 2 digits are major version
360
361 is_torch_cuda_compatible = (
362 torch_cuda_major == ort_cuda_major if (torch_cuda_major and ort_cuda_major) else False
363 )
364
365 if "torch" in sys.modules:
366 is_cuda_cudnn_imported_by_torch = is_torch_cuda_compatible
367 if torch_cuda_major and ort_cuda_major and torch_cuda_major != ort_cuda_major:
368 print(
369 f"\033[33mWARNING: The installed PyTorch {torch_version} uses CUDA {torch_cuda_major}.x, "
370 f"but {package_name} is built with CUDA {ort_cuda_major}.x. "
371 f"Please install PyTorch for CUDA {ort_cuda_major}.x to be compatible.\033[0m"
372 )
373
374 if is_torch_cuda_compatible and directory is None:
375 torch_root = _get_package_root("torch", "torch")
376 if torch_root:
377 directory = os.path.join(torch_root, "lib")
378
379 base_directory = directory or ".."
380 if not os.path.isabs(base_directory):
381 base_directory = os.path.join(os.path.dirname(__file__), base_directory)
382 base_directory = os.path.normpath(base_directory)
383 if not os.path.isdir(base_directory):
384 raise RuntimeError(f"Invalid parameter of directory={directory}. The directory does not exist!")
385
386 if is_cuda_cudnn_imported_by_torch:
387 # In Windows, PyTorch has loaded CUDA and cuDNN DLLs during `import torch`, no need to load them again.
388 print("Skip loading CUDA and cuDNN DLLs since torch is imported.")
389 return
390
391 # Try load DLLs from nvidia site packages.
392 dll_paths = _get_nvidia_dll_paths(is_windows, cuda, cudnn)
393 loaded_dlls = []
394 for relative_path in dll_paths:
395 dll_path = (
396 os.path.join(base_directory, relative_path[-1])
397 if directory
398 else os.path.join(base_directory, *relative_path)
399 )
400 if os.path.isfile(dll_path):
401 try:
402 _ = ctypes.CDLL(dll_path)
403 loaded_dlls.append(relative_path[-1])
404 except Exception as e:
405 print(f"Failed to load {dll_path}: {e}")
406
407 # Try load DLLs with default path settings.
408 has_failure = False
409 for relative_path in dll_paths:
410 dll_filename = relative_path[-1]
411 if dll_filename not in loaded_dlls:
412 try:
413 _ = ctypes.CDLL(dll_filename)
414 except Exception as e:
415 has_failure = True
416 print(f"Failed to load {dll_filename}: {e}")
417
418 if has_failure:
419 print("Please follow https://onnxruntime.ai/docs/install/#cuda-and-cudnn to install CUDA and CuDNN.")
420