Use library_dirs for cython linking, link cudatoolkit libs, allow setting UCX install location (#1698)

This PR is a continuation of #1694. Similar to rapidsai/cuml#4015, this PR updates setup.py to:

* Use `library_dirs` instead of `runtime_library_dirs` when linking Cython.
* Allow overriding UCX lib and include dirs via a `UCX_HOME` envvar.
* Link `cudart`, `cusparse`, and `cusolver`.

These are necessary to compile the Cython via `pip` when not inside a conda environment and when UCX is installed to a location other than `/usr` or `/usr/local`.

Authors:
  - Paul Taylor (https://github.com/trxcllnt)

Approvers:
  - Rick Ratzel (https://github.com/rlratzel)

URL: #1698

Author: trxcllnt

cpp/cmake/thirdparty/get_faiss.cmake

@@ -40,7 +40,10 @@ function(find_and_configure_faiss)
 
     if(FAISS_ADDED)
       set(FAISS_GPU_HEADERS ${FAISS_SOURCE_DIR} PARENT_SCOPE)
-      add_library(FAISS::FAISS ALIAS faiss)
+    endif()
+
+    if(TARGET faiss AND NOT TARGET FAISS::FAISS)
+        add_library(FAISS::FAISS ALIAS faiss)
     endif()
 
 endfunction()

python/setup.py

@@ -13,6 +13,7 @@
 
 import os
 import sys
+import sysconfig
 import shutil
 
 from setuptools import setup, find_packages, Command
@@ -29,13 +30,25 @@
 
 
 INSTALL_REQUIRES = ['numba', 'cython']
+CYTHON_FILES = ['cugraph/**/*.pyx']
+
+UCX_HOME = get_environment_option("UCX_HOME")
+CUDA_HOME = get_environment_option('CUDA_HOME')
+CONDA_PREFIX = get_environment_option('CONDA_PREFIX')
 
 conda_lib_dir = os.path.normpath(sys.prefix) + '/lib'
 conda_include_dir = os.path.normpath(sys.prefix) + '/include'
 
-CYTHON_FILES = ['cugraph/**/*.pyx']
+if CONDA_PREFIX:
+    conda_include_dir = CONDA_PREFIX + '/include'
+    conda_lib_dir = CONDA_PREFIX + '/lib'
+
+if not UCX_HOME:
+    UCX_HOME = CONDA_PREFIX if CONDA_PREFIX else os.sys.prefix
+
+ucx_include_dir = os.path.join(UCX_HOME, "include")
+ucx_lib_dir = os.path.join(UCX_HOME, "lib")
 
-CUDA_HOME = os.environ.get("CUDA_HOME", False)
 if not CUDA_HOME:
     path_to_cuda_gdb = shutil.which("cuda-gdb")
     if path_to_cuda_gdb is None:
@@ -53,11 +66,7 @@
     )
 
 cuda_include_dir = os.path.join(CUDA_HOME, "include")
-
-if (os.environ.get('CONDA_PREFIX', None)):
-    conda_prefix = os.environ.get('CONDA_PREFIX')
-    conda_include_dir = conda_prefix + '/include'
-    conda_lib_dir = conda_prefix + '/lib'
+cuda_lib_dir = os.path.join(CUDA_HOME, "lib64")
 
 # Optional location of C++ build folder that can be configured by the user
 libcugraph_path = get_environment_option('CUGRAPH_BUILD_PATH')
@@ -69,6 +78,9 @@
 # https://github.com/rapidsai/raft/issues/83
 raft_include_dir = use_raft_package(raft_path, libcugraph_path)
 
+if not libcugraph_path:
+    libcugraph_path = conda_lib_dir
+
 
 class CleanCommand(Command):
     """Custom clean command to tidy up the project root."""
@@ -101,16 +113,25 @@ def run(self):
 EXTENSIONS = [
     Extension("*",
               sources=CYTHON_FILES,
-              include_dirs=[conda_include_dir,
-                            '../cpp/include',
-                            "../thirdparty/cub",
-                            raft_include_dir,
-                            os.path.join(
-                                conda_include_dir, "libcudacxx"),
-                            cuda_include_dir],
-              library_dirs=[get_python_lib()],
-              runtime_library_dirs=[conda_lib_dir],
-              libraries=['cugraph', 'nccl'],
+              include_dirs=[
+                  conda_include_dir,
+                  ucx_include_dir,
+                  '../cpp/include',
+                  "../thirdparty/cub",
+                  raft_include_dir,
+                  os.path.join(conda_include_dir, "libcudacxx"),
+                  cuda_include_dir,
+                  os.path.dirname(sysconfig.get_path("include"))
+              ],
+              library_dirs=[
+                  get_python_lib(),
+                  conda_lib_dir,
+                  libcugraph_path,
+                  ucx_lib_dir,
+                  cuda_lib_dir,
+                  os.path.join(os.sys.prefix, "lib")
+              ],
+              libraries=['cudart', 'cusparse', 'cusolver', 'cugraph', 'nccl'],
               language='c++',
               extra_compile_args=['-std=c++17'])
 ]