Skip to content

Commit 7dfad5a

Browse files
committed
fix: add cuVS detection and C++ priority to backend selection
Add CUVS_AVAILABLE and CPP_CUVS_AVAILABLE flags to detect.py. Update get_optimal_backend() priority chain: C++ cuVS > Python cuVS > FAISS GPU > MPS > FAISS CPU > NumPy Signed-off-by: Maxime Grenu <maxime.grenu@gmail.com>
1 parent d144ef0 commit 7dfad5a

2 files changed

Lines changed: 39 additions & 2 deletions

File tree

python/zvec/backends/__init__.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33
from __future__ import annotations
44

55
from zvec.backends.detect import (
6+
CPP_CUVS_AVAILABLE,
7+
CUVS_AVAILABLE,
68
FAISS_AVAILABLE,
79
FAISS_CPU_AVAILABLE,
810
FAISS_GPU_AVAILABLE,
@@ -18,6 +20,8 @@
1820
)
1921

2022
__all__ = [
23+
"CPP_CUVS_AVAILABLE",
24+
"CUVS_AVAILABLE",
2125
"FAISS_AVAILABLE",
2226
"FAISS_CPU_AVAILABLE",
2327
"FAISS_GPU_AVAILABLE",

python/zvec/backends/detect.py

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,27 @@
7272
except ImportError:
7373
pass
7474

75+
# Try to detect cuVS (NVIDIA RAPIDS)
76+
CUVS_AVAILABLE = False
77+
try:
78+
import cuvs # noqa: F401
79+
80+
CUVS_AVAILABLE = True
81+
logger.info("cuVS (NVIDIA RAPIDS) available")
82+
except ImportError:
83+
pass
84+
85+
# Try to detect C++ cuVS bindings (via _zvec pybind11)
86+
CPP_CUVS_AVAILABLE = False
87+
try:
88+
import _zvec
89+
90+
CPP_CUVS_AVAILABLE = hasattr(_zvec, "create_cagra_float")
91+
if CPP_CUVS_AVAILABLE:
92+
logger.info("C++ cuVS bindings available (preferred path)")
93+
except ImportError:
94+
pass
95+
7596

7697
def get_available_backends() -> dict[str, bool]:
7798
"""Return a dictionary of available backends.
@@ -80,6 +101,8 @@ def get_available_backends() -> dict[str, bool]:
80101
Dictionary with backend availability information.
81102
"""
82103
return {
104+
"cpp_cuvs": CPP_CUVS_AVAILABLE,
105+
"cuvs": CUVS_AVAILABLE,
83106
"faiss": FAISS_AVAILABLE,
84107
"faiss_gpu": FAISS_GPU_AVAILABLE,
85108
"faiss_cpu": FAISS_CPU_AVAILABLE,
@@ -93,9 +116,19 @@ def get_available_backends() -> dict[str, bool]:
93116
def get_optimal_backend() -> str:
94117
"""Determine the optimal backend for the current system.
95118
119+
Priority: C++ cuVS > Python cuVS > FAISS GPU > MPS > FAISS CPU > NumPy.
120+
96121
Returns:
97-
Name of the optimal backend: "faiss_gpu", "faiss_cpu", or "numpy".
122+
Name of the optimal backend.
98123
"""
124+
if CPP_CUVS_AVAILABLE:
125+
logger.info("Using C++ cuVS backend (native, preferred)")
126+
return "cpp_cuvs"
127+
128+
if CUVS_AVAILABLE:
129+
logger.info("Using Python cuVS backend")
130+
return "cuvs"
131+
99132
if FAISS_GPU_AVAILABLE and NVIDIA_GPU_DETECTED:
100133
logger.info("Using FAISS GPU backend")
101134
return "faiss_gpu"
@@ -118,7 +151,7 @@ def is_gpu_available() -> bool:
118151
Returns:
119152
True if GPU acceleration is available.
120153
"""
121-
return FAISS_GPU_AVAILABLE or MPS_AVAILABLE
154+
return CPP_CUVS_AVAILABLE or CUVS_AVAILABLE or FAISS_GPU_AVAILABLE or MPS_AVAILABLE
122155

123156

124157
def get_backend_info() -> dict:

0 commit comments

Comments
 (0)