rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-08-31 10:38:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(functions): simplify GPU requirements

Browse Source
This commit is contained in:
Xuanwo
2026-08-29 19:12:45 +08:00
parent 9e8140ff88
commit 1443bf703f
5 changed files with 248 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/__init__.py
-1
View File
@@ -27,7 +27,6 @@ from .functions import (
FunctionApplication as FunctionApplication,
FunctionBinding as FunctionBinding,
FunctionRegistrationRequest as FunctionRegistrationRequest,
FunctionResourceRequirements as FunctionResourceRequirements,
FunctionVersion as FunctionVersion,
PythonRuntimeSpec as PythonRuntimeSpec,
RefreshColumnResult as RefreshColumnResult,
python/python/lancedb/functions.py
+65 -33
View File
@@ -51,10 +51,33 @@ from pydantic import (
_Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1)
_UInt32 = conint(strict=True, ge=0, le=2**32 - 1)
_PositiveUInt32 = conint(strict=True, gt=0, le=2**32 - 1)
_UInt64 = conint(strict=True, ge=0, le=2**64 - 1)
def _validate_gpu_wire_requirement(value: Any) -> str:
if not isinstance(value, str):
raise ValueError("runtime.gpu must be a string")
if not value or value != value.strip():
raise ValueError("runtime.gpu must be non-empty and trimmed")
if value.isascii() and value.isdigit():
count = int(value)
if not 0 < count <= 2**32 - 1 or str(count) != value:
raise ValueError(
"a numeric runtime.gpu must be a canonical positive uint32"
)
return value
def _normalize_gpu_requirement(value: Optional[int | str]) -> Optional[str]:
if value is None:
return None
if isinstance(value, int) and not isinstance(value, bool):
if not 0 < value <= 2**32 - 1:
raise ValueError("gpu must be a whole number greater than zero")
return str(value)
return _validate_gpu_wire_requirement(value)
class _FrozenDict(dict):
def _immutable(self, *args, **kwargs):
raise TypeError("remote canonical values are immutable")
@@ -229,12 +252,6 @@ class PythonEnvironmentSpec(_RemoteValue):
image: Optional[str] = None
class FunctionResourceRequirements(_RemoteValue):
"""Immutable resources required every time a Function version executes."""
num_gpus: _PositiveUInt32
class PythonRuntimeSpec(_RemoteValue):
"""Remote runtime definition with environment values.
@@ -246,7 +263,23 @@ class PythonRuntimeSpec(_RemoteValue):
python_version: Optional[str] = None
environment: Optional[PythonEnvironmentSpec] = None
env: Optional[Mapping[str, str]] = None
resources: Optional[FunctionResourceRequirements] = None
gpu: Optional[str] = None
@model_validator(mode="before")
@classmethod
def _discard_unknown_runtime_payload(cls, value):
if isinstance(value, Mapping):
kind = value.get("kind")
if isinstance(kind, str) and kind not in {"python", "python_v2"}:
return {"kind": kind}
return value
@field_validator("gpu", mode="before")
@classmethod
def _validate_gpu_requirement(cls, value):
if value is None:
return None
return _validate_gpu_wire_requirement(value)
@model_validator(mode="after")
def _validate_runtime_kind(self):
@@ -255,30 +288,28 @@ class PythonRuntimeSpec(_RemoteValue):
raise ValueError("python runtime requires python_version")
if self.environment is None:
raise ValueError("python runtime requires environment")
if self.resources is not None:
raise ValueError(
"python runtime with resources requires kind='python_v2'"
)
if self.gpu is not None:
raise ValueError("python runtime with gpu requires kind='python_v2'")
elif self.kind == "python_v2":
if self.python_version is None:
raise ValueError("python_v2 runtime requires python_version")
if self.environment is None:
raise ValueError("python_v2 runtime requires environment")
if self.resources is None:
raise ValueError("python_v2 runtime requires resources")
if self.gpu is None:
raise ValueError("python_v2 runtime requires gpu")
else:
object.__setattr__(self, "python_version", None)
object.__setattr__(self, "environment", None)
object.__setattr__(self, "env", None)
object.__setattr__(self, "resources", None)
object.__setattr__(self, "gpu", None)
return self
class FunctionVersion(_RemoteValue):
"""An exact immutable Function version returned by Enterprise.
Required execution resources are part of this identity. Priority,
concurrency, and retry policy belong to the submitting Job.
The GPU requirement is part of this identity. CPU and memory sizing,
priority, concurrency, and retry policy belong to the execution platform.
"""
name: str
@@ -930,18 +961,13 @@ class UdfDefinition:
pip: tuple[str, ...],
env: Mapping[str, str],
python_version: Optional[str],
num_gpus: Optional[int],
gpu: Optional[int | str] = None,
conda: tuple[str, ...] = (),
conda_channels: tuple[str, ...] = (),
):
function_name = name or function.__name__
if not _FUNCTION_NAME.fullmatch(function_name):
raise ValueError(f"invalid Function name: {function_name!r}")
resources = (
FunctionResourceRequirements(num_gpus=num_gpus)
if num_gpus is not None
else None
)
if pip and conda:
raise ValueError("a Function environment is pip or conda, not both")
if conda_channels and not conda:
@@ -964,13 +990,14 @@ class UdfDefinition:
signature = _infer_signature(function, input_schema, output_schema)
source = _package_source(function)
digest = f"sha256:{hashlib.sha256(source).hexdigest()}"
gpu_requirement = _normalize_gpu_requirement(gpu)
runtime = PythonRuntimeSpec(
kind="python_v2" if num_gpus is not None else "python",
kind="python_v2" if gpu_requirement is not None else "python",
python_version=python_version
or f"{sys.version_info.major}.{sys.version_info.minor}",
environment=environment_spec,
env=environment,
resources=resources,
gpu=gpu_requirement,
)
self._function = function
self._request = FunctionRegistrationRequest(
@@ -1016,7 +1043,7 @@ def udf(
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
python_version: Optional[str] = None,
num_gpus: Optional[int] = None,
gpu: Optional[int | str] = None,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
) -> Callable[[Callable[..., Any]], UdfDefinition]: ...
@@ -1031,7 +1058,7 @@ def udf(
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
python_version: Optional[str] = None,
num_gpus: Optional[int] = None,
gpu: Optional[int | str] = None,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
):
@@ -1064,10 +1091,11 @@ def udf(
Environment variables included in the Function definition.
python_version : str, optional
Remote Python major/minor version. Defaults to the client version.
num_gpus : int, optional
Number of whole NVIDIA GPUs required for every remote execution. Must
be greater than zero. The requirement is part of the immutable
Function version.
gpu : int or str, optional
GPU requirement for every remote execution. A positive integer
requests that many compatible NVIDIA GPUs. A string can select a
platform-supported model and count, such as ``"H100:8"``. The
requirement is part of the immutable Function version.
The packaged artifact is a snapshot: the function source plus exactly
the module-level names it references (modules as imports, importable
@@ -1092,6 +1120,11 @@ def udf(
... return value * 2
>>> score(1.5)
3.0
>>> @udf(pip=["cupy-cuda12x"], gpu="H100:8")
... def gpu_score(value: int) -> int:
... return value * 2
>>> gpu_score.registration_request.runtime.gpu
'H100:8'
"""
def decorate(target: Callable[..., Any]) -> UdfDefinition:
@@ -1103,7 +1136,7 @@ def udf(
pip=tuple(pip),
env={} if env is None else env,
python_version=python_version,
num_gpus=num_gpus,
gpu=gpu,
conda=tuple(conda),
conda_channels=tuple(conda_channels),
)
@@ -1123,7 +1156,6 @@ __all__ = [
"FunctionOutput",
"FunctionParameter",
"FunctionRegistrationRequest",
"FunctionResourceRequirements",
"FunctionResultField",
"FunctionSignature",
"FunctionVersion",
python/python/tests/test_first_class_function_slice2.py
+43 -7
View File
@@ -19,7 +19,7 @@ import pyarrow as pa
import pytest
import lancedb
from lancedb.functions import UdfDefinition, udf
from lancedb.functions import PythonRuntimeSpec, UdfDefinition, udf
THRESHOLD = 20
_CACHE = None
@@ -89,14 +89,21 @@ def test_udf_conda_environment():
udf(name="channels", conda_channels=["conda-forge"])(lambda value: value)
def test_udf_gpu_requirement_uses_resource_aware_runtime():
@udf(pip=["cupy-cuda12x"], num_gpus=1)
def test_udf_gpu_requirement_uses_gpu_runtime():
@udf(pip=["cupy-cuda12x"], gpu=1)
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"]["resources"] == {"num_gpus": 1}
assert request["runtime"]["gpu"] == "1"
@udf(pip=["cupy-cuda12x"], gpu="H100:8")
def double_on_h100(value: int) -> int:
return value * 2
h100_request = json.loads(double_on_h100.registration_request.to_canonical_json())
assert h100_request["runtime"]["gpu"] == "H100:8"
@udf(pip=["pyarrow"])
def cpu_function(value: int) -> int:
@@ -106,11 +113,40 @@ def test_udf_gpu_requirement_uses_resource_aware_runtime():
"runtime"
]
assert cpu_runtime["kind"] == "python"
assert "resources" not in cpu_runtime
assert "gpu" not in cpu_runtime
for invalid in [0, -1, 1.5, True]:
def identity(value: int) -> int:
return value
for invalid in [0, -1, 1.5, True, "", "0", "01", " 1"]:
with pytest.raises(ValueError):
udf(name="invalid_gpu", num_gpus=invalid)(lambda value: value)
udf(name="invalid_gpu", gpu=invalid)(identity)
base_runtime = {
"kind": "python_v2",
"python_version": "3.12",
"environment": {"kind": "pip"},
}
for requirement in ["H100", "H100:8"]:
runtime = PythonRuntimeSpec.model_validate({**base_runtime, "gpu": requirement})
assert runtime.gpu == requirement
for invalid in [1, 0, "", "0", "01", " 1"]:
with pytest.raises(ValueError):
PythonRuntimeSpec.model_validate({**base_runtime, "gpu": invalid})
def test_unknown_runtime_discards_payload_before_known_field_validation():
for payload in [
{"kind": "python_v3", "gpu": {"model": "H100"}},
{"kind": "python_v3", "resources": []},
{
"kind": "python_v3",
"environment": {"kind": []},
"python_version": 3.15,
},
]:
runtime = PythonRuntimeSpec.model_validate(payload)
assert runtime.to_canonical_json() == '{"kind":"python_v3"}'
def test_udf_packages_attribute_access_and_body_imports():