From 62beb4aeff7f2953a6ceaccdf3fd25786de884fa Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Tue, 10 May 2022 17:01:45 -0700
Subject: [PATCH 01/12] First pass at create a scikit-build enabled build.

---
 python/raft/CMakeLists.txt                   |  59 ++++++
 python/raft/pyproject.toml                   |  24 +++
 python/raft/raft/common/CMakeLists.txt       |  20 ++
 python/raft/raft/dask/common/CMakeLists.txt  |  20 ++
 python/raft/raft/include_test/CMakeLists.txt |  20 ++
 python/raft/setup.py                         | 211 ++-----------------
 6 files changed, 161 insertions(+), 193 deletions(-)
 create mode 100644 python/raft/CMakeLists.txt
 create mode 100644 python/raft/pyproject.toml
 create mode 100644 python/raft/raft/common/CMakeLists.txt
 create mode 100644 python/raft/raft/dask/common/CMakeLists.txt
 create mode 100644 python/raft/raft/include_test/CMakeLists.txt

diff --git a/python/raft/CMakeLists.txt b/python/raft/CMakeLists.txt
new file mode 100644
index 0000000000..a0561a210b
--- /dev/null
+++ b/python/raft/CMakeLists.txt
@@ -0,0 +1,59 @@
+# =============================================================================
+# 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
+  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 COMPONENTS distance)
+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.
+  enable_language(CUDA)
+  set(BUILD_TESTS OFF)
+  set(BUILD_BENCHMARKS 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)
diff --git a/python/raft/pyproject.toml b/python/raft/pyproject.toml
new file mode 100644
index 0000000000..e3fe544bef
--- /dev/null
+++ b/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",
+]
diff --git a/python/raft/raft/common/CMakeLists.txt b/python/raft/raft/common/CMakeLists.txt
new file mode 100644
index 0000000000..59b336574e
--- /dev/null
+++ b/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)
diff --git a/python/raft/raft/dask/common/CMakeLists.txt b/python/raft/raft/dask/common/CMakeLists.txt
new file mode 100644
index 0000000000..a65a68ffa8
--- /dev/null
+++ b/python/raft/raft/dask/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 comms_utils.pyx nccl.pyx)
+set(linked_libraries raft::raft)
+rapids_cython_create_modules(
+    SOURCE_FILES "${cython_sources}"
+    LINKED_LIBRARIES "${linked_libraries}"
+    CXX)
diff --git a/python/raft/raft/include_test/CMakeLists.txt b/python/raft/raft/include_test/CMakeLists.txt
new file mode 100644
index 0000000000..5e417f42ee
--- /dev/null
+++ b/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)
diff --git a/python/raft/setup.py b/python/raft/setup.py
index 9fe22c096c..3d5b394bb7 100644
--- a/python/raft/setup.py
+++ b/python/raft/setup.py
@@ -14,192 +14,12 @@
 # limitations under the License.
 #
 
-import numpy
-import os
-import shutil
-import sys
-import sysconfig
-
-# Must import in this order:
-#   setuptools -> Cython.Distutils.build_ext -> setuptools.command.build_ext
-# Otherwise, setuptools.command.build_ext ends up inheriting from
-# Cython.Distutils.old_build_ext which we do not want
-import setuptools
-
-try:
-    from Cython.Distutils.build_ext import new_build_ext as _build_ext
-except ImportError:
-    from setuptools.command.build_ext import build_ext as _build_ext
-
-from distutils.sysconfig import get_python_lib
-
-import setuptools.command.build_ext
-from setuptools import find_packages, setup
-from setuptools.extension import Extension
-
-from setuputils import clean_folder
-from setuputils import get_environment_option
-from setuputils import get_cli_option
-
-from pathlib import Path
+from setuptools import find_packages
+from skbuild import setup
 
 import versioneer
 
 
-##############################################################################
-# - Dependencies include and lib folder setup --------------------------------
-
-install_requires = [
-    'cython'
-]
-
-cuda_home = get_environment_option("CUDA_HOME")
-
-if not cuda_home:
-    cuda_home = (
-        os.popen('echo "$(dirname $(dirname $(which nvcc)))"').read().strip()
-    )
-    print("-- Using nvcc to detect CUDA, found at " + str(cuda_home))
-cuda_include_dir = os.path.join(cuda_home, "include")
-cuda_lib_dir = os.path.join(cuda_home, "lib64")
-
-clean_artifacts = get_cli_option('clean')
-single_gpu_build = get_cli_option('--singlegpu')
-
-ucx_home = get_environment_option("UCX_HOME") or os.sys.prefix
-ucx_lib_dir = os.path.join(ucx_home, "lib")
-ucx_include_dir = os.path.join(ucx_home, "include")
-
-rmm_include_dir = (
-    get_environment_option("RMM_INCLUDE_DIR")
-    or os.path.join(os.sys.prefix, "include")
-)
-thrust_include_dir = (
-    get_environment_option("THRUST_INCLUDE_DIR")
-    or os.path.join(os.sys.prefix, "include")
-)
-spdlog_include_dir = (
-    get_environment_option("SPDLOG_INCLUDE_DIR")
-    or os.path.join(os.sys.prefix, "include")
-)
-
-##############################################################################
-# - Clean target -------------------------------------------------------------
-
-if clean_artifacts:
-    print("-- Cleaning all Python and Cython build artifacts...")
-
-    try:
-        setup_file_path = str(Path(__file__).parent.absolute())
-        shutil.rmtree(setup_file_path + '/.pytest_cache', ignore_errors=True)
-        shutil.rmtree(setup_file_path + '/_external_repositories',
-                      ignore_errors=True)
-        shutil.rmtree(setup_file_path + '/raft.egg-info', ignore_errors=True)
-        shutil.rmtree(setup_file_path + '/__pycache__', ignore_errors=True)
-
-        clean_folder(setup_file_path + '/raft')
-        shutil.rmtree(setup_file_path + '/build')
-
-    except IOError:
-        pass
-
-    # need to terminate script so cythonizing doesn't get triggered after
-    # cleanup unintendedly
-    sys.argv.remove("clean")
-
-    if "--all" in sys.argv:
-        sys.argv.remove("--all")
-
-    if len(sys.argv) == 1:
-        sys.exit(0)
-
-
-##############################################################################
-# - Cython extensions build and parameters -----------------------------------
-
-
-libs = ['cudart', "nccl", "cusolver", "cusparse", "cublas"]
-
-include_dirs = [ucx_include_dir,
-                rmm_include_dir,
-                thrust_include_dir,
-                spdlog_include_dir,
-                cuda_include_dir,
-                numpy.get_include(),
-                "../../cpp/include/",
-                os.path.dirname(sysconfig.get_path("include"))]
-
-library_dirs = [ucx_lib_dir, cuda_lib_dir]
-
-extensions = [
-    Extension("*",
-              sources=["raft/**/*.pyx"],
-              include_dirs=include_dirs,
-              library_dirs=(
-                  library_dirs + [get_python_lib()]
-              ),
-              runtime_library_dirs=(
-                  library_dirs + [os.path.join(os.sys.prefix, "lib")]
-              ),
-              libraries=libs,
-              language='c++',
-              extra_compile_args=['-std=c++17'])
-]
-
-
-class build_ext_no_debug(_build_ext):
-
-    def build_extensions(self):
-        def remove_flags(compiler, *flags):
-            for flag in flags:
-                try:
-                    compiler.compiler_so = list(
-                        filter((flag).__ne__, compiler.compiler_so)
-                    )
-                except Exception:
-                    pass
-
-        # Full optimization
-        self.compiler.compiler_so.append("-O3")
-
-        # Ignore deprecation declaration warnings
-        self.compiler.compiler_so.append("-Wno-deprecated-declarations")
-
-        # No debug symbols, full optimization, no '-Wstrict-prototypes' warning
-        remove_flags(
-            self.compiler, "-g", "-G", "-O1", "-O2", "-Wstrict-prototypes"
-        )
-        super().build_extensions()
-
-    def finalize_options(self):
-        if self.distribution.ext_modules:
-            # Delay import this to allow for Cython-less installs
-            from Cython.Build.Dependencies import cythonize
-
-            nthreads = getattr(self, "parallel", None)  # -j option in Py3.5+
-            nthreads = int(nthreads) if nthreads else None
-            self.distribution.ext_modules = cythonize(
-                self.distribution.ext_modules,
-                nthreads=nthreads,
-                force=self.force,
-                gdb_debug=False,
-                compiler_directives=dict(
-                    profile=False, language_level=3, embedsignature=True
-                ),
-            )
-        # Skip calling super() and jump straight to setuptools
-        setuptools.command.build_ext.build_ext.finalize_options(self)
-
-
-cmdclass = dict()
-cmdclass.update(versioneer.get_cmdclass())
-cmdclass["build_ext"] = build_ext_no_debug
-
-
-##############################################################################
-# - Python package generation ------------------------------------------------
-
-
 setup(name='raft',
       description="RAPIDS Analytics Frameworks Toolset",
       version=versioneer.get_version(),
@@ -210,18 +30,23 @@ def finalize_options(self):
         "Programming Language :: Python :: 3.7"
       ],
       author="NVIDIA Corporation",
-      setup_requires=['cython'],
-      ext_modules=extensions,
-      package_data=dict.fromkeys(
-                         find_packages(include=["raft.dask.common",
-                                                "raft.dask.common.includes",
-                                                "raft.common",
-                                                "raft.common.includes"]),
-                         ["*.hpp", "*.pxd"],
-      ),
+      package_data={
+          # Note: A dict comprehension with an explicit copy is necessary
+          # (rather than something simpler like a dict.fromkeys) because
+          # otherwise every package will refer to the same list and skbuild
+          # modifies it in place.
+          key: ["*.hpp", "*.pxd"]
+          for key in find_packages(
+              include=[
+                  "raft.dask.common",
+                  "raft.dask.common.includes",
+                  "raft.common",
+                  "raft.common.includes"
+              ]
+          )
+      },
       packages=find_packages(include=['raft', 'raft.*']),
-      install_requires=install_requires,
       license="Apache",
-      cmdclass=cmdclass,
+      cmdclass=versioneer.get_cmdclass(),
       zip_safe=False
       )

From 3ddc1ab1d7d08dc825d09ffffea92b835c5df16c Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Tue, 10 May 2022 17:17:58 -0700
Subject: [PATCH 02/12] Update all ancillary files.

---
 build.sh                       | 5 ++---
 ci/release/update-version.sh   | 2 ++
 conda/recipes/pyraft/build.sh  | 2 +-
 conda/recipes/pyraft/meta.yaml | 2 ++
 4 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/build.sh b/build.sh
index 0287872341..367034641e 100755
--- a/build.sh
+++ b/build.sh
@@ -270,10 +270,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
 
diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh
index 85c256c61a..0ee92a6e22 100755
--- a/ci/release/update-version.sh
+++ b/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};
diff --git a/conda/recipes/pyraft/build.sh b/conda/recipes/pyraft/build.sh
index 1462f365ff..0e74faa699 100644
--- a/conda/recipes/pyraft/build.sh
+++ b/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\"
diff --git a/conda/recipes/pyraft/meta.yaml b/conda/recipes/pyraft/meta.yaml
index ab9dc28fda..a640c4beb7 100644
--- a/conda/recipes/pyraft/meta.yaml
+++ b/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 }}.*

From f4074a07b40da59bee2031a41ffcb047e233c125 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Fri, 13 May 2022 15:52:42 -0700
Subject: [PATCH 03/12] Fix CUDA architectures.

---
 python/raft/CMakeLists.txt | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/python/raft/CMakeLists.txt b/python/raft/CMakeLists.txt
index a0561a210b..6e02312a89 100644
--- a/python/raft/CMakeLists.txt
+++ b/python/raft/CMakeLists.txt
@@ -35,7 +35,7 @@ option(FIND_RAFT_CPP "Search for existing RAFT C++ installations before defaulti
 
 # If the user requested it we attempt to find RAFT.
 if(FIND_RAFT_CPP)
-    find_package(raft ${pyraft_version} REQUIRED COMPONENTS distance)
+    find_package(raft ${pyraft_version} REQUIRED)
 else()
   set(raft_FOUND OFF)
 endif()
@@ -44,7 +44,13 @@ 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)
   add_subdirectory(../../cpp raft-cpp)

From 5478bff634b1cf4703c8520b32c519257d724c6c Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Fri, 13 May 2022 15:52:52 -0700
Subject: [PATCH 04/12] Stop compiling unnecessary components.

---
 python/raft/CMakeLists.txt | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/python/raft/CMakeLists.txt b/python/raft/CMakeLists.txt
index 6e02312a89..02089fa5e9 100644
--- a/python/raft/CMakeLists.txt
+++ b/python/raft/CMakeLists.txt
@@ -53,6 +53,9 @@ if(NOT raft_FOUND)
 
   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()
 

From cb3a5527a499f82a4f9607adef8814fbae6866cc Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Fri, 13 May 2022 17:35:15 -0700
Subject: [PATCH 05/12] Add search for nccl.

---
 python/raft/raft/dask/common/CMakeLists.txt | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/python/raft/raft/dask/common/CMakeLists.txt b/python/raft/raft/dask/common/CMakeLists.txt
index a65a68ffa8..2f9d9f32ab 100644
--- a/python/raft/raft/dask/common/CMakeLists.txt
+++ b/python/raft/raft/dask/common/CMakeLists.txt
@@ -18,3 +18,13 @@ rapids_cython_create_modules(
     SOURCE_FILES "${cython_sources}"
     LINKED_LIBRARIES "${linked_libraries}"
     CXX)
+
+find_library(NCCL_PATH nccl REQUIRED)
+foreach(cython_module IN LISTS RAPIDS_CYTHON_CREATED_TARGETS)
+  set(is_nccl "")
+  string(FIND "${cython_module}" nccl is_nccl)
+
+  if(NOT is_nccl EQUAL -1)
+    target_link_libraries("${cython_module}" "${NCCL_PATH}")
+  endif()
+endforeach()

From 3c54d290699896637e8542312f8001dbddd625e6 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Fri, 13 May 2022 17:42:15 -0700
Subject: [PATCH 06/12] Use a simpler approach to specify the library.

---
 python/raft/raft/dask/common/CMakeLists.txt | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/python/raft/raft/dask/common/CMakeLists.txt b/python/raft/raft/dask/common/CMakeLists.txt
index 2f9d9f32ab..7e6c4b2ffc 100644
--- a/python/raft/raft/dask/common/CMakeLists.txt
+++ b/python/raft/raft/dask/common/CMakeLists.txt
@@ -12,19 +12,23 @@
 # the License.
 # =============================================================================
 
-set(cython_sources comms_utils.pyx nccl.pyx)
+set(cython_sources comms_utils.pyx)
 set(linked_libraries raft::raft)
 rapids_cython_create_modules(
     SOURCE_FILES "${cython_sources}"
     LINKED_LIBRARIES "${linked_libraries}"
     CXX)
 
+
+# TODO: This is a pretty hacky way to handle the NCCL dependency. There are a
+# couple of reasons I'm doing it this way right now:
+# 1. The nccl target name ends up making it impossible to specify that as a target library.
+# 2. nccl doesn't ship with a config file or FindNCCL.cmake file.
 find_library(NCCL_PATH nccl REQUIRED)
-foreach(cython_module IN LISTS RAPIDS_CYTHON_CREATED_TARGETS)
-  set(is_nccl "")
-  string(FIND "${cython_module}" nccl is_nccl)
 
-  if(NOT is_nccl EQUAL -1)
-    target_link_libraries("${cython_module}" "${NCCL_PATH}")
-  endif()
-endforeach()
+set(cython_sources nccl.pyx)
+set(linked_libraries raft::raft "${NCCL_PATH}")
+rapids_cython_create_modules(
+    SOURCE_FILES "${cython_sources}"
+    LINKED_LIBRARIES "${linked_libraries}"
+    CXX)

From 13772922a7687e0ccfd3ee4c63d4961d2d98d4d6 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Fri, 13 May 2022 18:26:05 -0700
Subject: [PATCH 07/12] Link both comm_utils.pyx and nccl.pyx to NCCL.

---
 python/raft/raft/dask/common/CMakeLists.txt | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/python/raft/raft/dask/common/CMakeLists.txt b/python/raft/raft/dask/common/CMakeLists.txt
index 7e6c4b2ffc..6b1308add6 100644
--- a/python/raft/raft/dask/common/CMakeLists.txt
+++ b/python/raft/raft/dask/common/CMakeLists.txt
@@ -12,21 +12,13 @@
 # the License.
 # =============================================================================
 
-set(cython_sources comms_utils.pyx)
-set(linked_libraries raft::raft)
-rapids_cython_create_modules(
-    SOURCE_FILES "${cython_sources}"
-    LINKED_LIBRARIES "${linked_libraries}"
-    CXX)
-
-
 # TODO: This is a pretty hacky way to handle the NCCL dependency. There are a
 # couple of reasons I'm doing it this way right now:
 # 1. The nccl target name ends up making it impossible to specify that as a target library.
 # 2. nccl doesn't ship with a config file or FindNCCL.cmake file.
 find_library(NCCL_PATH nccl REQUIRED)
 
-set(cython_sources nccl.pyx)
+set(cython_sources comms_utils.pyx nccl.pyx)
 set(linked_libraries raft::raft "${NCCL_PATH}")
 rapids_cython_create_modules(
     SOURCE_FILES "${cython_sources}"

From 4819938fe97946da321b555812c8c7f6086c8098 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Tue, 17 May 2022 10:50:58 -0700
Subject: [PATCH 08/12] Move CMake discovery scripts.

---
 {cpp => python/raft}/cmake/thirdparty/get_nccl.cmake | 0
 {cpp => python/raft}/cmake/thirdparty/get_ucx.cmake  | 0
 2 files changed, 0 insertions(+), 0 deletions(-)
 rename {cpp => python/raft}/cmake/thirdparty/get_nccl.cmake (100%)
 rename {cpp => python/raft}/cmake/thirdparty/get_ucx.cmake (100%)

diff --git a/cpp/cmake/thirdparty/get_nccl.cmake b/python/raft/cmake/thirdparty/get_nccl.cmake
similarity index 100%
rename from cpp/cmake/thirdparty/get_nccl.cmake
rename to python/raft/cmake/thirdparty/get_nccl.cmake
diff --git a/cpp/cmake/thirdparty/get_ucx.cmake b/python/raft/cmake/thirdparty/get_ucx.cmake
similarity index 100%
rename from cpp/cmake/thirdparty/get_ucx.cmake
rename to python/raft/cmake/thirdparty/get_ucx.cmake

From 22cd35148efc7e6800d8b68e2fe1c7b9cfe3672d Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Tue, 17 May 2022 12:00:02 -0700
Subject: [PATCH 09/12] Handle UCX and NCCL dependencies appropriately.

---
 build.sh                                    | 10 ----------
 cpp/CMakeLists.txt                          |  4 ----
 python/raft/CMakeLists.txt                  |  2 +-
 python/raft/cmake/thirdparty/get_nccl.cmake |  1 +
 python/raft/cmake/thirdparty/get_ucx.cmake  |  1 +
 python/raft/raft/dask/common/CMakeLists.txt |  9 +++------
 6 files changed, 6 insertions(+), 21 deletions(-)

diff --git a/build.sh b/build.sh
index 367034641e..b82fa6606e 100755
--- a/build.sh
+++ b/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
@@ -252,8 +244,6 @@ if (( ${NUMARGS} == 0 )) || hasArg libraft || hasArg docs || hasArg tests || has
           -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}
 
   if [[ ${CMAKE_TARGET} != "" ]]; then
diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index 072ff3ac18..66811ae850 100644
--- a/cpp/CMakeLists.txt
+++ b/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)
diff --git a/python/raft/CMakeLists.txt b/python/raft/CMakeLists.txt
index 02089fa5e9..21ca270108 100644
--- a/python/raft/CMakeLists.txt
+++ b/python/raft/CMakeLists.txt
@@ -21,7 +21,7 @@ file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-22.
 include(${CMAKE_BINARY_DIR}/RAPIDS.cmake)
 
 project(
-  raft
+  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
diff --git a/python/raft/cmake/thirdparty/get_nccl.cmake b/python/raft/cmake/thirdparty/get_nccl.cmake
index 118ae37704..5450612df5 100644
--- a/python/raft/cmake/thirdparty/get_nccl.cmake
+++ b/python/raft/cmake/thirdparty/get_nccl.cmake
@@ -14,6 +14,7 @@
 # limitations under the License.
 #=============================================================================
 
+include(rapids-find)
 function(find_and_configure_nccl)
 
     if(TARGET NCCL::NCCL)
diff --git a/python/raft/cmake/thirdparty/get_ucx.cmake b/python/raft/cmake/thirdparty/get_ucx.cmake
index 0ff9dc8734..6b1e62870a 100644
--- a/python/raft/cmake/thirdparty/get_ucx.cmake
+++ b/python/raft/cmake/thirdparty/get_ucx.cmake
@@ -14,6 +14,7 @@
 # limitations under the License.
 #=============================================================================
 
+include(rapids-find)
 function(find_and_configure_ucx)
     set(oneValueArgs VERSION PINNED_TAG)
     cmake_parse_arguments(PKG "${options}" "${oneValueArgs}"
diff --git a/python/raft/raft/dask/common/CMakeLists.txt b/python/raft/raft/dask/common/CMakeLists.txt
index 6b1308add6..c7fa163179 100644
--- a/python/raft/raft/dask/common/CMakeLists.txt
+++ b/python/raft/raft/dask/common/CMakeLists.txt
@@ -12,14 +12,11 @@
 # the License.
 # =============================================================================
 
-# TODO: This is a pretty hacky way to handle the NCCL dependency. There are a
-# couple of reasons I'm doing it this way right now:
-# 1. The nccl target name ends up making it impossible to specify that as a target library.
-# 2. nccl doesn't ship with a config file or FindNCCL.cmake file.
-find_library(NCCL_PATH nccl REQUIRED)
+include(${raft-python_SOURCE_DIR}/cmake/thirdparty/get_nccl.cmake)
+include(${raft-python_SOURCE_DIR}/cmake/thirdparty/get_ucx.cmake)
 
 set(cython_sources comms_utils.pyx nccl.pyx)
-set(linked_libraries raft::raft "${NCCL_PATH}")
+set(linked_libraries raft::raft NCCL::NCCL ucx::ucp)
 rapids_cython_create_modules(
     SOURCE_FILES "${cython_sources}"
     LINKED_LIBRARIES "${linked_libraries}"

From 3823ac991df05938d1e9e48802b3cd62329a57ab Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Tue, 17 May 2022 12:04:31 -0700
Subject: [PATCH 10/12] Fix typo and update copyrights.

---
 build.sh                                    | 2 +-
 python/raft/cmake/thirdparty/get_nccl.cmake | 2 +-
 python/raft/cmake/thirdparty/get_ucx.cmake  | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/build.sh b/build.sh
index b82fa6606e..876ad5090a 100755
--- a/build.sh
+++ b/build.sh
@@ -243,7 +243,7 @@ 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_thrust_DEPENDENCY=${ENABLE_thrust_DEPENDENCY} \
           ${CMAKE_ARGS}
 
   if [[ ${CMAKE_TARGET} != "" ]]; then
diff --git a/python/raft/cmake/thirdparty/get_nccl.cmake b/python/raft/cmake/thirdparty/get_nccl.cmake
index 5450612df5..c2cc17b399 100644
--- a/python/raft/cmake/thirdparty/get_nccl.cmake
+++ b/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.
diff --git a/python/raft/cmake/thirdparty/get_ucx.cmake b/python/raft/cmake/thirdparty/get_ucx.cmake
index 6b1e62870a..b83ce15b24 100644
--- a/python/raft/cmake/thirdparty/get_ucx.cmake
+++ b/python/raft/cmake/thirdparty/get_ucx.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.

From a9cce6ecdf45d30eb97270333758124ebb46a686 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Tue, 17 May 2022 12:49:34 -0700
Subject: [PATCH 11/12] Remove link-time dependency on UCX.

---
 python/raft/cmake/thirdparty/get_ucx.cmake  | 34 ---------------------
 python/raft/raft/dask/common/CMakeLists.txt |  3 +-
 2 files changed, 1 insertion(+), 36 deletions(-)
 delete mode 100644 python/raft/cmake/thirdparty/get_ucx.cmake

diff --git a/python/raft/cmake/thirdparty/get_ucx.cmake b/python/raft/cmake/thirdparty/get_ucx.cmake
deleted file mode 100644
index b83ce15b24..0000000000
--- a/python/raft/cmake/thirdparty/get_ucx.cmake
+++ /dev/null
@@ -1,34 +0,0 @@
-#=============================================================================
-# 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.
-# 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(rapids-find)
-function(find_and_configure_ucx)
-    set(oneValueArgs VERSION PINNED_TAG)
-    cmake_parse_arguments(PKG "${options}" "${oneValueArgs}"
-                          "${multiValueArgs}" ${ARGN} )
-
-    rapids_find_generate_module(ucx
-        HEADER_NAMES  ucp/api/ucp.h
-        LIBRARY_NAMES ucp
-    )
-
-    # Currently UCX has no CMake build-system so we require
-    # it built and installed on the machine already
-    rapids_find_package(ucx REQUIRED)
-
-endfunction()
-
-find_and_configure_ucx()
diff --git a/python/raft/raft/dask/common/CMakeLists.txt b/python/raft/raft/dask/common/CMakeLists.txt
index c7fa163179..c603ef6ee2 100644
--- a/python/raft/raft/dask/common/CMakeLists.txt
+++ b/python/raft/raft/dask/common/CMakeLists.txt
@@ -13,10 +13,9 @@
 # =============================================================================
 
 include(${raft-python_SOURCE_DIR}/cmake/thirdparty/get_nccl.cmake)
-include(${raft-python_SOURCE_DIR}/cmake/thirdparty/get_ucx.cmake)
 
 set(cython_sources comms_utils.pyx nccl.pyx)
-set(linked_libraries raft::raft NCCL::NCCL ucx::ucp)
+set(linked_libraries raft::raft NCCL::NCCL)
 rapids_cython_create_modules(
     SOURCE_FILES "${cython_sources}"
     LINKED_LIBRARIES "${linked_libraries}"

From e01f0914feb44534be0357ca39d60834334a9b14 Mon Sep 17 00:00:00 2001
From: Vyas Ramasubramani <vyasr@nvidia.com>
Date: Tue, 17 May 2022 13:30:26 -0700
Subject: [PATCH 12/12] Add back raft link.

---
 python/raft/raft/dask/common/CMakeLists.txt | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/python/raft/raft/dask/common/CMakeLists.txt b/python/raft/raft/dask/common/CMakeLists.txt
index c603ef6ee2..1abddae947 100644
--- a/python/raft/raft/dask/common/CMakeLists.txt
+++ b/python/raft/raft/dask/common/CMakeLists.txt
@@ -13,9 +13,10 @@
 # =============================================================================
 
 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)
+set(linked_libraries raft::raft NCCL::NCCL ucx::ucp)
 rapids_cython_create_modules(
     SOURCE_FILES "${cython_sources}"
     LINKED_LIBRARIES "${linked_libraries}"