mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
feat(python): expose Function GPU classes and placement (#4791)
## 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:
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
+100
-10
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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>()?;
|
||||
|
||||
+27
-5
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user