rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-11-13 14:16:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(python): expose Function GPU classes and placement (#4791)

Browse Source 
## Summary

- Add sync/async `DBConnection.list_gpu_classes()` returning typed
`GpuClass` records for `GET /v1/gpu/classes`. Records include the class
name, Kubernetes resource, GPU model and memory, per-worker GPU maximum,
and worker CPU/memory/ephemeral-storage limits.
- Add `udf(gpu_class=...)` for immutable Function-level GPU placement.
It requires `gpu=True`; Function versions expose their registered
`gpu_class`.
- Add `gpu_class` and `gpu_count` to `refresh_column` /
`refresh_column_async` and Rust `RefreshColumnOptions`. Both are sent as
top-level `backfill_column` fields and validated before any request.
`gpu_count` is the aggregate GPU budget across Function workers.
- Validate GPU class names with the server lowercase 63-character
DNS-label rule.
- Document the discovery-to-registration flow, configured hardware
versus free capacity, System Usage permission, and per-refresh GPU
placement in the API and web guides.

Older servers without `/v1/gpu/classes` produce a clear
`NotImplementedError` instead of a bare 404.

Part of ENT-3470. The client follows the in-flight Sophon contract on
`xuanwo/ent-3470-fix` for GPU class discovery and backfill placement
options.

## Breaking change

`PythonRuntimeSpec::PythonV2` adds `gpu_class`, and
`RefreshColumnOptions` adds `gpu_class` and `gpu_count`. Rust callers
constructing either value with all fields explicitly must supply the new
fields.

## Validation

At `7685c968`: Rust GPU tests (5 passed), Rust/PyO3 clippy, Python
extension build, pyright, ruff, and Python Function/runtime, discovery,
refresh placement, and refresh regressions (147 passed, 6 skipped).

At `b22428fd`: Rust GPU tests (5 passed), rustfmt, ruff, pyright,
Function slice 2 (85 passed), and remote GPU/refresh tests (20 passed).

The final commit only updates documentation and Rust doc comments.
This commit is contained in:
Xuanwo
2026-11-02 08:50:36 +08:00
committed by GitHub
parent 6a255926d7
commit 63aa06ef7d
24 changed files with 921 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/src/python/python.md
+2
View File
@@ -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
docs/web/curation/create-function.mdx
+20 -4
View File
@@ -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:
<CodeGroup>
```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"))})
```
</CodeGroup>
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.
docs/web/feature-engineering/refresh-column.mdx
+58
View File
@@ -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:
<CodeGroup>
```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<dyn std::error::Error>> {
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(())
}
```
</CodeGroup>
`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.
<Note>
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.
</Note>
## When a refresh fails {#failures}
What a failure keeps depends on where the refresh ran. `refresh_column` in your
python/python/lancedb/__init__.py
+1
View File
@@ -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
python/python/lancedb/_lancedb.pyi
+23 -2
View File
@@ -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,
python/python/lancedb/db.py
+22
View File
@@ -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,
python/python/lancedb/functions.py
+49 -2
View File
@@ -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),
python/python/lancedb/remote/db.py
+8 -1
View File
@@ -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,
python/python/lancedb/remote/table.py
+30 -4
View File
@@ -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(
python/python/lancedb/table.py
+100 -10
View File
@@ -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,
)
python/python/tests/test_first_class_function_slice1.py
+11
View File
@@ -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"))
python/python/tests/test_first_class_function_slice2.py
+21 -2
View File
@@ -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",
python/python/tests/test_package_metadata.py
+9
View File
@@ -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()
python/python/tests/test_remote_db.py
+149 -1
View File
@@ -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."""
python/src/connection.rs
+65
View File
@@ -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<LanceConnection>,
}
/// 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<String>,
pub gpu_memory_bytes: Option<u64>,
pub max_gpu_count_per_worker: Option<u32>,
pub cpu_millis: Option<u64>,
pub memory_bytes: Option<u64>,
pub ephemeral_storage_bytes: Option<u64>,
}
impl From<LanceGpuClass> 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<u64>| {
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<Bound<'_, PyAny>> {
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::<Vec<_>>())
})
}
pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
python/src/lib.rs
+1
View File
@@ -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::<Connection>()?;
m.add_class::<connection::GpuClass>()?;
m.add_class::<catalog::Catalog>()?;
m.add_class::<catalog::DatabaseNames>()?;
m.add_class::<listing::NameListing>()?;
python/src/table.rs
+27 -5
View File
@@ -42,12 +42,30 @@ use pyo3::{
mod scannable;
fn refresh_column_options(batch_size: Option<i64>) -> PyResult<RefreshColumnOptions> {
fn refresh_column_options(
batch_size: Option<i64>,
gpu_class: Option<String>,
gpu_count: Option<i64>,
) -> PyResult<RefreshColumnOptions> {
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<i64>,
gpu_class: Option<String>,
gpu_count: Option<i64>,
) -> PyResult<Bound<'_, PyAny>> {
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<i64>,
gpu_class: Option<String>,
gpu_count: Option<i64>,
) -> PyResult<Bound<'_, PyAny>> {
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
rust/lancedb/src/connection.rs
+11
View File
@@ -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<Vec<GpuClass>> {
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`]
rust/lancedb/src/database.rs
+11
View File
@@ -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<T>() -> Result<T> {
})
}
fn gpu_class_catalog_not_supported<T>() -> Result<T> {
Err(crate::error::Error::NotSupported {
message: "GPU class discovery is supported only by LanceDB Distributed".to_string(),
})
}
fn secret_catalog_not_supported<T>() -> Result<T> {
Err(crate::error::Error::NotSupported {
message: "Secret operations are not supported by this database".to_string(),
@@ -377,6 +384,10 @@ pub trait Database:
) -> Result<crate::job::Job<crate::function::FunctionVersion>> {
function_catalog_not_supported()
}
/// List the operator-configured GPU classes available to Function workers.
async fn list_gpu_classes(&self) -> Result<Vec<GpuClass>> {
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.
rust/lancedb/src/function.rs
+96
View File
@@ -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<String>,
/// Usable memory of one GPU in the class, in bytes.
pub gpu_memory_bytes: Option<u64>,
/// Maximum GPUs one worker in this class may request.
pub max_gpu_count_per_worker: Option<u32>,
/// CPU limit for a worker in this class, in millicores.
pub cpu_millis: Option<u64>,
/// Memory limit for a worker in this class, in bytes.
pub memory_bytes: Option<u64>,
/// Ephemeral-storage limit for a worker in this class, in bytes.
pub ephemeral_storage_bytes: Option<u64>,
}
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<String, String>,
/// Operator-defined GPU class; absent uses the deployment default.
gpu_class: Option<String>,
},
/// 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<String, String>,
#[serde(default)]
gpu: Option<Value>,
#[serde(default)]
gpu_class: Option<String>,
}
#[derive(Deserialize)]
@@ -336,6 +390,8 @@ struct PythonRuntimeV2Wire {
#[serde(default)]
env: BTreeMap<String, String>,
gpu: bool,
#[serde(default)]
gpu_class: Option<String>,
}
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<String, String>,
#[serde(skip_serializing_if = "Option::is_none")]
gpu: Option<bool>,
#[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<String, String>,
disabled: bool,
/// Operator-defined GPU class declared at registration.
#[serde(default, skip_serializing_if = "Option::is_none")]
gpu_class: Option<String>,
}
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::<PythonRuntimeSpec>(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]
rust/lancedb/src/remote/db.rs
+87 -1
View File
@@ -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<S: HttpSend> Database for RemoteDatabase<S> {
}
}
async fn list_gpu_classes(&self) -> Result<Vec<GpuClass>> {
#[derive(serde::Deserialize)]
struct Response {
gpu_classes: Vec<GpuClass>,
}
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<ListingPage<JobInfo>> {
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| {
rust/lancedb/src/remote/table.rs
+69 -1
View File
@@ -3794,6 +3794,12 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
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<String> {
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;
rust/lancedb/src/table.rs
+18 -7
View File
@@ -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<dyn std::error::Error>> {
/// 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(())
rust/lancedb/src/table/refresh.rs
+33 -1
View File
@@ -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<usize>,
/// 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<String>,
/// Positive aggregate GPU budget for Function workers in LanceDB
/// Distributed.
///
/// `None` leaves the server's configured budget unchanged.
pub gpu_count: Option<u32>,
}
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<String>) -> 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(())
}
}