mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +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:
@@ -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