Add robust dependency check for Python package (#6436)

* Add robust dependency check for Python package

* Add version_info.py to .gitignore

* Fix Linux build

* Fix Windows CPU build

* Fix Windows 32-bit build

* Minor tweak

* Generate version_info.py earlier in onnxruntime_python.cmake

* Print a user-friendly message if cuDNN is not found in

* Relax version requirements for CUDA 11 - only the major version has to match

* Fix PATH environment variable to include CUDA 11 in 'Python packaging pipeline' (Windows/GPU)

* Fix the build with cuDNN 7

.gitignore

@@ -46,6 +46,7 @@ java/gradlew.bat
 java/gradle
 java/.gradle
 java/hs_*.log
+onnxruntime/python/version_info.py
 /tools/perf_util/target/classes/com/msft/send_perf_metrics
 /tools/perf_util/send_perf_metrics.iml
 /tools/perf_util/target/classes

cmake/onnxruntime_python.cmake

@@ -153,6 +153,36 @@ else()
   set_target_properties(onnxruntime_pybind11_state PROPERTIES SUFFIX ".so")
 endif()
 
+# Generate version_info.py in Windows build.
+# Has to be done before onnxruntime_python_srcs is set.
+if (WIN32)
+  set(VERSION_INFO_FILE "${ONNXRUNTIME_ROOT}/python/version_info.py")
+
+  if (onnxruntime_USE_CUDA)
+    file(WRITE "${VERSION_INFO_FILE}" "use_cuda = True\n")
+
+    file(GLOB CUDNN_DLL_PATH "${onnxruntime_CUDNN_HOME}/bin/cudnn64_*.dll")
+    if (NOT CUDNN_DLL_PATH)
+      message(FATAL_ERROR "cuDNN not found in ${onnxruntime_CUDNN_HOME}")
+    endif()
+    get_filename_component(CUDNN_DLL_NAME ${CUDNN_DLL_PATH} NAME_WE)
+    string(REPLACE "cudnn64_" "" CUDNN_VERSION "${CUDNN_DLL_NAME}")
+
+    file(APPEND "${VERSION_INFO_FILE}"
+      "cuda_version = \"${onnxruntime_CUDA_VERSION}\"\n"
+      "cudnn_version = \"${CUDNN_VERSION}\"\n"
+    )
+  else()
+    file(WRITE "${VERSION_INFO_FILE}" "use_cuda = False\n")
+  endif()
+
+  if ("${MSVC_TOOLSET_VERSION}" STREQUAL "142")
+    file(APPEND "${VERSION_INFO_FILE}" "vs2019 = True\n")
+  else()
+    file(APPEND "${VERSION_INFO_FILE}" "vs2019 = False\n")
+  endif()
+endif()
+
 file(GLOB onnxruntime_backend_srcs CONFIGURE_DEPENDS
     "${ONNXRUNTIME_ROOT}/python/backend/*.py"
 )

onnxruntime/__init__.py

@@ -10,17 +10,6 @@ or the `Github project <https://github.com/microsoft/onnxruntime/>`_.
 __version__ = "1.7.0"
 __author__ = "Microsoft"
 
-import os
-import platform
-import sys
-
-# Python 3.8 (and later) on Windows doesn't search system PATH when loading DLLs,
-# so CUDA location needs to be specified explicitly. This needs to be done before importing
-# onnxruntime.capi._pybind_state
-if "CUDA_PATH" in os.environ and platform.system() == "Windows" and sys.version_info >= (3, 8):
-    cuda_bin_dir = os.path.join(os.environ["CUDA_PATH"], "bin")
-    os.add_dll_directory(cuda_bin_dir)
-
 from onnxruntime.capi._pybind_state import get_all_providers, get_available_providers, get_device, set_seed, \
     RunOptions, SessionOptions, set_default_logger_severity, enable_telemetry_events, disable_telemetry_events, \
     NodeArg, ModelMetadata, GraphOptimizationLevel, ExecutionMode, ExecutionOrder, OrtDevice, SessionIOBinding, \

onnxruntime/python/_pybind_state.py

@@ -2,25 +2,53 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
-
+"""
+Ensure that dependencies are available and then load the extension module.
+"""
 import os
 import platform
-import warnings
-import onnxruntime.capi._ld_preload  # noqa: F401
+import sys
 
-try:
-    from onnxruntime.capi.onnxruntime_pybind11_state import *  # noqa
-except ImportError as e:
-    warnings.warn("Cannot load onnxruntime.capi. Error: '{0}'.".format(str(e)))
+from . import _ld_preload  # noqa: F401
 
-    # If on Windows, check if this import error is caused by the user not installing the 2019 VC Runtime
-    # The VC Redist installer usually puts the VC Runtime dlls in the System32 folder
-    # This may not always paint the true picture as anyone building from source using VS 2017 might hit this error
-    # because the machine might be missing the 2019 VC Runtime but it is not actually needed in that case and the
-    # import error might actually be due to some other reason.
-    # TODO: Add a guard against False Positive error message
-    # As a proxy for checking if the 2019 VC Runtime is installed,
-    # we look for a specific dll only shipped with the 2019 VC Runtime
-    if platform.system().lower() == 'windows' and not os.path.isfile('c:\\Windows\\System32\\vcruntime140_1.dll'):
-        warnings.warn("Unless you have built the wheel using VS 2017, "
-                      "please install the 2019 Visual C++ runtime and then try again")
+if platform.system() == "Windows":
+    from . import version_info
+
+    if version_info.use_cuda:
+        cuda_version_major, cuda_version_minor = version_info.cuda_version.split(".")
+        if int(cuda_version_major) < 11:
+            # Prior to CUDA 11 both major and minor version at build time/runtime have to match.
+            cuda_env_variable = f"CUDA_PATH_V{cuda_version_major}_{cuda_version_minor}"
+            if cuda_env_variable not in os.environ:
+                raise ImportError(f"CUDA Toolkit {version_info.cuda_version} not installed on the machine.")
+        else:
+            # With CUDA 11 and newer only the major version at build time/runtime has to match.
+            # Use the most recent minor version available.
+            cuda_env_variable = None
+            for i in range(9, -1, -1):
+                if f"CUDA_PATH_V{cuda_version_major}_{i}" in os.environ:
+                    cuda_env_variable = f"CUDA_PATH_V{cuda_version_major}_{i}"
+                    break
+            if not cuda_env_variable:
+                raise ImportError(f"CUDA Toolkit {cuda_version_major}.x not installed on the machine.")
+
+        cuda_bin_dir = os.path.join(os.environ[cuda_env_variable], "bin")
+        if not os.path.isfile(os.path.join(cuda_bin_dir, f"cudnn64_{version_info.cudnn_version}.dll")):
+            raise ImportError(f"cuDNN {version_info.cudnn_version} not installed in {cuda_bin_dir}.")
+
+        if sys.version_info >= (3, 8):
+            # Python 3.8 (and later) doesn't search system PATH when loading DLLs, so the CUDA location needs to be
+            # specified explicitly using the new API introduced in Python 3.8.
+            os.add_dll_directory(cuda_bin_dir)
+        else:
+            # Python 3.7 (and earlier) searches directories listed in PATH variable.
+            # Make sure that the target CUDA version is at the beginning (important if multiple CUDA versions are
+            # installed on the machine.)
+            os.environ["PATH"] += cuda_bin_dir + os.pathsep + os.environ["PATH"]
+
+    if version_info.vs2019 and platform.architecture()[0] == "64bit":
+        if not os.path.isfile("C:\\Windows\\System32\\vcruntime140_1.dll"):
+            raise ImportError(
+                "Microsoft Visual C++ Redistributable for Visual Studio 2019 not installed on the machine.")
+
+from .onnxruntime_pybind11_state import *  # noqa

tools/ci_build/build.py

@@ -631,9 +631,10 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
         "-Donnxruntime_DEV_MODE=" + use_dev_mode(args),
         "-DPYTHON_EXECUTABLE=" + sys.executable,
         "-Donnxruntime_USE_CUDA=" + ("ON" if args.use_cuda else "OFF"),
+        "-Donnxruntime_CUDA_VERSION=" + (args.cuda_version if args.use_cuda else ""),
+        "-Donnxruntime_CUDA_HOME=" + (cuda_home if args.use_cuda else ""),
         "-Donnxruntime_CUDNN_HOME=" + (cudnn_home if args.use_cuda else ""),
         "-Donnxruntime_USE_FEATURIZERS=" + ("ON" if args.use_featurizers else "OFF"),
-        "-Donnxruntime_CUDA_HOME=" + (cuda_home if args.use_cuda else ""),
         "-Donnxruntime_USE_MIMALLOC_STL_ALLOCATOR=" + (
             "ON" if args.use_mimalloc == "stl" or args.use_mimalloc == "all" else "OFF"),
         "-Donnxruntime_USE_MIMALLOC_ARENA_ALLOCATOR=" + (
@@ -1903,6 +1904,12 @@ def main():
             install_python_deps(args.numpy_version)
         if args.enable_onnx_tests:
             setup_test_data(build_dir, configs)
+        if args.use_cuda and args.cuda_version is None:
+            if is_windows():
+                # cuda_version is used while generating version_info.py on Windows.
+                raise BuildError("cuda_version must be specified on Windows.")
+            else:
+                args.cuda_version = ""
         generate_build_tree(
             cmake_path, source_dir, build_dir, cuda_home, cudnn_home, rocm_home, mpi_home, nccl_home,
             tensorrt_home, migraphx_home, acl_home, acl_libs, armnn_home, armnn_libs,