diff --git a/docs/src/python/python.md b/docs/src/python/python.md
index d249aaea8..7bf7cd167 100644
--- a/docs/src/python/python.md
+++ b/docs/src/python/python.md
@@ -419,6 +419,8 @@ scoped to each invocation and support `read`, `readinto`, `seek`, `tell`,
::: lancedb.functions.PythonRuntimeSpec
+::: lancedb.GpuClass
+
::: lancedb.functions.FunctionVersionRef
::: lancedb.functions.ApplicationInput
diff --git a/docs/web/curation/create-function.mdx b/docs/web/curation/create-function.mdx
index 2d161e84b..8e4bdc657 100644
--- a/docs/web/curation/create-function.mdx
+++ b/docs/web/curation/create-function.mdx
@@ -92,16 +92,26 @@ and remote-API Functions both build on below.
## GPU: batch inference, loaded once {#gpu}
-`gpu=True` requires a GPU-capable worker for every execution. Combined with
-the batch and stateful forms above, a realistic embedding model loads once
-in `__init__` and runs one forward pass per batch, not one per row:
+`db.list_gpu_classes()` lists the GPU classes configured for the deployment.
+Each record reports the operator's hardware description and limits; it does
+not mean a GPU is currently free. Listing requires the System Usage
+permission. Pass a returned class name to `gpu_class` to fix that placement
+in the Function version. `gpu=True` remains required:
```python Python icon="python"
import pyarrow as pa
from lancedb import udf
-@udf(gpu=True, pip=["torch==2.12.0", "sentence-transformers==5.7.0"])
+gpu_class = next(
+ gpu.name for gpu in db.list_gpu_classes() if gpu.gpu_model == "NVIDIA A10G"
+)
+
+@udf(
+ gpu=True,
+ gpu_class=gpu_class,
+ pip=["torch==2.12.0", "sentence-transformers==5.7.0"],
+)
class embed_text:
def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
import torch
@@ -120,6 +130,12 @@ table.add_columns({"embedding": version(text=col("caption"))})
```
+Older servers that do not implement `GET /v1/gpu/classes` raise
+`NotImplementedError`. Without `gpu_class`, the deployment's default GPU
+class is used. The class is immutable with the Function version; set a
+per-refresh class or aggregate GPU budget through
+[`refresh_column_async`](/feature-engineering/refresh-column#resources).
+
The model loads once per worker instance, in `__init__`; `__call__` only
encodes the batch it was handed. Registering the same logic as a plain
function instead would reload the model on every single row.
diff --git a/docs/web/feature-engineering/refresh-column.mdx b/docs/web/feature-engineering/refresh-column.mdx
index ba3934f1e..77d603770 100644
--- a/docs/web/feature-engineering/refresh-column.mdx
+++ b/docs/web/feature-engineering/refresh-column.mdx
@@ -182,6 +182,64 @@ alternatives rather than a pair. `batch_size` controls how many rows a worker
takes at a time, and `error_policy = 'skip'` continues past rows that fail
rather than failing the job.
+For a column filled by a registered Function, the Python and Rust job APIs can
+override the Function's GPU placement for this refresh:
+
+
+```python
+classes = db.list_gpu_classes()
+gpu_class = next(gpu.name for gpu in classes if gpu.gpu_model == "NVIDIA H100")
+
+job = table.refresh_column_async(
+ "embedding",
+ gpu_class=gpu_class,
+ gpu_count=4,
+ batch_size=32,
+)
+```
+
+```rust
+use lancedb::{Connection, Table, table::RefreshColumnOptions};
+
+async fn refresh(
+ conn: &Connection,
+ table: &Table,
+) -> Result<(), Box> {
+ let gpu_class = conn
+ .list_gpu_classes()
+ .await?
+ .into_iter()
+ .find(|class| class.gpu_model.as_deref() == Some("NVIDIA H100"))
+ .map(|class| class.name)
+ .ok_or("H100 GPU class is not configured")?;
+ let _job = table
+ .refresh_column_async_with_options(
+ "embedding",
+ RefreshColumnOptions::default()
+ .gpu_class(gpu_class)
+ .gpu_count(4)
+ .batch_size(32),
+ )
+ .await?;
+ Ok(())
+}
+```
+
+
+`gpu_count` is the aggregate GPU budget for the refresh's Function workers,
+not a request for that many GPUs per worker. It must be positive.
+`gpu_class` must name a class returned by `list_gpu_classes()`; class names
+are lowercase DNS labels. Omit either option to use the Function's registered
+class or the server's configured budget. `list_gpu_classes()` requires System
+Usage permission, and raises `NotImplementedError` when the server predates
+GPU class discovery.
+
+
+The Rust GPU discovery and refresh controls are not in the latest Rust
+release. TypeScript does not expose GPU class discovery or placement; a
+tracking issue is needed before adding equivalent controls.
+
+
## When a refresh fails {#failures}
What a failure keeps depends on where the refresh ran. `refresh_column` in your
diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py
index 9cc1db0d5..a93c32e63 100644
--- a/python/python/lancedb/__init__.py
+++ b/python/python/lancedb/__init__.py
@@ -19,6 +19,7 @@ from ._lancedb import FtsToken
from ._lancedb import FunctionErrorFragment as FunctionErrorFragment
from ._lancedb import FunctionErrorRecord as FunctionErrorRecord
from ._lancedb import FunctionErrors as FunctionErrors
+from ._lancedb import GpuClass as GpuClass
from ._lancedb import RTEWriteSpec
from ._lancedb import tokenize as _tokenize
from .common import URI, sanitize_uri
diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi
index 2d9de72c4..8730917d3 100644
--- a/python/python/lancedb/_lancedb.pyi
+++ b/python/python/lancedb/_lancedb.pyi
@@ -112,6 +112,16 @@ class Session:
@property
def approx_num_items(self) -> int: ...
+class GpuClass:
+ name: str
+ resource_name: str
+ gpu_model: Optional[str]
+ gpu_memory_bytes: Optional[int]
+ max_gpu_count_per_worker: Optional[int]
+ cpu_millis: Optional[int]
+ memory_bytes: Optional[int]
+ ephemeral_storage_bytes: Optional[int]
+
class Connection(object):
uri: str
async def is_open(self): ...
@@ -153,6 +163,7 @@ class Connection(object):
limit: Optional[int],
) -> list[str]: ... # Deprecated: Use list_tables instead
async def open_job(self, job_id: str) -> Job: ...
+ async def list_gpu_classes(self) -> list[GpuClass]: ...
async def create_function_async(
self, request_json: str, namespace_path: Optional[List[str]] = None
) -> Job: ...
@@ -561,10 +572,20 @@ class Table:
auto_refresh_interval_seconds: Optional[int] = None,
) -> AddColumnsResult: ...
async def refresh_column(
- self, column: str, *, batch_size: Optional[int] = None
+ self,
+ column: str,
+ *,
+ batch_size: Optional[int] = None,
+ gpu_class: Optional[str] = None,
+ gpu_count: Optional[int] = None,
) -> RefreshColumnResult: ...
async def refresh_column_async(
- self, column: str, *, batch_size: Optional[int] = None
+ self,
+ column: str,
+ *,
+ batch_size: Optional[int] = None,
+ gpu_class: Optional[str] = None,
+ gpu_count: Optional[int] = None,
) -> Job: ...
async def function_errors(
self,
diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py
index e41a79cd5..d85685f0b 100644
--- a/python/python/lancedb/db.py
+++ b/python/python/lancedb/db.py
@@ -104,6 +104,7 @@ if TYPE_CHECKING:
from .pydantic import LanceModel
from ._lancedb import Connection as LanceDbConnection
+ from ._lancedb import GpuClass
from ._lancedb import JobInfo
from .common import DATA, URI
from .embeddings import EmbeddingFunctionConfig
@@ -1348,6 +1349,21 @@ class DBConnection(EnforceOverrides):
"""
raise NotImplementedError("open_job is not supported for this connection type")
+ def list_gpu_classes(self) -> List[GpuClass]:
+ """List the operator-configured Function GPU classes.
+
+ Each record names the class and reports its configured GPU model,
+ memory, per-worker GPU limit, and CPU, memory, and ephemeral-storage
+ limits. This describes configured hardware, not current free capacity.
+
+ LanceDB Distributed only; the principal needs System Usage permission.
+ Raises ``NotImplementedError`` when the server predates
+ ``GET /v1/gpu/classes``.
+ """
+ raise NotImplementedError(
+ "GPU class discovery is not supported for this connection type"
+ )
+
def list_jobs(
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> Listing[JobInfo]:
@@ -3450,6 +3466,12 @@ class AsyncConnection(object):
"""
return AsyncJob(await self._inner.open_job(job_id))
+ async def list_gpu_classes(self) -> List[GpuClass]:
+ """List configured Function GPU classes. See
+ [DBConnection.list_gpu_classes][lancedb.db.DBConnection.list_gpu_classes].
+ """
+ return await self._inner.list_gpu_classes()
+
async def create_function_async(
self,
definition: UdfDefinition,
diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py
index 721ba5ced..dc1e5684e 100644
--- a/python/python/lancedb/functions.py
+++ b/python/python/lancedb/functions.py
@@ -61,6 +61,11 @@ from .secrets import EnvVarSecret
_Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1)
_UInt32 = conint(strict=True, ge=0, le=2**32 - 1)
_UInt64 = conint(strict=True, ge=0, le=2**64 - 1)
+_GPU_CLASS_NAME = re.compile(r"^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$")
+_GPU_CLASS_ERROR = (
+ "gpu_class must be a lowercase name of at most 63 letters, digits, or "
+ "hyphens, and cannot start or end with a hyphen"
+)
def _validate_gpu_wire_marker(value: Any) -> bool:
@@ -75,6 +80,12 @@ def _normalize_gpu_marker(value: bool) -> Optional[bool]:
return True if value else None
+def _validate_gpu_class_name(value: Any) -> str:
+ if not isinstance(value, str) or not _GPU_CLASS_NAME.fullmatch(value):
+ raise ValueError(_GPU_CLASS_ERROR)
+ return value
+
+
class _FrozenDict(dict):
def _immutable(self, *args, **kwargs):
raise TypeError("remote canonical values are immutable")
@@ -291,6 +302,8 @@ class PythonRuntimeSpec(_RemoteValue):
V1 supports ``kind="python"``. Newer runtime kinds remain readable, while
their unknown payload fields are intentionally not retained by the client.
+ ``kind="python_v2"`` requires ``gpu=True`` and may name an operator-defined
+ ``gpu_class``.
"""
kind: str
@@ -298,6 +311,7 @@ class PythonRuntimeSpec(_RemoteValue):
environment: Optional[PythonEnvironmentSpec] = None
env: Optional[Mapping[str, str]] = None
gpu: Optional[bool] = None
+ gpu_class: Optional[str] = None
@model_validator(mode="before")
@classmethod
@@ -315,6 +329,13 @@ class PythonRuntimeSpec(_RemoteValue):
return None
return _validate_gpu_wire_marker(value)
+ @field_validator("gpu_class", mode="before")
+ @classmethod
+ def _validate_gpu_class(cls, value):
+ if value is None:
+ return None
+ return _validate_gpu_class_name(value)
+
@model_validator(mode="after")
def _validate_runtime_kind(self):
if self.kind == "python":
@@ -324,6 +345,10 @@ class PythonRuntimeSpec(_RemoteValue):
raise ValueError("python runtime requires environment")
if self.gpu is not None:
raise ValueError("python runtime with gpu requires kind='python_v2'")
+ if self.gpu_class is not None:
+ raise ValueError(
+ "python runtime with gpu_class requires kind='python_v2'"
+ )
elif self.kind == "python_v2":
if self.python_version is None:
raise ValueError("python_v2 runtime requires python_version")
@@ -336,6 +361,7 @@ class PythonRuntimeSpec(_RemoteValue):
object.__setattr__(self, "environment", None)
object.__setattr__(self, "env", None)
object.__setattr__(self, "gpu", None)
+ object.__setattr__(self, "gpu_class", None)
return self
@@ -373,6 +399,9 @@ class FunctionVersion(_RemoteValue):
``udf_def.registration_request.artifact.digest`` to check whether a local
[UdfDefinition][lancedb.functions.UdfDefinition] has the same source.
It is ``None`` when the server does not provide a source digest.
+
+ ``gpu_class`` is the operator-defined GPU pool fixed into this Function
+ version, or ``None`` when registration used the deployment default.
"""
name: str
@@ -387,6 +416,7 @@ class FunctionVersion(_RemoteValue):
created_at_millis: int
metadata: Mapping[str, str]
disabled: bool
+ gpu_class: Optional[str] = None
def __call__(self, **arguments: Any) -> FunctionApplication:
"""Bind this exact version to named table columns.
@@ -1900,6 +1930,7 @@ class UdfDefinition:
env: Mapping[str, str],
python_version: Optional[str],
gpu: bool = False,
+ gpu_class: Optional[str] = None,
conda: tuple[str, ...] = (),
conda_channels: tuple[str, ...] = (),
apt: tuple[str, ...] = (),
@@ -1918,6 +1949,11 @@ class UdfDefinition:
raise ValueError("a Function environment is pip or conda, not both")
if conda_channels and not conda:
raise ValueError("conda_channels requires conda packages")
+ gpu_marker = _normalize_gpu_marker(gpu)
+ if gpu_class is not None:
+ _validate_gpu_class_name(gpu_class)
+ if gpu_marker is None:
+ raise ValueError("gpu_class requires gpu=True")
packages = tuple(sorted(set(conda if conda else pip)))
system_packages = tuple(sorted(set(apt)))
if any(
@@ -1988,7 +2024,6 @@ class UdfDefinition:
kind = _PYTHON_CALLABLE_ARTIFACT
source = entry_source.encode("utf-8")
digest = f"sha256:{hashlib.sha256(source).hexdigest()}"
- gpu_marker = _normalize_gpu_marker(gpu)
runtime = PythonRuntimeSpec(
kind="python_v2" if gpu_marker is not None else "python",
python_version=python_version
@@ -1996,6 +2031,7 @@ class UdfDefinition:
environment=environment_spec,
env=environment,
gpu=gpu_marker,
+ gpu_class=gpu_class,
)
self._function = function
self._request = FunctionRegistrationRequest(
@@ -2110,6 +2146,7 @@ def udf(
env: Optional[Mapping[str, str]] = None,
python_version: Optional[str] = None,
gpu: bool = False,
+ gpu_class: Optional[str] = None,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
apt: tuple[str, ...] | list[str] = (),
@@ -2128,6 +2165,7 @@ def udf(
env: Optional[Mapping[str, str]] = None,
python_version: Optional[str] = None,
gpu: bool = False,
+ gpu_class: Optional[str] = None,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
apt: tuple[str, ...] | list[str] = (),
@@ -2251,6 +2289,12 @@ def udf(
Whether every remote execution requires a GPU. The execution platform
selects one compatible GPU for each worker. The requirement is part of
the immutable Function version.
+ gpu_class : str, optional
+ Operator-defined GPU class for every execution of this Function, such
+ as ``"a10g"`` or ``"h100"``. Requires ``gpu=True`` and follows the
+ lowercase 63-character DNS-label rule. The class is part of the
+ immutable Function version; per-refresh aggregate GPU budgets are
+ selected through ``refresh_column`` instead of registration.
code : sequence of module, optional
Top-level modules or packages whose importable ``.py`` files ship
with the Function and are importable by their own names.
@@ -2280,11 +2324,13 @@ def udf(
... return value
>>> batch_identity(pa.array([1, 2, 3])).to_pylist()
[1, 2, 3]
- >>> @udf(pip=["cupy-cuda12x"], gpu=True)
+ >>> @udf(pip=["cupy-cuda12x"], gpu=True, gpu_class="a10g")
... def gpu_score(value: int) -> int:
... return value * 2
>>> gpu_score.registration_request.runtime.gpu
True
+ >>> gpu_score.registration_request.runtime.gpu_class
+ 'a10g'
>>> @udf
... class scale:
... def __init__(self, factor: float, offset: float = 0.0):
@@ -2319,6 +2365,7 @@ def udf(
env={} if env is None else env,
python_version=python_version,
gpu=gpu,
+ gpu_class=gpu_class,
conda=tuple(conda),
conda_channels=tuple(conda_channels),
apt=tuple(apt),
diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py
index 0c1f127c8..8c87d55ed 100644
--- a/python/python/lancedb/remote/db.py
+++ b/python/python/lancedb/remote/db.py
@@ -49,7 +49,7 @@ from ..storage_integration import SecretReference, StorageIntegrationInfo
from ..view import ViewDescription
if TYPE_CHECKING:
- from .._lancedb import JobInfo, Session
+ from .._lancedb import GpuClass, JobInfo, Session
from ..embeddings import EmbeddingFunctionConfig
from lance_namespace import (
LanceNamespace,
@@ -958,6 +958,13 @@ class RemoteDBConnection(DBConnection):
"""
return Job(LOOP.run(self._conn.open_job(job_id)))
+ @override
+ def list_gpu_classes(self) -> List[GpuClass]:
+ """List configured Function GPU classes. See
+ [DBConnection.list_gpu_classes][lancedb.db.DBConnection.list_gpu_classes].
+ """
+ return LOOP.run(self._conn.list_gpu_classes())
+
@override
def create_function_async(
self,
diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py
index 67c58936f..b31325424 100644
--- a/python/python/lancedb/remote/table.py
+++ b/python/python/lancedb/remote/table.py
@@ -1108,14 +1108,40 @@ class RemoteTable(Table):
)
)
- def refresh_column(self, column: str, *, batch_size: Optional[int] = None):
- return LOOP.run(self._table.refresh_column(column, batch_size=batch_size))
+ def refresh_column(
+ self,
+ column: str,
+ *,
+ batch_size: Optional[int] = None,
+ gpu_class: Optional[str] = None,
+ gpu_count: Optional[int] = None,
+ ):
+ return LOOP.run(
+ self._table.refresh_column(
+ column,
+ batch_size=batch_size,
+ gpu_class=gpu_class,
+ gpu_count=gpu_count,
+ )
+ )
def refresh_column_async(
- self, column: str, *, batch_size: Optional[int] = None
+ self,
+ column: str,
+ *,
+ batch_size: Optional[int] = None,
+ gpu_class: Optional[str] = None,
+ gpu_count: Optional[int] = None,
) -> Job[RefreshColumnResult]:
return Job(
- LOOP.run(self._table.refresh_column_async(column, batch_size=batch_size))
+ LOOP.run(
+ self._table.refresh_column_async(
+ column,
+ batch_size=batch_size,
+ gpu_class=gpu_class,
+ gpu_count=gpu_count,
+ )
+ )
)
def function_errors(
diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py
index e16a29df1..11efea348 100644
--- a/python/python/lancedb/table.py
+++ b/python/python/lancedb/table.py
@@ -3078,7 +3078,12 @@ class Table(ABC):
@abstractmethod
def refresh_column(
- self, column: str, *, batch_size: Optional[int] = None
+ self,
+ column: str,
+ *,
+ batch_size: Optional[int] = None,
+ gpu_class: Optional[str] = None,
+ gpu_count: Optional[int] = None,
) -> "RefreshColumnResult":
"""
Fill the rows of a computed column that hold no value yet.
@@ -3103,6 +3108,15 @@ class Table(ABC):
Maximum rows per Function call on the remote server. Must be
greater than zero. None uses the server default. Has no effect on
local expression-backed columns.
+ gpu_class: str, optional
+ GPU class for Function workers. Must be a lowercase name of 1-63
+ letters, digits, or hyphens, without a leading or trailing hyphen.
+ None keeps the Function's registered class or the deployment
+ default. LanceDB Distributed only.
+ gpu_count: int, optional
+ Positive aggregate GPU budget for Function workers. None leaves
+ the server's configured budget unchanged. LanceDB Distributed
+ only.
Returns
-------
@@ -3113,7 +3127,12 @@ class Table(ABC):
@abstractmethod
def refresh_column_async(
- self, column: str, *, batch_size: Optional[int] = None
+ self,
+ column: str,
+ *,
+ batch_size: Optional[int] = None,
+ gpu_class: Optional[str] = None,
+ gpu_count: Optional[int] = None,
) -> Job[RefreshColumnJobResult]:
"""
Like :meth:`refresh_column`, but returns a handle to the refresh job
@@ -3137,6 +3156,15 @@ class Table(ABC):
greater than zero. None uses the server default. Lower this for
slow Functions to stay within the executor's per-call deadline.
Has no effect on local expression-backed columns.
+ gpu_class: str, optional
+ GPU class for Function workers. Must be a lowercase name of 1-63
+ letters, digits, or hyphens, without a leading or trailing hyphen.
+ None keeps the Function's registered class or the deployment
+ default. LanceDB Distributed only.
+ gpu_count: int, optional
+ Positive aggregate GPU budget for Function workers. None leaves
+ the server's configured budget unchanged. LanceDB Distributed
+ only.
Returns
-------
@@ -5200,22 +5228,46 @@ class LanceTable(Table):
)
def refresh_column(
- self, column: str, *, batch_size: Optional[int] = None
+ self,
+ column: str,
+ *,
+ batch_size: Optional[int] = None,
+ gpu_class: Optional[str] = None,
+ gpu_count: Optional[int] = None,
) -> "RefreshColumnResult":
"""Fill a computed column's unfilled rows and recompute those whose
inputs changed. See
[`AsyncTable.refresh_column`][lancedb.AsyncTable.refresh_column]."""
- return LOOP.run(self._table.refresh_column(column, batch_size=batch_size))
+ return LOOP.run(
+ self._table.refresh_column(
+ column,
+ batch_size=batch_size,
+ gpu_class=gpu_class,
+ gpu_count=gpu_count,
+ )
+ )
def refresh_column_async(
- self, column: str, *, batch_size: Optional[int] = None
+ self,
+ column: str,
+ *,
+ batch_size: Optional[int] = None,
+ gpu_class: Optional[str] = None,
+ gpu_count: Optional[int] = None,
) -> Job[RefreshColumnJobResult]:
"""Fill a computed column's unfilled rows and recompute those whose
inputs changed, returning a handle to the refresh job. See
[`Table.refresh_column_async`][lancedb.table.Table.refresh_column_async].
"""
return Job(
- LOOP.run(self._table.refresh_column_async(column, batch_size=batch_size))
+ LOOP.run(
+ self._table.refresh_column_async(
+ column,
+ batch_size=batch_size,
+ gpu_class=gpu_class,
+ gpu_count=gpu_count,
+ )
+ )
)
def function_errors(
@@ -7475,7 +7527,12 @@ class AsyncTable:
return await self._inner.add_columns(list(transforms.items()))
async def refresh_column(
- self, column: str, *, batch_size: Optional[int] = None
+ self,
+ column: str,
+ *,
+ batch_size: Optional[int] = None,
+ gpu_class: Optional[str] = None,
+ gpu_count: Optional[int] = None,
) -> RefreshColumnResult:
"""
Fill the rows of a computed column that hold no value yet.
@@ -7500,13 +7557,27 @@ class AsyncTable:
Maximum rows per Function call on the remote server. Must be
greater than zero. None uses the server default. Has no effect on
local expression-backed columns.
+ gpu_class: str, optional
+ GPU class for Function workers. Must be a lowercase name of 1-63
+ letters, digits, or hyphens, without a leading or trailing hyphen.
+ None keeps the Function's registered class or the deployment
+ default. LanceDB Distributed only.
+ gpu_count: int, optional
+ Positive aggregate GPU budget for Function workers. None leaves
+ the server's configured budget unchanged. LanceDB Distributed
+ only.
Returns
-------
RefreshColumnResult
The number of rows filled and the new version of the table.
"""
- return await self._inner.refresh_column(column, batch_size=batch_size)
+ return await self._inner.refresh_column(
+ column,
+ batch_size=batch_size,
+ gpu_class=gpu_class,
+ gpu_count=gpu_count,
+ )
async def function_errors(
self,
@@ -7576,7 +7647,12 @@ class AsyncTable:
)
async def refresh_column_async(
- self, column: str, *, batch_size: Optional[int] = None
+ self,
+ column: str,
+ *,
+ batch_size: Optional[int] = None,
+ gpu_class: Optional[str] = None,
+ gpu_count: Optional[int] = None,
) -> AsyncJob[RefreshColumnJobResult]:
"""
Like :meth:`refresh_column`, but returns a handle to the refresh job
@@ -7600,6 +7676,15 @@ class AsyncTable:
greater than zero. None uses the server default. Lower this for
slow Functions to stay within the executor's per-call deadline.
Has no effect on local expression-backed columns.
+ gpu_class: str, optional
+ GPU class for Function workers. Must be a lowercase name of 1-63
+ letters, digits, or hyphens, without a leading or trailing hyphen.
+ None keeps the Function's registered class or the deployment
+ default. LanceDB Distributed only.
+ gpu_count: int, optional
+ Positive aggregate GPU budget for Function workers. None leaves
+ the server's configured budget unchanged. LanceDB Distributed
+ only.
Returns
-------
@@ -7623,7 +7708,12 @@ class AsyncTable:
'finished'
"""
return _typed_job(
- await self._inner.refresh_column_async(column, batch_size=batch_size),
+ await self._inner.refresh_column_async(
+ column,
+ batch_size=batch_size,
+ gpu_class=gpu_class,
+ gpu_count=gpu_count,
+ ),
RefreshColumnJobResult.from_json,
)
diff --git a/python/python/tests/test_first_class_function_slice1.py b/python/python/tests/test_first_class_function_slice1.py
index 64a610b53..14c35d688 100644
--- a/python/python/tests/test_first_class_function_slice1.py
+++ b/python/python/tests/test_first_class_function_slice1.py
@@ -167,6 +167,17 @@ def test_function_version_preserves_source_digest(source_digest):
version.source_digest = "sha256:changed"
+def test_function_version_exposes_registered_gpu_class():
+ value = job_result("remote_function_job.json")
+ legacy = FunctionVersion.from_json(json.dumps(value))
+ assert legacy.gpu_class is None
+
+ value["gpu_class"] = "a10g"
+ version = FunctionVersion.from_json(json.dumps(value))
+ assert version.gpu_class == "a10g"
+ assert json.loads(version.to_canonical_json())["gpu_class"] == "a10g"
+
+
def test_function_version_binds_named_columns_as_one_immutable_application():
version = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py
index 4816c7555..b17e3208b 100644
--- a/python/python/tests/test_first_class_function_slice2.py
+++ b/python/python/tests/test_first_class_function_slice2.py
@@ -534,13 +534,14 @@ def test_udf_image_requires_digest_reference(image):
def test_udf_gpu_marker_uses_gpu_runtime():
- @udf(pip=["cupy-cuda12x"], gpu=True)
+ @udf(pip=["cupy-cuda12x"], gpu=True, gpu_class="a10g")
def double_on_gpu(value: int) -> int:
return value * 2
request = json.loads(double_on_gpu.registration_request.to_canonical_json())
assert request["runtime"]["kind"] == "python_v2"
assert request["runtime"]["gpu"] is True
+ assert request["runtime"]["gpu_class"] == "a10g"
@udf(pip=["pyarrow"])
def cpu_function(value: int) -> int:
@@ -551,6 +552,7 @@ def test_udf_gpu_marker_uses_gpu_runtime():
]
assert cpu_runtime["kind"] == "python"
assert "gpu" not in cpu_runtime
+ assert "gpu_class" not in cpu_runtime
def identity(value: int) -> int:
return value
@@ -558,6 +560,11 @@ def test_udf_gpu_marker_uses_gpu_runtime():
for invalid in [None, 0, 1, -1, 1.5, "", "true", "1", "H100"]:
with pytest.raises(ValueError, match="gpu must be a boolean"):
udf(name="invalid_gpu", gpu=invalid)(identity)
+ for invalid in ["H100", "a_b", "-a10g", "a10g-", "a" * 64]:
+ with pytest.raises(ValueError, match="gpu_class"):
+ udf(name="invalid_gpu_class", gpu=True, gpu_class=invalid)(identity)
+ with pytest.raises(ValueError, match="requires gpu=True"):
+ udf(name="gpu_class_without_gpu", gpu_class="a10g")(identity)
base_runtime = {
"kind": "python_v2",
@@ -566,14 +573,26 @@ def test_udf_gpu_marker_uses_gpu_runtime():
}
runtime = PythonRuntimeSpec.model_validate({**base_runtime, "gpu": True})
assert runtime.gpu is True
+ runtime = PythonRuntimeSpec.model_validate(
+ {**base_runtime, "gpu": True, "gpu_class": "a10g"}
+ )
+ assert runtime.gpu_class == "a10g"
for invalid in [False, 1, 0, "", "true", "1", "H100"]:
with pytest.raises(ValueError, match="runtime.gpu must be true"):
PythonRuntimeSpec.model_validate({**base_runtime, "gpu": invalid})
+ with pytest.raises(ValueError, match="gpu_class"):
+ PythonRuntimeSpec.model_validate(
+ {**base_runtime, "gpu": True, "gpu_class": "H100"}
+ )
def test_unknown_runtime_discards_payload_before_known_field_validation():
for payload in [
- {"kind": "python_v3", "gpu": {"model": "H100"}},
+ {
+ "kind": "python_v3",
+ "gpu": {"model": "H100"},
+ "gpu_class": "a10g",
+ },
{"kind": "python_v3", "resources": []},
{
"kind": "python_v3",
diff --git a/python/python/tests/test_package_metadata.py b/python/python/tests/test_package_metadata.py
index 9bdcbb3b0..1b4d3a711 100644
--- a/python/python/tests/test_package_metadata.py
+++ b/python/python/tests/test_package_metadata.py
@@ -31,6 +31,15 @@ def test_job_metadata_types_have_resolvable_modules(name):
assert getattr(defining_module, public_type.__name__, None) is public_type
+def test_gpu_class_reexport_has_resolvable_module():
+ import lancedb
+
+ public_type = lancedb.GpuClass
+ defining_module = importlib.import_module(public_type.__module__)
+
+ assert getattr(defining_module, public_type.__name__, None) is public_type
+
+
def test_pyo3_abi_matches_minimum_supported_python():
project_dir = Path(__file__).parents[2]
pyproject = (project_dir / "pyproject.toml").read_text()
diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py
index b2350104b..a7daeeb43 100644
--- a/python/python/tests/test_remote_db.py
+++ b/python/python/tests/test_remote_db.py
@@ -1700,7 +1700,14 @@ def test_remote_create_index_async_returns_job():
@pytest.mark.asyncio
@pytest.mark.parametrize("async_table", [False, True])
@pytest.mark.parametrize(
- "refresh_kwargs", [{}, {"batch_size": None}, {"batch_size": 1}, {"batch_size": 8}]
+ "refresh_kwargs",
+ [
+ {},
+ {"batch_size": None},
+ {"batch_size": 1},
+ {"batch_size": 8},
+ {"gpu_class": "a10g", "gpu_count": 2},
+ ],
)
async def test_remote_refresh_async_returns_typed_terminal_result(
async_table, refresh_kwargs
@@ -1720,6 +1727,10 @@ async def test_remote_refresh_async_returns_typed_terminal_result(
expected = {"column": "derived"}
if refresh_kwargs.get("batch_size") is not None:
expected["batch_size"] = refresh_kwargs["batch_size"]
+ if refresh_kwargs.get("gpu_class") is not None:
+ expected["gpu_class"] = refresh_kwargs["gpu_class"]
+ if refresh_kwargs.get("gpu_count") is not None:
+ expected["gpu_count"] = refresh_kwargs["gpu_count"]
assert json.loads(body) == expected
request.send_response(202)
request.send_header("Content-Type", "application/json")
@@ -1788,6 +1799,66 @@ async def test_remote_refresh_async_returns_typed_terminal_result(
assert result.version == 8
+@pytest.mark.parametrize(
+ "refresh_kwargs,message",
+ [
+ ({"gpu_class": "H100"}, "gpu_class"),
+ ({"gpu_class": "-a10g"}, "gpu_class"),
+ ({"gpu_class": "a_b"}, "gpu_class"),
+ ({"gpu_count": 0}, "gpu_count"),
+ ({"gpu_count": -1}, "gpu_count"),
+ ],
+)
+@pytest.mark.asyncio
+async def test_remote_refresh_rejects_invalid_gpu_options_before_request(
+ refresh_kwargs, message
+):
+ backfill_requests = []
+
+ def handler(request):
+ if request.path == "/v1/table/test/describe/":
+ request.send_response(200)
+ request.send_header("Content-Type", "application/json")
+ request.end_headers()
+ request.wfile.write(
+ json.dumps(
+ {
+ "version": 1,
+ "schema": {
+ "fields": [
+ {
+ "name": "id",
+ "type": {"type": "int64"},
+ "nullable": False,
+ }
+ ]
+ },
+ }
+ ).encode()
+ )
+ elif request.path == "/v1/table/test/backfill_column":
+ backfill_requests.append(request.path)
+ request.send_response(202)
+ request.send_header("Content-Type", "application/json")
+ request.end_headers()
+ request.wfile.write(b'{"job_id": "unexpected"}')
+ else:
+ request.send_response(404)
+ request.end_headers()
+
+ with mock_lancedb_connection(handler) as db:
+ table = db.open_table("test")
+ with pytest.raises(ValueError, match=message):
+ table.refresh_column_async("derived", **refresh_kwargs)
+
+ async with mock_lancedb_connection_async(handler) as db:
+ table = await db.open_table("test")
+ with pytest.raises(ValueError, match=message):
+ await table.refresh_column_async("derived", **refresh_kwargs)
+
+ assert backfill_requests == []
+
+
def test_remote_function_errors_lists_the_rows_a_refresh_skipped():
listing = {
"records": [
@@ -4914,6 +4985,83 @@ def test_remote_job_handle_reports_its_own_detail():
}
+@pytest.mark.parametrize("is_async", [False, True])
+@pytest.mark.asyncio
+async def test_remote_list_gpu_classes(is_async):
+ requests = []
+
+ def handler(request):
+ requests.append((request.command, request.path))
+ request.send_response(200)
+ request.send_header("Content-Type", "application/json")
+ request.end_headers()
+ request.wfile.write(
+ json.dumps(
+ {
+ "gpu_classes": [
+ {
+ "name": "default",
+ "resource_name": "nvidia.com/gpu",
+ "gpu_model": "NVIDIA A10G",
+ "gpu_memory_bytes": 25769803776,
+ "max_gpu_count_per_worker": 1,
+ "cpu_millis": 4000,
+ "memory_bytes": 25769803776,
+ "ephemeral_storage_bytes": 1073741824,
+ },
+ {
+ "name": "h100",
+ "resource_name": "nvidia.com/gpu",
+ "gpu_model": "NVIDIA H100",
+ "gpu_memory_bytes": 85899345920,
+ "max_gpu_count_per_worker": 8,
+ "cpu_millis": None,
+ "memory_bytes": None,
+ "ephemeral_storage_bytes": None,
+ },
+ ]
+ }
+ ).encode()
+ )
+
+ if is_async:
+ async with mock_lancedb_connection_async(handler) as db:
+ classes = await db.list_gpu_classes()
+ else:
+ with mock_lancedb_connection(handler) as db:
+ classes = db.list_gpu_classes()
+
+ assert requests == [("GET", "/v1/gpu/classes")]
+ assert len(classes) == 2
+ assert isinstance(classes[0], lancedb.GpuClass)
+ assert classes[0].name == "default"
+ assert classes[0].gpu_model == "NVIDIA A10G"
+ assert classes[0].gpu_memory_bytes == 25769803776
+ assert classes[1].max_gpu_count_per_worker == 8
+ assert classes[1].cpu_millis is None
+
+
+@pytest.mark.parametrize("is_async", [False, True])
+@pytest.mark.asyncio
+async def test_remote_list_gpu_classes_older_server(is_async):
+ requests = []
+
+ def handler(request):
+ requests.append(request.path)
+ request.send_response(404)
+ request.end_headers()
+
+ with pytest.raises(NotImplementedError, match="/v1/gpu/classes"):
+ if is_async:
+ async with mock_lancedb_connection_async(handler) as db:
+ await db.list_gpu_classes()
+ else:
+ with mock_lancedb_connection(handler) as db:
+ db.list_gpu_classes()
+
+ assert requests == ["/v1/gpu/classes"]
+
+
def _property_graph_handler(calls, respond):
"""A mock route handler recording ``(method, path, body)`` and answering
with the ``(status, body)`` that ``respond(path, body)`` returns, as JSON."""
diff --git a/python/src/connection.rs b/python/src/connection.rs
index d4516d405..d88991235 100644
--- a/python/src/connection.rs
+++ b/python/src/connection.rs
@@ -24,6 +24,7 @@ use lancedb::{
connection::NamespaceClientPushdownOperation,
database::namespace::LanceNamespaceDatabase,
database::{CreateTableMode, Database, ReadConsistency},
+ function::GpuClass as LanceGpuClass,
graph::PropertyGraphDefinition,
};
use pyo3::{
@@ -38,6 +39,62 @@ pub struct Connection {
inner: Option,
}
+/// Configured hardware and worker limits of a Function GPU class.
+#[pyclass(module = "lancedb._lancedb", get_all, from_py_object)]
+#[derive(Clone, Debug)]
+pub struct GpuClass {
+ pub name: String,
+ pub resource_name: String,
+ pub gpu_model: Option,
+ pub gpu_memory_bytes: Option,
+ pub max_gpu_count_per_worker: Option,
+ pub cpu_millis: Option,
+ pub memory_bytes: Option,
+ pub ephemeral_storage_bytes: Option,
+}
+
+impl From for GpuClass {
+ fn from(class: LanceGpuClass) -> Self {
+ Self {
+ name: class.name,
+ resource_name: class.resource_name,
+ gpu_model: class.gpu_model,
+ gpu_memory_bytes: class.gpu_memory_bytes,
+ max_gpu_count_per_worker: class.max_gpu_count_per_worker,
+ cpu_millis: class.cpu_millis,
+ memory_bytes: class.memory_bytes,
+ ephemeral_storage_bytes: class.ephemeral_storage_bytes,
+ }
+ }
+}
+
+#[pymethods]
+impl GpuClass {
+ fn __repr__(&self) -> String {
+ let gpu_model = self
+ .gpu_model
+ .as_ref()
+ .map_or_else(|| "None".to_string(), |value| format!("{value:?}"));
+ let optional = |value: Option| {
+ value.map_or_else(|| "None".to_string(), |value| value.to_string())
+ };
+ let max_gpu_count = self
+ .max_gpu_count_per_worker
+ .map_or_else(|| "None".to_string(), |value| value.to_string());
+ format!(
+ "GpuClass(name={:?}, resource_name={:?}, gpu_model={gpu_model}, \
+ gpu_memory_bytes={}, max_gpu_count_per_worker={max_gpu_count}, \
+ cpu_millis={}, memory_bytes={}, ephemeral_storage_bytes={})",
+ self.name,
+ self.resource_name,
+ optional(self.gpu_memory_bytes),
+ optional(self.cpu_millis),
+ optional(self.memory_bytes),
+ optional(self.ephemeral_storage_bytes),
+ )
+ }
+}
+
impl Connection {
pub(crate) fn new(inner: LanceConnection) -> Self {
Self { inner: Some(inner) }
@@ -1202,6 +1259,14 @@ impl Connection {
))
}
+ pub fn list_gpu_classes(self_: PyRef<'_, Self>) -> PyResult> {
+ let inner = self_.get_inner()?.clone();
+ future_into_py(self_.py(), async move {
+ let classes = inner.list_gpu_classes().await.infer_error()?;
+ Ok(classes.into_iter().map(GpuClass::from).collect::>())
+ })
+ }
+
pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
diff --git a/python/src/lib.rs b/python/src/lib.rs
index 71662add1..00cd233d0 100644
--- a/python/src/lib.rs
+++ b/python/src/lib.rs
@@ -65,6 +65,7 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
.write_style("LANCEDB_LOG_STYLE");
env_logger::init_from_env(env);
m.add_class::()?;
+ m.add_class::()?;
m.add_class::()?;
m.add_class::()?;
m.add_class::()?;
diff --git a/python/src/table.rs b/python/src/table.rs
index d9e006a38..8d758697c 100644
--- a/python/src/table.rs
+++ b/python/src/table.rs
@@ -42,12 +42,30 @@ use pyo3::{
mod scannable;
-fn refresh_column_options(batch_size: Option) -> PyResult {
+fn refresh_column_options(
+ batch_size: Option,
+ gpu_class: Option,
+ gpu_count: Option,
+) -> PyResult {
let mut options = RefreshColumnOptions::default();
options.batch_size = batch_size
.map(usize::try_from)
.transpose()
.map_err(|_| PyValueError::new_err("batch_size must be greater than zero"))?;
+ options.gpu_class = gpu_class;
+ options.gpu_count = gpu_count
+ .map(|count| {
+ u32::try_from(count)
+ .ok()
+ .filter(|count| *count > 0)
+ .ok_or_else(|| {
+ PyValueError::new_err(
+ "gpu_count must be a positive integer no greater than u32::MAX",
+ )
+ })
+ })
+ .transpose()?;
+ options.validate().infer_error()?;
Ok(options)
}
@@ -2080,13 +2098,15 @@ impl Table {
})
}
- #[pyo3(signature = (column, *, batch_size=None))]
+ #[pyo3(signature = (column, *, batch_size=None, gpu_class=None, gpu_count=None))]
pub fn refresh_column(
self_: PyRef<'_, Self>,
column: String,
batch_size: Option,
+ gpu_class: Option,
+ gpu_count: Option,
) -> PyResult> {
- let options = refresh_column_options(batch_size)?;
+ let options = refresh_column_options(batch_size, gpu_class, gpu_count)?;
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let result = inner
@@ -2097,13 +2117,15 @@ impl Table {
})
}
- #[pyo3(signature = (column, *, batch_size=None))]
+ #[pyo3(signature = (column, *, batch_size=None, gpu_class=None, gpu_count=None))]
pub fn refresh_column_async(
self_: PyRef<'_, Self>,
column: String,
batch_size: Option,
+ gpu_class: Option,
+ gpu_count: Option,
) -> PyResult> {
- let options = refresh_column_options(batch_size)?;
+ let options = refresh_column_options(batch_size, gpu_class, gpu_count)?;
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let job = inner
diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs
index c0c1890b9..ef77d3729 100644
--- a/rust/lancedb/src/connection.rs
+++ b/rust/lancedb/src/connection.rs
@@ -28,6 +28,7 @@ use crate::database::{
};
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};
use crate::error::{Error, Result};
+use crate::function::GpuClass;
use crate::graph::{PropertyGraphDefinition, PropertyGraphDescription};
use crate::listing::{Listing, ListingOptions, ListingPage};
#[cfg(feature = "remote")]
@@ -1244,6 +1245,16 @@ impl Connection {
self.internal.open_job(job_id.as_ref()).await
}
+ /// List the operator-configured GPU classes available to Function workers.
+ ///
+ /// The result describes configured hardware, not current free capacity.
+ /// The principal needs System Usage permission. LanceDB Embedded and
+ /// older Distributed servers without `GET /v1/gpu/classes` return
+ /// [`Error::NotSupported`].
+ pub async fn list_gpu_classes(&self) -> Result> {
+ self.database().list_gpu_classes().await
+ }
+
/// List server-side jobs across the database's tables.
///
/// Results are fetched lazily, one page at a time. See [`Listing`]
diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs
index 3dfdf444c..fda5fc911 100644
--- a/rust/lancedb/src/database.rs
+++ b/rust/lancedb/src/database.rs
@@ -29,6 +29,7 @@ use lance_namespace::models::{
use crate::data::scannable::Scannable;
use crate::error::Result;
+use crate::function::GpuClass;
use crate::graph::{PropertyGraphDefinition, PropertyGraphDescription};
use crate::job::Job;
use crate::listing::{ListingOptions, ListingPage};
@@ -305,6 +306,12 @@ fn function_catalog_not_supported() -> Result {
})
}
+fn gpu_class_catalog_not_supported() -> Result {
+ Err(crate::error::Error::NotSupported {
+ message: "GPU class discovery is supported only by LanceDB Distributed".to_string(),
+ })
+}
+
fn secret_catalog_not_supported() -> Result {
Err(crate::error::Error::NotSupported {
message: "Secret operations are not supported by this database".to_string(),
@@ -377,6 +384,10 @@ pub trait Database:
) -> Result> {
function_catalog_not_supported()
}
+ /// List the operator-configured GPU classes available to Function workers.
+ async fn list_gpu_classes(&self) -> Result> {
+ gpu_class_catalog_not_supported()
+ }
/// Create a materialized view through a remote catalog and return its
/// initial-population job. Local connections use the native declaration
/// path directly.
diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs
index a9797254e..d5086698c 100644
--- a/rust/lancedb/src/function.rs
+++ b/rust/lancedb/src/function.rs
@@ -19,6 +19,47 @@ use crate::{Error, Result};
/// Semantic Function type for a Blob v2 value.
pub const FUNCTION_BLOB_V2_TYPE: &str = "blob_v2";
+/// Configured placement-independent hardware and worker limits of a Function
+/// GPU class.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[non_exhaustive]
+pub struct GpuClass {
+ /// Stable class name accepted by Function registration and refresh options.
+ pub name: String,
+ /// Kubernetes extended resource requested by this class.
+ pub resource_name: String,
+ /// Human-readable GPU model, when configured by the operator.
+ pub gpu_model: Option,
+ /// Usable memory of one GPU in the class, in bytes.
+ pub gpu_memory_bytes: Option,
+ /// Maximum GPUs one worker in this class may request.
+ pub max_gpu_count_per_worker: Option,
+ /// CPU limit for a worker in this class, in millicores.
+ pub cpu_millis: Option,
+ /// Memory limit for a worker in this class, in bytes.
+ pub memory_bytes: Option,
+ /// Ephemeral-storage limit for a worker in this class, in bytes.
+ pub ephemeral_storage_bytes: Option,
+}
+
+pub(crate) const GPU_CLASS_NAME_ERROR: &str = "gpu_class must be a lowercase name of at most 63 letters, digits, or hyphens, and cannot start or end with a hyphen";
+
+pub(crate) fn validate_gpu_class_name(name: &str) -> Result<()> {
+ if name.is_empty()
+ || name.len() > 63
+ || name.starts_with('-')
+ || name.ends_with('-')
+ || !name
+ .bytes()
+ .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-')
+ {
+ return Err(Error::InvalidInput {
+ message: GPU_CLASS_NAME_ERROR.to_string(),
+ });
+ }
+ Ok(())
+}
+
fn invalid_json(error: impl std::fmt::Display) -> Error {
Error::InvalidInput {
message: format!("invalid remote Function JSON: {error}"),
@@ -260,6 +301,7 @@ pub enum PythonRuntimeSpec {
/// system_packages: Vec::new(),
/// },
/// env: BTreeMap::new(),
+ /// gpu_class: Some("a10g".to_string()),
/// };
/// assert!(runtime.requires_gpu());
/// ```
@@ -267,6 +309,8 @@ pub enum PythonRuntimeSpec {
python_version: String,
environment: PythonEnvironmentSpec,
env: BTreeMap,
+ /// Operator-defined GPU class; absent uses the deployment default.
+ gpu_class: Option,
},
/// A runtime kind introduced by a newer server.
///
@@ -317,6 +361,14 @@ impl PythonRuntimeSpec {
pub fn requires_gpu(&self) -> bool {
matches!(self, Self::PythonV2 { .. })
}
+
+ /// Operator-defined GPU class, or `None` to use the deployment default.
+ pub fn gpu_class(&self) -> Option<&str> {
+ match self {
+ Self::PythonV2 { gpu_class, .. } => gpu_class.as_deref(),
+ _ => None,
+ }
+ }
}
#[derive(Deserialize)]
@@ -327,6 +379,8 @@ struct PythonRuntimeV1Wire {
env: BTreeMap,
#[serde(default)]
gpu: Option,
+ #[serde(default)]
+ gpu_class: Option,
}
#[derive(Deserialize)]
@@ -336,6 +390,8 @@ struct PythonRuntimeV2Wire {
#[serde(default)]
env: BTreeMap,
gpu: bool,
+ #[serde(default)]
+ gpu_class: Option,
}
impl<'de> Deserialize<'de> for PythonRuntimeSpec {
@@ -356,6 +412,11 @@ impl<'de> Deserialize<'de> for PythonRuntimeSpec {
"python runtime with gpu requires kind='python_v2'",
));
}
+ if wire.gpu_class.is_some() {
+ return Err(de::Error::custom(
+ "python runtime with gpu_class requires kind='python_v2'",
+ ));
+ }
Ok(Self::Python {
python_version: wire.python_version,
environment: wire.environment,
@@ -368,10 +429,14 @@ impl<'de> Deserialize<'de> for PythonRuntimeSpec {
if !wire.gpu {
return Err(de::Error::custom("runtime.gpu must be true"));
}
+ if let Some(gpu_class) = wire.gpu_class.as_deref() {
+ validate_gpu_class_name(gpu_class).map_err(de::Error::custom)?;
+ }
Ok(Self::PythonV2 {
python_version: wire.python_version,
environment: wire.environment,
env: wire.env,
+ gpu_class: wire.gpu_class,
})
}
_ => Ok(Self::Unrecognized { kind }),
@@ -390,6 +455,8 @@ impl Serialize for PythonRuntimeSpec {
env: &'a BTreeMap,
#[serde(skip_serializing_if = "Option::is_none")]
gpu: Option,
+ #[serde(skip_serializing_if = "Option::is_none")]
+ gpu_class: Option<&'a str>,
}
#[derive(Serialize)]
@@ -408,18 +475,21 @@ impl Serialize for PythonRuntimeSpec {
environment,
env,
gpu: None,
+ gpu_class: None,
}
.serialize(serializer),
Self::PythonV2 {
python_version,
environment,
env,
+ gpu_class,
} => PythonRuntimeRef {
kind: "python_v2",
python_version,
environment,
env,
gpu: Some(true),
+ gpu_class: gpu_class.as_deref(),
}
.serialize(serializer),
Self::Unrecognized { kind } => UnrecognizedRuntimeRef { kind }.serialize(serializer),
@@ -473,6 +543,9 @@ pub struct FunctionVersion {
created_at_millis: i64,
metadata: BTreeMap,
disabled: bool,
+ /// Operator-defined GPU class declared at registration.
+ #[serde(default, skip_serializing_if = "Option::is_none")]
+ gpu_class: Option,
}
impl FunctionVersion {
@@ -527,6 +600,11 @@ impl FunctionVersion {
self.source_digest.as_deref()
}
+ /// Operator-defined GPU class declared at registration.
+ pub fn gpu_class(&self) -> Option<&str> {
+ self.gpu_class.as_deref()
+ }
+
pub fn signature(&self) -> &FunctionSignature {
&self.signature
}
@@ -985,14 +1063,26 @@ mod conda_environment_tests {
super::canonical_json(&runtime).unwrap(),
r#"{"environment":{"kind":"pip"},"gpu":true,"kind":"python_v2","python_version":"3.12"}"#
);
+ let runtime: PythonRuntimeSpec = serde_json::from_str(
+ r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":true,"gpu_class":"a10g"}"#,
+ )
+ .unwrap();
+ assert_eq!(runtime.gpu_class(), Some("a10g"));
+ assert_eq!(
+ super::canonical_json(&runtime).unwrap(),
+ r#"{"environment":{"kind":"pip"},"gpu":true,"gpu_class":"a10g","kind":"python_v2","python_version":"3.12"}"#
+ );
for invalid in [
r#"{"kind":"python","python_version":"3.12","environment":{"kind":"pip"},"gpu":true}"#,
+ r#"{"kind":"python","python_version":"3.12","environment":{"kind":"pip"},"gpu_class":"a10g"}"#,
r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"}}"#,
r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":1}"#,
r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":false}"#,
r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":"true"}"#,
r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":"H100"}"#,
+ r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":true,"gpu_class":"H100"}"#,
+ r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":true,"gpu_class":"-bad"}"#,
] {
assert!(serde_json::from_str::(invalid).is_err());
}
@@ -1043,6 +1133,12 @@ mod tests {
}
assert_eq!(FunctionVersion::from_json(&encoded).unwrap(), version);
}
+
+ document["gpu_class"] = serde_json::json!("a10g");
+ let version = FunctionVersion::from_json(&document.to_string()).unwrap();
+ assert_eq!(version.gpu_class(), Some("a10g"));
+ let encoded: Value = serde_json::from_str(&version.to_canonical_json().unwrap()).unwrap();
+ assert_eq!(encoded["gpu_class"], "a10g");
}
#[test]
diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs
index 99961986b..871460a99 100644
--- a/rust/lancedb/src/remote/db.rs
+++ b/rust/lancedb/src/remote/db.rs
@@ -25,7 +25,7 @@ use crate::database::{
ExecuteQueryRequest, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest,
};
use crate::error::Result;
-use crate::function::{FunctionRegistrationRequest, FunctionVersion};
+use crate::function::{FunctionRegistrationRequest, FunctionVersion, GpuClass};
use crate::graph::{PropertyGraphDefinition, PropertyGraphDescription};
use crate::job::Job;
use crate::listing::{ListingOptions, ListingPage};
@@ -1164,6 +1164,32 @@ impl Database for RemoteDatabase {
}
}
+ async fn list_gpu_classes(&self) -> Result> {
+ #[derive(serde::Deserialize)]
+ struct Response {
+ gpu_classes: Vec,
+ }
+
+ let req = self.client.get("/v1/gpu/classes");
+ let (request_id, response) = self.client.send(req).await?;
+ let response = match self.client.check_response(&request_id, response).await {
+ Ok(response) => response,
+ Err(Error::Http {
+ status_code: Some(StatusCode::NOT_FOUND),
+ ..
+ }) => {
+ return Err(Error::NotSupported {
+ message: "GPU class discovery requires a newer LanceDB server; \
+ GET /v1/gpu/classes returned 404 Not Found"
+ .to_string(),
+ });
+ }
+ Err(error) => return Err(error),
+ };
+ let response: Response = response.json().await.err_to_http(request_id)?;
+ Ok(response.gpu_classes)
+ }
+
async fn list_jobs(&self, options: ListingOptions) -> Result> {
let mut body = serde_json::json!({});
if let Some(token) = options.page_token {
@@ -3845,6 +3871,66 @@ mod tests {
}
}
+ #[tokio::test]
+ async fn test_list_gpu_classes_reports_configured_hardware() {
+ let conn = Connection::new_with_handler(|request| {
+ assert_eq!(request.method(), &reqwest::Method::GET);
+ assert_eq!(request.url().path(), "/v1/gpu/classes");
+ http::Response::builder()
+ .status(200)
+ .body(
+ serde_json::json!({
+ "gpu_classes": [
+ {
+ "name": "default",
+ "resource_name": "nvidia.com/gpu",
+ "gpu_model": "NVIDIA A10G",
+ "gpu_memory_bytes": 25_769_803_776_u64,
+ "max_gpu_count_per_worker": 1,
+ "cpu_millis": 4000,
+ "memory_bytes": 25_769_803_776_u64,
+ "ephemeral_storage_bytes": 1_073_741_824_u64,
+ },
+ {
+ "name": "h100",
+ "resource_name": "nvidia.com/gpu",
+ "gpu_model": "NVIDIA H100",
+ "gpu_memory_bytes": 85_899_345_920_u64,
+ "max_gpu_count_per_worker": 8,
+ }
+ ]
+ })
+ .to_string(),
+ )
+ .unwrap()
+ });
+
+ let classes = conn.list_gpu_classes().await.unwrap();
+ assert_eq!(classes.len(), 2);
+ assert_eq!(classes[0].name, "default");
+ assert_eq!(classes[0].gpu_model.as_deref(), Some("NVIDIA A10G"));
+ assert_eq!(classes[0].gpu_memory_bytes, Some(25_769_803_776));
+ assert_eq!(classes[1].max_gpu_count_per_worker, Some(8));
+ assert_eq!(classes[1].cpu_millis, None);
+ }
+
+ #[tokio::test]
+ async fn test_list_gpu_classes_explains_older_server_404() {
+ let conn = Connection::new_with_handler(|_| {
+ http::Response::builder()
+ .status(404)
+ .body("not found")
+ .unwrap()
+ });
+ let error = conn.list_gpu_classes().await.unwrap_err();
+ assert!(matches!(
+ error,
+ Error::NotSupported { message }
+ if message.contains("newer LanceDB server")
+ && message.contains("/v1/gpu/classes")
+ ));
+ }
+
#[tokio::test]
async fn test_list_jobs_does_not_truncate_after_100_pages() {
let conn = Connection::new_with_handler(|request| {
diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs
index adc0c1bfc..3a9cd567c 100644
--- a/rust/lancedb/src/remote/table.rs
+++ b/rust/lancedb/src/remote/table.rs
@@ -3794,6 +3794,12 @@ impl BaseTable for RemoteTable {
if let Some(batch_size) = options.batch_size {
body["batch_size"] = serde_json::json!(batch_size);
}
+ if let Some(gpu_class) = options.gpu_class {
+ body["gpu_class"] = serde_json::json!(gpu_class);
+ }
+ if let Some(gpu_count) = options.gpu_count {
+ body["gpu_count"] = serde_json::json!(gpu_count);
+ }
self.apply_branch_body(&mut body);
let request = self
.client
@@ -10558,7 +10564,10 @@ mod tests {
.create_branch("exp", Ref::Version(None, None))
.await
.unwrap();
- let options = RefreshColumnOptions { batch_size };
+ let options = RefreshColumnOptions {
+ batch_size,
+ ..Default::default()
+ };
let job = branch
.refresh_column_async_with_options("quality", options)
.await
@@ -10566,6 +10575,65 @@ mod tests {
assert_eq!(job.id(), Some("j-42"));
}
+ #[tokio::test]
+ async fn test_refresh_column_async_gpu_placement_options() {
+ use crate::table::RefreshColumnOptions;
+
+ let table = Table::new_with_handler("my_table", |request| {
+ assert_eq!(request.method(), "POST");
+ assert_eq!(request.url().path(), "/v1/table/my_table/backfill_column");
+ assert_eq!(
+ request_body_json(&request),
+ serde_json::json!({
+ "column": "quality",
+ "gpu_class": "a10g",
+ "gpu_count": 2,
+ })
+ );
+ http::Response::builder()
+ .status(202)
+ .body(r#"{"job_id": "j-gpu"}"#)
+ .unwrap()
+ });
+ let options = RefreshColumnOptions::default()
+ .gpu_class("a10g")
+ .gpu_count(2);
+ let job = table
+ .refresh_column_async_with_options("quality", options)
+ .await
+ .unwrap();
+ assert_eq!(job.id(), Some("j-gpu"));
+ }
+
+ #[tokio::test]
+ async fn test_refresh_column_rejects_invalid_gpu_options_before_request() {
+ use crate::table::RefreshColumnOptions;
+
+ let table = Table::new_with_handler("my_table", |_| -> http::Response {
+ panic!("invalid GPU options must not send a request")
+ });
+ for (options, expected) in [
+ (
+ RefreshColumnOptions::default().gpu_class("A10G"),
+ "gpu_class",
+ ),
+ (
+ RefreshColumnOptions::default().gpu_class("-a10g"),
+ "gpu_class",
+ ),
+ (RefreshColumnOptions::default().gpu_count(0), "gpu_count"),
+ ] {
+ let error = table
+ .refresh_column_async_with_options("quality", options)
+ .await
+ .unwrap_err();
+ assert!(
+ matches!(&error, Error::InvalidInput { message } if message.contains(expected)),
+ "{error:?}"
+ );
+ }
+ }
+
#[tokio::test]
async fn test_refresh_column_rejects_zero_batch_size_before_request() {
use crate::table::RefreshColumnOptions;
diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs
index ca1ee32e4..bd34427e6 100644
--- a/rust/lancedb/src/table.rs
+++ b/rust/lancedb/src/table.rs
@@ -1958,9 +1958,12 @@ impl Table {
/// Like [`Table::refresh_column`], with explicit refresh options.
///
/// [`RefreshColumnOptions::batch_size()`] caps rows per Function call in
- /// LanceDB Distributed and must be greater than zero when set. It has no
- /// effect on expression-backed columns in LanceDB Embedded. Remote refreshes
- /// require [`Table::refresh_column_async_with_options`].
+ /// LanceDB Distributed and must be greater than zero when set.
+ /// [`RefreshColumnOptions::gpu_class()`] selects a GPU pool and
+ /// [`RefreshColumnOptions::gpu_count()`] sets the aggregate GPU budget.
+ /// These options have no effect on expression-backed columns in LanceDB
+ /// Embedded. Remote refreshes require
+ /// [`Table::refresh_column_async_with_options`].
///
/// ```
/// # use lancedb::{Table, table::RefreshColumnOptions};
@@ -2015,15 +2018,23 @@ impl Table {
///
/// Set [`RefreshColumnOptions::batch_size()`] to cap the rows each remote
/// Function call receives, for example to keep slow Functions within the
- /// executor's per-call deadline. An unset batch size uses the server default;
- /// zero is rejected before submitting a job. Embedded expression refreshes
- /// ignore the batch size.
+ /// executor's per-call deadline. Set
+ /// [`RefreshColumnOptions::gpu_class()`] and
+ /// [`RefreshColumnOptions::gpu_count()`] to override the registered GPU
+ /// class or set an aggregate GPU budget for Function workers. Unset values
+ /// use the Function and server defaults; zero counts and malformed class
+ /// names are rejected before submitting a job. Embedded expression
+ /// refreshes ignore all three options.
///
/// ```
/// # use lancedb::{Table, table::RefreshColumnOptions};
/// # async fn refresh(table: &Table) -> Result<(), Box> {
/// let job = table.refresh_column_async_with_options(
- /// "quality", RefreshColumnOptions::default().batch_size(8),
+ /// "quality",
+ /// RefreshColumnOptions::default()
+ /// .batch_size(8)
+ /// .gpu_class("a10g")
+ /// .gpu_count(2),
/// ).await?;
/// let result = job.wait().await?;
/// # Ok(())
diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs
index 7000be879..18dd832eb 100644
--- a/rust/lancedb/src/table/refresh.rs
+++ b/rust/lancedb/src/table/refresh.rs
@@ -51,6 +51,7 @@ use serde::{Deserialize, Serialize};
use super::computed_columns::{BoundExpression, ComputedColumnKind, computed_column_from_field};
use super::freshness::{self, SignatureMap, StalenessPlan};
use super::{BaseTable, NativeTable};
+use crate::function::validate_gpu_class_name;
use crate::job::Job;
use crate::{Error, Result};
@@ -63,6 +64,16 @@ pub struct RefreshColumnOptions {
/// Must be greater than zero. `None` uses the server default. This has no
/// effect on expression-backed columns in LanceDB Embedded.
pub batch_size: Option,
+ /// GPU class for Function workers in LanceDB Distributed.
+ ///
+ /// `None` keeps the Function's registered class, or the deployment default
+ /// when the Function does not declare one.
+ pub gpu_class: Option,
+ /// Positive aggregate GPU budget for Function workers in LanceDB
+ /// Distributed.
+ ///
+ /// `None` leaves the server's configured budget unchanged.
+ pub gpu_count: Option,
}
impl RefreshColumnOptions {
@@ -72,12 +83,33 @@ impl RefreshColumnOptions {
self
}
- pub(crate) fn validate(&self) -> Result<()> {
+ /// Request Function workers from `gpu_class`.
+ pub fn gpu_class(mut self, gpu_class: impl Into) -> Self {
+ self.gpu_class = Some(gpu_class.into());
+ self
+ }
+
+ /// Request an aggregate GPU budget of `gpu_count`.
+ pub fn gpu_count(mut self, gpu_count: u32) -> Self {
+ self.gpu_count = Some(gpu_count);
+ self
+ }
+
+ /// Validate client-side constraints before any request is sent.
+ pub fn validate(&self) -> Result<()> {
if self.batch_size == Some(0) {
return Err(Error::InvalidInput {
message: "batch_size must be greater than zero".into(),
});
}
+ if let Some(gpu_class) = &self.gpu_class {
+ validate_gpu_class_name(gpu_class)?;
+ }
+ if self.gpu_count == Some(0) {
+ return Err(Error::InvalidInput {
+ message: "gpu_count must be greater than zero".into(),
+ });
+ }
Ok(())
}
}