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