pylibraft builds mostly working

ci/build_wheel_pylibraft.sh

@@ -5,17 +5,8 @@ set -euo pipefail
 
 package_dir="python/pylibraft"
 
-case "${RAPIDS_CUDA_VERSION}" in
-  12.*)
-    EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=ON"
-  ;;
-  11.*)
-    EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=OFF"
-  ;;
-esac
-
 # Set up skbuild options. Enable sccache in skbuild config options
-export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF${EXTRA_CMAKE_ARGS}"
+export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;"
 
 ci/build_wheel.sh pylibraft ${package_dir} python
 ci/validate_wheel.sh ${package_dir} final_dist pylibraft

dependencies.yaml

@@ -65,6 +65,13 @@ files:
       - build_pylibraft
       - depends_on_librmm
       - rapids_build
+  py_run_libraft:
+    output: pyproject
+    pyproject_dir: python/libraft
+    extras:
+      table: project
+    includes:
+      - cuda_wheels
   py_build_pylibraft:
     output: pyproject
     pyproject_dir: python/pylibraft
@@ -79,6 +86,7 @@ files:
       table: tool.rapids-build-backend
       key: requires
     includes:
+      - depends_on_libraft
       - rapids_build
       - build_pylibraft
   py_run_pylibraft:
@@ -88,6 +96,7 @@ files:
       table: project
     includes:
       - cuda_wheels
+      - depends_on_libraft
       - run_pylibraft
   py_test_pylibraft:
     output: pyproject
@@ -113,6 +122,7 @@ files:
       table: tool.rapids-build-backend
       key: requires
     includes:
+      - depends_on_libraft
       - rapids_build
       - depends_on_ucx_build
   py_run_raft_dask:
@@ -121,6 +131,7 @@ files:
     extras:
       table: project
     includes:
+      - depends_on_libraft
       - run_raft_dask
       - depends_on_distributed_ucxx
   py_test_raft_dask:
@@ -554,6 +565,29 @@ dependencies:
             packages:
               - distributed-ucxx-cu11==0.42.*,>=0.0.0a0
           - {matrix: null, packages: [*distributed_ucxx_unsuffixed]}
+  depends_on_libraft:
+    common:
+      - output_types: requirements
+        packages:
+          # pip recognizes the index as a global option for the requirements.txt file
+          - --extra-index-url=https://pypi.nvidia.com
+          - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
+    specific:
+      - output_types: [requirements, pyproject]
+        matrices:
+          - matrix:
+              cuda: "12.*"
+              cuda_suffixed: "true"
+            packages:
+              - libraft-cu12==25.2.*,>=0.0.0a0
+          - matrix:
+              cuda: "11.*"
+              cuda_suffixed: "true"
+            packages:
+              - libraft-cu11==25.2.*,>=0.0.0a0
+          - matrix:
+            packages:
+              - libraft==25.2.*,>=0.0.0a0
   depends_on_librmm:
     common:
       - output_types: conda

python/libraft/CMakeLists.txt

@@ -52,6 +52,9 @@ elseif(USE_CUDA_MATH_WHEELS)
   message(FATAL_ERROR "Cannot use CUDA math wheels with CUDA < 12.0")
 endif()
 
+# ensure installs are inside the libraft/ folder in site-packages/
+set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${PROJECT_BINARY_DIR}/lib)
+
 add_subdirectory(../../cpp raft-cpp)
 
 if(NOT CUDA_STATIC_MATH_LIBRARIES AND USE_CUDA_MATH_WHEELS)

python/libraft/pyproject.toml

@@ -31,13 +31,10 @@ authors = [
 license = { text = "Apache 2.0" }
 requires-python = ">=3.10"
 dependencies = [
-    "cuda-python",
-    "numpy>=1.23,<3.0a0",
     "nvidia-cublas",
     "nvidia-curand",
     "nvidia-cusolver",
     "nvidia-cusparse",
-    "rmm==25.2.*,>=0.0.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 classifiers = [
     "Intended Audience :: Developers",
@@ -56,6 +53,9 @@ test = [
 Homepage = "https://github.com/rapidsai/raft"
 Documentation = "https://docs.rapids.ai/api/raft/stable/"
 
+[project.entry-points."cmake.prefix"]
+libraft = "libraft"
+
 [tool.isort]
 line_length = 79
 multi_line_output = 3
@@ -100,6 +100,8 @@ ninja.make-fallback = true
 sdist.exclude = ["*tests*"]
 sdist.reproducible = true
 wheel.packages = ["libraft"]
+wheel.install-dir = "libraft"
+wheel.py-api = "py3"
 
 [tool.scikit-build.metadata.version]
 provider = "scikit_build_core.metadata.regex"
@@ -112,6 +114,7 @@ requires = [
     "cmake>=3.26.4,!=3.30.0",
     "cuda-python",
     "cython>=3.0.0,<3.1.0a0",
+    "librmm==25.2.*,>=0.0.0a0",
     "ninja",
     "rmm==25.2.*,>=0.0.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.

python/pylibraft/CMakeLists.txt

@@ -27,21 +27,7 @@ project(
   LANGUAGES CXX CUDA
 )
 
-option(USE_CUDA_MATH_WHEELS "Use the CUDA math wheels instead of the system libraries" OFF)
-
-# If the user requested it we attempt to find RAFT.
-if(FIND_RAFT_CPP)
-  find_package(raft "${RAPIDS_VERSION}" REQUIRED COMPONENTS compiled)
-  if(NOT TARGET raft::raft_lib)
-    message(
-      FATAL_ERROR
-        "Building against a preexisting libraft library requires the compiled libraft to have been built!"
-    )
-
-  endif()
-else()
-  set(raft_FOUND OFF)
-endif()
+find_package(raft "${RAPIDS_VERSION}" REQUIRED COMPONENTS raft compiled)
 
 include(rapids-cython-core)
 

python/pylibraft/pylibraft/__init__.py

@@ -13,4 +13,15 @@
 # limitations under the License.
 #
 
+# If libraft was installed as a wheel, we must request it to load the library
+# symbols. Otherwise, we assume that the library was installed in a system path that ld
+# can find.
+try:
+    import libraft
+except ModuleNotFoundError:
+    pass
+else:
+    libraft.load_library()
+    del libraft
+
 from pylibraft._version import __git_commit__, __version__

python/pylibraft/pylibraft/common/CMakeLists.txt

@@ -20,5 +20,5 @@ set(linked_libraries raft::raft)
 rapids_cython_create_modules(
   CXX
   SOURCE_FILES "${cython_sources}"
-  LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS raft MODULE_PREFIX common_
+  LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX common_
 )

python/pylibraft/pylibraft/random/CMakeLists.txt

@@ -23,5 +23,5 @@ set(linked_libraries raft::raft raft::compiled)
 rapids_cython_create_modules(
   CXX
   SOURCE_FILES "${cython_sources}"
-  LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS raft MODULE_PREFIX random_
+  LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX random_
 )

python/pylibraft/pylibraft/sparse/linalg/CMakeLists.txt

@@ -23,5 +23,5 @@ set(linked_libraries raft::raft raft::compiled)
 rapids_cython_create_modules(
   CXX
   SOURCE_FILES "${cython_sources}"
-  LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS raft MODULE_PREFIX sparse_
+  LINKED_LIBRARIES "${linked_libraries}" MODULE_PREFIX sparse_
 )

python/pylibraft/pyproject.toml

@@ -32,6 +32,7 @@ license = { text = "Apache 2.0" }
 requires-python = ">=3.10"
 dependencies = [
     "cuda-python",
+    "libraft==25.2.*,>=0.0.0a0",
     "numpy>=1.23,<3.0a0",
     "nvidia-cublas",
     "nvidia-curand",
@@ -124,11 +125,12 @@ requires = [
     "cmake>=3.26.4,!=3.30.0",
     "cuda-python",
     "cython>=3.0.0,<3.1.0a0",
+    "libraft==25.2.*,>=0.0.0a0",
     "ninja",
     "rmm==25.2.*,>=0.0.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 dependencies-file = "../../dependencies.yaml"
-matrix-entry = "cuda_suffixed=true;use_cuda_wheels=true"
+matrix-entry = "cuda_suffixed=true"
 
 [tool.pydistcheck]
 select = [

python/raft-dask/pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
     "dask-cuda==25.2.*,>=0.0.0a0",
     "distributed-ucxx==0.42.*,>=0.0.0a0",
     "joblib>=0.11",
+    "libraft==25.2.*,>=0.0.0a0",
     "numba>=0.57",
     "pylibraft==25.2.*,>=0.0.0a0",
     "rapids-dask-dependency==25.2.*,>=0.0.0a0",
@@ -121,6 +122,7 @@ build-backend = "scikit_build_core.build"
 requires = [
     "cmake>=3.26.4,!=3.30.0",
     "cython>=3.0.0,<3.1.0a0",
+    "libraft==25.2.*,>=0.0.0a0",
     "libucx==1.15.0",
     "ninja",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.

python/raft-dask/raft_dask/__init__.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 #
 
-from raft_dask._version import __git_commit__, __version__
-
 # If libucx was installed as a wheel, we must request it to load the library symbols.
 # Otherwise, we assume that the library was installed in a system path that ld can find.
 try:
@@ -24,3 +22,16 @@
 else:
     libucx.load_library()
     del libucx
+
+# If libraft was installed as a wheel, we must request it to load the library
+# symbols. Otherwise, we assume that the library was installed in a system path that ld
+# can find.
+try:
+    import libraft
+except ModuleNotFoundError:
+    pass
+else:
+    libraft.load_library()
+    del libraft
+
+from raft_dask._version import __git_commit__, __version__

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

@@ -15,6 +15,6 @@
 set(cython_sources comms_utils.pyx nccl.pyx)
 set(linked_libraries raft::raft raft::distributed)
 rapids_cython_create_modules(
-  SOURCE_FILES "${cython_sources}" ASSOCIATED_TARGETS raft LINKED_LIBRARIES "${linked_libraries}"
+  SOURCE_FILES "${cython_sources}" LINKED_LIBRARIES "${linked_libraries}"
                                                                             CXX
 )

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

@@ -15,6 +15,6 @@
 set(cython_sources raft_include_test.pyx)
 set(linked_libraries raft::raft)
 rapids_cython_create_modules(
-  SOURCE_FILES "${cython_sources}" ASSOCIATED_TARGETS raft LINKED_LIBRARIES "${linked_libraries}"
+  SOURCE_FILES "${cython_sources}" LINKED_LIBRARIES "${linked_libraries}"
                                                                             CXX
 )