Repository: cupy/cupy
Pull request: Updating pylibraft `pairwise_distance` to cuvs
This PR updates cupy so that it can make use of the cuvs library for pairwise distance, which the vector search APIs from pylibraft have recently moved. This PR maintains the backwards compatibility with pylibraft while will exposing double precision float support. 

Closes https://github.com/rapidsai/cuvs/issues/542

Problem statement:
[BUG] Update cuPy to use `cuvs.distance`
cuPy is currently using pylibraft for cuvs.distance, but the distance APIs were recently removed from `pylibraft`, which means users who have installed the latest `pylibraft` get an error. We need to update `cupy` to use `cuvs.distance`.

Hints from issue discussion:
cc @kmaehashi for vis

Commit message: Fixing style

Patch to review:
diff --git a/.pfnci/linux/tests/cuda-rapids.sh b/.pfnci/linux/tests/cuda-rapids.sh
index 80eb9c53f39..dca5dcffa72 100644
--- a/.pfnci/linux/tests/cuda-rapids.sh
+++ b/.pfnci/linux/tests/cuda-rapids.sh
@@ -14,6 +14,6 @@ export NVCC="ccache nvcc"
 CXX=/usr/lib/ccache/g++ CUDA_PATH="/usr/local/cuda" pip install -v ".[test]"
 "$ACTIONS/unittest.sh" "not slow" \
     "cupyx_tests/scipy_tests/sparse_tests/csgraph_tests" \
-    "cupyx_tests/scipy_tests/spatial_tests/test_kdtree_pylibraft.py" \
+    "cupyx_tests/scipy_tests/spatial_tests/test_kdtree_cuvs.py" \
     "cupyx_tests/scipy_tests/spatial_tests/test_distance.py"
 "$ACTIONS/cleanup.sh"
diff --git a/cupyx/scipy/spatial/distance.py b/cupyx/scipy/spatial/distance.py
index b5352000f08..425f8b09274 100644
--- a/cupyx/scipy/spatial/distance.py
+++ b/cupyx/scipy/spatial/distance.py
@@ -1,9 +1,18 @@
 import cupy
+
+cuvs_available = False
+pylibraft_available = False
 try:
-    from pylibraft.distance import pairwise_distance
-    pylibraft_available = True
-except ModuleNotFoundError:
-    pylibraft_available = False
+    from cuvs.distance import pairwise_distance
+    cuvs_available = True
+except ImportError:
+    try:
+        # cuVS distance primitives were previously in pylibraft
+        from pylibraft.distance import pairwise_distance
+        pylibraft_available = True
+    except ImportError:
+        cuvs_available = False
+        pylibraft_available = False
 
 
 def _convert_to_type(X, out_type):
@@ -102,6 +111,13 @@ def __init__(self, canonical_name=None, aka=None,
 _METRICS_NAMES = list(_METRICS.keys())
 
 
+def check_soft_dependencies():
+    if not cuvs_available:
+        if not pylibraft_available:
+            raise RuntimeError('cuVS >= 24.12 or pylibraft < '
+                               '24.12 should be installed to use this feature')
+
+
 def minkowski(u, v, p):
     """Compute the Minkowski distance between two 1-D arrays.
 
@@ -116,8 +132,9 @@ def minkowski(u, v, p):
     Returns:
         minkowski (double): The Minkowski distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -141,8 +158,8 @@ def canberra(u, v):
     Returns:
         canberra (double): The Canberra distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -166,8 +183,8 @@ def chebyshev(u, v):
     Returns:
         chebyshev (double): The Chebyshev distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -192,8 +209,8 @@ def cityblock(u, v):
         cityblock (double): The City Block distance between
         vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -222,8 +239,8 @@ def correlation(u, v):
         correlation (double): The correlation distance between
         vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -249,8 +266,8 @@ def cosine(u, v):
     Returns:
         cosine (double): The Cosine distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -278,8 +295,8 @@ def hamming(u, v):
     Returns:
         hamming (double): The Hamming distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -303,8 +320,8 @@ def euclidean(u, v):
     Returns:
         euclidean (double): The Euclidean distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -332,8 +349,8 @@ def jensenshannon(u, v):
         jensenshannon (double): The Jensen-Shannon distance between
         vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -361,8 +378,8 @@ def russellrao(u, v):
     Returns:
         hamming (double): The Hamming distance between vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -387,8 +404,8 @@ def sqeuclidean(u, v):
         sqeuclidean (double): The squared Euclidean distance between
         vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed.')
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -414,8 +431,8 @@ def hellinger(u, v):
         hellinger (double): The Hellinger distance between
         vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -440,8 +457,8 @@ def kl_divergence(u, v):
         kl_divergence (double): The Kullback-Leibler divergence between
         vectors `u` and `v`.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
+    check_soft_dependencies()
+
     u = cupy.asarray(u)
     v = cupy.asarray(v)
     output_arr = cupy.zeros((1,), dtype=u.dtype)
@@ -479,10 +496,15 @@ def cdist(XA, XB, metric='euclidean', out=None, **kwargs):
             ``dist(u=XA[i], v=XB[j])`` is computed and stored in the
             :math:`ij` th entry.
     """
-    if not pylibraft_available:
-        raise RuntimeError('pylibraft is not installed')
-    XA = cupy.asarray(XA, dtype='float32')
-    XB = cupy.asarray(XB, dtype='float32')
+    check_soft_dependencies()
+
+    if pylibraft_available or \
+            (cuvs_available and XA.dtype not in ['float32', 'float64']):
+        XA = cupy.asarray(XA, dtype='float32')
+
+    if pylibraft_available or \
+            (cuvs_available and XB.dtype not in ['float32', 'float64']):
+        XB = cupy.asarray(XB, dtype='float32')
 
     s = XA.shape
     sB = XB.shape
@@ -501,7 +523,8 @@ def cdist(XA, XB, metric='euclidean', out=None, **kwargs):
     p = kwargs["p"] if "p" in kwargs else 2.0
 
     if out is not None:
-        if out.dtype != 'float32':
+        if (pylibraft_available and out.dtype != 'float32') or \
+                (cuvs_available and out.dtype not in ['float32', 'float64']):
             out = out.astype('float32', copy=False)
         if out.shape != (mA, mB):
             cupy.resize(out, (mA, mB))
diff --git a/docs/source/reference/scipy_spatial.rst b/docs/source/reference/scipy_spatial.rst
index e34ead0915d..7b38ee923a4 100644
--- a/docs/source/reference/scipy_spatial.rst
+++ b/docs/source/reference/scipy_spatial.rst
@@ -7,8 +7,8 @@ Spatial algorithms and data structures  (:mod:`cupyx.scipy.spatial`)
 
 .. note::
 
-   The ``spatial`` module uses ``pylibraft`` as a backend.
-   You need to install `pylibraft package <https://anaconda.org/rapidsai/pylibraft>` from ``rapidsai`` Conda channel to use features listed on this page.
+   The ``spatial`` module uses the ``cuVS`` library as a backend.
+   You need to install `cuVS package <https://anaconda.org/rapidsai/cuvs>` from ``rapidsai`` Conda channel to use features listed on this page.
 
 .. note::
    Currently, the ``spatial`` module is not supported on AMD ROCm platforms.
diff --git a/docs/source/reference/scipy_spatial_distance.rst b/docs/source/reference/scipy_spatial_distance.rst
index aa7b696bcd3..5d119036489 100644
--- a/docs/source/reference/scipy_spatial_distance.rst
+++ b/docs/source/reference/scipy_spatial_distance.rst
@@ -5,8 +5,8 @@ Distance computations (:mod:`cupyx.scipy.spatial.distance`)
 
 .. note::
 
-   The ``distance`` module uses ``pylibraft`` as a backend.
-   You need to install `pylibraft package <https://anaconda.org/rapidsai/pylibraft>` from ``rapidsai`` Conda channel to use features listed on this page.
+   The ``distance`` module uses the ``cuVS`` library as a backend.
+   You need to install `cuVS package <https://anaconda.org/rapidsai/cuvs>` from ``rapidsai`` Conda channel to use features listed on this page.
 
 .. note::
    Currently, the ``distance`` module is not supported on AMD ROCm platforms.
diff --git a/tests/cupyx_tests/scipy_tests/spatial_tests/test_distance.py b/tests/cupyx_tests/scipy_tests/spatial_tests/test_distance.py
index 941e73eb609..be3af3fd1bc 100644
--- a/tests/cupyx_tests/scipy_tests/spatial_tests/test_distance.py
+++ b/tests/cupyx_tests/scipy_tests/spatial_tests/test_distance.py
@@ -11,10 +11,10 @@
     scipy_available = False
 import cupyx.scipy.spatial.distance  # NOQA
 try:
-    import pylibraft  # NOQA
-    pylibraft_available = True
+    import cuvs  # NOQA
+    cuvs_available = True
 except ModuleNotFoundError:
-    pylibraft_available = False
+    cuvs_available = False
 from cupy import testing
 
 
@@ -30,7 +30,7 @@
     'order': ["C", "F"]
 }))
 @pytest.mark.skipif(cupy.cuda.runtime.is_hip, reason="tests for CUDA only")
-@pytest.mark.skipif(not scipy_available or not pylibraft_available,
+@pytest.mark.skipif(not scipy_available or not cuvs_available,
                     reason='requires scipy and pylibraft')
 class TestCdist(unittest.TestCase):
 
@@ -92,7 +92,7 @@ def test_cdist_(self, xp, scp):
     'order': ["C", "F"]
 }))
 @pytest.mark.skipif(cupy.cuda.runtime.is_hip, reason="tests for CUDA only")
-@pytest.mark.skipif(not scipy_available or not pylibraft_available,
+@pytest.mark.skipif(not scipy_available or not cuvs_available,
                     reason='requires scipy and pylibraft')
 class TestPdist:
 
@@ -118,7 +118,7 @@ def test_pdist_(self, xp, scp):
     'order': ["C", "F"]
 }))
 @pytest.mark.skipif(cupy.cuda.runtime.is_hip, reason="tests for CUDA only")
-@pytest.mark.skipif(not scipy_available or not pylibraft_available,
+@pytest.mark.skipif(not scipy_available or not cuvs_available,
                     reason='requires scipy and pylibraft')
 class TestDistanceMatrix(unittest.TestCase):
 
@@ -143,7 +143,7 @@ def test_distance_matrix_(self, xp, scp):
     'order': ["C", "F"]
 }))
 @pytest.mark.skipif(cupy.cuda.runtime.is_hip, reason="tests for CUDA only")
-@pytest.mark.skipif(not scipy_available or not pylibraft_available,
+@pytest.mark.skipif(not scipy_available or not cuvs_available,
                     reason='requires scipy and pylibraft')
 class TestDistanceFunction(unittest.TestCase):
 
diff --git a/tests/cupyx_tests/scipy_tests/spatial_tests/test_kdtree_pylibraft.py b/tests/cupyx_tests/scipy_tests/spatial_tests/test_kdtree_cuvs.py
similarity index 85%
rename from tests/cupyx_tests/scipy_tests/spatial_tests/test_kdtree_pylibraft.py
rename to tests/cupyx_tests/scipy_tests/spatial_tests/test_kdtree_cuvs.py
index 8e92e1ec882..508e32c325d 100644
--- a/tests/cupyx_tests/scipy_tests/spatial_tests/test_kdtree_pylibraft.py
+++ b/tests/cupyx_tests/scipy_tests/spatial_tests/test_kdtree_cuvs.py
@@ -3,10 +3,10 @@
 from cupy import testing
 
 try:
-    import pylibraft  # NOQA
-    pylibraft_available = True
+    import cuvs  # NOQA
+    cuvs_available = True
 except ModuleNotFoundError:
-    pylibraft_available = False
+    cuvs_available = False
 
 
 def create_random_kd_tree(xp, scp, n, m, n_points=1, x_offset=0, scale=10,
@@ -19,7 +19,7 @@ def create_random_kd_tree(xp, scp, n, m, n_points=1, x_offset=0, scale=10,
 
 
 @testing.with_requires('scipy')
-@pytest.mark.skipif(not pylibraft_available, reason='pylibraft is required')
+@pytest.mark.skipif(not cuvs_available, reason='pylibraft is required')
 class TestDistance:
     @pytest.mark.parametrize('args', [
         (100, 4, 1, 0),


Code review comments:
diff_hunk: @@ -19,7 +19,7 @@ def create_random_kd_tree(xp, scp, n, m, n_points=1, x_offset=0, scale=10,
 
 
 @testing.with_requires('scipy')
-@pytest.mark.skipif(not pylibraft_available, reason='pylibraft is required')
+@pytest.mark.skipif(not cuvs_available, reason='pylibraft is required')
path: tests/cupyx_tests/scipy_tests/spatial_tests/test_kdtree_cuvs.py
text: @user2:
```suggestion
@user1.mark.skipif(not cuvs_available, reason='cuVS is required')
```