finalize dependencies, verify tests locally, fix script

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

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

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

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

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

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

python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py

@@ -68,19 +68,34 @@ 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 "
         " generating groundtruths using cuVS."
     )
 
 
+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",

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",