Add wheel builds

.github/workflows/wheels.yml

@@ -0,0 +1,72 @@
+name: RAFT wheels
+
+on:
+  workflow_call:
+    inputs:
+      versioneer-override:
+        type: string
+        default: ''
+      build-tag:
+        type: string
+        default: ''
+      branch:
+        required: true
+        type: string
+      date:
+        required: true
+        type: string
+      sha:
+        required: true
+        type: string
+      build-type:
+        type: string
+        default: nightly
+
+concurrency:
+  group: "raft-${{ github.workflow }}-${{ github.ref }}"
+  cancel-in-progress: true
+
+jobs:
+  pylibraft-wheel:
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux.yml@main
+    with:
+      repo: rapidsai/raft
+
+      build-type: ${{ inputs.build-type }}
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+
+      package-dir: python/pylibraft
+      package-name: pylibraft
+
+      python-package-versioneer-override: ${{ inputs.versioneer-override }}
+      python-package-build-tag: ${{ inputs.build-tag }}
+
+      skbuild-configure-options: "-DRAFT_BUILD_WHEELS=ON -DDETECT_CONDA_ENV=OFF -DFIND_RAFT_CPP=OFF"
+
+      test-extras: test
+      test-unittest: "python -m pytest -v ./python/pylibraft/pylibraft/test"
+    secrets: inherit
+  raft-dask-wheel:
+    needs: pylibraft-wheel
+    uses: rapidsai/shared-action-workflows/.github/workflows/wheels-manylinux.yml@main
+    with:
+      repo: rapidsai/raft
+
+      build-type: ${{ inputs.build-type }}
+      branch: ${{ inputs.branch }}
+      sha: ${{ inputs.sha }}
+      date: ${{ inputs.date }}
+
+      package-dir: python/raft-dask
+      package-name: raft_dask
+
+      python-package-versioneer-override: ${{ inputs.versioneer-override }}
+      python-package-build-tag: ${{ inputs.build-tag }}
+
+      skbuild-configure-options: "-DRAFT_BUILD_WHEELS=ON -DDETECT_CONDA_ENV=OFF -DFIND_RAFT_CPP=OFF"
+
+      test-extras: test
+      test-unittest: "python -m pytest -v ./python/raft-dask/raft_dask/test"
+    secrets: inherit

.gitignore

@@ -49,3 +49,6 @@ _skbuild
 
 ## doxygen build check inside ci/checks/style.sh
 doxygen_check/
+
+## cibuildwheel
+/wheelhouse

cpp/CMakeLists.txt

@@ -30,11 +30,6 @@ project(
   LANGUAGES CXX CUDA
 )
 
-# Needed because GoogleBenchmark changes the state of FindThreads.cmake, causing subsequent runs to
-# have different values for the `Threads::Threads` target. Setting this flag ensures
-# `Threads::Threads` is the same value in first run and subsequent runs.
-set(THREADS_PREFER_PTHREAD_FLAG ON)
-
 # Write the version header
 rapids_cmake_write_version_file(include/raft/version_config.hpp)
 
@@ -77,6 +72,13 @@ option(RAFT_ENABLE_NN_DEPENDENCIES "Search for raft::nn dependencies like faiss"
 
 option(RAFT_ENABLE_thrust_DEPENDENCY "Enable Thrust dependency" ON)
 
+if(BUILD_TESTS OR BUILD_BENCH)
+  # Needed because GoogleBenchmark changes the state of FindThreads.cmake, causing subsequent runs
+  # to have different values for the `Threads::Threads` target. Setting this flag ensures
+  # `Threads::Threads` is the same value in first run and subsequent runs.
+  set(THREADS_PREFER_PTHREAD_FLAG ON)
+endif()
+
 if(BUILD_TESTS AND NOT RAFT_ENABLE_thrust_DEPENDENCY)
   message(VERBOSE "RAFT: BUILD_TESTS is enabled, overriding RAFT_ENABLE_thrust_DEPENDENCY")
   set(RAFT_ENABLE_thrust_DEPENDENCY ON)
@@ -134,12 +136,7 @@ endif()
 
 set(_ctk_static_suffix "")
 if(CUDA_STATIC_RUNTIME)
-  # If we're statically linking CTK cuBLAS, we also want to statically link BLAS
-  set(BLA_STATIC ON)
   set(_ctk_static_suffix "_static")
-  # Control legacy FindCUDA.cmake behavior too Remove this after we push it into rapids-cmake:
-  # https://github.com/rapidsai/rapids-cmake/pull/259
-  set(CUDA_USE_STATIC_CUDA_RUNTIME ON)
 endif()
 
 # CUDA runtime

python/pylibraft/CMakeLists.txt

@@ -32,6 +32,8 @@ option(FIND_RAFT_CPP "Search for existing RAFT C++ installations before defaulti
        ON
 )
 
+option(RAFT_BUILD_WHEELS "Whether this build is generating a Python wheel." OFF)
+
 # If the user requested it we attempt to find RAFT.
 if(FIND_RAFT_CPP)
   find_package(raft ${pylibraft_version} REQUIRED COMPONENTS distance)
@@ -62,7 +64,16 @@ if(NOT raft_FOUND)
   set(BUILD_BENCH OFF)
   set(RAFT_COMPILE_LIBRARIES OFF)
   set(RAFT_COMPILE_DIST_LIBRARY ON)
-  add_subdirectory(../../cpp raft-cpp)
+
+  set(_exclude_from_all "")
+  if(RAFT_BUILD_WHEELS)
+    # Statically link dependencies if building wheels
+    set(CUDA_STATIC_RUNTIME ON)
+    # Don't install the raft C++ targets into wheels
+    set(_exclude_from_all EXCLUDE_FROM_ALL)
+  endif()
+
+  add_subdirectory(../../cpp raft-cpp ${_exclude_from_all})
 
   # When building the C++ libraries from source we must copy libraft_distance.so alongside the
   # pairwise_distance and random Cython libraries TODO: when we have a single 'compiled' raft

python/pylibraft/LICENSE

@@ -0,0 +1 @@
+../../LICENSE

python/pylibraft/_custom_build/backend.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.
+
+"""Custom build backend for pylibraft to get versioned requirements.
+
+Based on https://setuptools.pypa.io/en/latest/build_meta.html
+"""
+import os
+from functools import wraps
+
+from setuptools import build_meta as _orig
+
+# Alias the required bits
+build_wheel = _orig.build_wheel
+build_sdist = _orig.build_sdist
+
+
+def replace_requirements(func):
+    @wraps(func)
+    def wrapper(config_settings=None):
+        orig_list = getattr(_orig, func.__name__)(config_settings)
+        append_list = [
+            f"rmm{os.getenv('RAPIDS_PY_WHEEL_CUDA_SUFFIX', default='')}"
+        ]
+        return orig_list + append_list
+
+    return wrapper
+
+
+get_requires_for_build_wheel = replace_requirements(
+    _orig.get_requires_for_build_wheel
+)
+get_requires_for_build_sdist = replace_requirements(
+    _orig.get_requires_for_build_sdist
+)
+get_requires_for_build_editable = replace_requirements(
+    _orig.get_requires_for_build_editable
+)

python/pylibraft/pylibraft/cluster/CMakeLists.txt

@@ -14,15 +14,11 @@
 
 # Set the list of Cython files to build
 set(cython_sources kmeans.pyx)
-set(linked_libraries raft::raft raft::distance)
+set(linked_libraries raft::distance)
 
 # Build all of the Cython targets
 rapids_cython_create_modules(
   CXX
   SOURCE_FILES "${cython_sources}"
-  LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX cluster_
+  LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS raft MODULE_PREFIX cluster_
 )
-
-foreach(cython_module IN LISTS RAPIDS_CYTHON_CREATED_TARGETS)
-  set_target_properties(${cython_module} PROPERTIES INSTALL_RPATH "\$ORIGIN;\$ORIGIN/../library")
-endforeach()

python/pylibraft/pylibraft/neighbors/CMakeLists.txt

@@ -19,11 +19,7 @@ set(linked_libraries raft::raft raft::distance)
 rapids_cython_create_modules(
   CXX
   SOURCE_FILES ""
-  LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX neighbors_
+  LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS raft MODULE_PREFIX neighbors_
 )
 
-foreach(cython_module IN LISTS RAPIDS_CYTHON_CREATED_TARGETS)
-  set_target_properties(${cython_module} PROPERTIES INSTALL_RPATH "\$ORIGIN;\$ORIGIN/../library")
-endforeach()
-
 add_subdirectory(ivf_pq)

python/pylibraft/pylibraft/neighbors/ivf_pq/CMakeLists.txt

@@ -20,9 +20,5 @@ set(linked_libraries raft::raft raft::distance)
 rapids_cython_create_modules(
   CXX
   SOURCE_FILES "${cython_sources}"
-  LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX neighbors_ivfpq_
+  LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS raft MODULE_PREFIX neighbors_ivfpq_
 )
-
-foreach(cython_module IN LISTS RAPIDS_CYTHON_CREATED_TARGETS)
-  set_target_properties(${cython_module} PROPERTIES INSTALL_RPATH "\$ORIGIN;\$ORIGIN/../library")
-endforeach()

python/pylibraft/pylibraft/random/CMakeLists.txt

@@ -14,6 +14,7 @@
 
 # Set the list of Cython files to build
 set(cython_sources rmat_rectangular_generator.pyx)
+
 # TODO: should finally be replaced with 'compiled' library to be more generic, when that is
 # available
 set(linked_libraries raft::raft raft::distance)

python/pylibraft/pyproject.toml

@@ -18,7 +18,11 @@ requires = [
     "wheel",
     "setuptools",
     "cython>=0.29,<0.30",
+    "cuda-python>=11.7.1,<12.0",
     "scikit-build>=0.13.1",
     "cmake>=3.23.1,!=3.25.0",
+    "versioneer",
     "ninja"
 ]
+build-backend = "backend"
+backend-path = ["_custom_build"]

python/pylibraft/setup.cfg

@@ -44,8 +44,3 @@ skip=
     build
     dist
     __init__.py
-
-[options]
-packages = find:
-install_requires = numpy
-python_requires = >=3.7,<3.9

python/pylibraft/setup.py

@@ -14,12 +14,54 @@
 # limitations under the License.
 #
 
+import os
+
 import versioneer
 from setuptools import find_packages
 from skbuild import setup
 
+cuda_suffix = os.getenv("RAPIDS_PY_WHEEL_CUDA_SUFFIX", default="")
+
+install_requires = [
+    "numpy",
+    "cuda-python>=11.7.1,<12.0",
+    f"rmm{cuda_suffix}",
+]
+
+extras_require = {
+    "test": [
+        "pytest",
+        "scipy",
+        "scikit-learn",
+    ]
+}
+
+
+def exclude_libcxx_symlink(cmake_manifest):
+    return list(
+        filter(
+            lambda name: not ("include/rapids/libcxx/include" in name),
+            cmake_manifest,
+        )
+    )
+
+
+# Make versioneer produce PyPI-compatible nightly versions for wheels.
+if "RAPIDS_PY_WHEEL_VERSIONEER_OVERRIDE" in os.environ:
+    orig_get_versions = versioneer.get_versions
+
+    version_override = os.environ["RAPIDS_PY_WHEEL_VERSIONEER_OVERRIDE"]
+
+    def get_versions():
+        data = orig_get_versions()
+        data["version"] = version_override
+        return data
+
+    versioneer.get_versions = get_versions
+
+
 setup(
-    name="pylibraft",
+    name=f"pylibraft{cuda_suffix}",
     description="RAFT: Reusable Algorithms Functions and other Tools",
     version=versioneer.get_version(),
     classifiers=[
@@ -29,6 +71,7 @@
         "Programming Language :: Python :: 3.9",
     ],
     author="NVIDIA Corporation",
+    include_package_data=True,
     package_data={
         # Note: A dict comprehension with an explicit copy is necessary
         # (rather than something simpler like a dict.fromkeys) because
@@ -46,8 +89,12 @@
             ]
         )
     },
+    install_requires=install_requires,
+    extras_require=extras_require,
+    # Don't want libcxx getting pulled into wheel builds.
+    cmake_process_manifest_hook=exclude_libcxx_symlink,
     packages=find_packages(include=["pylibraft", "pylibraft.*"]),
-    license="Apache",
+    license="Apache 2.0",
     cmdclass=versioneer.get_cmdclass(),
     zip_safe=False,
 )

python/raft-dask/CMakeLists.txt

@@ -32,6 +32,8 @@ option(FIND_RAFT_CPP "Search for existing RAFT C++ installations before defaulti
        OFF
 )
 
+option(RAFT_BUILD_WHEELS "Whether this build is generating a Python wheel." OFF)
+
 # If the user requested it we attempt to find RAFT.
 if(FIND_RAFT_CPP)
   find_package(raft ${raft_dask_version} REQUIRED COMPONENTS distributed)
@@ -57,7 +59,16 @@ if(NOT raft_FOUND)
   set(RAFT_COMPILE_LIBRARIES OFF)
   set(RAFT_COMPILE_DIST_LIBRARY OFF)
   set(RAFT_COMPILE_NN_LIBRARY OFF)
-  add_subdirectory(../../cpp raft-cpp)
+
+  set(_exclude_from_all "")
+  if(RAFT_BUILD_WHEELS)
+    # Statically link dependencies if building wheels
+    set(CUDA_STATIC_RUNTIME ON)
+    # Don't install the raft C++ targets into wheels
+    set(_exclude_from_all EXCLUDE_FROM_ALL)
+  endif()
+
+  add_subdirectory(../../cpp raft-cpp ${_exclude_from_all})
 endif()
 
 include(rapids-cython)

python/raft-dask/LICENSE

@@ -0,0 +1 @@
+../../LICENSE

python/raft-dask/pyproject.toml

@@ -21,5 +21,4 @@ requires = [
     "scikit-build>=0.13.1",
     "cmake>=3.23.1,!=3.25.0",
     "ninja",
-    "pylibraft"
 ]

python/raft-dask/pytest.ini

@@ -5,4 +5,4 @@ markers =
   stress: marks stress tests
   mg: marks a test as multi-GPU
   memleak: marks a test as a memory leak test
-
+  nccl: marks a test as using NCCL

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

@@ -17,5 +17,6 @@ include(${raft-dask-python_SOURCE_DIR}/cmake/thirdparty/get_nccl.cmake)
 set(cython_sources comms_utils.pyx nccl.pyx)
 set(linked_libraries raft::raft raft::distributed NCCL::NCCL)
 rapids_cython_create_modules(
-  SOURCE_FILES "${cython_sources}" LINKED_LIBRARIES "${linked_libraries}" CXX
+  SOURCE_FILES "${cython_sources}" ASSOCIATED_TARGETS raft LINKED_LIBRARIES "${linked_libraries}"
+                                                                            CXX
 )

python/raft-dask/raft_dask/include_test/CMakeLists.txt

@@ -15,5 +15,6 @@
 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
+  SOURCE_FILES "${cython_sources}" ASSOCIATED_TARGETS raft LINKED_LIBRARIES "${linked_libraries}"
+                                                                            CXX
 )