-
Notifications
You must be signed in to change notification settings - Fork 539
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Migrate to use cuVS for vector search #6085
Merged
Merged
Changes from 51 commits
Commits
Show all changes
52 commits
Select commit
Hold shift + click to select a range
2f2c55c
Migrate from raft to cuvs for pairwise_distance and bfknn
benfred 9116ae8
Merge branch 'branch-24.08' into cuvs
benfred 84bc77a
.
benfred cbb79ec
use cuvs::distance::DistanceType where possible
benfred f653059
Revert "use cuvs::distance::DistanceType where possible"
benfred a6f2c2a
use stats from raft
benfred c2c9c04
Merge remote-tracking branch 'origin/branch-24.10' into cuvs
benfred 95b9c14
use ivf-* from cuvs
benfred 1934d40
Merge remote-tracking branch 'origin/branch-24.10' into cuvs
benfred 0601412
Merge branch 'rapidsai:branch-24.10' into cuvs
benfred f30f933
add libcuvs to dependencies.yaml
benfred 69d2398
Merge branch 'cuvs' of https://github.com/benfred/cuml into cuvs
benfred 6c46624
attempt to fix build error in CI
benfred 3fb5479
Merge branch 'branch-24.10' into cuvs
benfred 5380c17
fix tsne
benfred dc1f7a6
Merge branch 'cuvs' of https://github.com/benfred/cuml into cuvs
benfred fdd18c5
fix ivf-fla
benfred 590865d
fix test_nearest_neighbors_rbc test for haversine distance
benfred c7d1b0e
re-add MetricProcessor code
benfred e8c1b18
suggestions from code review
benfred bd58347
fix dask pytests
benfred adb450a
attempt to fix python build errors in CI
benfred df47d3c
use raft in header only mode
benfred c72173c
Use kmeans/mutual reachability code from pending cuvs PR's
benfred 6139db0
remove comment
benfred 3e1c465
cmake fixes
benfred 732ea10
.
benfred 4942bb0
pick up right cuvs version
benfred fef9920
.
benfred 3c0c47e
empty commit for ci
benfred d8e6b6d
Exclude libcuvs.so in auditwheel.
bdice 257898e
add cuvs to python dependencies
benfred 9f805cd
Merge branch 'cuvs' of https://github.com/benfred/cuml into cuvs
benfred 176c9a9
use l2expanded distance in kmeans transform
benfred 2122120
Merge branch 'branch-24.10' into cuvs
benfred c12f1d5
Set rpath for cuvs
KyleFromNVIDIA f3edb8f
Don't link Python modules against cuvs directly
KyleFromNVIDIA f906e79
Remove superfluous cuvs::cuvs references
KyleFromNVIDIA 585acd4
Add cuvs rpath
KyleFromNVIDIA 1964a94
remove cuvs pin
benfred f1db388
Merge branch 'cuvs' of https://github.com/benfred/cuml into cuvs
benfred 259256a
updates to handle bfknn api changes
benfred 4efd97a
link cuvs statically in python wheels
benfred 41daf01
empty commit for ci
benfred 31331d7
empty commit for ci
benfred 6bef472
empty commit for ci
benfred 3e88b8e
remove pin
benfred 064cced
Merge branch 'branch-24.10' into cuvs
cjnolet 3959c58
re-add pin + suggestions from code review
benfred dc6de84
Merge branch 'cuvs' of https://github.com/benfred/cuml into cuvs
benfred 461a271
fix
benfred f50b9d2
remove pin
benfred File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,77 @@ | ||
#============================================================================= | ||
# Copyright (c) 2024, NVIDIA CORPORATION. | ||
# | ||
# Licensed under the Apache License, Version 2.0 (the "License"); | ||
# you may not use this file except in compliance with the License. | ||
# You may obtain a copy of the License at | ||
# | ||
# http://www.apache.org/licenses/LICENSE-2.0 | ||
# | ||
# Unless required by applicable law or agreed to in writing, software | ||
# distributed under the License is distributed on an "AS IS" BASIS, | ||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
# See the License for the specific language governing permissions and | ||
# limitations under the License. | ||
#============================================================================= | ||
|
||
set(CUML_MIN_VERSION_cuvs "${CUML_VERSION_MAJOR}.${CUML_VERSION_MINOR}.00") | ||
set(CUML_BRANCH_VERSION_cuvs "${CUML_VERSION_MAJOR}.${CUML_VERSION_MINOR}") | ||
|
||
function(find_and_configure_cuvs) | ||
set(oneValueArgs VERSION FORK PINNED_TAG EXCLUDE_FROM_ALL USE_CUVS_STATIC COMPILE_LIBRARY 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_cuvs}") | ||
message(STATUS "CUML: CUVS pinned tag found: ${PKG_PINNED_TAG}. Cloning cuvs locally.") | ||
set(CPM_DOWNLOAD_cuvs ON) | ||
elseif(PKG_USE_CUVS_STATIC AND (NOT CPM_cuvs_SOURCE)) | ||
message(STATUS "CUML: Cloning cuvs locally to build static libraries.") | ||
set(CPM_DOWNLOAD_cuvs ON) | ||
else() | ||
message(STATUS "Not cloning cuvs locally") | ||
endif() | ||
|
||
if(PKG_USE_CUVS_STATIC) | ||
set(CUVS_LIB cuvs::cuvs_static PARENT_SCOPE) | ||
else() | ||
set(CUVS_LIB cuvs::cuvs PARENT_SCOPE) | ||
endif() | ||
|
||
rapids_cpm_find(cuvs ${PKG_VERSION} | ||
GLOBAL_TARGETS cuvs::cuvs | ||
BUILD_EXPORT_SET cuml-exports | ||
INSTALL_EXPORT_SET cuml-exports | ||
CPM_ARGS | ||
GIT_REPOSITORY https://github.com/${PKG_FORK}/cuvs.git | ||
GIT_TAG ${PKG_PINNED_TAG} | ||
SOURCE_SUBDIR cpp | ||
EXCLUDE_FROM_ALL ${PKG_EXCLUDE_FROM_ALL} | ||
OPTIONS | ||
"BUILD_TESTS OFF" | ||
"BUILD_BENCH OFF" | ||
) | ||
|
||
if(cuvs_ADDED) | ||
message(VERBOSE "CUML: Using CUVS located in ${cuvs_SOURCE_DIR}") | ||
else() | ||
message(VERBOSE "CUML: Using CUVS located in ${cuvs_DIR}") | ||
endif() | ||
|
||
|
||
endfunction() | ||
|
||
# Change pinned tag here to test a commit in CI | ||
# To use a different CUVS locally, set the CMake variable | ||
# CPM_cuvs_SOURCE=/path/to/local/cuvs | ||
find_and_configure_cuvs(VERSION ${CUML_MIN_VERSION_cuvs} | ||
FORK benfred | ||
PINNED_TAG static_lib2 | ||
benfred marked this conversation as resolved.
Show resolved
Hide resolved
|
||
EXCLUDE_FROM_ALL ${CUML_EXCLUDE_CUVS_FROM_ALL} | ||
# When PINNED_TAG above doesn't match cuml, | ||
# force local cuvs clone in build directory | ||
# even if it's already installed. | ||
CLONE_ON_PIN ${CUML_CUVS_CLONE_ON_PIN} | ||
COMPILE_LIBRARY ${CUML_CUVS_COMPILED} | ||
USE_CUVS_STATIC ${CUML_USE_CUVS_STATIC} | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -18,7 +18,7 @@ | |
|
||
#include <cuml/common/log_levels.hpp> | ||
|
||
#include <raft/cluster/kmeans_types.hpp> | ||
#include <cuvs/cluster/kmeans.hpp> | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Looks great! |
||
|
||
namespace raft { | ||
class handle_t; | ||
|
@@ -28,7 +28,7 @@ namespace ML { | |
|
||
namespace kmeans { | ||
|
||
using KMeansParams = raft::cluster::KMeansParams; | ||
using KMeansParams = cuvs::cluster::kmeans::params; | ||
|
||
/** | ||
* @brief Compute k-means clustering and predicts cluster index for each sample | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't think cuml needs a dependency on cuvs python today. We could keep this for the release, though, just to avoid additional changes.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think you are correct here (we only need
cuvs
for the requirements/pyproject outputs for wheels). Let's file a follow-up PR once this one merges so we don't lose track. It can target 24.12.Thankfully,
cuvs
is only ~350 kB: https://anaconda.org/rapidsai-nightly/cuvs/files