ENH Re-enable wheels CI

rapidsai · Mar 1, 2024 · 2b3500c · 2b3500c
1 parent 3136d55
commit 2b3500c
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 40 deletions.
diff --git a/ci/build_wheel_cuvs.sh b/ci/build_wheel_cuvs.sh
@@ -6,4 +6,4 @@ set -euo pipefail
 # Set up skbuild options. Enable sccache in skbuild config options
 export SKBUILD_CONFIGURE_OPTIONS="-DCUVS_BUILD_WHEELS=ON -DDETECT_CONDA_ENV=OFF -DFIND_CUVS_CPP=OFF"
 
-#ci/build_wheel.sh cuvs python/cuvs
+ci/build_wheel.sh cuvs python/cuvs
diff --git a/ci/test_wheel_cuvs.sh b/ci/test_wheel_cuvs.sh
@@ -4,15 +4,15 @@
 set -euo pipefail
 
 mkdir -p ./dist
-#RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
-#RAPIDS_PY_WHEEL_NAME="cuvs_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist
+RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen ${RAPIDS_CUDA_VERSION})"
+RAPIDS_PY_WHEEL_NAME="cuvs_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./dist
 
-## echo to expand wildcard before adding `[extra]` requires for pip
-#python -m pip install $(echo ./dist/cuvs*.whl)[test]
-#
-## Run smoke tests for aarch64 pull requests
-#if [[ "$(arch)" == "aarch64" && "${RAPIDS_BUILD_TYPE}" == "pull-request" ]]; then
-#    python ./ci/wheel_smoke_test_cuvs.py
-#else
-#    python -m pytest ./python/cuvs/cuvs/test
-#fi
+# echo to expand wildcard before adding `[extra]` requires for pip
+python -m pip install $(echo ./dist/cuvs*.whl)[test]
+
+# Run smoke tests for aarch64 pull requests
+if [[ "$(arch)" == "aarch64" && "${RAPIDS_BUILD_TYPE}" == "pull-request" ]]; then
+   python ./ci/wheel_smoke_test_cuvs.py
+else
+   python -m pytest ./python/cuvs/cuvs/test
+fi
diff --git a/ci/wheel_smoke_test_cuvs.py b/ci/wheel_smoke_test_cuvs.py
@@ -13,41 +13,30 @@
 # limitations under the License.
 #
 
+import cupy as cp
 import numpy as np
-from scipy.spatial.distance import cdist
 
-from pylibrat.common import Handle, Stream, device_ndarray
-from cuvs.distance import pairwise_distance
+from cuvs.neighbors import cagra
+from pylibraft.common import Stream, DeviceResources
 
 
 if __name__ == "__main__":
-    metric = "euclidean"
-    n_rows = 1337
-    n_cols = 1337
+    n_samples = 1000
+    n_features = 50
+    n_queries = 1000
+    k = 10
 
-    input1 = np.random.random_sample((n_rows, n_cols))
-    input1 = np.asarray(input1, order="C").astype(np.float64)
+    dataset = cp.random.random_sample((n_samples,
+                                       n_features)).astype(cp.float32)
 
-    output = np.zeros((n_rows, n_rows), dtype=np.float64)
+    build_params = cagra.IndexParams(metric="sqeuclidean",
+                                     build_algo="nn_descent")
 
-    expected = cdist(input1, input1, metric)
+    index = cagra.build_index(build_params, dataset)
 
-    expected[expected <= 1e-5] = 0.0
+    distances, neighbors = cagra.search(cagra.SearchParams(),
+                                          index, dataset,
+                                          k)
 
-    input1_device = device_ndarray(input1)
-    output_device = None
-
-    s2 = Stream()
-    handle = Handle(stream=s2)
-    ret_output = pairwise_distance(
-        input1_device, input1_device, output_device, metric, handle=handle
-    )
-    handle.sync()
-
-    output_device = ret_output
-
-    actual = output_device.copy_to_host()
-
-    actual[actual <= 1e-5] = 0.0
-
-    assert np.allclose(expected, actual, rtol=1e-4)
+    distances = cp.asarray(distances)
+    neighbors = cp.asarray(neighbors)