Renaming pyraft -> raft-dask

BUILD.md

@@ -34,8 +34,8 @@ In addition to the libraries included with cudatoolkit 11.0+, there are some oth
 - [cuCollections](https://github.com/NVIDIA/cuCollections) - Used in `raft::sparse::distance` API.
 - [Libcu++](https://github.com/NVIDIA/libcudacxx) v1.7.0
 - [FAISS](https://github.com/facebookresearch/faiss) v1.7.0 - Used in `raft::spatial::knn` API and needed to build tests.
-- [NCCL](https://github.com/NVIDIA/nccl) - Used in `raft::comms` API and needed to build `Pyraft`
-- [UCX](https://github.com/openucx/ucx) - Used in `raft::comms` API and needed to build `Pyraft`
+- [NCCL](https://github.com/NVIDIA/nccl) - Used in `raft::comms` API and needed to build `raft-dask`
+- [UCX](https://github.com/openucx/ucx) - Used in `raft::comms` API and needed to build `raft-dask`
 - [Googletest](https://github.com/google/googletest) - Needed to build tests
 - [Googlebench](https://github.com/google/benchmark) - Needed to build benchmarks
 - [Doxygen](https://github.com/doxygen/doxygen) - Needed to build docs
@@ -128,9 +128,6 @@ RAFT's cmake has the following configurable flags available:.
 | RAFT_COMPILE_DIST_LIBRARY | ON, OFF | OFF | Compiles the `libraft-distance` shared library |
 | RAFT_ENABLE_NN_DEPENDENCIES | ON, OFF | OFF | Searches for dependencies of nearest neighbors API, such as FAISS, and compiles them if not found. Needed for `raft::spatial::knn` |
 | RAFT_ENABLE_thrust_DEPENDENCY | ON, OFF | ON | Enables the Thrust dependency. This can be disabled when using many simple utilities or to override with a different Thrust version. |
-| RAFT_ENABLE_mdspan_DEPENDENCY | ON, OFF | ON | Enables the std::mdspan dependency. This can be disabled when using many simple utilities. |
-| RAFT_ENABLE_nccl_DEPENDENCY | ON, OFF | OFF | Enables NCCL dependency used by `raft::comms` and needed to build `pyraft` |
-| RAFT_ENABLE_ucx_DEPENDENCY | ON, OFF | OFF | Enables UCX dependency used by `raft::comms` and needed to build `pyraft` |
 | RAFT_USE_FAISS_STATIC | ON, OFF | OFF | Statically link FAISS into `libraft-nn` | 
 | RAFT_STATIC_LINK_LIBRARIES | ON, OFF | ON | Build static link libraries instead of shared libraries |
 | DETECT_CONDA_ENV | ON, OFF | ON | Enable detection of conda environment for dependencies |
@@ -153,12 +150,12 @@ mamba activate raft_env_name
 The Python APIs can be built using the `build.sh` script:
 
 ```bash
-./build.sh pyraft pylibraft
+./build.sh raft-dask pylibraft
 ```
 
 `setup.py` can also be used to build the Python APIs manually:
 ```bash
-cd python/raft
+cd python/raft-dask
 python setup.py build_ext --inplace
 python setup.py install
 
@@ -169,7 +166,7 @@ python setup.py install
 
 To run the Python tests:
 ```bash
-cd python/raft
+cd python/raft-dask
 py.test -s -v raft
 
 cd python pylibraft
@@ -187,7 +184,7 @@ When the needed [build dependencies](#required_depenencies) are already satisfie
 set(RAFT_GIT_DIR ${CMAKE_CURRENT_BINARY_DIR}/raft CACHE STRING "Path to RAFT repo")
 ExternalProject_Add(raft
   GIT_REPOSITORY    [email protected]:rapidsai/raft.git
-  GIT_TAG           branch-22.04
+  GIT_TAG           branch-22.10
   PREFIX            ${RAFT_GIT_DIR}
   CONFIGURE_COMMAND ""
   BUILD_COMMAND     ""
@@ -219,15 +216,15 @@ The following `cmake` snippet enables a flexible configuration of RAFT:
 
 ```cmake
 
-set(RAFT_VERSION "22.04")
+set(RAFT_VERSION "22.10")
 set(RAFT_FORK "rapidsai")
 set(RAFT_PINNED_TAG "branch-${RAFT_VERSION}")
 
 function(find_and_configure_raft)
   set(oneValueArgs VERSION FORK PINNED_TAG USE_FAISS_STATIC
           COMPILE_LIBRARIES ENABLE_NN_DEPENDENCIES CLONE_ON_PIN
           USE_NN_LIBRARY USE_DISTANCE_LIBRARY 
-          ENABLE_thrust_DEPENDENCY ENABLE_mdspan_DEPENDENCY)
+          ENABLE_thrust_DEPENDENCY)
   cmake_parse_arguments(PKG "${options}" "${oneValueArgs}"
                             "${multiValueArgs}" ${ARGN} )
 
@@ -272,7 +269,6 @@ function(find_and_configure_raft)
           "RAFT_USE_FAISS_STATIC ${PKG_USE_FAISS_STATIC}"
           "RAFT_COMPILE_LIBRARIES ${PKG_COMPILE_LIBRARIES}"
           "RAFT_ENABLE_thrust_DEPENDENCY ${PKG_ENABLE_thrust_DEPENDENCY}"
-          "RAFT_ENABLE_mdspan_DEPENDENCY ${PKG_ENABLE_mdspan_DEPENDENCY}"
   )
 
 endfunction()
@@ -295,7 +291,6 @@ find_and_configure_raft(VERSION    ${RAFT_VERSION}.00
         ENABLE_NN_DEPENDENCIES   NO  # This builds FAISS if not installed
         USE_FAISS_STATIC         NO
         ENABLE_thrust_DEPENDENCY YES
-        ENABLE_mdspan_DEPENDENCY YES
 )
 ```
 

README.md

@@ -24,7 +24,7 @@ RAFT provides a header-only C++ library and pre-compiled shared libraries that c
 
 RAFT also provides 2 Python libraries:
 - `pylibraft` - low-level Python wrappers around RAFT algorithms and primitives.
-- `pyraft` - reusable infrastructure for building analytics, including tools for building both single-GPU and multi-node multi-GPU algorithms.
+- `raft-dask` - reusable infrastructure for building analytics, including tools for building both single-GPU and multi-node multi-GPU algorithms.
 
 ## Getting started
 
@@ -108,11 +108,11 @@ The easiest way to install RAFT is through conda and several packages are provid
 - `libraft-nn` (optional) contains shared libraries for the nearest neighbors primitives.
 - `libraft-distance` (optional) contains shared libraries for distance primitives.
 - `pylibraft` (optional) Python wrappers around RAFT algorithms and primitives
-- `pyraft` (optional) contains reusable Python infrastructure and tools to accelerate Python algorithm development.
+- `raft-dask` (optional) enables deployment of multi-node multi-GPU algorithms that use RAFT `raft::comms` in Dask clusters.
 
 Use the following command to install RAFT with conda (replace `rapidsai` with `rapidsai-nightly` to install more up-to-date but less stable nightly packages). `mamba` is preferred over the `conda` command.
 ```bash
-mamba install -c rapidsai libraft-headers libraft-nn libraft-distance pyraft pylibraft
+mamba install -c rapidsai libraft-headers libraft-nn libraft-distance raft-dask pylibraft
 ```
 
 After installing RAFT, `find_package(raft COMPONENTS nn distance)` can be used in your CUDA/C++ cmake build to compile and/or link against needed dependencies in your raft target. `COMPONENTS` are optional and will depend on the packages installed.
@@ -181,7 +181,7 @@ mamba env create --name raft_dev_env -f conda/environments/raft_dev_cuda11.5.yml
 mamba activate raft_dev_env
 ```
 ```
-./build.sh pyraft pylibraft libraft tests bench --compile-libs
+./build.sh raft-dask pylibraft libraft tests bench --compile-libs
 ```
 
 The [build](BUILD.md) instructions contain more details on building RAFT from source and including it in downstream projects. You can also find a more comprehensive version of the above CPM code snippet the [Building RAFT C++ from source](BUILD.md#build_cxx_source) section of the build instructions.

build.sh

@@ -18,14 +18,14 @@ ARGS=$*
 # script, and that this script resides in the repo dir!
 REPODIR=$(cd $(dirname $0); pwd)
 
-VALIDARGS="clean libraft pyraft pylibraft docs tests bench clean -v -g --install --compile-libs --compile-nn --compile-dist --allgpuarch --no-nvtx --show_depr_warn -h --buildfaiss --minimal-deps"
+VALIDARGS="clean libraft pylibraft raft-dask docs tests bench clean -v -g --install --compile-libs --compile-nn --compile-dist --allgpuarch --no-nvtx --show_depr_warn -h --buildfaiss --minimal-deps"
 HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<tool>]
  where <target> is:
    clean            - remove all existing build artifacts and configuration (start over)
    libraft          - build the raft C++ code only. Also builds the C-wrapper library
                       around the C++ code.
-   pyraft           - build the pyraft Python package
    pylibraft        - build the pylibraft Python package
+   raft-dask        - build the raft-dask Python package
    docs             - build the documentation
    tests            - build the tests
    bench            - build the benchmarks
@@ -50,7 +50,7 @@ HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<to
                                  to speedup the build process.
    -h                          - print this text
 
- default action (no args) is to build both libraft and pyraft targets
+ default action (no args) is to build both libraft and raft-dask targets
 "
 LIBRAFT_BUILD_DIR=${LIBRAFT_BUILD_DIR:=${REPODIR}/cpp/build}
 SPHINX_BUILD_DIR=${REPODIR}/docs
@@ -241,7 +241,7 @@ if (( ${CLEAN} == 1 )); then
       fi
     done
 
-    cd ${REPODIR}/python/raft
+    cd ${REPODIR}/python/raft-dask
     python setup.py clean --all
     cd ${REPODIR}
 
@@ -290,10 +290,10 @@ if (( ${NUMARGS} == 0 )) || hasArg libraft || hasArg pylibraft || hasArg docs ||
   fi
 fi
 
-# Build and (optionally) install the pyraft Python package
-if (( ${NUMARGS} == 0 )) || hasArg pyraft || hasArg docs; then
+# Build and (optionally) install the raft-dask Python package
+if (( ${NUMARGS} == 0 )) || hasArg raft-dask || hasArg docs; then
 
-    cd ${REPODIR}/python/raft
+    cd ${REPODIR}/python/raft-dask
     python setup.py build_ext --inplace -- -DCMAKE_PREFIX_PATH="${LIBRAFT_BUILD_DIR};${INSTALL_PREFIX}" -DCMAKE_LIBRARY_PATH=${LIBRAFT_BUILD_DIR} ${EXTRA_CMAKE_ARGS} -- -j${PARALLEL_LEVEL:-1}
     if [[ ${INSTALL_TARGET} != "" ]]; then
         python setup.py install --single-version-externally-managed --record=record.txt -- -DCMAKE_PREFIX_PATH=${INSTALL_PREFIX} ${EXTRA_CMAKE_ARGS}

ci/cpu/build.sh

@@ -99,12 +99,12 @@ fi
 if [ "$BUILD_RAFT" == '1' ]; then
   gpuci_logger "Building Python conda packages for raft"
   if [[ -z "$PROJECT_FLASH" || "$PROJECT_FLASH" == "0" ]]; then
-    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/pyraft --python=$PYTHON
+    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/raft-dask --python=$PYTHON
     gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/pylibraft --python=$PYTHON
   else
-    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/pyraft -c ${CONDA_LOCAL_CHANNEL} --dirty --no-remove-work-dir --python=$PYTHON
-    mkdir -p ${CONDA_BLD_DIR}/pyraft/work
-    mv ${CONDA_BLD_DIR}/work ${CONDA_BLD_DIR}/pyraft/work
+    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/raft-dask -c ${CONDA_LOCAL_CHANNEL} --dirty --no-remove-work-dir --python=$PYTHON
+    mkdir -p ${CONDA_BLD_DIR}/raft-dask/work
+    mv ${CONDA_BLD_DIR}/work ${CONDA_BLD_DIR}/raft-dask/work
 
     gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/pylibraft -c ${CONDA_LOCAL_CHANNEL} --dirty --no-remove-work-dir --python=$PYTHON
     mkdir -p ${CONDA_BLD_DIR}/pylibraft/work

ci/cpu/upload.sh

@@ -38,12 +38,12 @@ if [[ "$BUILD_LIBRAFT" == "1" && "$UPLOAD_LIBRAFT" == "1" ]]; then
 fi
 
 if [[ "$BUILD_RAFT" == "1" && "$UPLOAD_RAFT" == "1" ]]; then
-  PYRAFT_FILE=$(conda build --no-build-id --croot ${CONDA_BLD_DIR} -c ${CONDA_LOCAL_CHANNEL} conda/recipes/pyraft --python=$PYTHON --output)
+  DASKRAFT_FILE=$(conda build --no-build-id --croot ${CONDA_BLD_DIR} -c ${CONDA_LOCAL_CHANNEL} conda/recipes/raft-dask --python=$PYTHON --output)
   PYLIBRAFT_FILE=$(conda build --no-build-id --croot ${CONDA_BLD_DIR} -c ${CONDA_LOCAL_CHANNEL} conda/recipes/pylibraft --python=$PYTHON --output)
-  test -e ${PYRAFT_FILE}
-  echo "Upload pyraft"
-  echo ${PYRAFT_FILE}
-  gpuci_retry anaconda -t ${MY_UPLOAD_KEY} upload -u ${CONDA_USERNAME:-rapidsai} ${LABEL_OPTION} --skip-existing ${PYRAFT_FILE} --no-progress
+  test -e ${DASKRAFT_FILE}
+  echo "Upload raft-dask"
+  echo ${DASKRAFT_FILE}
+  gpuci_retry anaconda -t ${MY_UPLOAD_KEY} upload -u ${CONDA_USERNAME:-rapidsai} ${LABEL_OPTION} --skip-existing ${DASKRAFT_FILE} --no-progress
 
   test -e ${PYLIBRAFT_FILE}
   echo "Upload pylibraft"

ci/gpu/build.sh

@@ -63,9 +63,13 @@ conda list --show-channel-urls
 gpuci_logger "Build and install Python targets"
 CONDA_BLD_DIR="$WORKSPACE/.conda-bld"
 gpuci_mamba_retry install boa
-gpuci_conda_retry mambabuild --no-build-id --croot "${CONDA_BLD_DIR}" conda/recipes/pyraft -c "${CONDA_ARTIFACT_PATH}" --python="${PYTHON}"
+
+# Install pylibraft first since it's a dependency of raft-dask
 gpuci_conda_retry mambabuild --no-build-id --croot "${CONDA_BLD_DIR}" conda/recipes/pylibraft -c "${CONDA_ARTIFACT_PATH}" --python="${PYTHON}"
-gpuci_mamba_retry install -y -c "${CONDA_BLD_DIR}" -c "${CONDA_ARTIFACT_PATH}" pyraft pylibraft
+gpuci_mamba_retry install -y -c "${CONDA_BLD_DIR}" -c "${CONDA_ARTIFACT_PATH}" pylibraft
+
+gpuci_conda_retry mambabuild --no-build-id --croot "${CONDA_BLD_DIR}" conda/recipes/raft-dask -c "${CONDA_ARTIFACT_PATH}" --python="${PYTHON}"
+gpuci_mamba_retry install -y -c "${CONDA_BLD_DIR}" -c "${CONDA_ARTIFACT_PATH}" raft-dask
 
 ################################################################################
 # TEST - Run GoogleTest and py.tests for RAFT
@@ -86,12 +90,12 @@ set +x
 gpuci_logger "Check GPU usage"
 nvidia-smi
 
-gpuci_logger "GoogleTest for raft"
+gpuci_logger "GoogleTest for libraft"
 GTEST_OUTPUT="xml:${WORKSPACE}/test-results/raft_cpp/" $CONDA_PREFIX/bin/libraft/gtests/test_raft
 
-gpuci_logger "Python pytest for pyraft"
-cd "$WORKSPACE/python/raft/raft/test"
-pytest --cache-clear --junitxml="$WORKSPACE/junit-pyraft.xml" -v -s
+gpuci_logger "Python pytest for raft-dask"
+cd "$WORKSPACE/python/raft-dask/raft_dask/test"
+pytest --cache-clear --junitxml="$WORKSPACE/junit-raft-dask.xml" -v -s
 
 gpuci_logger "Python pytest for pylibraft"
 cd "$WORKSPACE/python/pylibraft/pylibraft/test"

conda/recipes/pyraft/build.sh → conda/recipes/raft-dask/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 raft-dask --install --no-nvtx

conda/recipes/pyraft/meta.yaml → conda/recipes/raft-dask/meta.yaml

@@ -10,7 +10,7 @@
 {% set ucx_py_version=environ.get('UCX_PY_VERSION') %}
 
 package:
-  name: pyraft
+  name: raft-dask
   version: {{ version }}
 
 source:
@@ -35,7 +35,7 @@ requirements:
     - cython>=0.29,<0.30
     - scikit-build>=0.13.1
     - rmm {{ minor_version }}
-    - libraft-headers {{ version }}
+    - pylibraft {{ version }}
     - cudatoolkit {{ cuda_version }}.*
     - cuda-python >=11.5,<11.7.1
     - nccl>=2.9.9
@@ -45,7 +45,7 @@ requirements:
   run:
     - python x.x
     - dask-cuda {{ minor_version }}
-    - libraft-headers {{ version }}
+    - pylibraft {{ version }}
     - nccl>=2.9.9
     - rmm {{ minor_version }}
     - ucx >={{ ucx_version }}
@@ -61,10 +61,10 @@ tests:                                 # [linux64]
   requirements:                        # [linux64]
     - cudatoolkit {{ cuda_version }}.* # [linux64]
   imports:                             # [linux64]
-    - raft                             # [linux64]
+    - raft_dask                        # [linux64]
 
 about:
   home: http://rapids.ai/
   license: Apache-2.0
   # license_file: LICENSE
-  summary: pyraft library
+  summary: raft-dask library

docs/source/conf.py

@@ -24,7 +24,7 @@
 # is relative to the documentation root, use os.path.abspath to make it
 # absolute, like shown here.
 sys.path.insert(0, os.path.abspath("sphinxext"))
-sys.path.insert(0, os.path.abspath("../../python/raft"))
+sys.path.insert(0, os.path.abspath("../../python/raft-dask"))
 sys.path.insert(0, os.path.abspath("../../python/pylibraft"))
 
 from github_link import make_linkcode_resolve  # noqa

docs/source/index.rst

@@ -8,7 +8,7 @@ RAFT contains fundamental widely-used algorithms and primitives for data science
    :caption: Contents:
 
    cpp_api.rst
-   pyraft_api.rst
+   raft_dask_api.rst
    pylibraft_api.rst
 
 

docs/source/python.rst

@@ -5,6 +5,6 @@ RAFT Python APIs
    :maxdepth: 2
    :caption: Contents:
 
-   pyraft_api.rst
    pylibraft_api.rst
+   raft_dask_api.rst
 

docs/source/raft_dask_api.rst

@@ -0,0 +1,13 @@
+~~~~~~~~~~~~~~~~~~~~~~~
+RAFT Dask API Reference
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. role:: py(code)
+   :language: python
+   :class: highlight
+
+Dask-based Multi-Node Multi-GPU Communicator
+--------------------------------------------
+
+.. autoclass:: raft_dask.common.Comms
+    :members:

python/pylibraft/CMakeLists.txt

@@ -68,5 +68,6 @@ endif()
 include(rapids-cython)
 rapids_cython_init()
 
+add_subdirectory(pylibraft/common)
 add_subdirectory(pylibraft/distance)
 add_subdirectory(pylibraft/random)

python/raft/raft/common/CMakeLists.txt → python/pylibraft/pylibraft/common/CMakeLists.txt

@@ -12,9 +12,17 @@
 # the License.
 # =============================================================================
 
+# Set the list of Cython files to build
 set(cython_sources cuda.pyx handle.pyx interruptible.pyx)
 set(linked_libraries raft::raft)
+
+# Build all of the Cython targets
 rapids_cython_create_modules(
-    SOURCE_FILES "${cython_sources}"
-    LINKED_LIBRARIES "${linked_libraries}"
-    CXX)
+        CXX
+        SOURCE_FILES "${cython_sources}"
+        LINKED_LIBRARIES "${linked_libraries}"
+        MODULE_PREFIX common_)
+
+foreach(cython_module IN LISTS RAPIDS_CYTHON_CREATED_TARGETS)
+    set_target_properties(${cython_module} PROPERTIES INSTALL_RPATH "\$ORIGIN;\$ORIGIN/../library")
+endforeach()

python/pylibraft/pylibraft/common/handle.pxd

@@ -19,6 +19,7 @@
 # cython: embedsignature = True
 # cython: language_level = 3
 
+
 from libcpp.memory cimport shared_ptr
 from rmm._lib.cuda_stream_view cimport cuda_stream_view
 from rmm._lib.cuda_stream_pool cimport cuda_stream_pool
@@ -33,3 +34,8 @@ cdef extern from "raft/core/handle.hpp" namespace "raft" nogil:
                  shared_ptr[cuda_stream_pool] stream_pool) except +
         cuda_stream_view get_stream() except +
         void sync_stream() except +
+
+cdef class Handle:
+    cdef unique_ptr[handle_t] c_obj
+    cdef shared_ptr[cuda_stream_pool] stream_pool
+    cdef int n_streams