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:
lancedb-gatefixer[bot]
2026-10-23 01:26:13 +08:00
committed by GitHub
parent d37a6b76a4
commit 66b5b5bbdc
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
@@ -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"),
+46 -1
View File
@@ -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