Build pylibraft with scikit-build (#633)

This PR changes the build system for pylibraft to use scikit-build and CMake, allowing it to standardize more logic across all of RAPIDS and leverage the benefits of CMake to simplify building the C++ libraries along with the Python packages.

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - AJ Schmidt (https://github.com/ajschmidt8)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #633

.gitignore

@@ -24,6 +24,10 @@ log
 .DS_Store
 dask-worker-space/
 *.egg-info/
+
+## scikit-build
+_skbuild
+
 ## eclipse
 .project
 .cproject

build.sh

@@ -19,7 +19,7 @@ ARGS=$*
 REPODIR=$(cd $(dirname $0); pwd)
 
 VALIDARGS="clean libraft pyraft pylibraft docs tests bench clean -v -g --install --compile-libs --compile-nn --compile-dist --allgpuarch --no-nvtx --show_depr_warn -h --buildfaiss --minimal-deps"
-HELP="$0 [<target> ...] [<flag> ...]
+HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"]
  where <target> is:
    clean            - remove all existing build artifacts and configuration (start over)
    libraft          - build the raft C++ code only. Also builds the C-wrapper library
@@ -31,26 +31,27 @@ HELP="$0 [<target> ...] [<flag> ...]
    bench            - build the benchmarks
 
  and <flag> is:
-   -v               - verbose build mode
-   -g               - build for debug
-   --compile-libs   - compile shared libraries for all components
-   --compile-nn     - compile shared library for nn component
-   --compile-dist   - compile shared library for distance component
-   --minimal-deps   - disables dependencies like thrust so they can be overridden.
-                      can be useful for a pure header-only install
-   --allgpuarch     - build for all supported GPU architectures
-   --buildfaiss     - build faiss statically into raft
-   --install        - install cmake targets
-   --no-nvtx        - disable nvtx (profiling markers), but allow enabling it in downstream projects
-   --show_depr_warn - show cmake deprecation warnings
-   -h               - print this text
+   -v                          - verbose build mode
+   -g                          - build for debug
+   --compile-libs              - compile shared libraries for all components
+   --compile-nn                - compile shared library for nn component
+   --compile-dist              - compile shared library for distance component
+   --minimal-deps              - disables dependencies like thrust so they can be overridden.
+                                 can be useful for a pure header-only install
+   --allgpuarch                - build for all supported GPU architectures
+   --buildfaiss                - build faiss statically into raft
+   --install                   - install cmake targets
+   --no-nvtx                   - disable nvtx (profiling markers), but allow enabling it in downstream projects
+   --show_depr_warn            - show cmake deprecation warnings
+   --cmake-args=\\\"<args>\\\" - pass arbitrary list of CMake configuration options (escape all quotes in argument)
+   -h                          - print this text
 
  default action (no args) is to build both libraft and pyraft targets
 "
 LIBRAFT_BUILD_DIR=${LIBRAFT_BUILD_DIR:=${REPODIR}/cpp/build}
 SPHINX_BUILD_DIR=${REPODIR}/docs
 PY_RAFT_BUILD_DIR=${REPODIR}/python/raft/build
-PY_LIBRAFT_BUILD_DIR=${REPODIR}/python/pylibraft/build
+PY_LIBRAFT_BUILD_DIR=${REPODIR}/python/pylibraft/_skbuild
 BUILD_DIRS="${LIBRAFT_BUILD_DIR} ${PY_RAFT_BUILD_DIR} ${PY_LIBRAFT_BUILD_DIR}"
 
 # Set defaults for vars modified by flags to this script
@@ -91,28 +92,51 @@ function hasArg {
     (( ${NUMARGS} != 0 )) && (echo " ${ARGS} " | grep -q " $1 ")
 }
 
+function cmakeArgs {
+    # Check for multiple cmake args options
+    if [[ $(echo $ARGS | { grep -Eo "\-\-cmake\-args" || true; } | wc -l ) -gt 1 ]]; then
+        echo "Multiple --cmake-args options were provided, please provide only one: ${ARGS}"
+        exit 1
+    fi
+
+    # Check for cmake args option
+    if [[ -n $(echo $ARGS | { grep -E "\-\-cmake\-args" || true; } ) ]]; then
+        # There are possible weird edge cases that may cause this regex filter to output nothing and fail silently
+        # the true pipe will catch any weird edge cases that may happen and will cause the program to fall back
+        # on the invalid option error
+        CMAKE_ARGS=$(echo $ARGS | { grep -Eo "\-\-cmake\-args=\".+\"" || true; })
+        if [[ -n ${CMAKE_ARGS} ]]; then
+            # Remove the full  CMAKE_ARGS argument from list of args so that it passes validArgs function
+            ARGS=${ARGS//$CMAKE_ARGS/}
+            # Filter the full argument down to just the extra string that will be added to cmake call
+            CMAKE_ARGS=$(echo $CMAKE_ARGS | grep -Eo "\".+\"" | sed -e 's/^"//' -e 's/"$//')
+        fi
+    fi
+}
+
 if hasArg -h || hasArg --help; then
     echo "${HELP}"
     exit 0
 fi
 
 # Check for valid usage
 if (( ${NUMARGS} != 0 )); then
+    cmakeArgs
     for a in ${ARGS}; do
-  if ! (echo " ${VALIDARGS} " | grep -q " ${a} "); then
-      echo "Invalid option: ${a}"
-      exit 1
-  fi
+        if ! (echo " ${VALIDARGS} " | grep -q " ${a} "); then
+            echo "Invalid option: ${a}"
+            exit 1
+        fi
     done
 fi
 
 # Process flags
 if hasArg --install; then
-  INSTALL_TARGET="install"
+    INSTALL_TARGET="install"
 fi
 
 if hasArg --minimal-deps; then
-  ENABLE_thrust_DEPENDENCY=OFF
+    ENABLE_thrust_DEPENDENCY=OFF
 fi
 
 if hasArg -v; then
@@ -128,7 +152,7 @@ if hasArg --allgpuarch; then
 fi
 
 if hasArg --compile-libs || (( ${NUMARGS} == 0 )); then
-  COMPILE_LIBRARIES=ON
+    COMPILE_LIBRARIES=ON
 fi
 
 if hasArg --compile-nn || hasArg --compile-libs || (( ${NUMARGS} == 0 )); then
@@ -167,11 +191,11 @@ if hasArg clean; then
     CLEAN=1
 fi
 if hasArg uninstall; then
-  UNINSTALL=1
+    UNINSTALL=1
 fi
 
 if [[ ${CMAKE_TARGET} == "" ]]; then
-  CMAKE_TARGET="all"
+    CMAKE_TARGET="all"
 fi
 
 # If clean given, run it prior to any other steps
@@ -198,8 +222,8 @@ fi
 
 # Pyraft requires ucx + nccl
 if (( ${NUMARGS} == 0 )) || hasArg pyraft || hasArg docs; then
-  ENABLE_nccl_DEPENDENCY=ON
-  ENABLE_ucx_DEPENDENCY=ON
+    ENABLE_nccl_DEPENDENCY=ON
+    ENABLE_ucx_DEPENDENCY=ON
 fi
 ################################################################################
 # Configure for building all C++ targets
@@ -229,7 +253,8 @@ if (( ${NUMARGS} == 0 )) || hasArg libraft || hasArg docs || hasArg tests || has
           -DRAFT_USE_FAISS_STATIC=${BUILD_STATIC_FAISS} \
           -DRAFT_ENABLE_nccl_DEPENDENCY=${ENABLE_nccl_DEPENDENCY} \
           -DRAFT_ENABLE_ucx_DEPENDENCY=${ENABLE_ucx_DEPENDENCY} \
-          -DRAFT_ENABLE_thrust_DEPENDENCY=${ENABLE_thrust_DEPENDENCY}
+          -DRAFT_ENABLE_thrust_DEPENDENCY=${ENABLE_thrust_DEPENDENCY} \
+          ${CMAKE_ARGS}
 
   if [[ ${CMAKE_TARGET} != "" ]]; then
       echo "-- Compiling targets: ${CMAKE_TARGET}, verbose=${VERBOSE_FLAG}"
@@ -256,10 +281,9 @@ fi
 if (( ${NUMARGS} == 0 )) || hasArg pylibraft; then
 
     cd ${REPODIR}/python/pylibraft
+    python setup.py build_ext -j${PARALLEL_LEVEL:-1} --inplace -- -DCMAKE_PREFIX_PATH=${INSTALL_PREFIX} -DCMAKE_LIBRARY_PATH=${LIBRAFT_BUILD_DIR} ${CMAKE_ARGS}
     if [[ ${INSTALL_TARGET} != "" ]]; then
-        python setup.py build_ext -j${PARALLEL_LEVEL:-1} --inplace --library-dir=${LIBRAFT_BUILD_DIR} install --single-version-externally-managed --record=record.txt
-    else
-        python setup.py build_ext -j${PARALLEL_LEVEL:-1} --inplace --library-dir=${LIBRAFT_BUILD_DIR}
+        python setup.py install --single-version-externally-managed --record=record.txt -- -DCMAKE_PREFIX_PATH=${INSTALL_PREFIX} ${CMAKE_ARGS}
     fi
 fi
 

ci/gpu/build.sh

@@ -103,9 +103,9 @@ export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH
 gpuci_logger "Build C++ and Python targets"
 # These should link against the existing shared libs
 if hasArg --skip-tests; then
-  "$WORKSPACE/build.sh" pyraft pylibraft libraft -v
+  "$WORKSPACE/build.sh" pyraft pylibraft libraft -v --cmake-args=\"-DFIND_RAFT_CPP=ON\"
 else
-  "$WORKSPACE/build.sh" pyraft pylibraft libraft tests bench  -v
+  "$WORKSPACE/build.sh" pyraft pylibraft libraft tests bench -v --cmake-args=\"-DFIND_RAFT_CPP=ON\"
 fi
 
 gpuci_logger "sccache stats"

ci/release/update-version.sh

@@ -32,7 +32,9 @@ function sed_runner() {
 }
 
 sed_runner 's/'"RAFT VERSION .* LANGUAGES"'/'"RAFT VERSION ${NEXT_FULL_TAG} LANGUAGES"'/g' cpp/CMakeLists.txt
+sed_runner 's/'"pylibraft_version .*)"'/'"pylibraft_version ${NEXT_FULL_TAG})"'/g' python/pylibraft/CMakeLists.txt
 sed_runner 's/'"branch-.*\/RAPIDS.cmake"'/'"branch-${NEXT_SHORT_TAG}\/RAPIDS.cmake"'/g' cpp/CMakeLists.txt
+sed_runner 's/'"branch-.*\/RAPIDS.cmake"'/'"branch-${NEXT_SHORT_TAG}\/RAPIDS.cmake"'/g' python/pylibraft/CMakeLists.txt
 
 for FILE in conda/environments/*.yml; do
   sed_runner "s/ucx-py=.*/ucx-py=${NEXT_UCX_PY_VERSION}/g" ${FILE};

conda/environments/raft_dev_cuda11.0.yml

@@ -8,6 +8,9 @@ dependencies:
 - cudatoolkit=11.0
 - clang=11.1.0
 - clang-tools=11.1.0
+- cython>=0.29,<0.30
+- cmake>=3.20.1,!=3.23.0
+- scikit-build>=0.13.1
 - rapids-build-env=22.02.*
 - rapids-notebook-env=22.02.*
 - rapids-doc-env=22.02.*

conda/environments/raft_dev_cuda11.2.yml

@@ -8,6 +8,9 @@ dependencies:
 - cudatoolkit=11.2
 - clang=11.1.0
 - clang-tools=11.1.0
+- cython>=0.29,<0.30
+- cmake>=3.20.1,!=3.23.0
+- scikit-build>=0.13.1
 - rapids-build-env=22.02.*
 - rapids-notebook-env=22.02.*
 - rapids-doc-env=22.02.*

conda/environments/raft_dev_cuda11.4.yml

@@ -8,6 +8,9 @@ dependencies:
 - cudatoolkit=11.4
 - clang=11.1.0
 - clang-tools=11.1.0
+- cython>=0.29,<0.30
+- cmake>=3.20.1,!=3.23.0
+- scikit-build>=0.13.1
 - rapids-build-env=22.02.*
 - rapids-notebook-env=22.02.*
 - rapids-doc-env=22.02.*

conda/environments/raft_dev_cuda11.5.yml

@@ -9,6 +9,9 @@ dependencies:
 - cuda-python >=11.5,<12.0
 - clang=11.1.0
 - clang-tools=11.1.0
+- cython>=0.29,<0.30
+- cmake>=3.20.1,!=3.23.0
+- scikit-build>=0.13.1
 - rapids-build-env=22.02.*
 - rapids-notebook-env=22.02.*
 - rapids-doc-env=22.02.*

conda/recipes/pylibraft/build.sh

@@ -2,4 +2,4 @@
 #!/usr/bin/env bash
 
 # This assumes the script is executed from the root of the repo directory
-./build.sh pylibraft --install --no-nvtx
+./build.sh pylibraft --install --no-nvtx --cmake-args=\"-DFIND_RAFT_CPP=ON\"

conda/recipes/pylibraft/meta.yaml

@@ -28,6 +28,8 @@ requirements:
     - python x.x
     - setuptools
     - cython>=0.29,<0.30
+    - cmake>=3.20.1,!=3.23.0
+    - scikit-build>=0.13.1
     - rmm {{ minor_version }}
     - libraft-headers {{ version }}
     - libraft-distance {{ version }}

cpp/CMakeLists.txt

@@ -62,8 +62,8 @@ option(DISABLE_OPENMP "Disable OpenMP" OFF)
 option(RAFT_NVTX "Enable nvtx markers" OFF)
 
 option(RAFT_COMPILE_LIBRARIES "Enable building raft shared library instantiations" ${BUILD_TESTS})
-option(RAFT_COMPILE_NN_LIBRARY "Enable building raft nearest neighbors shared library instantiations" OFF)
-option(RAFT_COMPILE_DIST_LIBRARY "Enable building raft distant shared library instantiations" OFF)
+option(RAFT_COMPILE_NN_LIBRARY "Enable building raft nearest neighbors shared library instantiations" ${RAFT_COMPILE_LIBRARIES})
+option(RAFT_COMPILE_DIST_LIBRARY "Enable building raft distant shared library instantiations" ${RAFT_COMPILE_LIBRARIES})
 option(RAFT_ENABLE_NN_DEPENDENCIES "Search for raft::nn dependencies like faiss" ${RAFT_COMPILE_LIBRARIES})
 
 option(RAFT_ENABLE_thrust_DEPENDENCY "Enable Thrust dependency" ON)
@@ -131,7 +131,12 @@ include(cmake/modules/ConfigureCUDA.cmake)
 ##############################################################################
 # - Requirements -------------------------------------------------------------
 
-if(distance IN_LIST raft_FIND_COMPONENTS OR RAFT_COMPILE_LIBRARIES OR RAFT_COMPILE_DIST_LIBRARY)
+if(RAFT_COMPILE_LIBRARIES)
+  set(RAFT_COMPILE_DIST_LIBRARY ON)
+  set(RAFT_COMPILE_NN_LIBRARY ON)
+endif()
+
+if(RAFT_COMPILE_DIST_LIBRARY OR distance IN_LIST raft_FIND_COMPONENTS)
   set(RAFT_ENABLE_cuco_DEPENDENCY ON)
 endif()
 
@@ -176,7 +181,7 @@ target_link_libraries(raft INTERFACE
 
 target_compile_features(raft INTERFACE cxx_std_17 $<BUILD_INTERFACE:cuda_std_17>)
 
-if(RAFT_COMPILE_LIBRARIES OR RAFT_COMPILE_DIST_LIBRARY OR RAFT_COMPILE_NN_LIBRARY)
+if(RAFT_COMPILE_DIST_LIBRARY OR RAFT_COMPILE_NN_LIBRARY)
   file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld"
           [=[
 SECTIONS
@@ -219,7 +224,7 @@ endif()
 
 set_target_properties(raft_distance PROPERTIES EXPORT_NAME distance)
 
-if(RAFT_COMPILE_LIBRARIES OR RAFT_COMPILE_DIST_LIBRARY)
+if(RAFT_COMPILE_DIST_LIBRARY)
   add_library(raft_distance_lib
     src/distance/pairwise_distance.cu
     src/distance/specializations/detail/canberra.cu
@@ -295,7 +300,7 @@ endif()
 
 set_target_properties(raft_nn PROPERTIES EXPORT_NAME nn)
 
-if(RAFT_COMPILE_LIBRARIES OR RAFT_COMPILE_NN_LIBRARY)
+if(RAFT_COMPILE_NN_LIBRARY)
   add_library(raft_nn_lib
           src/nn/specializations/ball_cover.cu
           src/nn/specializations/detail/ball_cover_lowdim_pass_one_2d.cu
@@ -364,6 +369,8 @@ if(TARGET raft_distance_lib)
           DESTINATION ${lib_dir}
           COMPONENT raft_distance
           EXPORT raft-distance-lib-exports)
+  install(DIRECTORY include/raft_distance
+          DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})
 endif()
 
 if(TARGET raft_nn_lib)

python/pylibraft/CMakeLists.txt

@@ -0,0 +1,70 @@
+# =============================================================================
+# Copyright (c) 2022, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under the License
+# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+# or implied. See the License for the specific language governing permissions and limitations under
+# the License.
+# =============================================================================
+
+cmake_minimum_required(VERSION 3.20.1 FATAL_ERROR)
+
+set(pylibraft_version 22.06.00)
+
+file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-22.06/RAPIDS.cmake
+        ${CMAKE_BINARY_DIR}/RAPIDS.cmake)
+include(${CMAKE_BINARY_DIR}/RAPIDS.cmake)
+
+project(
+  pylibraft
+  VERSION ${pylibraft_version}
+  LANGUAGES # TODO: Building Python extension modules via the python_extension_module requires the C
+            # language to be enabled here. The test project that is built in scikit-build to verify
+            # various linking options for the python library is hardcoded to build with C, so until
+            # that is fixed we need to keep C.
+            C
+            CXX)
+
+option(FIND_RAFT_CPP "Search for existing RAFT C++ installations before defaulting to local files"
+       OFF)
+
+# If the user requested it we attempt to find RAFT.
+if(FIND_RAFT_CPP)
+    find_package(raft ${pylibraft_version} REQUIRED COMPONENTS distance)
+    if(NOT TARGET raft::raft_distance_lib)
+        message(FATAL_ERROR "Building against a preexisting libraft library requires the distance component of that library to have been built!")
+    endif()
+else()
+  set(raft_FOUND OFF)
+endif()
+
+if(NOT raft_FOUND)
+  # TODO: This will not be necessary once we upgrade to CMake 3.22, which will
+  # pull in the required languages for the C++ project even if this project
+  # does not require those languges.
+  include(rapids-cuda)
+  rapids_cuda_init_architectures(pylibraft)
+  enable_language(CUDA)
+  # Since pylibraft only enables CUDA optionally we need to manually include the file that
+  # rapids_cuda_init_architectures relies on `project` including.
+  include("${CMAKE_PROJECT_pylibraft_INCLUDE}")
+
+  set(BUILD_TESTS OFF)
+  set(BUILD_BENCHMARKS OFF)
+  set(RAFT_COMPILE_DIST_LIBRARY ON)
+  add_subdirectory(../../cpp raft-cpp)
+
+  # When building the C++ libraries from source we must copy
+  # libraft_distance.so alongside the pairwise_distance Cython library.
+  install(TARGETS raft_distance_lib DESTINATION pylibraft/distance)
+endif()
+
+include(rapids-cython)
+rapids_cython_init()
+
+add_subdirectory(pylibraft/distance)