Triage NN Descent failing CI wheel tests

cpp/src/neighbors/detail/nn_descent.cuh

@@ -16,22 +16,27 @@
 
 #pragma once
 
+#include "cagra/device_common.hpp"
+
 #include <cuvs/neighbors/nn_descent.hpp>
 
-#include "ann_utils.cuh"
-#include "cagra/device_common.hpp"
 #include <raft/core/device_mdarray.hpp>
+#include <raft/core/device_mdspan.hpp>
 #include <raft/core/error.hpp>
 #include <raft/core/host_mdarray.hpp>
+#include <raft/core/mdspan.hpp>
+#include <raft/core/operators.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resources.hpp>
-
+#include <raft/spatial/knn/detail/ann_utils.cuh>
 #include <raft/util/arch.cuh>  // raft::util::arch::SM_*
 #include <raft/util/cuda_dev_essentials.cuh>
 #include <raft/util/cuda_rt_essentials.hpp>
 #include <raft/util/cudart_utils.hpp>
 #include <raft/util/pow2_utils.cuh>
 
+#include <rmm/device_uvector.hpp>
+
 #include <cub/cub.cuh>
 #include <cuda_runtime.h>
 #include <thrust/execution_policy.h>
@@ -44,12 +49,13 @@
 #include <omp.h>
 
 #include <limits>
+#include <optional>
 #include <queue>
 #include <random>
 
 namespace cuvs::neighbors::nn_descent::detail {
-static const std::string RAFT_NAME = "raft";
-using pinned_memory_resource       = thrust::universal_host_pinned_memory_resource;
+
+using pinned_memory_resource = thrust::universal_host_pinned_memory_resource;
 template <typename T>
 using pinned_memory_allocator = thrust::mr::stateless_resource_allocator<T, pinned_memory_resource>;
 
@@ -1217,13 +1223,14 @@ void GNND<Data_t, Index_t>::build(Data_t* data, const Index_t nrow, Index_t* out
 
   cudaStream_t stream = raft::resource::get_cuda_stream(res);
   nrow_               = nrow;
+  graph_.nrow         = nrow;
   graph_.h_graph      = (InternalID_t<Index_t>*)output_graph;
 
   cudaPointerAttributes data_ptr_attr;
   RAFT_CUDA_TRY(cudaPointerGetAttributes(&data_ptr_attr, data));
   size_t batch_size = (data_ptr_attr.devicePointer == nullptr) ? 100000 : nrow_;
 
-  cuvs::spatial::knn::detail::utils::batch_load_iterator vec_batches{
+  raft::spatial::knn::detail::utils::batch_load_iterator vec_batches{
     data, static_cast<size_t>(nrow_), build_config_.dataset_dim, batch_size, stream};
   for (auto const& batch : vec_batches) {
     preprocess_data_kernel<<<batch.size(),
@@ -1452,4 +1459,4 @@ index<IdxT> build(
   return idx;
 }
 
-}  // namespace  cuvs::neighbors::nn_descent::detail
+}  // namespace cuvs::neighbors::nn_descent::detail

python/cuvs/cuvs/test/ann_utils.py

@@ -16,6 +16,7 @@
 
 
 def generate_data(shape, dtype):
+    np.random.seed(0)
     if dtype == np.byte:
         x = np.random.randint(-127, 128, size=shape, dtype=np.byte)
     elif dtype == np.ubyte:

python/cuvs/cuvs/test/test_cagra.py

@@ -19,6 +19,8 @@
 from sklearn.neighbors import NearestNeighbors
 from sklearn.preprocessing import normalize
 
+import rmm
+
 from cuvs.neighbors import cagra
 from cuvs.test.ann_utils import calc_recall, generate_data
 
@@ -37,85 +39,91 @@ def run_cagra_build_search_test(
     compare=True,
     inplace=True,
     add_data_on_build=True,
-    search_params={},
+    search_params1={},
     compression=None,
 ):
-    dataset = generate_data((n_rows, n_cols), dtype)
-    if metric == "inner_product":
-        dataset = normalize(dataset, norm="l2", axis=1)
-    dataset_device = device_ndarray(dataset)
-
-    build_params = cagra.IndexParams(
-        metric=metric,
-        intermediate_graph_degree=intermediate_graph_degree,
-        graph_degree=graph_degree,
-        build_algo=build_algo,
-        compression=compression,
-    )
+    for i in range(2):
+        rmm.reinitialize(
+            managed_memory=False,
+            pool_allocator=False,
+            initial_pool_size=2 << 30,
+        )
+        dataset = generate_data((n_rows, n_cols), dtype)
+        if metric == "inner_product":
+            dataset = normalize(dataset, norm="l2", axis=1)
+        dataset_device = device_ndarray(dataset)
+
+        build_params = cagra.IndexParams(
+            metric=metric,
+            intermediate_graph_degree=intermediate_graph_degree,
+            graph_degree=graph_degree,
+            build_algo=build_algo,
+            compression=compression,
+        )
 
-    if array_type == "device":
-        index = cagra.build(build_params, dataset_device)
-    else:
-        index = cagra.build(build_params, dataset)
-
-    if not add_data_on_build:
-        dataset_1 = dataset[: n_rows // 2, :]
-        dataset_2 = dataset[n_rows // 2 :, :]
-        indices_1 = np.arange(n_rows // 2, dtype=np.uint32)
-        indices_2 = np.arange(n_rows // 2, n_rows, dtype=np.uint32)
         if array_type == "device":
-            dataset_1_device = device_ndarray(dataset_1)
-            dataset_2_device = device_ndarray(dataset_2)
-            indices_1_device = device_ndarray(indices_1)
-            indices_2_device = device_ndarray(indices_2)
-            index = cagra.extend(index, dataset_1_device, indices_1_device)
-            index = cagra.extend(index, dataset_2_device, indices_2_device)
+            index = cagra.build(build_params, dataset_device)
         else:
-            index = cagra.extend(index, dataset_1, indices_1)
-            index = cagra.extend(index, dataset_2, indices_2)
-
-    queries = generate_data((n_queries, n_cols), dtype)
-    out_idx = np.zeros((n_queries, k), dtype=np.uint32)
-    out_dist = np.zeros((n_queries, k), dtype=np.float32)
-
-    queries_device = device_ndarray(queries)
-    out_idx_device = device_ndarray(out_idx) if inplace else None
-    out_dist_device = device_ndarray(out_dist) if inplace else None
-
-    search_params = cagra.SearchParams(**search_params)
-
-    ret_output = cagra.search(
-        search_params,
-        index,
-        queries_device,
-        k,
-        neighbors=out_idx_device,
-        distances=out_dist_device,
-    )
-
-    if not inplace:
-        out_dist_device, out_idx_device = ret_output
-
-    if not compare:
-        return
-
-    out_idx = out_idx_device.copy_to_host()
-    out_dist = out_dist_device.copy_to_host()
-
-    # Calculate reference values with sklearn
-    skl_metric = {
-        "sqeuclidean": "sqeuclidean",
-        "inner_product": "cosine",
-        "euclidean": "euclidean",
-    }[metric]
-    nn_skl = NearestNeighbors(
-        n_neighbors=k, algorithm="brute", metric=skl_metric
-    )
-    nn_skl.fit(dataset)
-    skl_idx = nn_skl.kneighbors(queries, return_distance=False)
-
-    recall = calc_recall(out_idx, skl_idx)
-    assert recall > 0.7
+            index = cagra.build(build_params, dataset)
+
+            # if not add_data_on_build:
+            dataset_1 = dataset[: n_rows // 2, :]
+            dataset_2 = dataset[n_rows // 2 :, :]
+            indices_1 = np.arange(n_rows // 2, dtype=np.uint32)
+            indices_2 = np.arange(n_rows // 2, n_rows, dtype=np.uint32)
+            if array_type == "device":
+                dataset_1_device = device_ndarray(dataset_1)
+                dataset_2_device = device_ndarray(dataset_2)
+                indices_1_device = device_ndarray(indices_1)
+                indices_2_device = device_ndarray(indices_2)
+                index = cagra.extend(index, dataset_1_device, indices_1_device)
+                index = cagra.extend(index, dataset_2_device, indices_2_device)
+            else:
+                index = cagra.extend(index, dataset_1, indices_1)
+                index = cagra.extend(index, dataset_2, indices_2)
+
+        queries = generate_data((n_queries, n_cols), dtype)
+        out_idx = np.zeros((n_queries, k), dtype=np.uint32)
+        out_dist = np.zeros((n_queries, k), dtype=np.float32)
+
+        queries_device = device_ndarray(queries)
+        out_idx_device = device_ndarray(out_idx) if inplace else None
+        out_dist_device = device_ndarray(out_dist) if inplace else None
+
+        search_params = cagra.SearchParams(**search_params1)
+
+        ret_output = cagra.search(
+            search_params,
+            index,
+            queries_device,
+            k,
+            neighbors=out_idx_device,
+            distances=out_dist_device,
+        )
+
+        if not inplace:
+            out_dist_device, out_idx_device = ret_output
+
+        if not compare:
+            return
+
+        out_idx = out_idx_device.copy_to_host()
+        out_dist = out_dist_device.copy_to_host()
+
+        # Calculate reference values with sklearn
+        skl_metric = {
+            "sqeuclidean": "sqeuclidean",
+            "inner_product": "cosine",
+            "euclidean": "euclidean",
+        }[metric]
+        nn_skl = NearestNeighbors(
+            n_neighbors=k, algorithm="brute", metric=skl_metric
+        )
+        nn_skl.fit(dataset)
+        skl_idx = nn_skl.kneighbors(queries, return_distance=False)
+
+        recall = calc_recall(out_idx, skl_idx)
+        assert recall > 0.7
 
 
 @pytest.mark.parametrize("inplace", [True, False])

python/cuvs/cuvs/test/test_hnsw.py

@@ -13,27 +13,36 @@
 # limitations under the License.
 #
 
+import cupy
 import numpy as np
 import pytest
+from ann_utils import calc_recall, generate_data
 from sklearn.neighbors import NearestNeighbors
 from sklearn.preprocessing import normalize
 
 from cuvs.neighbors import cagra, hnsw
-from cuvs.test.ann_utils import calc_recall, generate_data
+
+# import rmm
 
 
 def run_hnsw_build_search_test(
-    n_rows=1000,
-    n_cols=10,
-    n_queries=100,
+    n_rows=10000,
+    n_cols=1,
+    n_queries=5,
     k=10,
     dtype=np.float32,
     metric="sqeuclidean",
     build_algo="ivf_pq",
     intermediate_graph_degree=128,
     graph_degree=64,
-    search_params={},
+    search_params1={},
 ):
+    # for i in range(2):
+    # rmm.reinitialize(
+    #     managed_memory=False,
+    #     pool_allocator=False,
+    #     initial_pool_size=2 << 30,
+    # )
     dataset = generate_data((n_rows, n_cols), dtype)
     if metric == "inner_product":
         dataset = normalize(dataset, norm="l2", axis=1)
@@ -50,18 +59,20 @@ def run_hnsw_build_search_test(
         graph_degree=graph_degree,
         build_algo=build_algo,
     )
-
     index = cagra.build(build_params, dataset)
 
     assert index.trained
 
     hnsw_index = hnsw.from_cagra(index)
 
     queries = generate_data((n_queries, n_cols), dtype)
+    print("queries", queries)
 
-    search_params = hnsw.SearchParams(**search_params)
+    search_params = hnsw.SearchParams(**search_params1)
 
     out_dist, out_idx = hnsw.search(search_params, hnsw_index, queries, k)
+    print("out_idx", cupy.asarray(out_idx))
+    print("out_dist", cupy.asarray(out_dist))
 
     # Calculate reference values with sklearn
     skl_metric = {
@@ -74,13 +85,17 @@ def run_hnsw_build_search_test(
     )
     nn_skl.fit(dataset)
     skl_dist, skl_idx = nn_skl.kneighbors(queries, return_distance=True)
+    print("skl_idx", skl_idx)
+    print("skl_dist", skl_dist)
 
     recall = calc_recall(out_idx, skl_idx)
+    print("recall", recall)
     assert recall > 0.95
+    del index
 
 
-@pytest.mark.parametrize("dtype", [np.float32, np.int8, np.uint8])
-@pytest.mark.parametrize("k", [10, 20])
+@pytest.mark.parametrize("dtype", [np.float32])
+@pytest.mark.parametrize("k", [3])
 @pytest.mark.parametrize("ef", [30, 40])
 @pytest.mark.parametrize("num_threads", [2, 4])
 @pytest.mark.parametrize("metric", ["sqeuclidean"])
@@ -93,5 +108,5 @@ def test_hnsw(dtype, k, ef, num_threads, metric, build_algo):
         k=k,
         metric=metric,
         build_algo=build_algo,
-        search_params={"ef": ef, "num_threads": num_threads},
+        search_params1={"ef": ef, "num_threads": num_threads},
     )