Using Pyraft and automatically cloning when raft pin changes (#4593)

This PR depends on rapidsai/raft#520

Authors:
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - AJ Schmidt (https://github.com/ajschmidt8)
  - Dante Gama Dessavre (https://github.com/dantegd)

URL: #4593

ci/gpu/build.sh

@@ -61,6 +61,7 @@ gpuci_mamba_retry install -c conda-forge -c rapidsai -c rapidsai-nightly -c nvid
       "libraft-headers=${MINOR_VERSION}" \
       "libraft-distance=${MINOR_VERSION}" \
       "libraft-nn=${MINOR_VERSION}" \
+      "pyraft=${MINOR_VERSION}" \
       "dask-cudf=${MINOR_VERSION}" \
       "dask-cuda=${MINOR_VERSION}" \
       "ucx-py=${UCX_PY_VERSION}" \

conda/environments/cuml_dev_cuda11.0.yml

@@ -16,6 +16,7 @@ dependencies:
 - libraft-headers=22.04.*
 - libraft-distance=22.04.*
 - libraft-nn=22.04.*
+- pyraft=22.04.*
 - dask-cudf=22.04.*
 - dask-cuda=22.04.*
 - ucx-py=0.25.*

conda/environments/cuml_dev_cuda11.2.yml

@@ -16,6 +16,7 @@ dependencies:
 - libraft-headers=22.04.*
 - libraft-distance=22.04.*
 - libraft-nn=22.04.*
+- pyraft=22.04.*
 - dask-cudf=22.04.*
 - dask-cuda=22.04.*
 - ucx-py=0.25.*

conda/environments/cuml_dev_cuda11.4.yml

@@ -16,6 +16,7 @@ dependencies:
 - libraft-headers=22.04.*
 - libraft-distance=22.04.*
 - libraft-nn=22.04.*
+- pyraft=22.04.*
 - dask-cudf=22.04.*
 - dask-cuda=22.04.*
 - ucx-py=0.25.*

conda/environments/cuml_dev_cuda11.5.yml

@@ -16,6 +16,7 @@ dependencies:
 - libraft-headers=22.04.*
 - libraft-distance=22.04.*
 - libraft-nn=22.04.*
+- pyraft=22.04.*
 - dask-cudf=22.04.*
 - dask-cuda=22.04.*
 - ucx-py=0.25.*

conda/recipes/cuml/meta.yaml

@@ -34,6 +34,7 @@ requirements:
     - cudf {{ minor_version }}
     - libcuml={{ version }}
     - libcumlprims {{ minor_version }}
+    - pyraft {{ minor_version }}
     - cudatoolkit {{ cuda_version }}.*
     - ucx-py {{ ucx_py_version }}
     - ucx-proc=*=gpu
@@ -44,6 +45,7 @@ requirements:
     - dask-cudf {{ minor_version }}
     - libcuml={{ version }}
     - libcumlprims {{ minor_version }}
+    - pyraft {{ minor_version }}
     - cupy>=7.8.0,<11.0.0a0
     - treelite=2.3.0
     - nccl>=2.9.9

cpp/cmake/thirdparty/get_raft.cmake

@@ -14,12 +14,20 @@
 # limitations under the License.
 #=============================================================================
 
+set(CUML_MIN_VERSION_raft "${CUML_VERSION_MAJOR}.${CUML_VERSION_MINOR}.00")
+set(CUML_BRANCH_VERSION_raft "${CUML_VERSION_MAJOR}.${CUML_VERSION_MINOR}")
+
 function(find_and_configure_raft)
 
-    set(oneValueArgs VERSION FORK PINNED_TAG USE_RAFT_NN USE_FAISS_STATIC)
+    set(oneValueArgs VERSION FORK PINNED_TAG USE_RAFT_NN USE_FAISS_STATIC CLONE_ON_PIN)
     cmake_parse_arguments(PKG "${options}" "${oneValueArgs}"
             "${multiValueArgs}" ${ARGN} )
 
+    if(PKG_CLONE_ON_PIN AND NOT PKG_PINNED_TAG STREQUAL "branch-${CUML_BRANCH_VERSION_raft}")
+        message("Pinned tag found: ${PKG_PINNED_TAG}. Cloning raft locally.")
+        set(CPM_DOWNLOAD_raft ON)
+    endif()
+
     string(APPEND RAFT_COMPONENTS "distance")
     if(PKG_USE_RAFT_NN)
         string(APPEND RAFT_COMPONENTS " nn")
@@ -48,17 +56,20 @@ function(find_and_configure_raft)
         message(VERBOSE "CUML: Using RAFT located in ${raft_DIR}")
     endif()
 
-endfunction()
 
-set(CUML_MIN_VERSION_raft "${CUML_VERSION_MAJOR}.${CUML_VERSION_MINOR}.00")
-set(CUML_BRANCH_VERSION_raft "${CUML_VERSION_MAJOR}.${CUML_VERSION_MINOR}")
+endfunction()
 
 # Change pinned tag here to test a commit in CI
 # To use a different RAFT locally, set the CMake variable
 # CPM_raft_SOURCE=/path/to/local/raft
 find_and_configure_raft(VERSION          ${CUML_MIN_VERSION_raft}
                         FORK             rapidsai
                         PINNED_TAG       branch-${CUML_BRANCH_VERSION_raft}
+
+                        # When PINNED_TAG above doesn't match cuml,
+                        # force local raft clone in build directory
+                        # even if it's already installed.
+                        CLONE_ON_PIN     ON
                         USE_RAFT_NN      ${CUML_USE_RAFT_NN}
                         USE_FAISS_STATIC ${CUML_USE_FAISS_STATIC}
                         )

python/cuml/__init__.py

@@ -91,10 +91,6 @@
 from cuml.common.pointer_utils import device_of_gpu_matrix
 from cuml.common.memory_utils import set_global_output_type, using_output_type
 
-# RAFT
-
-from cuml.raft import raft_include_test
-
 # Import verion. Remove at end of file
 from ._version import get_versions
 

python/cuml/cluster/agglomerative.pyx

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2019-2021, NVIDIA CORPORATION.
+# Copyright (c) 2019-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.
@@ -23,7 +23,7 @@ import numpy as np
 from cuml.common.array import CumlArray
 from cuml.common.base import Base
 from cuml.common.doc_utils import generate_docstring
-from cuml.raft.common.handle cimport handle_t
+from raft.common.handle cimport handle_t
 from cuml.common import input_to_cuml_array
 from cuml.common.array_descriptor import CumlArrayDescriptor
 from cuml.common.mixins import ClusterMixin

python/cuml/cluster/dbscan.pyx

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2019-2021, NVIDIA CORPORATION.
+# Copyright (c) 2019-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.
@@ -28,7 +28,7 @@ from libc.stdlib cimport calloc, malloc, free
 from cuml.common.array import CumlArray
 from cuml.common.base import Base
 from cuml.common.doc_utils import generate_docstring
-from cuml.raft.common.handle cimport handle_t
+from raft.common.handle cimport handle_t
 from cuml.common import input_to_cuml_array
 from cuml.common import using_output_type
 from cuml.common.array_descriptor import CumlArrayDescriptor

python/cuml/cluster/hdbscan.pyx

@@ -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.
@@ -27,7 +27,9 @@ import cupy as cp
 from cuml.common.array import CumlArray
 from cuml.common.base import Base
 from cuml.common.doc_utils import generate_docstring
-from cuml.raft.common.handle cimport handle_t
+from raft.common.handle cimport handle_t
+
+from raft.common.handle import Handle
 from cuml.common import input_to_cuml_array
 from cuml.common.array_descriptor import CumlArrayDescriptor
 from cuml.common.mixins import ClusterMixin
@@ -192,7 +194,7 @@ def condense_hierarchy(dendrogram,
     condensed_tree : hdbscan.plots.CondensedTree object
     """
 
-    handle = cuml.raft.common.handle.Handle()
+    handle = Handle()
     cdef handle_t *handle_ = <handle_t*> <size_t> handle.getHandle()
 
     n_leaves = dendrogram.shape[0]+1

python/cuml/cluster/kmeans.pyx

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2019-2021, NVIDIA CORPORATION.
+# Copyright (c) 2019-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.
@@ -35,7 +35,7 @@ from cuml.common.mixins import ClusterMixin
 from cuml.common.mixins import CMajorInputTagMixin
 from cuml.common import input_to_cuml_array
 from cuml.cluster.kmeans_utils cimport *
-from cuml.raft.common.handle cimport handle_t
+from raft.common.handle cimport handle_t
 
 cdef extern from "cuml/cluster/kmeans.hpp" namespace "ML::kmeans":
 

python/cuml/cluster/kmeans_mg.pyx

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2019-2021, NVIDIA CORPORATION.
+# Copyright (c) 2019-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.
@@ -29,7 +29,7 @@ from libc.stdlib cimport calloc, malloc, free
 
 from cuml.common.array import CumlArray
 from cuml.common.base import Base
-from cuml.raft.common.handle cimport handle_t
+from raft.common.handle cimport handle_t
 from cuml.common import input_to_cuml_array
 
 from cuml.cluster import KMeans

python/cuml/common/base.pyx

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2019-2021, NVIDIA CORPORATION.
+# Copyright (c) 2019-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.
@@ -24,7 +24,7 @@ import cuml.common
 import cuml.common.cuda
 import cuml.common.logger as logger
 import cuml.internals
-import cuml.raft.common.handle
+import raft.common.handle
 import cuml.common.input_utils
 
 from cuml.common.doc_utils import generate_docstring
@@ -164,7 +164,7 @@ class Base(TagsMixin,
         Constructor. All children must call init method of this base class.
 
         """
-        self.handle = cuml.raft.common.handle.Handle() if handle is None \
+        self.handle = raft.common.handle.Handle() if handle is None \
             else handle
 
         # Internally, self.verbose follows the spdlog/c++ standard of

python/cuml/common/handle.pyx

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2020, NVIDIA CORPORATION.
+# Copyright (c) 2020-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.
@@ -16,6 +16,6 @@
 
 # distutils: language = c++
 
-from cuml.raft.common.handle import Handle as raftHandle
+from raft.common.handle import Handle as raftHandle
 
 Handle = raftHandle

python/cuml/dask/cluster/dbscan.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2021, NVIDIA CORPORATION.
+# Copyright (c) 2020-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.
@@ -20,8 +20,8 @@
 from cuml.dask.common.base import DelayedTransformMixin
 from cuml.dask.common.base import mnmg_import
 
-from cuml.raft.dask.common.comms import Comms
-from cuml.raft.dask.common.comms import get_raft_comm_state
+from raft.dask.common.comms import Comms
+from raft.dask.common.comms import get_raft_comm_state
 
 from cuml.dask.common.utils import wait_and_raise_from_futures
 

python/cuml/dask/cluster/kmeans.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2021, NVIDIA CORPORATION.
+# Copyright (c) 2019-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.
@@ -23,8 +23,8 @@
 from cuml.dask.common.input_utils import concatenate
 from cuml.dask.common.input_utils import DistributedDataHandler
 
-from cuml.raft.dask.common.comms import Comms
-from cuml.raft.dask.common.comms import get_raft_comm_state
+from raft.dask.common.comms import Comms
+from raft.dask.common.comms import get_raft_comm_state
 
 from cuml.dask.common.utils import wait_and_raise_from_futures
 

python/cuml/dask/common/base.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2021, NVIDIA CORPORATION.
+# Copyright (c) 2020-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.
@@ -25,7 +25,7 @@
 from cuml import Base
 from cuml.common.array import CumlArray
 from cuml.dask.common.utils import wait_and_raise_from_futures
-from cuml.raft.dask.common.comms import Comms
+from raft.dask.common.comms import Comms
 from cuml.dask.common.input_utils import DistributedDataHandler
 from cuml.dask.common import parts_to_ranks
 from cuml.internals import BaseMetaClass

python/cuml/dask/decomposition/base.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2021, NVIDIA CORPORATION.
+# Copyright (c) 2019-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.
@@ -14,8 +14,8 @@
 #
 
 from cuml.dask.common import raise_exception_from_futures
-from cuml.raft.dask.common.comms import get_raft_comm_state
-from cuml.raft.dask.common.comms import Comms
+from raft.dask.common.comms import get_raft_comm_state
+from raft.dask.common.comms import Comms
 
 from cuml.dask.common.input_utils import to_output
 from cuml.dask.common import parts_to_ranks

python/cuml/dask/linear_model/linear_regression.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2021, NVIDIA CORPORATION.
+# Copyright (c) 2019-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.
@@ -17,7 +17,7 @@
 from cuml.dask.common.base import DelayedPredictionMixin
 from cuml.dask.common.base import mnmg_import
 from cuml.dask.common.base import SyncFitMixinLinearModel
-from cuml.raft.dask.common.comms import get_raft_comm_state
+from raft.dask.common.comms import get_raft_comm_state
 
 
 class LinearRegression(BaseEstimator,

python/cuml/dask/linear_model/ridge.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2021, NVIDIA CORPORATION.
+# Copyright (c) 2019-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.
@@ -17,7 +17,7 @@
 from cuml.dask.common.base import DelayedPredictionMixin
 from cuml.dask.common.base import mnmg_import
 from cuml.dask.common.base import SyncFitMixinLinearModel
-from cuml.raft.dask.common.comms import get_raft_comm_state
+from raft.dask.common.comms import get_raft_comm_state
 
 
 class Ridge(BaseEstimator,

python/cuml/dask/neighbors/kneighbors_classifier.py

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2020-2021, NVIDIA CORPORATION.
+# Copyright (c) 2020-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.
@@ -20,7 +20,7 @@
 from cuml.dask.common import flatten_grouped_results
 from cuml.dask.common.utils import raise_mg_import_exception
 from cuml.dask.common.utils import wait_and_raise_from_futures
-from cuml.raft.dask.common.comms import get_raft_comm_state
+from raft.dask.common.comms import get_raft_comm_state
 from cuml.dask.neighbors import NearestNeighbors
 from dask.dataframe import Series as DaskSeries
 import dask.array as da

python/cuml/dask/neighbors/kneighbors_regressor.py

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2020-2021, NVIDIA CORPORATION.
+# Copyright (c) 2020-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.
@@ -20,7 +20,7 @@
 from cuml.dask.common import flatten_grouped_results
 from cuml.dask.common.utils import raise_mg_import_exception
 from cuml.dask.common.utils import wait_and_raise_from_futures
-from cuml.raft.dask.common.comms import get_raft_comm_state
+from raft.dask.common.comms import get_raft_comm_state
 from cuml.dask.neighbors import NearestNeighbors
 import dask.array as da
 from uuid import uuid1

python/cuml/dask/neighbors/nearest_neighbors.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2021, NVIDIA CORPORATION.
+# Copyright (c) 2019-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.
@@ -19,8 +19,8 @@
 from cuml.dask.common import raise_mg_import_exception
 from cuml.dask.common.base import BaseEstimator
 
-from cuml.raft.dask.common.comms import get_raft_comm_state
-from cuml.raft.dask.common.comms import Comms
+from raft.dask.common.comms import get_raft_comm_state
+from raft.dask.common.comms import Comms
 from cuml.dask.common.input_utils import to_output
 from cuml.dask.common.input_utils import DistributedDataHandler