diff --git a/.devcontainer/cuda11.8-conda/devcontainer.json b/.devcontainer/cuda11.8-conda/devcontainer.json index 203f52f1a2..4c2161cab7 100644 --- a/.devcontainer/cuda11.8-conda/devcontainer.json +++ b/.devcontainer/cuda11.8-conda/devcontainer.json @@ -5,12 +5,12 @@ "args": { "CUDA": "11.8", "PYTHON_PACKAGE_MANAGER": "conda", - "BASE": "rapidsai/devcontainers:23.12-cpp-llvm16-cuda11.8-mambaforge-ubuntu22.04" + "BASE": "rapidsai/devcontainers:24.02-cpp-llvm16-cuda11.8-mambaforge-ubuntu22.04" } }, "hostRequirements": {"gpu": "optional"}, "features": { - "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:23.12": {} + "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:24.2": {} }, "overrideFeatureInstallOrder": [ "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils" diff --git a/.devcontainer/cuda11.8-pip/devcontainer.json b/.devcontainer/cuda11.8-pip/devcontainer.json index 080ece996e..8bbc016620 100644 --- a/.devcontainer/cuda11.8-pip/devcontainer.json +++ b/.devcontainer/cuda11.8-pip/devcontainer.json @@ -5,13 +5,13 @@ "args": { "CUDA": "11.8", "PYTHON_PACKAGE_MANAGER": "pip", - "BASE": "rapidsai/devcontainers:23.12-cpp-llvm16-cuda11.8-ubuntu22.04" + "BASE": "rapidsai/devcontainers:24.02-cpp-llvm16-cuda11.8-ubuntu22.04" } }, "hostRequirements": {"gpu": "optional"}, "features": { - "ghcr.io/rapidsai/devcontainers/features/ucx:23.12": {"version": "1.14.1"}, - "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:23.12": {} + "ghcr.io/rapidsai/devcontainers/features/ucx:24.2": {"version": "1.14.1"}, + "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:24.2": {} }, "overrideFeatureInstallOrder": [ "ghcr.io/rapidsai/devcontainers/features/ucx", diff --git a/.devcontainer/cuda12.0-conda/devcontainer.json b/.devcontainer/cuda12.0-conda/devcontainer.json index da8bfb4db9..ebb101d9d5 100644 --- a/.devcontainer/cuda12.0-conda/devcontainer.json +++ b/.devcontainer/cuda12.0-conda/devcontainer.json @@ -5,12 +5,12 @@ "args": { "CUDA": "12.0", "PYTHON_PACKAGE_MANAGER": "conda", - "BASE": "rapidsai/devcontainers:23.12-cpp-mambaforge-ubuntu22.04" + "BASE": "rapidsai/devcontainers:24.02-cpp-mambaforge-ubuntu22.04" } }, "hostRequirements": {"gpu": "optional"}, "features": { - "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:23.12": {} + "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:24.2": {} }, "overrideFeatureInstallOrder": [ "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils" diff --git a/.devcontainer/cuda12.0-pip/devcontainer.json b/.devcontainer/cuda12.0-pip/devcontainer.json index e2bee94f8a..3e2efb2f37 100644 --- a/.devcontainer/cuda12.0-pip/devcontainer.json +++ b/.devcontainer/cuda12.0-pip/devcontainer.json @@ -5,13 +5,13 @@ "args": { "CUDA": "12.0", "PYTHON_PACKAGE_MANAGER": "pip", - "BASE": "rapidsai/devcontainers:23.12-cpp-llvm16-cuda12.0-ubuntu22.04" + "BASE": "rapidsai/devcontainers:24.02-cpp-llvm16-cuda12.0-ubuntu22.04" } }, "hostRequirements": {"gpu": "optional"}, "features": { - "ghcr.io/rapidsai/devcontainers/features/ucx:23.12": {"version": "1.14.1"}, - "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:23.12": {} + "ghcr.io/rapidsai/devcontainers/features/ucx:24.2": {"version": "1.14.1"}, + "ghcr.io/rapidsai/devcontainers/features/rapids-build-utils:24.2": {} }, "overrideFeatureInstallOrder": [ "ghcr.io/rapidsai/devcontainers/features/ucx", diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 2435c477ca..4c6e3d0ed4 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -28,7 +28,7 @@ concurrency: jobs: cpp-build: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-24.02 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -37,7 +37,7 @@ jobs: python-build: needs: [cpp-build] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-24.02 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -46,7 +46,7 @@ jobs: upload-conda: needs: [cpp-build, python-build] secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-upload-packages.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/conda-upload-packages.yaml@branch-24.02 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -57,7 +57,7 @@ jobs: if: github.ref_type == 'branch' needs: python-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-24.02 with: arch: "amd64" branch: ${{ inputs.branch }} @@ -69,7 +69,7 @@ jobs: sha: ${{ inputs.sha }} wheel-build-pylibraft: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-24.02 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -79,7 +79,7 @@ jobs: wheel-publish-pylibraft: needs: wheel-build-pylibraft secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-24.02 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -89,7 +89,7 @@ jobs: wheel-build-raft-dask: needs: wheel-publish-pylibraft secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-24.02 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} @@ -99,7 +99,7 @@ jobs: wheel-publish-raft-dask: needs: wheel-build-raft-dask secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-24.02 with: build_type: ${{ inputs.build_type || 'branch' }} branch: ${{ inputs.branch }} diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index bd37b8ac01..19b990c8e9 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -25,29 +25,29 @@ jobs: - wheel-tests-raft-dask - devcontainer secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-24.02 checks: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@branch-24.02 with: enable_check_generated_files: false conda-cpp-build: needs: checks secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-24.02 with: build_type: pull-request node_type: cpu16 conda-cpp-tests: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-24.02 with: build_type: pull-request conda-cpp-checks: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-24.02 with: build_type: pull-request enable_check_symbols: true @@ -55,19 +55,19 @@ jobs: conda-python-build: needs: conda-cpp-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-24.02 with: build_type: pull-request conda-python-tests: needs: conda-python-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-24.02 with: build_type: pull-request docs-build: needs: conda-python-build secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-24.02 with: build_type: pull-request node_type: "gpu-v100-latest-1" @@ -77,34 +77,34 @@ jobs: wheel-build-pylibraft: needs: checks secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-24.02 with: build_type: pull-request script: ci/build_wheel_pylibraft.sh wheel-tests-pylibraft: needs: wheel-build-pylibraft secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-24.02 with: build_type: pull-request script: ci/test_wheel_pylibraft.sh wheel-build-raft-dask: needs: wheel-tests-pylibraft secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-24.02 with: build_type: pull-request script: "ci/build_wheel_raft_dask.sh" wheel-tests-raft-dask: needs: wheel-build-raft-dask secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-24.02 with: build_type: pull-request script: ci/test_wheel_raft_dask.sh devcontainer: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@branch-24.02 with: build_command: | sccache -z; diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 0c46d83cf9..138fd0e300 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -16,7 +16,7 @@ on: jobs: conda-cpp-checks: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-post-build-checks.yaml@branch-24.02 with: build_type: nightly branch: ${{ inputs.branch }} @@ -26,7 +26,7 @@ jobs: symbol_exclusions: (void (thrust::|cub::)|_ZN\d+raft_cutlass) conda-cpp-tests: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-24.02 with: build_type: nightly branch: ${{ inputs.branch }} @@ -34,7 +34,7 @@ jobs: sha: ${{ inputs.sha }} conda-python-tests: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-24.02 with: build_type: nightly branch: ${{ inputs.branch }} @@ -42,7 +42,7 @@ jobs: sha: ${{ inputs.sha }} wheel-tests-pylibraft: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-24.02 with: build_type: nightly branch: ${{ inputs.branch }} @@ -51,7 +51,7 @@ jobs: script: ci/test_wheel_pylibraft.sh wheel-tests-raft-dask: secrets: inherit - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-23.12 + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-24.02 with: build_type: nightly branch: ${{ inputs.branch }} diff --git a/README.md b/README.md index 5b1297b63c..9ab1168bdb 100755 --- a/README.md +++ b/README.md @@ -287,7 +287,7 @@ You can also install the conda packages individually using the `mamba` command a mamba install -c rapidsai -c conda-forge -c nvidia libraft libraft-headers cuda-version=12.0 ``` -If installing the C++ APIs please see [using libraft](https://docs.rapids.ai/api/raft/nightly/using_libraft/) for more information on using the pre-compiled shared library. You can also refer to the [example C++ template project](https://github.com/rapidsai/raft/tree/branch-23.12/cpp/template) for a ready-to-go CMake configuration that you can drop into your project and build against installed RAFT development artifacts above. +If installing the C++ APIs please see [using libraft](https://docs.rapids.ai/api/raft/nightly/using_libraft/) for more information on using the pre-compiled shared library. You can also refer to the [example C++ template project](https://github.com/rapidsai/raft/tree/branch-24.02/cpp/template) for a ready-to-go CMake configuration that you can drop into your project and build against installed RAFT development artifacts above. ### Installing Python through Pip diff --git a/VERSION b/VERSION index a193fff41e..3c6c5e2b70 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -23.12.00 +24.02.00 diff --git a/ci/build_docs.sh b/ci/build_docs.sh index 2eb9f7da6d..b7c8c7a3e0 100755 --- a/ci/build_docs.sh +++ b/ci/build_docs.sh @@ -28,7 +28,7 @@ rapids-mamba-retry install \ pylibraft \ raft-dask -export RAPIDS_VERSION_NUMBER="23.12" +export RAPIDS_VERSION_NUMBER="24.02" export RAPIDS_DOCS_DIR="$(mktemp -d)" rapids-logger "Build CPP docs" diff --git a/ci/checks/copyright.py b/ci/checks/copyright.py index 123aeba87b..2af8b1b8ff 100644 --- a/ci/checks/copyright.py +++ b/ci/checks/copyright.py @@ -20,6 +20,8 @@ import os import sys +import git + SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.expanduser(__file__))) # Add the scripts dir for gitutils @@ -37,7 +39,12 @@ re.compile(r"setup[.]cfg$"), re.compile(r"meta[.]yaml$") ] -ExemptFiles = ["cpp/include/raft/neighbors/detail/faiss_select/"] +ExemptFiles = [ + re.compile("cpp/include/raft/neighbors/detail/faiss_select/"), + re.compile("cpp/include/raft/thirdparty/"), + re.compile("docs/source/sphinxext/github_link.py"), + re.compile("cpp/cmake/modules/FindAVX.cmake") +] # this will break starting at year 10000, which is probably OK :) CheckSimple = re.compile( @@ -48,10 +55,12 @@ def checkThisFile(f): - # This check covers things like symlinks which point to files that DNE - if not (os.path.exists(f)): - return False - if gitutils and gitutils.isFileEmpty(f): + if isinstance(f, git.Diff): + if f.deleted_file or f.b_blob.size == 0: + return False + f = f.b_path + elif not os.path.exists(f) or os.stat(f).st_size == 0: + # This check covers things like symlinks which point to files that DNE return False for exempt in ExemptFiles: if exempt.search(f): @@ -62,36 +71,90 @@ def checkThisFile(f): return False +def modifiedFiles(): + """Get a set of all modified files, as Diff objects. + + The files returned have been modified in git since the merge base of HEAD + and the upstream of the target branch. We return the Diff objects so that + we can read only the staged changes. + """ + repo = git.Repo() + # Use the environment variable TARGET_BRANCH or RAPIDS_BASE_BRANCH (defined in CI) if possible + target_branch = os.environ.get("TARGET_BRANCH", os.environ.get("RAPIDS_BASE_BRANCH")) + if target_branch is None: + # Fall back to the closest branch if not on CI + target_branch = repo.git.describe( + all=True, tags=True, match="branch-*", abbrev=0 + ).lstrip("heads/") + + upstream_target_branch = None + if target_branch in repo.heads: + # Use the tracking branch of the local reference if it exists. This + # returns None if no tracking branch is set. + upstream_target_branch = repo.heads[target_branch].tracking_branch() + if upstream_target_branch is None: + # Fall back to the remote with the newest target_branch. This code + # path is used on CI because the only local branch reference is + # current-pr-branch, and thus target_branch is not in repo.heads. + # This also happens if no tracking branch is defined for the local + # target_branch. We use the remote with the latest commit if + # multiple remotes are defined. + candidate_branches = [ + remote.refs[target_branch] for remote in repo.remotes + if target_branch in remote.refs + ] + if len(candidate_branches) > 0: + upstream_target_branch = sorted( + candidate_branches, + key=lambda branch: branch.commit.committed_datetime, + )[-1] + else: + # If no remotes are defined, try to use the local version of the + # target_branch. If this fails, the repo configuration must be very + # strange and we can fix this script on a case-by-case basis. + upstream_target_branch = repo.heads[target_branch] + merge_base = repo.merge_base("HEAD", upstream_target_branch.commit)[0] + diff = merge_base.diff() + changed_files = {f for f in diff if f.b_path is not None} + return changed_files + + def getCopyrightYears(line): res = CheckSimple.search(line) if res: - return (int(res.group(1)), int(res.group(1))) + return int(res.group(1)), int(res.group(1)) res = CheckDouble.search(line) if res: - return (int(res.group(1)), int(res.group(2))) - return (None, None) + return int(res.group(1)), int(res.group(2)) + return None, None def replaceCurrentYear(line, start, end): # first turn a simple regex into double (if applicable). then update years res = CheckSimple.sub(r"Copyright (c) \1-\1, NVIDIA CORPORATION", line) res = CheckDouble.sub( - r"Copyright (c) {:04d}-{:04d}, NVIDIA CORPORATION".format(start, end), - res) + rf"Copyright (c) {start:04d}-{end:04d}, NVIDIA CORPORATION", + res, + ) return res def checkCopyright(f, update_current_year): - """ - Checks for copyright headers and their years - """ + """Checks for copyright headers and their years.""" errs = [] thisYear = datetime.datetime.now().year lineNum = 0 crFound = False yearMatched = False - with io.open(f, "r", encoding="utf-8") as fp: - lines = fp.readlines() + + if isinstance(f, git.Diff): + path = f.b_path + lines = f.b_blob.data_stream.read().decode().splitlines(keepends=True) + else: + path = f + with open(f, encoding="utf-8") as fp: + lines = fp.readlines() + for line in lines: lineNum += 1 start, end = getCopyrightYears(line) @@ -100,20 +163,19 @@ def checkCopyright(f, update_current_year): crFound = True if start > end: e = [ - f, + path, lineNum, "First year after second year in the copyright " "header (manual fix required)", - None + None, ] errs.append(e) - if thisYear < start or thisYear > end: + elif thisYear < start or thisYear > end: e = [ - f, + path, lineNum, - "Current year not included in the " - "copyright header", - None + "Current year not included in the copyright header", + None, ] if thisYear < start: e[-1] = replaceCurrentYear(line, thisYear, end) @@ -122,15 +184,14 @@ def checkCopyright(f, update_current_year): errs.append(e) else: yearMatched = True - fp.close() # copyright header itself not found if not crFound: e = [ - f, + path, 0, "Copyright header missing or formatted incorrectly " "(manual fix required)", - None + None, ] errs.append(e) # even if the year matches a copyright header, make the check pass @@ -140,21 +201,19 @@ def checkCopyright(f, update_current_year): if update_current_year: errs_update = [x for x in errs if x[-1] is not None] if len(errs_update) > 0: - print("File: {}. Changing line(s) {}".format( - f, ', '.join(str(x[1]) for x in errs if x[-1] is not None))) + lines_changed = ", ".join(str(x[1]) for x in errs_update) + print(f"File: {path}. Changing line(s) {lines_changed}") for _, lineNum, __, replacement in errs_update: lines[lineNum - 1] = replacement - with io.open(f, "w", encoding="utf-8") as out_file: - for new_line in lines: - out_file.write(new_line) - errs = [x for x in errs if x[-1] is None] + with open(path, "w", encoding="utf-8") as out_file: + out_file.writelines(lines) return errs def getAllFilesUnderDir(root, pathFilter=None): retList = [] - for (dirpath, dirnames, filenames) in os.walk(root): + for dirpath, dirnames, filenames in os.walk(root): for fn in filenames: filePath = os.path.join(dirpath, fn) if pathFilter(filePath): @@ -169,49 +228,37 @@ def checkCopyright_main(): it compares between branches "$PR_TARGET_BRANCH" and "current-pr-branch" """ retVal = 0 - global ExemptFiles argparser = argparse.ArgumentParser( - "Checks for a consistent copyright header in git's modified files") - argparser.add_argument("--update-current-year", - dest='update_current_year', - action="store_true", - required=False, - help="If set, " - "update the current year if a header " - "is already present and well formatted.") - argparser.add_argument("--git-modified-only", - dest='git_modified_only', - action="store_true", - required=False, - help="If set, " - "only files seen as modified by git will be " - "processed.") - argparser.add_argument("--exclude", - dest='exclude', - action="append", - required=False, - default=["python/cuml/_thirdparty/", - "cpp/include/raft/thirdparty/", - "cpp/cmake/modules/FindAVX.cmake"], - help=("Exclude the paths specified (regexp). " - "Can be specified multiple times.")) - - (args, dirs) = argparser.parse_known_args() - try: - ExemptFiles = ExemptFiles + [pathName for pathName in args.exclude] - ExemptFiles = [re.compile(file) for file in ExemptFiles] - except re.error as reException: - print("Regular expression error:") - print(reException) - return 1 + "Checks for a consistent copyright header in git's modified files" + ) + argparser.add_argument( + "--update-current-year", + dest="update_current_year", + action="store_true", + required=False, + help="If set, " + "update the current year if a header is already " + "present and well formatted.", + ) + argparser.add_argument( + "--git-modified-only", + dest="git_modified_only", + action="store_true", + required=False, + help="If set, " + "only files seen as modified by git will be " + "processed.", + ) + + args, dirs = argparser.parse_known_args() if args.git_modified_only: - files = gitutils.modifiedFiles(pathFilter=checkThisFile) + files = [f for f in modifiedFiles() if checkThisFile(f)] else: files = [] for d in [os.path.abspath(d) for d in dirs]: - if not (os.path.isdir(d)): + if not os.path.isdir(d): raise ValueError(f"{d} is not a directory.") files += getAllFilesUnderDir(d, pathFilter=checkThisFile) @@ -220,24 +267,24 @@ def checkCopyright_main(): errors += checkCopyright(f, args.update_current_year) if len(errors) > 0: - print("Copyright headers incomplete in some of the files!") + if any(e[-1] is None for e in errors): + print("Copyright headers incomplete in some of the files!") for e in errors: print(" %s:%d Issue: %s" % (e[0], e[1], e[2])) print("") n_fixable = sum(1 for e in errors if e[-1] is not None) path_parts = os.path.abspath(__file__).split(os.sep) - file_from_repo = os.sep.join(path_parts[path_parts.index("ci"):]) - if n_fixable > 0: - print(("You can run `python {} --git-modified-only " - "--update-current-year` to fix {} of these " - "errors.\n").format(file_from_repo, n_fixable)) + file_from_repo = os.sep.join(path_parts[path_parts.index("ci") :]) + if n_fixable > 0 and not args.update_current_year: + print( + f"You can run `python {file_from_repo} --git-modified-only " + "--update-current-year` and stage the results in git to " + f"fix {n_fixable} of these errors.\n" + ) retVal = 1 - else: - print("Copyright check passed") return retVal if __name__ == "__main__": - import sys sys.exit(checkCopyright_main()) diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index d5bc17be56..7aee8726c9 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -74,7 +74,6 @@ for FILE in python/*/pyproject.toml; do for DEP in "${DEPENDENCIES[@]}"; do sed_runner "/\"${DEP}==/ s/==.*\"/==${NEXT_SHORT_TAG_PEP440}.*\"/g" ${FILE} done - sed_runner "s/^version = .*/version = \"${NEXT_FULL_TAG}\"/g" "${FILE}" sed_runner "/\"ucx-py==/ s/==.*\"/==${NEXT_UCX_PY_SHORT_TAG_PEP440}.*\"/g" ${FILE} done diff --git a/ci/wheel_smoke_test_pylibraft.py b/ci/wheel_smoke_test_pylibraft.py index 7fee674691..c0df2fe45c 100644 --- a/ci/wheel_smoke_test_pylibraft.py +++ b/ci/wheel_smoke_test_pylibraft.py @@ -1,3 +1,18 @@ +# Copyright (c) 2023, 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. +# + import numpy as np from scipy.spatial.distance import cdist diff --git a/conda/environments/all_cuda-118_arch-aarch64.yaml b/conda/environments/all_cuda-118_arch-aarch64.yaml index c28f1961e6..ce1fed94f7 100644 --- a/conda/environments/all_cuda-118_arch-aarch64.yaml +++ b/conda/environments/all_cuda-118_arch-aarch64.yaml @@ -20,7 +20,7 @@ dependencies: - cupy>=12.0.0 - cxx-compiler - cython>=3.0.0 -- dask-cuda==23.12.* +- dask-cuda==24.2.* - doxygen>=1.8.20 - gcc_linux-aarch64=11.* - gmock>=1.13.0 @@ -46,9 +46,9 @@ dependencies: - pydata-sphinx-theme - pytest - pytest-cov -- rapids-dask-dependency==23.12.* +- rapids-dask-dependency==24.2.* - recommonmark -- rmm==23.12.* +- rmm==24.2.* - scikit-build>=0.13.1 - scikit-learn - scipy @@ -56,6 +56,6 @@ dependencies: - sphinx-markdown-tables - sysroot_linux-aarch64==2.17 - ucx-proc=*=gpu -- ucx-py==0.35.* +- ucx-py==0.36.* - ucx>=1.13.0 name: all_cuda-118_arch-aarch64 diff --git a/conda/environments/all_cuda-118_arch-x86_64.yaml b/conda/environments/all_cuda-118_arch-x86_64.yaml index 9b7c110bc3..16504fb4ca 100644 --- a/conda/environments/all_cuda-118_arch-x86_64.yaml +++ b/conda/environments/all_cuda-118_arch-x86_64.yaml @@ -20,7 +20,7 @@ dependencies: - cupy>=12.0.0 - cxx-compiler - cython>=3.0.0 -- dask-cuda==23.12.* +- dask-cuda==24.2.* - doxygen>=1.8.20 - gcc_linux-64=11.* - gmock>=1.13.0 @@ -46,9 +46,9 @@ dependencies: - pydata-sphinx-theme - pytest - pytest-cov -- rapids-dask-dependency==23.12.* +- rapids-dask-dependency==24.2.* - recommonmark -- rmm==23.12.* +- rmm==24.2.* - scikit-build>=0.13.1 - scikit-learn - scipy @@ -56,6 +56,6 @@ dependencies: - sphinx-markdown-tables - sysroot_linux-64==2.17 - ucx-proc=*=gpu -- ucx-py==0.35.* +- ucx-py==0.36.* - ucx>=1.13.0 name: all_cuda-118_arch-x86_64 diff --git a/conda/environments/all_cuda-120_arch-aarch64.yaml b/conda/environments/all_cuda-120_arch-aarch64.yaml index 8d614d3c2c..b2d4a50df8 100644 --- a/conda/environments/all_cuda-120_arch-aarch64.yaml +++ b/conda/environments/all_cuda-120_arch-aarch64.yaml @@ -21,7 +21,7 @@ dependencies: - cupy>=12.0.0 - cxx-compiler - cython>=3.0.0 -- dask-cuda==23.12.* +- dask-cuda==24.2.* - doxygen>=1.8.20 - gcc_linux-aarch64=11.* - gmock>=1.13.0 @@ -42,9 +42,9 @@ dependencies: - pydata-sphinx-theme - pytest - pytest-cov -- rapids-dask-dependency==23.12.* +- rapids-dask-dependency==24.2.* - recommonmark -- rmm==23.12.* +- rmm==24.2.* - scikit-build>=0.13.1 - scikit-learn - scipy @@ -52,6 +52,6 @@ dependencies: - sphinx-markdown-tables - sysroot_linux-aarch64==2.17 - ucx-proc=*=gpu -- ucx-py==0.35.* +- ucx-py==0.36.* - ucx>=1.13.0 name: all_cuda-120_arch-aarch64 diff --git a/conda/environments/all_cuda-120_arch-x86_64.yaml b/conda/environments/all_cuda-120_arch-x86_64.yaml index f9d65cee39..adb1fd74b7 100644 --- a/conda/environments/all_cuda-120_arch-x86_64.yaml +++ b/conda/environments/all_cuda-120_arch-x86_64.yaml @@ -21,7 +21,7 @@ dependencies: - cupy>=12.0.0 - cxx-compiler - cython>=3.0.0 -- dask-cuda==23.12.* +- dask-cuda==24.2.* - doxygen>=1.8.20 - gcc_linux-64=11.* - gmock>=1.13.0 @@ -42,9 +42,9 @@ dependencies: - pydata-sphinx-theme - pytest - pytest-cov -- rapids-dask-dependency==23.12.* +- rapids-dask-dependency==24.2.* - recommonmark -- rmm==23.12.* +- rmm==24.2.* - scikit-build>=0.13.1 - scikit-learn - scipy @@ -52,6 +52,6 @@ dependencies: - sphinx-markdown-tables - sysroot_linux-64==2.17 - ucx-proc=*=gpu -- ucx-py==0.35.* +- ucx-py==0.36.* - ucx>=1.13.0 name: all_cuda-120_arch-x86_64 diff --git a/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml index b5fc4e3bd5..39c6d24889 100644 --- a/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml @@ -38,7 +38,7 @@ dependencies: - openblas - pandas - pyyaml -- rmm==23.12.* +- rmm==24.2.* - scikit-build>=0.13.1 - sysroot_linux-aarch64==2.17 name: bench_ann_cuda-118_arch-aarch64 diff --git a/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml index b868f26e15..e0f46085d3 100644 --- a/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml @@ -38,7 +38,7 @@ dependencies: - openblas - pandas - pyyaml -- rmm==23.12.* +- rmm==24.2.* - scikit-build>=0.13.1 - sysroot_linux-64==2.17 name: bench_ann_cuda-118_arch-x86_64 diff --git a/conda/environments/bench_ann_cuda-120_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-120_arch-aarch64.yaml index 4a3818fe5d..2c69bc5325 100644 --- a/conda/environments/bench_ann_cuda-120_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-120_arch-aarch64.yaml @@ -34,7 +34,7 @@ dependencies: - openblas - pandas - pyyaml -- rmm==23.12.* +- rmm==24.2.* - scikit-build>=0.13.1 - sysroot_linux-aarch64==2.17 name: bench_ann_cuda-120_arch-aarch64 diff --git a/conda/environments/bench_ann_cuda-120_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-120_arch-x86_64.yaml index 3d6f8c4ec1..4f4b0d5019 100644 --- a/conda/environments/bench_ann_cuda-120_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-120_arch-x86_64.yaml @@ -34,7 +34,7 @@ dependencies: - openblas - pandas - pyyaml -- rmm==23.12.* +- rmm==24.2.* - scikit-build>=0.13.1 - sysroot_linux-64==2.17 name: bench_ann_cuda-120_arch-x86_64 diff --git a/conda/recipes/libraft/conda_build_config.yaml b/conda/recipes/libraft/conda_build_config.yaml index d156f2609b..25493a34fa 100644 --- a/conda/recipes/libraft/conda_build_config.yaml +++ b/conda/recipes/libraft/conda_build_config.yaml @@ -71,3 +71,9 @@ cuda11_cuda_profiler_api_host_version: cuda11_cuda_profiler_api_run_version: - ">=11.4.240,<12" + +spdlog_version: + - ">=1.11.0,<1.12" + +fmt_version: + - ">=9.1.0,<10" diff --git a/conda/recipes/libraft/meta.yaml b/conda/recipes/libraft/meta.yaml index 5203f7b3bd..116f8d9e6e 100644 --- a/conda/recipes/libraft/meta.yaml +++ b/conda/recipes/libraft/meta.yaml @@ -63,12 +63,16 @@ outputs: {% endif %} - cuda-version ={{ cuda_version }} - librmm ={{ minor_version }} + - spdlog {{ spdlog_version }} + - fmt {{ fmt_version }} run: - {{ pin_compatible('cuda-version', max_pin='x', min_pin='x') }} {% if cuda_major == "11" %} - cudatoolkit {% endif %} - librmm ={{ minor_version }} + - spdlog {{ spdlog_version }} + - fmt {{ fmt_version }} about: home: https://rapids.ai/ license: Apache-2.0 diff --git a/conda/recipes/raft-ann-bench-cpu/conda_build_config.yaml b/conda/recipes/raft-ann-bench-cpu/conda_build_config.yaml index 0bd424f85b..fda3e4e53d 100644 --- a/conda/recipes/raft-ann-bench-cpu/conda_build_config.yaml +++ b/conda/recipes/raft-ann-bench-cpu/conda_build_config.yaml @@ -18,3 +18,9 @@ h5py_version: nlohmann_json_version: - ">=3.11.2" + +spdlog_version: + - ">=1.11.0,<1.12" + +fmt_version: + - ">=9.1.0,<10" diff --git a/conda/recipes/raft-ann-bench-cpu/meta.yaml b/conda/recipes/raft-ann-bench-cpu/meta.yaml index 6db9400d97..fce85d5ffc 100644 --- a/conda/recipes/raft-ann-bench-cpu/meta.yaml +++ b/conda/recipes/raft-ann-bench-cpu/meta.yaml @@ -48,6 +48,8 @@ requirements: - glog {{ glog_version }} - matplotlib - nlohmann_json {{ nlohmann_json_version }} + - spdlog {{ spdlog_version }} + - fmt {{ fmt_version }} - python - pyyaml - pandas diff --git a/conda/recipes/raft-dask/conda_build_config.yaml b/conda/recipes/raft-dask/conda_build_config.yaml index d89dbae4df..c12c35af3b 100644 --- a/conda/recipes/raft-dask/conda_build_config.yaml +++ b/conda/recipes/raft-dask/conda_build_config.yaml @@ -17,7 +17,7 @@ ucx_version: - ">=1.14.1,<1.16.0" ucx_py_version: - - "0.35.*" + - "0.36.*" cmake_version: - ">=3.26.4" diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 5d2864e2e0..acb77ec8c7 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -10,8 +10,8 @@ # 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(RAPIDS_VERSION "23.12") -set(RAFT_VERSION "23.12.00") +set(RAPIDS_VERSION "24.02") +set(RAFT_VERSION "24.02.00") cmake_minimum_required(VERSION 3.26.4 FATAL_ERROR) include(../fetch_rapids.cmake) diff --git a/cpp/bench/ann/src/raft/raft_ivf_flat_wrapper.h b/cpp/bench/ann/src/raft/raft_ivf_flat_wrapper.h index 24b3c69bb6..13ea20d483 100644 --- a/cpp/bench/ann/src/raft/raft_ivf_flat_wrapper.h +++ b/cpp/bench/ann/src/raft/raft_ivf_flat_wrapper.h @@ -78,7 +78,7 @@ class RaftIvfFlatGpu : public ANN { AlgoProperty get_preference() const override { AlgoProperty property; - property.dataset_memory_type = MemoryType::Device; + property.dataset_memory_type = MemoryType::HostMmap; property.query_memory_type = MemoryType::Device; return property; } diff --git a/cpp/doxygen/Doxyfile b/cpp/doxygen/Doxyfile index f1dcfa9db7..3eb0763eaf 100644 --- a/cpp/doxygen/Doxyfile +++ b/cpp/doxygen/Doxyfile @@ -38,7 +38,7 @@ PROJECT_NAME = "RAFT C++ API" # could be handy for archiving the generated documentation or if some version # control system is used. -PROJECT_NUMBER = "23.12" +PROJECT_NUMBER = "24.02" # Using the PROJECT_BRIEF tag one can provide an optional one line description # for a project that appears at the top of each page and should give viewer a diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh index a9a6ac025f..a35cb9e1f1 100644 --- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh +++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh @@ -120,7 +120,8 @@ RAFT_KERNEL build_index_kernel(const LabelT* labels, uint32_t* list_sizes_ptr, IdxT n_rows, uint32_t dim, - uint32_t veclen) + uint32_t veclen, + IdxT batch_offset = 0) { const IdxT i = IdxT(blockDim.x) * IdxT(blockIdx.x) + threadIdx.x; if (i >= n_rows) { return; } @@ -131,7 +132,7 @@ RAFT_KERNEL build_index_kernel(const LabelT* labels, auto* list_data = list_data_ptrs[list_id]; // Record the source vector id in the index - list_index[inlist_id] = source_ixs == nullptr ? i : source_ixs[i]; + list_index[inlist_id] = source_ixs == nullptr ? i + batch_offset : source_ixs[i]; // The data is written in interleaved groups of `index::kGroupSize` vectors using interleaved_group = Pow2; @@ -180,16 +181,33 @@ void extend(raft::resources const& handle, auto new_labels = raft::make_device_vector(handle, n_rows); raft::cluster::kmeans_balanced_params kmeans_params; - kmeans_params.metric = index->metric(); - auto new_vectors_view = raft::make_device_matrix_view(new_vectors, n_rows, dim); + kmeans_params.metric = index->metric(); auto orig_centroids_view = raft::make_device_matrix_view(index->centers().data_handle(), n_lists, dim); - raft::cluster::kmeans_balanced::predict(handle, - kmeans_params, - new_vectors_view, - orig_centroids_view, - new_labels.view(), - utils::mapping{}); + // Calculate the batch size for the input data if it's not accessible directly from the device + constexpr size_t kReasonableMaxBatchSize = 65536; + size_t max_batch_size = std::min(n_rows, kReasonableMaxBatchSize); + + // Predict the cluster labels for the new data, in batches if necessary + utils::batch_load_iterator vec_batches(new_vectors, + n_rows, + index->dim(), + max_batch_size, + stream, + resource::get_workspace_resource(handle)); + + for (const auto& batch : vec_batches) { + auto batch_data_view = + raft::make_device_matrix_view(batch.data(), batch.size(), index->dim()); + auto batch_labels_view = raft::make_device_vector_view( + new_labels.data_handle() + batch.offset(), batch.size()); + raft::cluster::kmeans_balanced::predict(handle, + kmeans_params, + batch_data_view, + orig_centroids_view, + batch_labels_view, + utils::mapping{}); + } auto* list_sizes_ptr = index->list_sizes().data_handle(); auto old_list_sizes_dev = raft::make_device_vector(handle, n_lists); @@ -202,14 +220,19 @@ void extend(raft::resources const& handle, auto list_sizes_view = raft::make_device_vector_view, IdxT>( list_sizes_ptr, n_lists); - auto const_labels_view = make_const_mdspan(new_labels.view()); - raft::cluster::kmeans_balanced::helpers::calc_centers_and_sizes(handle, - new_vectors_view, - const_labels_view, - centroids_view, - list_sizes_view, - false, - utils::mapping{}); + for (const auto& batch : vec_batches) { + auto batch_data_view = + raft::make_device_matrix_view(batch.data(), batch.size(), index->dim()); + auto batch_labels_view = raft::make_device_vector_view( + new_labels.data_handle() + batch.offset(), batch.size()); + raft::cluster::kmeans_balanced::helpers::calc_centers_and_sizes(handle, + batch_data_view, + batch_labels_view, + centroids_view, + list_sizes_view, + false, + utils::mapping{}); + } } else { raft::stats::histogram(raft::stats::HistTypeAuto, reinterpret_cast(list_sizes_ptr), @@ -244,20 +267,39 @@ void extend(raft::resources const& handle, // we'll rebuild the `list_sizes_ptr` in the following kernel, using it as an atomic counter. raft::copy(list_sizes_ptr, old_list_sizes_dev.data_handle(), n_lists, stream); - // Kernel to insert the new vectors - const dim3 block_dim(256); - const dim3 grid_dim(raft::ceildiv(n_rows, block_dim.x)); - build_index_kernel<<>>(new_labels.data_handle(), - new_vectors, - new_indices, - index->data_ptrs().data_handle(), - index->inds_ptrs().data_handle(), - list_sizes_ptr, - n_rows, - dim, - index->veclen()); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - + utils::batch_load_iterator vec_indices( + new_indices, n_rows, 1, max_batch_size, stream, resource::get_workspace_resource(handle)); + utils::batch_load_iterator idx_batch = vec_indices.begin(); + size_t next_report_offset = 0; + size_t d_report_offset = n_rows * 5 / 100; + for (const auto& batch : vec_batches) { + auto batch_data_view = + raft::make_device_matrix_view(batch.data(), batch.size(), index->dim()); + // Kernel to insert the new vectors + const dim3 block_dim(256); + const dim3 grid_dim(raft::ceildiv(batch.size(), block_dim.x)); + build_index_kernel + <<>>(new_labels.data_handle() + batch.offset(), + batch_data_view.data_handle(), + idx_batch->data(), + index->data_ptrs().data_handle(), + index->inds_ptrs().data_handle(), + list_sizes_ptr, + batch.size(), + dim, + index->veclen(), + batch.offset()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + if (batch.offset() > next_report_offset) { + float progress = batch.offset() * 100.0f / n_rows; + RAFT_LOG_DEBUG("ivf_flat::extend added vectors %zu, %6.1f%% complete", + static_cast(batch.offset()), + progress); + next_report_offset += d_report_offset; + } + ++idx_batch; + } // Precompute the centers vector norms for L2Expanded distance if (!index->center_norms().has_value()) { index->allocate_center_norms(handle); diff --git a/cpp/include/raft/neighbors/ivf_flat-ext.cuh b/cpp/include/raft/neighbors/ivf_flat-ext.cuh index 8dbe7587ff..063105cf46 100644 --- a/cpp/include/raft/neighbors/ivf_flat-ext.cuh +++ b/cpp/include/raft/neighbors/ivf_flat-ext.cuh @@ -48,6 +48,18 @@ void build(raft::resources const& handle, raft::device_matrix_view dataset, raft::neighbors::ivf_flat::index& idx) RAFT_EXPLICIT; +template +auto build(raft::resources const& handle, + const index_params& params, + raft::host_matrix_view dataset) + -> index RAFT_EXPLICIT; + +template +void build(raft::resources const& handle, + const index_params& params, + raft::host_matrix_view dataset, + raft::neighbors::ivf_flat::index& idx) RAFT_EXPLICIT; + template auto extend(raft::resources const& handle, const index& orig_index, @@ -74,6 +86,19 @@ void extend(raft::resources const& handle, std::optional> new_indices, index* index) RAFT_EXPLICIT; +template +auto extend(raft::resources const& handle, + raft::host_matrix_view new_vectors, + std::optional> new_indices, + const raft::neighbors::ivf_flat::index& orig_index) + -> raft::neighbors::ivf_flat::index RAFT_EXPLICIT; + +template +void extend(raft::resources const& handle, + raft::host_matrix_view new_vectors, + std::optional> new_indices, + index* index) RAFT_EXPLICIT; + template void search_with_filtering(raft::resources const& handle, const search_params& params, @@ -137,6 +162,18 @@ void search(raft::resources const& handle, raft::resources const& handle, \ const raft::neighbors::ivf_flat::index_params& params, \ raft::device_matrix_view dataset, \ + raft::neighbors::ivf_flat::index& idx); \ + \ + extern template auto raft::neighbors::ivf_flat::build( \ + raft::resources const& handle, \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::host_matrix_view dataset) \ + ->raft::neighbors::ivf_flat::index; \ + \ + extern template void raft::neighbors::ivf_flat::build( \ + raft::resources const& handle, \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::host_matrix_view dataset, \ raft::neighbors::ivf_flat::index& idx); instantiate_raft_neighbors_ivf_flat_build(float, int64_t); @@ -171,7 +208,20 @@ instantiate_raft_neighbors_ivf_flat_build(uint8_t, int64_t); raft::resources const& handle, \ raft::device_matrix_view new_vectors, \ std::optional> new_indices, \ - raft::neighbors::ivf_flat::index* index); + raft::neighbors::ivf_flat::index* index); \ + \ + extern template void raft::neighbors::ivf_flat::extend( \ + raft::resources const& handle, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices, \ + raft::neighbors::ivf_flat::index* index); \ + \ + extern template auto raft::neighbors::ivf_flat::extend( \ + const raft::resources& handle, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices, \ + const raft::neighbors::ivf_flat::index& idx) \ + ->raft::neighbors::ivf_flat::index; instantiate_raft_neighbors_ivf_flat_extend(float, int64_t); instantiate_raft_neighbors_ivf_flat_extend(int8_t, int64_t); diff --git a/cpp/include/raft/neighbors/ivf_flat-inl.cuh b/cpp/include/raft/neighbors/ivf_flat-inl.cuh index 692fb08810..b540de7f14 100644 --- a/cpp/include/raft/neighbors/ivf_flat-inl.cuh +++ b/cpp/include/raft/neighbors/ivf_flat-inl.cuh @@ -55,7 +55,7 @@ namespace raft::neighbors::ivf_flat { * * @param[in] handle * @param[in] params configure the index building - * @param[in] dataset a device pointer to a row-major matrix [n_rows, dim] + * @param[in] dataset a host or device pointer to a row-major matrix [n_rows, dim] * @param[in] n_rows the number of samples * @param[in] dim the dimensionality of the data * @@ -102,7 +102,7 @@ auto build(raft::resources const& handle, * * @param[in] handle * @param[in] params configure the index building - * @param[in] dataset a device pointer to a row-major matrix [n_rows, dim] + * @param[in] dataset a device matrix [n_rows, dim] * * @return the constructed ivf-flat index */ @@ -118,6 +118,20 @@ auto build(raft::resources const& handle, static_cast(dataset.extent(1))); } +/** + * @brief Build the index from a dataset in host memory. + */ +template +auto build(raft::resources const& handle, + const index_params& params, + raft::host_matrix_view dataset) -> index +{ + return raft::neighbors::ivf_flat::detail::build(handle, + params, + dataset.data_handle(), + static_cast(dataset.extent(0)), + static_cast(dataset.extent(1))); +} /** * @brief Build the index from the dataset for efficient search. * @@ -162,6 +176,21 @@ void build(raft::resources const& handle, static_cast(dataset.extent(1))); } +/** + * @brief Build the index from a dataset in host memory. + */ +template +void build(raft::resources const& handle, + const index_params& params, + raft::host_matrix_view dataset, + raft::neighbors::ivf_flat::index& idx) +{ + idx = raft::neighbors::ivf_flat::detail::build(handle, + params, + dataset.data_handle(), + static_cast(dataset.extent(0)), + static_cast(dataset.extent(1))); +} /** @} */ /** @@ -188,8 +217,8 @@ void build(raft::resources const& handle, * * @param[in] handle * @param[in] orig_index original index - * @param[in] new_vectors a device pointer to a row-major matrix [n_rows, index.dim()] - * @param[in] new_indices a device pointer to a vector of indices [n_rows]. + * @param[in] new_vectors a device/host pointer to a row-major matrix [n_rows, index.dim()] + * @param[in] new_indices a device/host pointer to a vector of indices [n_rows]. * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` * here to imply a continuous range `[0...n_rows)`. * @param[in] n_rows number of rows in `new_vectors` @@ -257,6 +286,23 @@ auto extend(raft::resources const& handle, new_vectors.extent(0)); } +/** + * @brief Extend the index with additional vectors. + * + * This overloads takes input data in host memory. + */ +template +auto extend(raft::resources const& handle, + raft::host_matrix_view new_vectors, + std::optional> new_indices, + const index& orig_index) -> index +{ + return extend(handle, + orig_index, + new_vectors.data_handle(), + new_indices.has_value() ? new_indices.value().data_handle() : nullptr, + new_vectors.extent(0)); +} /** @} */ /** @@ -279,8 +325,8 @@ auto extend(raft::resources const& handle, * * @param handle * @param[inout] index - * @param[in] new_vectors a device pointer to a row-major matrix [n_rows, index.dim()] - * @param[in] new_indices a device pointer to a vector of indices [n_rows]. + * @param[in] new_vectors a device/host pointer to a row-major matrix [n_rows, index.dim()] + * @param[in] new_indices a device/host pointer to a vector of indices [n_rows]. * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` * here to imply a continuous range `[0...n_rows)`. * @param[in] n_rows the number of samples @@ -339,6 +385,23 @@ void extend(raft::resources const& handle, static_cast(new_vectors.extent(0))); } +/** + * @brief Extend the index with additional vectors. + * + * This overloads takes input data in host memory. + */ +template +void extend(raft::resources const& handle, + raft::host_matrix_view new_vectors, + std::optional> new_indices, + index* index) +{ + extend(handle, + index, + new_vectors.data_handle(), + new_indices.has_value() ? new_indices.value().data_handle() : nullptr, + static_cast(new_vectors.extent(0))); +} /** @} */ /** diff --git a/cpp/src/neighbors/ivf_flat_00_generate.py b/cpp/src/neighbors/ivf_flat_00_generate.py index b02606a23e..d987a4e17d 100644 --- a/cpp/src/neighbors/ivf_flat_00_generate.py +++ b/cpp/src/neighbors/ivf_flat_00_generate.py @@ -41,63 +41,88 @@ """ types = dict( - float_int64_t= ("float", "int64_t"), + float_int64_t=("float", "int64_t"), int8_t_int64_t=("int8_t", "int64_t"), uint8_t_int64_t=("uint8_t", "int64_t"), ) build_macro = """ #define instantiate_raft_neighbors_ivf_flat_build(T, IdxT) \\ - template auto raft::neighbors::ivf_flat::build( \\ - raft::resources const& handle, \\ + template auto raft::neighbors::ivf_flat::build( \\ + raft::resources const& handle, \\ const raft::neighbors::ivf_flat::index_params& params, \\ const T* dataset, \\ IdxT n_rows, \\ uint32_t dim) \\ ->raft::neighbors::ivf_flat::index; \\ \\ - template auto raft::neighbors::ivf_flat::build( \\ - raft::resources const& handle, \\ + template auto raft::neighbors::ivf_flat::build( \\ + raft::resources const& handle, \\ const raft::neighbors::ivf_flat::index_params& params, \\ raft::device_matrix_view dataset) \\ ->raft::neighbors::ivf_flat::index; \\ \\ - template void raft::neighbors::ivf_flat::build( \\ - raft::resources const& handle, \\ + template void raft::neighbors::ivf_flat::build( \\ + raft::resources const& handle, \\ const raft::neighbors::ivf_flat::index_params& params, \\ raft::device_matrix_view dataset, \\ + raft::neighbors::ivf_flat::index& idx); \\ + \\ + template auto raft::neighbors::ivf_flat::build( \\ + raft::resources const& handle, \\ + const raft::neighbors::ivf_flat::index_params& params, \\ + raft::host_matrix_view dataset) \\ + ->raft::neighbors::ivf_flat::index; \\ + \\ + template void raft::neighbors::ivf_flat::build( \\ + raft::resources const& handle, \\ + const raft::neighbors::ivf_flat::index_params& params, \\ + raft::host_matrix_view dataset, \\ raft::neighbors::ivf_flat::index& idx); """ extend_macro = """ #define instantiate_raft_neighbors_ivf_flat_extend(T, IdxT) \\ - template auto raft::neighbors::ivf_flat::extend( \\ - raft::resources const& handle, \\ + template auto raft::neighbors::ivf_flat::extend( \\ + raft::resources const& handle, \\ const raft::neighbors::ivf_flat::index& orig_index, \\ const T* new_vectors, \\ const IdxT* new_indices, \\ IdxT n_rows) \\ ->raft::neighbors::ivf_flat::index; \\ \\ - template auto raft::neighbors::ivf_flat::extend( \\ - raft::resources const& handle, \\ + template auto raft::neighbors::ivf_flat::extend( \\ + raft::resources const& handle, \\ raft::device_matrix_view new_vectors, \\ std::optional> new_indices, \\ const raft::neighbors::ivf_flat::index& orig_index) \\ ->raft::neighbors::ivf_flat::index; \\ \\ - template void raft::neighbors::ivf_flat::extend( \\ - raft::resources const& handle, \\ + template void raft::neighbors::ivf_flat::extend( \\ + raft::resources const& handle, \\ raft::neighbors::ivf_flat::index* index, \\ const T* new_vectors, \\ const IdxT* new_indices, \\ IdxT n_rows); \\ \\ - template void raft::neighbors::ivf_flat::extend( \\ - raft::resources const& handle, \\ + template void raft::neighbors::ivf_flat::extend( \\ + raft::resources const& handle, \\ raft::device_matrix_view new_vectors, \\ std::optional> new_indices, \\ - raft::neighbors::ivf_flat::index* index); + raft::neighbors::ivf_flat::index* index); \\ + \\ + template auto raft::neighbors::ivf_flat::extend( \\ + const raft::resources& handle, \\ + raft::host_matrix_view new_vectors, \\ + std::optional> new_indices, \\ + const raft::neighbors::ivf_flat::index& idx) \\ + -> raft::neighbors::ivf_flat::index; \\ + \\ + template void raft::neighbors::ivf_flat::extend( \\ + raft::resources const& handle, \\ + raft::host_matrix_view new_vectors, \\ + std::optional> new_indices, \\ + raft::neighbors::ivf_flat::index* index); """ search_macro = """ @@ -125,13 +150,16 @@ macros = dict( build=dict( definition=build_macro, - name="instantiate_raft_neighbors_ivf_flat_build"), + name="instantiate_raft_neighbors_ivf_flat_build", + ), extend=dict( definition=extend_macro, - name="instantiate_raft_neighbors_ivf_flat_extend"), + name="instantiate_raft_neighbors_ivf_flat_extend", + ), search=dict( definition=search_macro, - name="instantiate_raft_neighbors_ivf_flat_search"), + name="instantiate_raft_neighbors_ivf_flat_search", + ), ) for type_path, (T, IdxT) in types.items(): @@ -139,8 +167,7 @@ path = f"ivf_flat_{macro_path}_{type_path}.cu" with open(path, "w") as f: f.write(header) - f.write(macro['definition']) - + f.write(macro["definition"]) f.write(f"{macro['name']}({T}, {IdxT});\n\n") f.write(f"#undef {macro['name']}\n") diff --git a/cpp/src/neighbors/ivf_flat_build_float_int64_t.cu b/cpp/src/neighbors/ivf_flat_build_float_int64_t.cu index 2ae795db56..cf3cb6b1b2 100644 --- a/cpp/src/neighbors/ivf_flat_build_float_int64_t.cu +++ b/cpp/src/neighbors/ivf_flat_build_float_int64_t.cu @@ -44,6 +44,18 @@ raft::resources const& handle, \ const raft::neighbors::ivf_flat::index_params& params, \ raft::device_matrix_view dataset, \ + raft::neighbors::ivf_flat::index& idx); \ + \ + template auto raft::neighbors::ivf_flat::build( \ + raft::resources const& handle, \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::host_matrix_view dataset) \ + ->raft::neighbors::ivf_flat::index; \ + \ + template void raft::neighbors::ivf_flat::build( \ + raft::resources const& handle, \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::host_matrix_view dataset, \ raft::neighbors::ivf_flat::index& idx); instantiate_raft_neighbors_ivf_flat_build(float, int64_t); diff --git a/cpp/src/neighbors/ivf_flat_build_int8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_build_int8_t_int64_t.cu index deb31bf441..e1cf64907e 100644 --- a/cpp/src/neighbors/ivf_flat_build_int8_t_int64_t.cu +++ b/cpp/src/neighbors/ivf_flat_build_int8_t_int64_t.cu @@ -44,6 +44,18 @@ raft::resources const& handle, \ const raft::neighbors::ivf_flat::index_params& params, \ raft::device_matrix_view dataset, \ + raft::neighbors::ivf_flat::index& idx); \ + \ + template auto raft::neighbors::ivf_flat::build( \ + raft::resources const& handle, \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::host_matrix_view dataset) \ + ->raft::neighbors::ivf_flat::index; \ + \ + template void raft::neighbors::ivf_flat::build( \ + raft::resources const& handle, \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::host_matrix_view dataset, \ raft::neighbors::ivf_flat::index& idx); instantiate_raft_neighbors_ivf_flat_build(int8_t, int64_t); diff --git a/cpp/src/neighbors/ivf_flat_build_uint8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_build_uint8_t_int64_t.cu index 402fdbab97..26d1647954 100644 --- a/cpp/src/neighbors/ivf_flat_build_uint8_t_int64_t.cu +++ b/cpp/src/neighbors/ivf_flat_build_uint8_t_int64_t.cu @@ -44,6 +44,18 @@ raft::resources const& handle, \ const raft::neighbors::ivf_flat::index_params& params, \ raft::device_matrix_view dataset, \ + raft::neighbors::ivf_flat::index& idx); \ + \ + template auto raft::neighbors::ivf_flat::build( \ + raft::resources const& handle, \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::host_matrix_view dataset) \ + ->raft::neighbors::ivf_flat::index; \ + \ + template void raft::neighbors::ivf_flat::build( \ + raft::resources const& handle, \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::host_matrix_view dataset, \ raft::neighbors::ivf_flat::index& idx); instantiate_raft_neighbors_ivf_flat_build(uint8_t, int64_t); diff --git a/cpp/src/neighbors/ivf_flat_extend_float_int64_t.cu b/cpp/src/neighbors/ivf_flat_extend_float_int64_t.cu index 9e7701f773..16472c6692 100644 --- a/cpp/src/neighbors/ivf_flat_extend_float_int64_t.cu +++ b/cpp/src/neighbors/ivf_flat_extend_float_int64_t.cu @@ -52,6 +52,19 @@ raft::resources const& handle, \ raft::device_matrix_view new_vectors, \ std::optional> new_indices, \ + raft::neighbors::ivf_flat::index* index); \ + \ + template auto raft::neighbors::ivf_flat::extend( \ + const raft::resources& handle, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices, \ + const raft::neighbors::ivf_flat::index& idx) \ + ->raft::neighbors::ivf_flat::index; \ + \ + template void raft::neighbors::ivf_flat::extend( \ + raft::resources const& handle, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices, \ raft::neighbors::ivf_flat::index* index); instantiate_raft_neighbors_ivf_flat_extend(float, int64_t); diff --git a/cpp/src/neighbors/ivf_flat_extend_int8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_extend_int8_t_int64_t.cu index 5d3d23c3ab..d98b5225c3 100644 --- a/cpp/src/neighbors/ivf_flat_extend_int8_t_int64_t.cu +++ b/cpp/src/neighbors/ivf_flat_extend_int8_t_int64_t.cu @@ -52,6 +52,19 @@ raft::resources const& handle, \ raft::device_matrix_view new_vectors, \ std::optional> new_indices, \ + raft::neighbors::ivf_flat::index* index); \ + \ + template auto raft::neighbors::ivf_flat::extend( \ + const raft::resources& handle, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices, \ + const raft::neighbors::ivf_flat::index& idx) \ + ->raft::neighbors::ivf_flat::index; \ + \ + template void raft::neighbors::ivf_flat::extend( \ + raft::resources const& handle, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices, \ raft::neighbors::ivf_flat::index* index); instantiate_raft_neighbors_ivf_flat_extend(int8_t, int64_t); diff --git a/cpp/src/neighbors/ivf_flat_extend_uint8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_extend_uint8_t_int64_t.cu index 3150a676eb..520c3be536 100644 --- a/cpp/src/neighbors/ivf_flat_extend_uint8_t_int64_t.cu +++ b/cpp/src/neighbors/ivf_flat_extend_uint8_t_int64_t.cu @@ -52,6 +52,19 @@ raft::resources const& handle, \ raft::device_matrix_view new_vectors, \ std::optional> new_indices, \ + raft::neighbors::ivf_flat::index* index); \ + \ + template auto raft::neighbors::ivf_flat::extend( \ + const raft::resources& handle, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices, \ + const raft::neighbors::ivf_flat::index& idx) \ + ->raft::neighbors::ivf_flat::index; \ + \ + template void raft::neighbors::ivf_flat::extend( \ + raft::resources const& handle, \ + raft::host_matrix_view new_vectors, \ + std::optional> new_indices, \ raft::neighbors::ivf_flat::index* index); instantiate_raft_neighbors_ivf_flat_extend(uint8_t, int64_t); diff --git a/cpp/template/cmake/thirdparty/fetch_rapids.cmake b/cpp/template/cmake/thirdparty/fetch_rapids.cmake index 22afb1a3f5..15b6c43a6f 100644 --- a/cpp/template/cmake/thirdparty/fetch_rapids.cmake +++ b/cpp/template/cmake/thirdparty/fetch_rapids.cmake @@ -12,7 +12,7 @@ # the License. # Use this variable to update RAPIDS and RAFT versions -set(RAPIDS_VERSION "23.12") +set(RAPIDS_VERSION "24.02") if(NOT EXISTS ${CMAKE_CURRENT_BINARY_DIR}/RAFT_RAPIDS.cmake) file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-${RAPIDS_VERSION}/RAPIDS.cmake diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh index a9fd696f1f..39439d392d 100644 --- a/cpp/test/neighbors/ann_ivf_flat.cuh +++ b/cpp/test/neighbors/ann_ivf_flat.cuh @@ -72,6 +72,7 @@ struct AnnIvfFlatInputs { IdxT nlist; raft::distance::DistanceType metric; bool adaptive_centers; + bool host_dataset; }; template @@ -79,7 +80,7 @@ template { os << "{ " << p.num_queries << ", " << p.num_db_vecs << ", " << p.dim << ", " << p.k << ", " << p.nprobe << ", " << p.nlist << ", " << static_cast(p.metric) << ", " - << p.adaptive_centers << '}' << std::endl; + << p.adaptive_centers << ", " << p.host_dataset << '}' << std::endl; return os; } @@ -178,36 +179,69 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { index_params.kmeans_trainset_fraction = 0.5; index_params.metric_arg = 0; - auto database_view = raft::make_device_matrix_view( - (const DataT*)database.data(), ps.num_db_vecs, ps.dim); - - auto idx = ivf_flat::build(handle_, index_params, database_view); + ivf_flat::index idx(handle_, index_params, ps.dim); + ivf_flat::index index_2(handle_, index_params, ps.dim); + + if (!ps.host_dataset) { + auto database_view = raft::make_device_matrix_view( + (const DataT*)database.data(), ps.num_db_vecs, ps.dim); + idx = ivf_flat::build(handle_, index_params, database_view); + rmm::device_uvector vector_indices(ps.num_db_vecs, stream_); + thrust::sequence(resource::get_thrust_policy(handle_), + thrust::device_pointer_cast(vector_indices.data()), + thrust::device_pointer_cast(vector_indices.data() + ps.num_db_vecs)); + resource::sync_stream(handle_); - rmm::device_uvector vector_indices(ps.num_db_vecs, stream_); - thrust::sequence(resource::get_thrust_policy(handle_), - thrust::device_pointer_cast(vector_indices.data()), - thrust::device_pointer_cast(vector_indices.data() + ps.num_db_vecs)); - resource::sync_stream(handle_); + IdxT half_of_data = ps.num_db_vecs / 2; - IdxT half_of_data = ps.num_db_vecs / 2; + auto half_of_data_view = raft::make_device_matrix_view( + (const DataT*)database.data(), half_of_data, ps.dim); - auto half_of_data_view = raft::make_device_matrix_view( - (const DataT*)database.data(), half_of_data, ps.dim); + const std::optional> no_opt = std::nullopt; + index_2 = ivf_flat::extend(handle_, half_of_data_view, no_opt, idx); - const std::optional> no_opt = std::nullopt; - index index_2 = ivf_flat::extend(handle_, half_of_data_view, no_opt, idx); + auto new_half_of_data_view = raft::make_device_matrix_view( + database.data() + half_of_data * ps.dim, IdxT(ps.num_db_vecs) - half_of_data, ps.dim); - auto new_half_of_data_view = raft::make_device_matrix_view( - database.data() + half_of_data * ps.dim, IdxT(ps.num_db_vecs) - half_of_data, ps.dim); + auto new_half_of_data_indices_view = raft::make_device_vector_view( + vector_indices.data() + half_of_data, IdxT(ps.num_db_vecs) - half_of_data); - auto new_half_of_data_indices_view = raft::make_device_vector_view( - vector_indices.data() + half_of_data, IdxT(ps.num_db_vecs) - half_of_data); + ivf_flat::extend(handle_, + new_half_of_data_view, + std::make_optional>( + new_half_of_data_indices_view), + &index_2); - ivf_flat::extend(handle_, - new_half_of_data_view, - std::make_optional>( - new_half_of_data_indices_view), - &index_2); + } else { + auto host_database = raft::make_host_matrix(ps.num_db_vecs, ps.dim); + raft::copy( + host_database.data_handle(), database.data(), ps.num_db_vecs * ps.dim, stream_); + idx = + ivf_flat::build(handle_, index_params, raft::make_const_mdspan(host_database.view())); + + auto vector_indices = raft::make_host_vector(handle_, ps.num_db_vecs); + std::iota(vector_indices.data_handle(), vector_indices.data_handle() + ps.num_db_vecs, 0); + + IdxT half_of_data = ps.num_db_vecs / 2; + + auto half_of_data_view = raft::make_host_matrix_view( + (const DataT*)host_database.data_handle(), half_of_data, ps.dim); + + const std::optional> no_opt = std::nullopt; + index_2 = ivf_flat::extend(handle_, half_of_data_view, no_opt, idx); + + auto new_half_of_data_view = raft::make_host_matrix_view( + host_database.data_handle() + half_of_data * ps.dim, + IdxT(ps.num_db_vecs) - half_of_data, + ps.dim); + auto new_half_of_data_indices_view = raft::make_host_vector_view( + vector_indices.data_handle() + half_of_data, IdxT(ps.num_db_vecs) - half_of_data); + ivf_flat::extend(handle_, + new_half_of_data_view, + std::make_optional>( + new_half_of_data_indices_view), + &index_2); + } auto search_queries_view = raft::make_device_matrix_view( search_queries.data(), ps.num_queries, ps.dim); @@ -574,6 +608,15 @@ const std::vector> inputs = { {1000, 100000, 16, 10, 20, 1024, raft::distance::DistanceType::L2Expanded, true}, {10000, 131072, 8, 10, 20, 1024, raft::distance::DistanceType::L2Expanded, false}, + // host input data + {1000, 10000, 16, 10, 40, 1024, raft::distance::DistanceType::L2Expanded, false, true}, + {1000, 10000, 16, 10, 50, 1024, raft::distance::DistanceType::L2Expanded, false, true}, + {1000, 10000, 16, 10, 70, 1024, raft::distance::DistanceType::L2Expanded, false, true}, + {100, 10000, 16, 10, 20, 512, raft::distance::DistanceType::L2Expanded, false, true}, + {20, 100000, 16, 10, 20, 1024, raft::distance::DistanceType::L2Expanded, false, true}, + {1000, 100000, 16, 10, 20, 1024, raft::distance::DistanceType::L2Expanded, false, true}, + {10000, 131072, 8, 10, 20, 1024, raft::distance::DistanceType::L2Expanded, false, true}, + {1000, 10000, 16, 10, 40, 1024, raft::distance::DistanceType::InnerProduct, true}, {1000, 10000, 16, 10, 50, 1024, raft::distance::DistanceType::InnerProduct, true}, {1000, 10000, 16, 10, 70, 1024, raft::distance::DistanceType::InnerProduct, false}, diff --git a/dependencies.yaml b/dependencies.yaml index 419cd0c22e..9f113d0969 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -177,7 +177,7 @@ dependencies: common: - output_types: [conda] packages: - - &rmm_conda rmm==23.12.* + - &rmm_conda rmm==24.2.* - output_types: requirements packages: # pip recognizes the index as a global option for the requirements.txt file @@ -197,12 +197,12 @@ dependencies: matrices: - matrix: {cuda: "12.2"} packages: &build_pylibraft_packages_cu12 - - &rmm_cu12 rmm-cu12==23.12.* + - &rmm_cu12 rmm-cu12==24.2.* - {matrix: {cuda: "12.1"}, packages: *build_pylibraft_packages_cu12} - {matrix: {cuda: "12.0"}, packages: *build_pylibraft_packages_cu12} - matrix: {cuda: "11.8"} packages: &build_pylibraft_packages_cu11 - - &rmm_cu11 rmm-cu11==23.12.* + - &rmm_cu11 rmm-cu11==24.2.* - {matrix: {cuda: "11.5"}, packages: *build_pylibraft_packages_cu11} - {matrix: {cuda: "11.4"}, packages: *build_pylibraft_packages_cu11} - {matrix: {cuda: "11.2"}, packages: *build_pylibraft_packages_cu11} @@ -436,20 +436,20 @@ dependencies: common: - output_types: [conda, pyproject] packages: - - dask-cuda==23.12.* + - dask-cuda==24.2.* - joblib>=0.11 - numba>=0.57 - *numpy - - rapids-dask-dependency==23.12.* - - ucx-py==0.35.* + - rapids-dask-dependency==24.2.* + - ucx-py==0.36.* - output_types: conda packages: - ucx>=1.13.0 - ucx-proc=*=gpu - - &ucx_py_conda ucx-py==0.35.* + - &ucx_py_conda ucx-py==0.36.* - output_types: pyproject packages: - - &pylibraft_conda pylibraft==23.12.* + - &pylibraft_conda pylibraft==24.2.* - output_types: requirements packages: # pip recognizes the index as a global option for the requirements.txt file @@ -460,13 +460,13 @@ dependencies: matrices: - matrix: {cuda: "12.2"} packages: &run_raft_dask_packages_cu12 - - &pylibraft_cu12 pylibraft-cu12==23.12.* + - &pylibraft_cu12 pylibraft-cu12==24.2.* - &ucx_py_cu12 ucx-py-cu12==0.35.* - {matrix: {cuda: "12.1"}, packages: *run_raft_dask_packages_cu12} - {matrix: {cuda: "12.0"}, packages: *run_raft_dask_packages_cu12} - matrix: {cuda: "11.8"} packages: &run_raft_dask_packages_cu11 - - &pylibraft_cu11 pylibraft-cu11==23.12.* + - &pylibraft_cu11 pylibraft-cu11==24.2.* - &ucx_py_cu11 ucx-py-cu11==0.35.* - {matrix: {cuda: "11.5"}, packages: *run_raft_dask_packages_cu11} - {matrix: {cuda: "11.4"}, packages: *run_raft_dask_packages_cu11} diff --git a/docs/source/build.md b/docs/source/build.md index 4be0a84090..ae7734d0ed 100644 --- a/docs/source/build.md +++ b/docs/source/build.md @@ -56,7 +56,7 @@ You can also install the conda packages individually using the `mamba` command a mamba install -c rapidsai -c conda-forge -c nvidia libraft libraft-headers cuda-version=12.0 ``` -If installing the C++ APIs Please see [using libraft](https://docs.rapids.ai/api/raft/nightly/using_libraft/) for more information on using the pre-compiled shared library. You can also refer to the [example C++ template project](https://github.com/rapidsai/raft/tree/branch-23.12/cpp/template) for a ready-to-go CMake configuration that you can drop into your project and build against installed RAFT development artifacts above. +If installing the C++ APIs Please see [using libraft](https://docs.rapids.ai/api/raft/nightly/using_libraft/) for more information on using the pre-compiled shared library. You can also refer to the [example C++ template project](https://github.com/rapidsai/raft/tree/branch-24.02/cpp/template) for a ready-to-go CMake configuration that you can drop into your project and build against installed RAFT development artifacts above. ## Installing Python through Pip diff --git a/docs/source/conf.py b/docs/source/conf.py index 6c523cceb5..2a2c700926 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -67,9 +67,9 @@ # built documents. # # The short X.Y version. -version = '23.12' +version = '24.02' # The full version, including alpha/beta/rc tags. -release = '23.12.00' +release = '24.02.00' # The language for content autogenerated by Sphinx. Refer to documentation # for a list of supported languages. diff --git a/docs/source/developer_guide.md b/docs/source/developer_guide.md index c45cec1666..c5bcd03f69 100644 --- a/docs/source/developer_guide.md +++ b/docs/source/developer_guide.md @@ -187,7 +187,7 @@ RAFT relies on `clang-format` to enforce code style across all C++ and CUDA sour 1. Do not split empty functions/records/namespaces. 2. Two-space indentation everywhere, including the line continuations. 3. Disable reflowing of comments. - The reasons behind these deviations from the Google style guide are given in comments [here](https://github.com/rapidsai/raft/blob/branch-23.12/cpp/.clang-format). + The reasons behind these deviations from the Google style guide are given in comments [here](https://github.com/rapidsai/raft/blob/branch-24.02/cpp/.clang-format). [`doxygen`](https://doxygen.nl/) is used as documentation generator and also as a documentation linter. In order to run doxygen as a linter on C++/CUDA code, run @@ -205,7 +205,7 @@ you can run `codespell -i 3 -w .` from the repository root directory. This will bring up an interactive prompt to select which spelling fixes to apply. ### #include style -[include_checker.py](https://github.com/rapidsai/raft/blob/branch-23.12/cpp/scripts/include_checker.py) is used to enforce the include style as follows: +[include_checker.py](https://github.com/rapidsai/raft/blob/branch-24.02/cpp/scripts/include_checker.py) is used to enforce the include style as follows: 1. `#include "..."` should be used for referencing local files only. It is acceptable to be used for referencing files in a sub-folder/parent-folder of the same algorithm, but should never be used to include files in other algorithms or between algorithms and the primitives or other dependencies. 2. `#include <...>` should be used for referencing everything else @@ -215,7 +215,7 @@ python ./cpp/scripts/include_checker.py --inplace [cpp/include cpp/test ... list ``` ### Copyright header -[copyright.py](https://github.com/rapidsai/raft/blob/branch-23.12/ci/checks/copyright.py) checks the Copyright header for all git-modified files +[copyright.py](https://github.com/rapidsai/raft/blob/branch-24.02/ci/checks/copyright.py) checks the Copyright header for all git-modified files Manually, you can run the following to bulk-fix the header if only the years need to be updated: ```bash @@ -229,7 +229,7 @@ Call CUDA APIs via the provided helper macros `RAFT_CUDA_TRY`, `RAFT_CUBLAS_TRY` ## Logging ### Introduction -Anything and everything about logging is defined inside [logger.hpp](https://github.com/rapidsai/raft/blob/branch-23.12/cpp/include/raft/core/logger.hpp). It uses [spdlog](https://github.com/gabime/spdlog) underneath, but this information is transparent to all. +Anything and everything about logging is defined inside [logger.hpp](https://github.com/rapidsai/raft/blob/branch-24.02/cpp/include/raft/core/logger.hpp). It uses [spdlog](https://github.com/gabime/spdlog) underneath, but this information is transparent to all. ### Usage ```cpp diff --git a/docs/source/raft_ann_benchmarks.md b/docs/source/raft_ann_benchmarks.md index 24fc3801d9..dcdfc2cec9 100644 --- a/docs/source/raft_ann_benchmarks.md +++ b/docs/source/raft_ann_benchmarks.md @@ -62,7 +62,7 @@ Nightly images are located in [dockerhub](https://hub.docker.com/r/rapidsai/raft - The following command pulls the nightly container for python version 10, cuda version 12, and RAFT version 23.10: ```bash -docker pull rapidsai/raft-ann-bench:23.12a-cuda12.0-py3.10 #substitute raft-ann-bench for the exact desired container. +docker pull rapidsai/raft-ann-bench:24.02a-cuda12.0-py3.10 #substitute raft-ann-bench for the exact desired container. ``` The CUDA and python versions can be changed for the supported values: @@ -83,7 +83,7 @@ You can see the exact versions as well in the dockerhub site: [//]: # () [//]: # (```bash) -[//]: # (docker pull nvcr.io/nvidia/rapidsai/raft-ann-bench:23.12-cuda11.8-py3.10 #substitute raft-ann-bench for the exact desired container.) +[//]: # (docker pull nvcr.io/nvidia/rapidsai/raft-ann-bench:24.02-cuda11.8-py3.10 #substitute raft-ann-bench for the exact desired container.) [//]: # (```) @@ -198,27 +198,33 @@ options: --dataset-path DATASET_PATH path to dataset folder (default: ${RAPIDS_DATASET_ROOT_DIR}) ``` -Build statistics CSV file is stored in `/result/build/` -and index search statistics CSV file in `/result/search/`. +Build statistics CSV file is stored in `/result/build/` +and index search statistics CSV file in `/result/search/`, where suffix has three values: +1. `raw`: All search results are exported +2. `throughput`: Pareto frontier of throughput results is exported +3. `latency`: Pareto frontier of latency results is exported + ### Step 4: Plot Results The script `raft-ann-bench.plot` will plot results for all algorithms found in index search statistics -CSV file in `/result/search/<-k{k}-batch_size{batch_size}>.csv`. +CSV files `/result/search/*.csv`. The usage of this script is: ```bash -usage: __main__.py [-h] [--dataset DATASET] [--dataset-path DATASET_PATH] [--output-filepath OUTPUT_FILEPATH] [--algorithms ALGORITHMS] [--groups GROUPS] [--algo-groups ALGO_GROUPS] [-k COUNT] - [-bs BATCH_SIZE] [--build] [--search] [--x-scale X_SCALE] [--y-scale {linear,log,symlog,logit}] [--raw] +usage: [-h] [--dataset DATASET] [--dataset-path DATASET_PATH] [--output-filepath OUTPUT_FILEPATH] [--algorithms ALGORITHMS] [--groups GROUPS] [--algo-groups ALGO_GROUPS] + [-k COUNT] [-bs BATCH_SIZE] [--build] [--search] [--x-scale X_SCALE] [--y-scale {linear,log,symlog,logit}] [--mode {throughput,latency}] [--time-unit {s,ms,us}] + [--raw] options: -h, --help show this help message and exit --dataset DATASET dataset to plot (default: glove-100-inner) --dataset-path DATASET_PATH - path to dataset folder (default: os.getcwd()/datasets/) + path to dataset folder (default: /home/coder/raft/datasets/) --output-filepath OUTPUT_FILEPATH - directory for PNG to be saved (default: os.getcwd()) + directory for PNG to be saved (default: /home/coder/raft) --algorithms ALGORITHMS - plot only comma separated list of named algorithms. If parameters `groups` and `algo-groups are both undefined, then group `base` is plot by default (default: None) + plot only comma separated list of named algorithms. If parameters `groups` and `algo-groups are both undefined, then group `base` is plot by default + (default: None) --groups GROUPS plot only comma separated groups of parameters (default: base) --algo-groups ALGO_GROUPS, --algo-groups ALGO_GROUPS add comma separated . to plot. Example usage: "--algo-groups=raft_cagra.large,hnswlib.large" (default: None) @@ -231,8 +237,14 @@ options: --x-scale X_SCALE Scale to use when drawing the X-axis. Typically linear, logit or a2 (default: linear) --y-scale {linear,log,symlog,logit} Scale to use when drawing the Y-axis (default: linear) - --raw Show raw results (not just Pareto frontier) in faded colours (default: False) + --mode {throughput,latency} + search mode whose Pareto frontier is used on the y-axis (default: throughput) + --time-unit {s,ms,us} + time unit to plot when mode is latency (default: ms) + --raw Show raw results (not just Pareto frontier) of mode arg (default: False) ``` +`mode`: plots pareto frontier of `throughput` or `latency` results exported in the previous step + `algorithms`: plots all algorithms that it can find results for the specified `dataset`. By default, only `base` group will be plotted. `groups`: plot only specific groups of parameters configurations for an algorithm. Groups are defined in YAML configs (see `configuration`), and by default run `base` group @@ -331,7 +343,7 @@ For GPU-enabled systems, the `DATA_FOLDER` variable should be a local folder whe export DATA_FOLDER=path/to/store/datasets/and/results docker run --gpus all --rm -it -u $(id -u) \ -v $DATA_FOLDER:/data/benchmarks \ - rapidsai/raft-ann-bench:23.12a-cuda11.8-py3.10 \ + rapidsai/raft-ann-bench:24.02a-cuda11.8-py3.10 \ "--dataset deep-image-96-angular" \ "--normalize" \ "--algorithms raft_cagra,raft_ivf_pq --batch-size 10 -k 10" \ @@ -342,7 +354,7 @@ Usage of the above command is as follows: | Argument | Description | |-----------------------------------------------------------|----------------------------------------------------------------------------------------------------| -| `rapidsai/raft-ann-bench:23.12a-cuda11.8-py3.10` | Image to use. Can be either `raft-ann-bench` or `raft-ann-bench-datasets` | +| `rapidsai/raft-ann-bench:24.02a-cuda11.8-py3.10` | Image to use. Can be either `raft-ann-bench` or `raft-ann-bench-datasets` | | `"--dataset deep-image-96-angular"` | Dataset name | | `"--normalize"` | Whether to normalize the dataset | | `"--algorithms raft_cagra,hnswlib --batch-size 10 -k 10"` | Arguments passed to the `run` script, such as the algorithms to benchmark, the batch size, and `k` | @@ -359,7 +371,7 @@ The container arguments in the above section also be used for the CPU-only conta export DATA_FOLDER=path/to/store/datasets/and/results docker run --rm -it -u $(id -u) \ -v $DATA_FOLDER:/data/benchmarks \ - rapidsai/raft-ann-bench-cpu:23.12a-py3.10 \ + rapidsai/raft-ann-bench-cpu:24.02a-py3.10 \ "--dataset deep-image-96-angular" \ "--normalize" \ "--algorithms hnswlib --batch-size 10 -k 10" \ @@ -376,7 +388,7 @@ docker run --gpus all --rm -it -u $(id -u) \ --entrypoint /bin/bash \ --workdir /data/benchmarks \ -v $DATA_FOLDER:/data/benchmarks \ - rapidsai/raft-ann-bench:23.12a-cuda11.8-py3.10 + rapidsai/raft-ann-bench:24.02a-cuda11.8-py3.10 ``` This will drop you into a command line in the container, with the `raft-ann-bench` python package ready to use, as described in the [Running the benchmarks](#running-the-benchmarks) section above: diff --git a/fetch_rapids.cmake b/fetch_rapids.cmake index c74c09fb26..ca871c5759 100644 --- a/fetch_rapids.cmake +++ b/fetch_rapids.cmake @@ -12,7 +12,7 @@ # the License. # ============================================================================= if(NOT EXISTS ${CMAKE_CURRENT_BINARY_DIR}/RAFT_RAPIDS.cmake) - file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-23.12/RAPIDS.cmake + file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-24.02/RAPIDS.cmake ${CMAKE_CURRENT_BINARY_DIR}/RAFT_RAPIDS.cmake ) endif() diff --git a/python/pylibraft/CMakeLists.txt b/python/pylibraft/CMakeLists.txt index 96489ea045..1109eb79cc 100644 --- a/python/pylibraft/CMakeLists.txt +++ b/python/pylibraft/CMakeLists.txt @@ -16,7 +16,7 @@ cmake_minimum_required(VERSION 3.26.4 FATAL_ERROR) include(../../fetch_rapids.cmake) -set(pylibraft_version 23.12.00) +set(pylibraft_version 24.02.00) # We always need CUDA for pylibraft because the raft dependency brings in a header-only cuco # dependency that enables CUDA unconditionally. diff --git a/python/pylibraft/pyproject.toml b/python/pylibraft/pyproject.toml index b946a3b942..f811afa55f 100644 --- a/python/pylibraft/pyproject.toml +++ b/python/pylibraft/pyproject.toml @@ -19,7 +19,7 @@ requires = [ "cuda-python>=11.7.1,<12.0a0", "cython>=3.0.0", "ninja", - "rmm==23.12.*", + "rmm==24.2.*", "scikit-build>=0.13.1", "setuptools", "wheel", @@ -39,7 +39,7 @@ requires-python = ">=3.9" dependencies = [ "cuda-python>=11.7.1,<12.0a0", "numpy>=1.21", - "rmm==23.12.*", + "rmm==24.2.*", ] # 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", diff --git a/python/raft-ann-bench/pyproject.toml b/python/raft-ann-bench/pyproject.toml index 0cb36deaa7..e3ae36ef62 100644 --- a/python/raft-ann-bench/pyproject.toml +++ b/python/raft-ann-bench/pyproject.toml @@ -9,7 +9,7 @@ requires = [ [project] name = "raft-ann-bench" -version = "23.12.00" +version = "24.02.00" description = "RAFT ANN benchmarks" authors = [ { name = "NVIDIA Corporation" }, diff --git a/python/raft-ann-bench/src/raft-ann-bench/constraints/__init__.py b/python/raft-ann-bench/src/raft-ann-bench/constraints/__init__.py index d827c3d1e1..2b7b2728fe 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/constraints/__init__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/constraints/__init__.py @@ -16,6 +16,12 @@ DTYPE_SIZES = {"float": 4, "half": 2, "fp8": 1} +def raft_cagra_build_constraints(params, dims): + if "graph_degree" in params and "intermediate_graph_degree" in params: + return params["graph_degree"] <= params["intermediate_graph_degree"] + return True + + def raft_ivf_pq_build_constraints(params, dims): if "pq_dim" in params: return params["pq_dim"] <= dims @@ -36,8 +42,10 @@ def raft_ivf_pq_search_constraints(params, build_params, k, batch_size): def raft_cagra_search_constraints(params, build_params, k, batch_size): + ret = True if "itopk" in params: - return params["itopk"] >= k + ret = ret and params["itopk"] >= k + return ret def hnswlib_search_constraints(params, build_params, k, batch_size): diff --git a/python/raft-ann-bench/src/raft-ann-bench/data_export/__main__.py b/python/raft-ann-bench/src/raft-ann-bench/data_export/__main__.py index 4978c99d60..5cb06c573f 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/data_export/__main__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/data_export/__main__.py @@ -43,15 +43,34 @@ ) skip_search_cols = ( - set(["recall", "qps", "items_per_second", "Recall"]) | skip_build_cols + set(["recall", "qps", "latency", "items_per_second", "Recall", "Latency"]) + | skip_build_cols ) +metrics = { + "k-nn": { + "description": "Recall", + "worst": float("-inf"), + "lim": [0.0, 1.03], + }, + "throughput": { + "description": "Queries per second (1/s)", + "worst": float("-inf"), + }, + "latency": { + "description": "Search Latency (s)", + "worst": float("inf"), + }, +} + def read_file(dataset, dataset_path, method): dir = os.path.join(dataset_path, dataset, "result", method) for file in os.listdir(dir): if file.endswith(".json"): - with open(os.path.join(dir, file), "r") as f: + with open( + os.path.join(dir, file), "r", encoding="ISO-8859-1" + ) as f: try: data = json.load(f) df = pd.DataFrame(data["benchmarks"]) @@ -92,6 +111,31 @@ def convert_json_to_csv_build(dataset, dataset_path): traceback.print_exc() +def create_pointset(data, xn, yn): + xm, ym = (metrics[xn], metrics[yn]) + rev_y = -1 if ym["worst"] < 0 else 1 + rev_x = -1 if xm["worst"] < 0 else 1 + + y_idx = 3 if yn == "throughput" else 4 + data.sort(key=lambda t: (rev_y * t[y_idx], rev_x * t[2])) + + lines = [] + last_x = xm["worst"] + comparator = ( + (lambda xv, lx: xv > lx) if last_x < 0 else (lambda xv, lx: xv < lx) + ) + for d in data: + if comparator(d[2], last_x): + last_x = d[2] + lines.append(d) + return lines + + +def get_frontier(df, metric): + lines = create_pointset(df.values.tolist(), "k-nn", metric) + return pd.DataFrame(lines, columns=df.columns) + + def convert_json_to_csv_search(dataset, dataset_path): for file, algo_name, df in read_file(dataset, dataset_path, "search"): try: @@ -100,14 +144,21 @@ def convert_json_to_csv_search(dataset, dataset_path): ) algo_name = algo_name.replace("_base", "") df["name"] = df["name"].str.split("/").str[0] - write = pd.DataFrame( - { - "algo_name": [algo_name] * len(df), - "index_name": df["name"], - "recall": df["Recall"], - "qps": df["items_per_second"], - } - ) + try: + write = pd.DataFrame( + { + "algo_name": [algo_name] * len(df), + "index_name": df["name"], + "recall": df["Recall"], + "throughput": df["items_per_second"], + "latency": df["Latency"], + } + ) + except Exception as e: + print( + "Search file %s (%s) missing a key. Skipping..." + % (file, e) + ) for name in df: if name not in skip_search_cols: write[name] = df[name] @@ -120,20 +171,29 @@ def convert_json_to_csv_search(dataset, dataset_path): write["build cpu_time"] = None write["build GPU"] = None - for col_idx in range(6, len(build_df.columns)): - col_name = build_df.columns[col_idx] - write[col_name] = None - - for s_index, search_row in write.iterrows(): - for b_index, build_row in build_df.iterrows(): - if search_row["index_name"] == build_row["index_name"]: - write.iloc[s_index, write_ncols] = build_df.iloc[ - b_index, 2 - ] - write.iloc[ - s_index, write_ncols + 1 : - ] = build_df.iloc[b_index, 3:] - break + try: + for col_idx in range(6, len(build_df.columns)): + col_name = build_df.columns[col_idx] + write[col_name] = None + + for s_index, search_row in write.iterrows(): + for b_index, build_row in build_df.iterrows(): + if ( + search_row["index_name"] + == build_row["index_name"] + ): + write.iloc[ + s_index, write_ncols + ] = build_df.iloc[b_index, 2] + write.iloc[ + s_index, write_ncols + 1 : + ] = build_df.iloc[b_index, 3:] + break + except Exception as e: + print( + "Build file %s (%s) missing a key. Skipping..." + % (build_file, e) + ) else: warnings.warn( f"Build CSV not found for {algo_name}, " @@ -141,7 +201,13 @@ def convert_json_to_csv_search(dataset, dataset_path): "appended in the Search CSV" ) - write.to_csv(file.replace(".json", ".csv"), index=False) + write.to_csv(file.replace(".json", "_raw.csv"), index=False) + throughput = get_frontier(write, "throughput") + throughput.to_csv( + file.replace(".json", "_throughput.csv"), index=False + ) + latency = get_frontier(write, "latency") + latency.to_csv(file.replace(".json", "_latency.csv"), index=False) except Exception as e: print( "An error occurred processing file %s (%s). Skipping..." diff --git a/python/raft-ann-bench/src/raft-ann-bench/plot/__main__.py b/python/raft-ann-bench/src/raft-ann-bench/plot/__main__.py index c45ff5b14e..8bd54170c9 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/plot/__main__.py +++ b/python/raft-ann-bench/src/raft-ann-bench/plot/__main__.py @@ -38,10 +38,14 @@ "worst": float("-inf"), "lim": [0.0, 1.03], }, - "qps": { + "throughput": { "description": "Queries per second (1/s)", "worst": float("-inf"), }, + "latency": { + "description": "Search Latency (s)", + "worst": float("inf"), + }, } @@ -98,53 +102,20 @@ def create_linestyles(unique_algorithms): ) -def get_up_down(metric): - if metric["worst"] == float("inf"): - return "down" - return "up" - - -def get_left_right(metric): - if metric["worst"] == float("inf"): - return "left" - return "right" - - -def create_pointset(data, xn, yn): - xm, ym = (metrics[xn], metrics[yn]) - rev_y = -1 if ym["worst"] < 0 else 1 - rev_x = -1 if xm["worst"] < 0 else 1 - data.sort(key=lambda t: (rev_y * t[-1], rev_x * t[-2])) - - axs, ays, als, aidxs = [], [], [], [] - # Generate Pareto frontier - xs, ys, ls, idxs = [], [], [], [] - last_x = xm["worst"] - comparator = ( - (lambda xv, lx: xv > lx) if last_x < 0 else (lambda xv, lx: xv < lx) - ) - for algo_name, index_name, xv, yv in data: - if not xv or not yv: - continue - axs.append(xv) - ays.append(yv) - als.append(algo_name) - aidxs.append(algo_name) - if comparator(xv, last_x): - last_x = xv - xs.append(xv) - ys.append(yv) - ls.append(algo_name) - idxs.append(index_name) - return xs, ys, ls, idxs, axs, ays, als, aidxs - - def create_plot_search( - all_data, raw, x_scale, y_scale, fn_out, linestyles, dataset, k, batch_size + all_data, + x_scale, + y_scale, + fn_out, + linestyles, + dataset, + k, + batch_size, + mode, + time_unit, ): xn = "k-nn" - yn = "qps" - xm, ym = (metrics[xn], metrics[yn]) + xm, ym = (metrics[xn], metrics[mode]) # Now generate each plot handles = [] labels = [] @@ -152,17 +123,15 @@ def create_plot_search( # Sorting by mean y-value helps aligning plots with labels def mean_y(algo): - xs, ys, ls, idxs, axs, ays, als, aidxs = create_pointset( - all_data[algo], xn, yn - ) - return -np.log(np.array(ys)).mean() + points = np.array(all_data[algo], dtype=object) + return -np.log(np.array(points[:, 3], dtype=np.float32)).mean() # Find range for logit x-scale min_x, max_x = 1, 0 for algo in sorted(all_data.keys(), key=mean_y): - xs, ys, ls, idxs, axs, ays, als, aidxs = create_pointset( - all_data[algo], xn, yn - ) + points = np.array(all_data[algo], dtype=object) + xs = points[:, 2] + ys = points[:, 3] min_x = min([min_x] + [x for x in xs if x > 0]) max_x = max([max_x] + [x for x in xs if x < 1]) color, faded, linestyle, marker = linestyles[algo] @@ -178,23 +147,15 @@ def mean_y(algo): marker=marker, ) handles.append(handle) - if raw: - (handle2,) = plt.plot( - axs, - ays, - "-", - label=algo, - color=faded, - ms=5, - mew=2, - lw=2, - marker=marker, - ) + labels.append(algo) ax = plt.gca() - ax.set_ylabel(ym["description"]) - ax.set_xlabel(xm["description"]) + y_description = ym["description"] + if mode == "latency": + y_description = y_description.replace("(s)", f"({time_unit})") + ax.set_ylabel(y_description) + ax.set_xlabel("Recall") # Custom scales of the type --x-scale a3 if x_scale[0] == "a": alpha = float(x_scale[1:]) @@ -250,10 +211,8 @@ def inv_fun(x): def create_plot_build( - build_results, search_results, linestyles, fn_out, dataset, k, batch_size + build_results, search_results, linestyles, fn_out, dataset ): - xn = "k-nn" - yn = "qps" qps_85 = [-1] * len(linestyles) bt_85 = [0] * len(linestyles) @@ -271,16 +230,17 @@ def create_plot_build( colors = OrderedDict() # Sorting by mean y-value helps aligning plots with labels + def mean_y(algo): - xs, ys, ls, idxs, axs, ays, als, aidxs = create_pointset( - search_results[algo], xn, yn - ) - return -np.log(np.array(ys)).mean() + points = np.array(search_results[algo], dtype=object) + return -np.log(np.array(points[:, 3], dtype=np.float32)).mean() for pos, algo in enumerate(sorted(search_results.keys(), key=mean_y)): - xs, ys, ls, idxs, axs, ays, als, aidxs = create_pointset( - search_results[algo], xn, yn - ) + points = np.array(search_results[algo], dtype=object) + xs = points[:, 2] + ys = points[:, 3] + ls = points[:, 0] + idxs = points[:, 1] # x is recall, y is qps, ls is algo_name, idxs is index_name for i in range(len(xs)): if xs[i] >= 0.85 and xs[i] < 0.9 and ys[i] > qps_85[pos]: @@ -311,11 +271,11 @@ def mean_y(algo): fig.savefig(fn_out) -def load_lines(results_path, result_files, method, index_key): +def load_lines(results_path, result_files, method, index_key, mode, time_unit): results = dict() for result_filename in result_files: - if result_filename.endswith(".csv"): + try: with open(os.path.join(results_path, result_filename), "r") as f: lines = f.readlines() lines = lines[:-1] if lines[-1] == "\n" else lines @@ -323,7 +283,8 @@ def load_lines(results_path, result_files, method, index_key): if method == "build": key_idx = [2] elif method == "search": - key_idx = [2, 3] + y_idx = 3 if mode == "throughput" else 4 + key_idx = [2, y_idx] for line in lines[1:]: split_lines = line.split(",") @@ -340,7 +301,22 @@ def load_lines(results_path, result_files, method, index_key): to_add = [algo_name, index_name] for key_i in key_idx: to_add.append(float(split_lines[key_i])) + if ( + mode == "latency" + and time_unit != "s" + and method == "search" + ): + to_add[-1] = ( + to_add[-1] * (10**3) + if time_unit == "ms" + else to_add[-1] * (10**6) + ) results[dict_key].append(to_add) + except Exception: + print( + f"An error occurred processing file {result_filename}. " + "Skipping..." + ) return results @@ -354,12 +330,31 @@ def load_all_results( batch_size, method, index_key, + raw, + mode, + time_unit, ): results_path = os.path.join(dataset_path, "result", method) result_files = os.listdir(results_path) - result_files = [ - result_file for result_file in result_files if ".csv" in result_file - ] + if method == "build": + result_files = [ + result_file + for result_file in result_files + if ".csv" in result_file + ] + elif method == "search": + if raw: + suffix = "_raw" + else: + suffix = f"_{mode}" + result_files = [ + result_file + for result_file in result_files + if f"{suffix}.csv" in result_file + ] + if len(result_files) == 0: + raise FileNotFoundError(f"No CSV result files found in {results_path}") + if method == "search": result_files = [ result_filename @@ -407,7 +402,9 @@ def load_all_results( final_results = final_results + final_algo_groups final_results = set(final_results) - results = load_lines(results_path, final_results, method, index_key) + results = load_lines( + results_path, final_results, method, index_key, mode, time_unit + ) return results @@ -481,9 +478,21 @@ def main(): choices=["linear", "log", "symlog", "logit"], default="linear", ) + parser.add_argument( + "--mode", + help="search mode whose Pareto frontier is used on the y-axis", + choices=["throughput", "latency"], + default="throughput", + ) + parser.add_argument( + "--time-unit", + help="time unit to plot when mode is latency", + choices=["s", "ms", "us"], + default="ms", + ) parser.add_argument( "--raw", - help="Show raw results (not just Pareto frontier) in faded colours", + help="Show raw results (not just Pareto frontier) of mode arg", action="store_true", ) @@ -528,12 +537,14 @@ def main(): batch_size, "search", "algo", + args.raw, + args.mode, + args.time_unit, ) linestyles = create_linestyles(sorted(search_results.keys())) if search: create_plot_search( search_results, - args.raw, args.x_scale, args.y_scale, search_output_filepath, @@ -541,6 +552,8 @@ def main(): args.dataset, k, batch_size, + args.mode, + args.time_unit, ) if build: build_results = load_all_results( @@ -552,6 +565,9 @@ def main(): batch_size, "build", "index", + args.raw, + args.mode, + args.time_unit, ) create_plot_build( build_results, @@ -559,8 +575,6 @@ def main(): linestyles, build_output_filepath, args.dataset, - k, - batch_size, ) diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml index d8015da5c6..374458989a 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/raft_cagra.yaml @@ -1,5 +1,6 @@ name: raft_cagra constraints: + build: raft-ann-bench.constraints.raft_cagra_build_constraints search: raft-ann-bench.constraints.raft_cagra_search_constraints groups: base: diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml index be63e1d535..188d24d20f 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/datasets.yaml @@ -112,7 +112,7 @@ groundtruth_neighbors_file: wiki_all_1M/groundtruth.1M.neighbors.ibin distance: euclidean -- name: wiki_all_10M, +- name: wiki_all_10M dims: 768 base_file: wiki_all_10M/base.10M.fbin query_file: wiki_all_10M/queries.fbin diff --git a/python/raft-dask/CMakeLists.txt b/python/raft-dask/CMakeLists.txt index d29997b4a3..f6ebd4c7ef 100644 --- a/python/raft-dask/CMakeLists.txt +++ b/python/raft-dask/CMakeLists.txt @@ -14,7 +14,7 @@ cmake_minimum_required(VERSION 3.26.4 FATAL_ERROR) -set(raft_dask_version 23.12.00) +set(raft_dask_version 24.02.00) include(../../fetch_rapids.cmake) include(rapids-cuda) diff --git a/python/raft-dask/pyproject.toml b/python/raft-dask/pyproject.toml index be030f839d..807c413160 100644 --- a/python/raft-dask/pyproject.toml +++ b/python/raft-dask/pyproject.toml @@ -34,13 +34,13 @@ authors = [ license = { text = "Apache 2.0" } requires-python = ">=3.9" dependencies = [ - "dask-cuda==23.12.*", + "dask-cuda==24.2.*", "joblib>=0.11", "numba>=0.57", "numpy>=1.21", - "pylibraft==23.12.*", - "rapids-dask-dependency==23.12.*", - "ucx-py==0.35.*", + "pylibraft==24.2.*", + "rapids-dask-dependency==24.2.*", + "ucx-py==0.36.*", ] # 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",