mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-12 23:02:28 +00:00
fix(python): clarify unavailable CUDA index acceleration
This commit is contained in:
@@ -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"),
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user