libraft minimally working

ci/build_wheel.sh

@@ -5,6 +5,7 @@ set -euo pipefail
 
 package_name=$1
 package_dir=$2
+package_type=$3
 underscore_package_name=$(echo "${package_name}" | tr "-" "_")
 
 # Clear out system ucx files to ensure that we're getting ucx from the wheel.
@@ -55,4 +56,4 @@ sccache --show-adv-stats
 mkdir -p final_dist
 python -m auditwheel repair -w final_dist "${EXCLUDE_ARGS[@]}" dist/*
 
-RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 python final_dist
+RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 ${package_type} final_dist

ci/build_wheel_libraft.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+package_name="libraft"
+package_dir="python/libraft"
+
+rapids-logger "Generating build requirements"
+matrix_selectors="cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch);py=${RAPIDS_PY_VERSION};cuda_suffixed=true"
+
+rapids-dependency-file-generator \
+  --output requirements \
+  --file-key "py_build_${package_name}" \
+  --file-key "py_rapids_build_${package_name}" \
+  --matrix "${matrix_selectors}" \
+| tee /tmp/requirements-build.txt
+
+rapids-logger "Installing build requirements"
+python -m pip install \
+    -v \
+    --prefer-binary \
+    -r /tmp/requirements-build.txt
+
+# build with '--no-build-isolation', for better sccache hit rate
+# 0 really means "add --no-build-isolation" (ref: https://github.com/pypa/pip/issues/5735)
+export PIP_NO_BUILD_ISOLATION=0
+
+RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
+
+case "${RAPIDS_CUDA_VERSION}" in
+  12.*)
+    EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=ON"
+  ;;
+  11.*)
+    EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=OFF"
+  ;;
+esac
+
+export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF${EXTRA_CMAKE_ARGS}"
+
+ci/build_wheel.sh libraft ${package_dir} cpp
+ci/validate_wheel.sh ${package_dir} final_dist libraft

ci/build_wheel_pylibraft.sh

@@ -17,5 +17,5 @@ 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}"
 
-ci/build_wheel.sh pylibraft ${package_dir}
+ci/build_wheel.sh pylibraft ${package_dir} python
 ci/validate_wheel.sh ${package_dir} final_dist pylibraft

ci/build_wheel_raft_dask.sh

@@ -8,5 +8,5 @@ package_dir="python/raft-dask"
 # Set up skbuild options. Enable sccache in skbuild config options
 export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF"
 
-ci/build_wheel.sh raft-dask ${package_dir}
+ci/build_wheel.sh raft-dask ${package_dir} python
 ci/validate_wheel.sh ${package_dir} final_dist raft-dask

dependencies.yaml

@@ -48,6 +48,23 @@ files:
       - docs
       - py_version
       - test_pylibraft
+  py_build_libraft:
+    output: pyproject
+    pyproject_dir: python/libraft
+    extras:
+      table: build-system
+    includes:
+      - rapids_build_skbuild
+  py_rapids_build_libraft:
+    output: pyproject
+    pyproject_dir: python/libraft
+    extras:
+      table: tool.rapids-build-backend
+      key: requires
+    includes:
+      - build_pylibraft
+      - depends_on_librmm
+      - rapids_build
   py_build_pylibraft:
     output: pyproject
     pyproject_dir: python/pylibraft
@@ -537,6 +554,32 @@ dependencies:
             packages:
               - distributed-ucxx-cu11==0.42.*,>=0.0.0a0
           - {matrix: null, packages: [*distributed_ucxx_unsuffixed]}
+  depends_on_librmm:
+    common:
+      - output_types: conda
+        packages:
+          - &librmm_unsuffixed librmm==25.2.*,>=0.0.0a0
+      - 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:
+              - librmm-cu12==25.2.*,>=0.0.0a0
+          - matrix:
+              cuda: "11.*"
+              cuda_suffixed: "true"
+            packages:
+              - librmm-cu11==25.2.*,>=0.0.0a0
+          - matrix:
+            packages:
+              - *librmm_unsuffixed
   depends_on_ucx_build:
     common:
       - output_types: conda

python/libraft/CMakeLists.txt

@@ -29,6 +29,10 @@ option(USE_CUDA_MATH_WHEELS "Use the CUDA math wheels instead of the system libr
 
 # Check if raft is already available. If so, it is the user's responsibility to ensure that the
 # CMake package is also available at build time of the Python raft package.
+
+find_package(CUDAToolkit REQUIRED)
+
+# TODO(jameslamb): make this work
 find_package(raft "${RAPIDS_VERSION}")
 
 if(raft_FOUND)

python/libraft/libraft/_version.py

@@ -15,7 +15,10 @@
 import importlib.resources
 
 __version__ = (
-    importlib.resources.files(__package__).joinpath("VERSION").read_text().strip()
+    importlib.resources.files(__package__)
+    .joinpath("VERSION")
+    .read_text()
+    .strip()
 )
 try:
     __git_commit__ = (

python/libraft/libraft/load.py

@@ -36,15 +36,18 @@ def _load_wheel_installation(soname: str):
     """Try to dlopen() the library indicated by ``soname``
     Returns ``None`` if the library cannot be loaded.
     """
-    if os.path.isfile(lib := os.path.join(os.path.dirname(__file__), "lib64", soname)):
+    if os.path.isfile(
+        lib := os.path.join(os.path.dirname(__file__), "lib64", soname)
+    ):
         return ctypes.CDLL(lib, PREFERRED_LOAD_FLAG)
     return None
 
 
 def load_library():
-    """Dynamically load libcugraph.so and its dependencies"""
+    """Dynamically load libraft.so and its dependencies"""
     prefer_system_installation = (
-        os.getenv("RAPIDS_LIBRAFT_PREFER_SYSTEM_LIBRARY", "false").lower() != "false"
+        os.getenv("RAPIDS_LIBRAFT_PREFER_SYSTEM_LIBRARY", "false").lower()
+        != "false"
     )
 
     soname = "libraft.so"
@@ -73,5 +76,5 @@ def load_library():
 
     # The caller almost never needs to do anything with this library, but no
     # harm in offering the option since this object at least provides a handle
-    # to inspect where libcugraph was loaded from.
+    # to inspect where libraft was loaded from.
     return libraft_lib