fix(python): clarify unavailable CUDA index acceleration

This commit is contained in:
Gatefixer
2026-09-23 00:54:11 +00:00
parent dd2539c2ca
commit bd83c6568a
2 changed files with 79 additions and 3 deletions
+33 -2
View File
@@ -119,6 +119,33 @@ def _should_push_down_query_table(
return namespace_client is not None and "QueryTable" in pushdown_operations
_CUDA_UNAVAILABLE_MESSAGE = (
"CUDA acceleration is unavailable in this Python environment. Install a "
"CUDA-enabled PyTorch build and verify that the CUDA device is accessible, or "
"omit `accelerator` (or set `accelerator=None`) to create the index on CPU."
)
def _validate_index_accelerator(accelerator: Any) -> None:
if not isinstance(accelerator, str) or accelerator.split(":", 1)[0] != "cuda":
return
try:
import torch
cuda_available = torch.cuda.is_available()
except (ImportError, RuntimeError, AssertionError) as err:
raise RuntimeError(_CUDA_UNAVAILABLE_MESSAGE) from err
if not cuda_available:
raise RuntimeError(_CUDA_UNAVAILABLE_MESSAGE)
try:
torch.cuda.init()
except (RuntimeError, AssertionError) as err:
raise RuntimeError(_CUDA_UNAVAILABLE_MESSAGE) from err
def _polars_predicate_pushdown_barrier(frame: Any) -> Any:
"""Return a Polars frame unchanged while blocking predicate pushdown."""
return frame
@@ -3287,7 +3314,9 @@ class LanceTable(Table):
# Handle accelerator through pylance
if accelerator is not None:
self.to_lance().create_index(
dataset = self.to_lance()
_validate_index_accelerator(accelerator)
dataset.create_index(
column=column,
index_type=index_type,
metric=metric,
@@ -3312,6 +3341,8 @@ class LanceTable(Table):
acc = getattr(config, "accelerator", None)
if acc is not None:
# Dispatch to pylance for GPU acceleration
dataset = self.to_lance()
_validate_index_accelerator(acc)
index_type_map = {
"IvfFlat": "IVF_FLAT",
"IvfSq": "IVF_SQ",
@@ -3323,7 +3354,7 @@ class LanceTable(Table):
cfg_type = type(config).__name__
lance_index_type = index_type_map.get(cfg_type, "IVF_PQ")
self.to_lance().create_index(
dataset.create_index(
column=column,
index_type=lance_index_type,
metric=getattr(config, "distance_type", "l2"),
+46 -1
View File
@@ -15,7 +15,7 @@ from datetime import date, datetime, timedelta
from decimal import Decimal
from time import sleep
from typing import List
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import lancedb
from lancedb import table as table_module
@@ -1642,6 +1642,51 @@ def test_create_index_method(mock_create_index, mem_db: DBConnection):
)
@pytest.mark.parametrize("api", ["legacy", "unified"])
@pytest.mark.parametrize("cuda_failure", ["unavailable", "init_error"])
def test_create_index_cuda_failure_suggests_cpu(
api, cuda_failure, mem_db: DBConnection
):
table = mem_db.create_table(
f"test_cuda_unavailable_{api}",
data=[{"vector": [3.1, 4.1]}, {"vector": [5.9, 26.5]}],
)
torch = MagicMock()
torch.cuda.is_available.return_value = cuda_failure == "init_error"
if cuda_failure == "init_error":
torch.cuda.init.side_effect = RuntimeError("CUDA driver initialization failed")
with (
patch.dict(sys.modules, {"torch": torch}),
patch.object(table, "to_lance") as to_lance,
):
if api == "legacy":
with (
pytest.warns(DeprecationWarning, match="create_index"),
pytest.raises(RuntimeError, match="accelerator=None.*CPU"),
):
table.create_index(
metric="dot",
num_partitions=1,
num_sub_vectors=1,
vector_column_name="vector",
accelerator="cuda",
)
else:
with pytest.raises(RuntimeError, match="accelerator=None.*CPU"):
table.create_index(
"vector",
config=IvfPq(
distance_type="dot",
num_partitions=1,
num_sub_vectors=1,
accelerator="cuda",
),
)
to_lance.return_value.create_index.assert_not_called()
@patch("lancedb.table.AsyncTable.create_index")
def test_create_index_name_and_train_parameters(
mock_create_index, mem_db: DBConnection