Implement CUDA_PYTHON_CUDA_HOME_PRIORITY first, last, with default first

Author: rwgk

cuda_bindings/cuda/bindings/_path_finder/find_nvidia_dynamic_library.py

@@ -53,15 +53,23 @@ def _find_dll_using_nvidia_bin_dirs(libname, lib_searched_for, error_messages, a
     return None
 
 
-def _get_cuda_home():
+def _get_cuda_home(priority):
+    supported_priorities = ("first", "last")
+    assert priority in supported_priorities
+    env_priority = os.environ.get("CUDA_PYTHON_CUDA_HOME_PRIORITY")
+    if env_priority:
+        if env_priority not in supported_priorities:
+            raise RuntimeError(f"Invalid CUDA_PYTHON_CUDA_HOME_PRIORITY {env_priority!r} ({supported_priorities=})")
+        if priority != env_priority:
+            return None
     cuda_home = os.environ.get("CUDA_HOME")
     if cuda_home is None:
         cuda_home = os.environ.get("CUDA_PATH")
     return cuda_home
 
 
-def _find_lib_dir_using_cuda_home(libname):
-    cuda_home = _get_cuda_home()
+def _find_lib_dir_using_cuda_home(libname, priority):
+    cuda_home = _get_cuda_home(priority)
     if cuda_home is None:
         return None
     if IS_WINDOWS:
@@ -126,7 +134,7 @@ def __init__(self, libname: str):
         self.attachments = []
         self.abs_path = None
 
-        cuda_home_lib_dir = _find_lib_dir_using_cuda_home(libname)
+        cuda_home_lib_dir = _find_lib_dir_using_cuda_home(libname, "first")
         if IS_WINDOWS:
             self.lib_searched_for = f"{libname}*.dll"
             if cuda_home_lib_dir is not None:
@@ -148,6 +156,18 @@ def __init__(self, libname: str):
                     libname, self.lib_searched_for, self.error_messages, self.attachments
                 )
 
+    def retry_with_cuda_home_priority_last(self):
+        cuda_home_lib_dir = _find_lib_dir_using_cuda_home(self.libname, "last")
+        if cuda_home_lib_dir is not None:
+            if IS_WINDOWS:
+                self.abs_path = _find_dll_using_lib_dir(
+                    cuda_home_lib_dir, self.libname, self.error_messages, self.attachments
+                )
+            else:
+                self.abs_path = _find_so_using_lib_dir(
+                    cuda_home_lib_dir, self.lib_searched_for, self.error_messages, self.attachments
+                )
+
     def retry_with_anchor_abs_path(self, anchor_abs_path):
         assert self.libname == "nvvm"
         nvvm_lib_dir = _find_nvvm_lib_dir_from_anchor_abs_path(anchor_abs_path)

cuda_bindings/cuda/bindings/_path_finder/load_nvidia_dynamic_library.py

@@ -62,14 +62,16 @@ def _load_nvidia_dynamic_library_no_cache(libname: str) -> LoadedDL:
         loaded = load_with_system_search(libname, found.lib_searched_for)
         if loaded is not None:
             return loaded
-        if libname == "nvvm":
+        if libname == "nvvm" and ANCHOR_LIBNAME is not None:
             loaded_anchor = check_if_already_loaded_from_elsewhere(ANCHOR_LIBNAME)
             if loaded_anchor is not None:
                 found.retry_with_anchor_abs_path(loaded_anchor.abs_path)
             else:
                 anchor_abs_path = _load_anchor_in_subprocess(ANCHOR_LIBNAME, found.error_messages)
                 if anchor_abs_path is not None:
                     found.retry_with_anchor_abs_path(anchor_abs_path)
+        if found.abs_path is None:
+            found.retry_with_cuda_home_priority_last()
         found.raise_if_abs_path_is_None()
 
     # Load the library from the found path