Integrate cuvs kmeans into training/assignments for ivf

python/python/lance/cuvs/kmeans.py

@@ -6,9 +6,11 @@
 
 from lance.dependencies import numpy as np
 from lance.dependencies import torch
-from lance.dependencies import cagra
-from lance.dependencies import device_ndarray
+#from lance.dependencies import cagra
+#from lance.dependencies import device_ndarray
 from lance.torch.kmeans import KMeans as KMeansTorch
+from pylibraft.common import device_ndarray
+from cuvs.neighbors import cagra
 
 __all__ = ["KMeans"]
 
@@ -70,7 +72,7 @@ def __init__(
             device=device,
         )
 
-        if self.device.type != "cuda":
+        if self.device.type != "cuda" or not torch.cuda.is_available():
             raise ValueError("KMeans::__init__: cuda is not enabled/available")
 
         self.itopk_size = itopk_size

python/python/lance/dependencies.py

@@ -168,7 +168,7 @@ def _lazy_import(module_name: str) -> tuple[ModuleType, bool]:
     import numpy
     import pandas
     import polars
-    import pylibraft.device_ndarray
+    import pylibraft.common.device_ndarray
     import ray
     import tensorflow
     import torch
@@ -179,7 +179,7 @@ def _lazy_import(module_name: str) -> tuple[ModuleType, bool]:
     polars, _POLARS_AVAILABLE = _lazy_import("polars")
     torch, _TORCH_AVAILABLE = _lazy_import("torch")
     cagra, _CAGRA_AVAILABLE = _lazy_import("cuvs.neighbors.cagra")
-    device_ndarray, _DEVICE_NDARRAY_AVAILABLE = _lazy_import("pylibraft.device_ndarray")
+    device_ndarray, _DEVICE_NDARRAY_AVAILABLE = _lazy_import("pylibraft.common.device_ndarray")
     datasets, _HUGGING_FACE_AVAILABLE = _lazy_import("datasets")
     tensorflow, _TENSORFLOW_AVAILABLE = _lazy_import("tensorflow")
     ray, _RAY_AVAILABLE = _lazy_import("ray")

python/python/lance/vector.py

@@ -139,7 +139,11 @@ def train_ivf_centroids_on_accelerator(
 ) -> (np.ndarray, str):
     """Use accelerator (GPU or MPS) to train kmeans."""
     if isinstance(accelerator, str) and (
-        not (CUDA_REGEX.match(accelerator) or accelerator == "mps")
+        not (
+            CUDA_REGEX.match(accelerator)
+            or accelerator == "mps"
+            or accelerator == "cuvs"
+        )
     ):
         raise ValueError(
             "Train ivf centroids on accelerator: "
@@ -168,14 +172,27 @@ def train_ivf_centroids_on_accelerator(
         cache=True,
     )
 
-    logging.info("Training IVF partitions using GPU(%s)", accelerator)
-    kmeans = KMeans(
-        k,
-        max_iters=max_iters,
-        metric=metric_type,
-        device=accelerator,
-        centroids=init_centroids,
-    )
+    if accelerator == "cuvs":
+        logging.info("Training IVF partitions using cuVS+GPU")
+        print("Training IVF partitions using cuVS+GPU")
+        from lance.cuvs.kmeans import KMeans as KMeansCuVS
+
+        kmeans = KMeansCuVS(
+            k,
+            max_iters=max_iters,
+            metric=metric_type,
+            device="cuda",
+            centroids=init_centroids,
+        )
+    else:
+        logging.info("Training IVF partitions using GPU(%s)", accelerator)
+        kmeans = KMeans(
+            k,
+            max_iters=max_iters,
+            metric=metric_type,
+            device=accelerator,
+            centroids=init_centroids,
+        )
     kmeans.fit(ds)
 
     centroids = kmeans.centroids.cpu().numpy()