mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 20:22:37 +00:00
fix(python): clarify unavailable CUDA index acceleration (#4249)
## Summary - preflight CUDA availability and initialization before accelerated index training - direct unavailable or misconfigured CUDA environments to the explicit CPU path - cover both legacy and unified Python index APIs ## Root cause LanceDB delegated accelerator requests directly to PyLance. When PyTorch lacked CUDA support or CUDA initialization failed, users received a low-level assertion or runtime error with no indication that omitting `accelerator` or setting `accelerator=None` creates the index on CPU. The repair validates CUDA before training and preserves explicit accelerator semantics without silently changing the requested backend. ## Scope This PR addresses only the separately reproduced CUDA availability and initialization error. It does not close #1766; the reported query failure after successful CUDA index creation remains open pending reproducible evidence. ## Validation - `uv run --extra tests pytest python/tests/test_table.py -k create_index -q` (10 passed) - `ruff format .` - `ruff check .` - CPU-only integration check confirms the actionable fallback message Related to #1766 <!-- lance-gatekeeper-fix:v1 agent=d05c37237b04c2c8f88395a82e00bc67 generation=1 --> Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
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
|
||||
@@ -3600,7 +3627,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,
|
||||
@@ -3625,6 +3654,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",
|
||||
@@ -3636,7 +3667,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, localcontext
|
||||
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
|
||||
@@ -1773,6 +1773,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