diff --git a/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml index bc86c5e86..59d471bda 100644 --- a/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-118_arch-aarch64.yaml @@ -19,6 +19,7 @@ dependencies: - cuda-version=11.8 - cudatoolkit - cupy>=12.0.0 +- cuvs==24.12.*,>=0.0.0a0 - cxx-compiler - cython>=3.0.0 - dlpack>=0.8,<1.0 @@ -33,6 +34,7 @@ dependencies: - libcusolver=11.4.1.48 - libcusparse-dev=11.7.5.86 - libcusparse=11.7.5.86 +- libcuvs==24.12.*,>=0.0.0a0 - librmm==24.12.*,>=0.0.0a0 - matplotlib - nccl>=2.19 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 607cbf8c5..31a416eb5 100644 --- a/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml @@ -19,6 +19,7 @@ dependencies: - cuda-version=11.8 - cudatoolkit - cupy>=12.0.0 +- cuvs==24.12.*,>=0.0.0a0 - cxx-compiler - cython>=3.0.0 - dlpack>=0.8,<1.0 @@ -33,6 +34,7 @@ dependencies: - libcusolver=11.4.1.48 - libcusparse-dev=11.7.5.86 - libcusparse=11.7.5.86 +- libcuvs==24.12.*,>=0.0.0a0 - librmm==24.12.*,>=0.0.0a0 - matplotlib - nccl>=2.19 diff --git a/conda/environments/bench_ann_cuda-125_arch-aarch64.yaml b/conda/environments/bench_ann_cuda-125_arch-aarch64.yaml index 6cea3684e..3efe9ebde 100644 --- a/conda/environments/bench_ann_cuda-125_arch-aarch64.yaml +++ b/conda/environments/bench_ann_cuda-125_arch-aarch64.yaml @@ -20,6 +20,7 @@ dependencies: - cuda-python>=12.0,<13.0a0,<=12.6.0 - cuda-version=12.5 - cupy>=12.0.0 +- cuvs==24.12.*,>=0.0.0a0 - cxx-compiler - cython>=3.0.0 - dlpack>=0.8,<1.0 @@ -30,6 +31,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev +- libcuvs==24.12.*,>=0.0.0a0 - librmm==24.12.*,>=0.0.0a0 - matplotlib - nccl>=2.19 diff --git a/conda/environments/bench_ann_cuda-125_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-125_arch-x86_64.yaml index ec0475402..7fbd77368 100644 --- a/conda/environments/bench_ann_cuda-125_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-125_arch-x86_64.yaml @@ -20,6 +20,7 @@ dependencies: - cuda-python>=12.0,<13.0a0,<=12.6.0 - cuda-version=12.5 - cupy>=12.0.0 +- cuvs==24.12.*,>=0.0.0a0 - cxx-compiler - cython>=3.0.0 - dlpack>=0.8,<1.0 @@ -30,6 +31,7 @@ dependencies: - libcurand-dev - libcusolver-dev - libcusparse-dev +- libcuvs==24.12.*,>=0.0.0a0 - librmm==24.12.*,>=0.0.0a0 - matplotlib - nccl>=2.19 diff --git a/conda/recipes/cuvs-bench/meta.yaml b/conda/recipes/cuvs-bench/meta.yaml index d85c12640..0681a1038 100644 --- a/conda/recipes/cuvs-bench/meta.yaml +++ b/conda/recipes/cuvs-bench/meta.yaml @@ -92,7 +92,7 @@ requirements: - libcublas {% endif %} - glog {{ glog_version }} - - libcuvs {{ version }} + - cuvs {{ version }} - h5py {{ h5py_version }} - matplotlib - pandas diff --git a/dependencies.yaml b/dependencies.yaml index 9977dcf46..98cac5300 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -476,11 +476,13 @@ dependencies: - h5py>=3.8.0 - benchmark>=1.8.2 - openblas + - libcuvs==24.12.*,>=0.0.0a0 bench_python: common: - output_types: [conda, pyproject, requirements] packages: - click + - cuvs==24.12.*,>=0.0.0a0 - matplotlib - pandas - pyyaml diff --git a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py index 43ce56092..88ec55dfa 100644 --- a/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py +++ b/python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py @@ -68,12 +68,12 @@ def import_with_fallback(primary_lib, secondary_lib=None, alias=None): xp = import_with_fallback("cupy", "numpy") rmm = import_with_fallback("rmm") +gpu_system = False -if rmm is not None: - gpu_system = True - from pylibraft.common import DeviceResources - from rmm.allocators.cupy import rmm_cupy_allocator -else: + +def force_fallback_to_numpy(): + global xp, gpu_system + xp = import_with_fallback("numpy") gpu_system = False warnings.warn( "Consider using a GPU-based system to greatly accelerate " @@ -81,6 +81,21 @@ def import_with_fallback(primary_lib, secondary_lib=None, alias=None): ) +if rmm is not None: + gpu_system = True + try: + from pylibraft.common import DeviceResources + from rmm.allocators.cupy import rmm_cupy_allocator + + from cuvs.neighbors.brute_force import build, search + except ImportError: + # RMM is available, cupy is available, but cuVS is not + force_fallback_to_numpy() +else: + # No RMM, no cuVS, but cupy is available + force_fallback_to_numpy() + + def generate_random_queries(n_queries, n_features, dtype=xp.float32): print("Generating random queries") if xp.issubdtype(dtype, xp.integer): @@ -173,8 +188,6 @@ def calc_truth(dataset, queries, k, metric="sqeuclidean"): queries = xp.asarray(queries, dtype=xp.float32) if gpu_system: - from cuvs.neighbors.brute_force import build, search - resources = DeviceResources() while i < n_samples: @@ -209,12 +222,15 @@ def calc_truth(dataset, queries, k, metric="sqeuclidean"): def main(): - if gpu_system: + if gpu_system and xp.__name__ == "cupy": pool = rmm.mr.PoolMemoryResource( rmm.mr.CudaMemoryResource(), initial_pool_size=2**30 ) rmm.mr.set_current_device_resource(pool) xp.cuda.set_allocator(rmm_cupy_allocator) + else: + # RMM is available, but cupy is not + force_fallback_to_numpy() parser = argparse.ArgumentParser( prog="generate_groundtruth", diff --git a/python/cuvs_bench/pyproject.toml b/python/cuvs_bench/pyproject.toml index 41ebad116..5b17f7228 100644 --- a/python/cuvs_bench/pyproject.toml +++ b/python/cuvs_bench/pyproject.toml @@ -19,6 +19,7 @@ license = { text = "Apache 2.0" } requires-python = ">=3.10" dependencies = [ "click", + "cuvs==24.12.*,>=0.0.0a0", "matplotlib", "pandas", "pyyaml",