Build pyraft with scikit-build (#644)

This PR changes the pyraft build system to use scikit-build and CMake. It depends on the changes in #633 and should be merged after that PR.

This PR is largely complete, except that I am not yet sure exactly which libraries/headers are entirely necessary for the build. This version builds locally for me, but not all tests pass since I don't have nccl. However, I am a bit surprised to see the build succeeding without me even searching for UCX (perhaps this is because it's already on my system include path) so I need to figure out what do with UCX, NCCL, and the various CUDA math libraries (cusparse, cublas, etc).

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

Approvers:
  - Sevag Hanssian (https://github.com/sevagh)
  - Ray Douglass (https://github.com/raydouglass)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #644

build.sh

@@ -68,9 +68,6 @@ ENABLE_NN_DEPENDENCIES=OFF
 
 ENABLE_thrust_DEPENDENCY=ON
 
-ENABLE_ucx_DEPENDENCY=OFF
-ENABLE_nccl_DEPENDENCY=OFF
-
 NVTX=ON
 CLEAN=0
 UNINSTALL=0
@@ -220,11 +217,6 @@ if (( ${CLEAN} == 1 )); then
     cd ${REPODIR}
 fi
 
-# Pyraft requires ucx + nccl
-if (( ${NUMARGS} == 0 )) || hasArg pyraft || hasArg docs; then
-    ENABLE_nccl_DEPENDENCY=ON
-    ENABLE_ucx_DEPENDENCY=ON
-fi
 ################################################################################
 # Configure for building all C++ targets
 if (( ${NUMARGS} == 0 )) || hasArg libraft || hasArg docs || hasArg tests || hasArg bench; then
@@ -251,8 +243,6 @@ if (( ${NUMARGS} == 0 )) || hasArg libraft || hasArg docs || hasArg tests || has
           -DRAFT_COMPILE_NN_LIBRARY=${COMPILE_NN_LIBRARY} \
           -DRAFT_COMPILE_DIST_LIBRARY=${COMPILE_DIST_LIBRARY} \
           -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} \
           ${CMAKE_ARGS}
 
@@ -270,10 +260,9 @@ fi
 if (( ${NUMARGS} == 0 )) || hasArg pyraft || hasArg docs; then
 
     cd ${REPODIR}/python/raft
+    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/release/update-version.sh

@@ -33,8 +33,10 @@ 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/'"pyraft_version .*)"'/'"pyraft_version ${NEXT_FULL_TAG})"'/g' python/raft/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
+sed_runner 's/'"branch-.*\/RAPIDS.cmake"'/'"branch-${NEXT_SHORT_TAG}\/RAPIDS.cmake"'/g' python/pyraft/CMakeLists.txt
 
 for FILE in conda/environments/*.yml; do
   sed_runner "s/ucx-py=.*/ucx-py=${NEXT_UCX_PY_VERSION}/g" ${FILE};

conda/recipes/pyraft/build.sh

@@ -3,4 +3,4 @@
 # Copyright (c) 2022, NVIDIA CORPORATION.
 
 # This assumes the script is executed from the root of the repo directory
-./build.sh pyraft --install --no-nvtx
+./build.sh pyraft --install --no-nvtx --cmake-args=\"-DFIND_RAFT_CPP=ON\"

conda/recipes/pyraft/meta.yaml

@@ -29,6 +29,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 }}
     - cudatoolkit {{ cuda_version }}.*

cpp/CMakeLists.txt

@@ -73,10 +73,6 @@ if(BUILD_TESTS AND NOT RAFT_ENABLE_thrust_DEPENDENCY)
   set(RAFT_ENABLE_thrust_DEPENDENCY ON)
 endif()
 
-# Currently, UCX and NCCL are only needed to build Pyraft and so a simple find_package() is sufficient
-option(RAFT_ENABLE_nccl_DEPENDENCY "Enable NCCL dependency" OFF)
-option(RAFT_ENABLE_ucx_DEPENDENCY "Enable ucx dependency" OFF)
-
 option(RAFT_EXCLUDE_FAISS_FROM_ALL "Exclude FAISS targets from RAFT's 'all' target" ON)
 
 include(CMakeDependentOption)

python/raft/CMakeLists.txt

@@ -0,0 +1,68 @@
+# =============================================================================
+# 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(pyraft_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(
+  raft-python
+  VERSION ${pyraft_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 ${pyraft_version} REQUIRED)
+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_LIBRARIES OFF)
+  set(RAFT_COMPILE_DIST_LIBRARY OFF)
+  set(RAFT_COMPILE_NN_LIBRARY OFF)
+  add_subdirectory(../../cpp raft-cpp)
+endif()
+
+include(rapids-cython)
+rapids_cython_init()
+
+# TODO: Figure out which of ucx, nccl, cusolver, cusparse, and cublas I need to add as include directories or linked libraries for which components
+add_subdirectory(raft/common)
+add_subdirectory(raft/dask/common)
+add_subdirectory(raft/include_test)

cpp/cmake/thirdparty/get_nccl.cmake → python/raft/cmake/thirdparty/get_nccl.cmake

@@ -1,5 +1,5 @@
 #=============================================================================
-# Copyright (c) 2021, NVIDIA CORPORATION.
+# Copyright (c) 2021-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.
@@ -14,6 +14,7 @@
 # limitations under the License.
 #=============================================================================
 
+include(rapids-find)
 function(find_and_configure_nccl)
 
     if(TARGET NCCL::NCCL)

python/raft/pyproject.toml

@@ -0,0 +1,24 @@
+# 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.
+
+[build-system]
+
+requires = [
+    "wheel",
+    "setuptools",
+    "cython>=0.29,<0.30",
+    "scikit-build>=0.13.1",
+    "cmake>=3.20.1,!=3.23.0",
+    "ninja",
+]

python/raft/raft/common/CMakeLists.txt

@@ -0,0 +1,20 @@
+# =============================================================================
+# 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.
+# =============================================================================
+
+set(cython_sources cuda.pyx handle.pyx interruptible.pyx)
+set(linked_libraries raft::raft)
+rapids_cython_create_modules(
+    SOURCE_FILES "${cython_sources}"
+    LINKED_LIBRARIES "${linked_libraries}"
+    CXX)

python/raft/raft/dask/common/CMakeLists.txt

@@ -0,0 +1,23 @@
+# =============================================================================
+# 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.
+# =============================================================================
+
+include(${raft-python_SOURCE_DIR}/cmake/thirdparty/get_nccl.cmake)
+find_package(ucx REQUIRED)
+
+set(cython_sources comms_utils.pyx nccl.pyx)
+set(linked_libraries raft::raft NCCL::NCCL ucx::ucp)
+rapids_cython_create_modules(
+    SOURCE_FILES "${cython_sources}"
+    LINKED_LIBRARIES "${linked_libraries}"
+    CXX)

python/raft/raft/include_test/CMakeLists.txt

@@ -0,0 +1,20 @@
+# =============================================================================
+# 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.
+# =============================================================================
+
+set(cython_sources raft_include_test.pyx)
+set(linked_libraries raft::raft)
+rapids_cython_create_modules(
+    SOURCE_FILES "${cython_sources}"
+    LINKED_LIBRARIES "${linked_libraries}"
+    CXX)