rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-09-04 04:28:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(functions): use a GPU capability marker

Browse Source
This commit is contained in:
Xuanwo
2026-08-30 20:06:14 +08:00
parent 1443bf703f
commit d8f321c1f4
3 changed files with 58 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/functions.py
+24 -37
View File
@@ -54,28 +54,16 @@ _UInt32 = conint(strict=True, ge=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 _validate_gpu_wire_marker(value: Any) -> bool:
if value is not True:
raise ValueError("runtime.gpu must be true")
return True
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)
def _normalize_gpu_marker(value: bool) -> Optional[bool]:
if not isinstance(value, bool):
raise ValueError("gpu must be a boolean")
return True if value else None
class _FrozenDict(dict):
@@ -263,7 +251,7 @@ class PythonRuntimeSpec(_RemoteValue):
python_version: Optional[str] = None
environment: Optional[PythonEnvironmentSpec] = None
env: Optional[Mapping[str, str]] = None
gpu: Optional[str] = None
gpu: Optional[bool] = None
@model_validator(mode="before")
@classmethod
@@ -276,10 +264,10 @@ class PythonRuntimeSpec(_RemoteValue):
@field_validator("gpu", mode="before")
@classmethod
def _validate_gpu_requirement(cls, value):
def _validate_gpu_marker(cls, value):
if value is None:
return None
return _validate_gpu_wire_requirement(value)
return _validate_gpu_wire_marker(value)
@model_validator(mode="after")
def _validate_runtime_kind(self):
@@ -308,7 +296,7 @@ class PythonRuntimeSpec(_RemoteValue):
class FunctionVersion(_RemoteValue):
"""An exact immutable Function version returned by Enterprise.
The GPU requirement is part of this identity. CPU and memory sizing,
The GPU execution requirement is part of this identity. CPU and memory sizing,
priority, concurrency, and retry policy belong to the execution platform.
"""
@@ -961,7 +949,7 @@ class UdfDefinition:
pip: tuple[str, ...],
env: Mapping[str, str],
python_version: Optional[str],
gpu: Optional[int | str] = None,
gpu: bool = False,
conda: tuple[str, ...] = (),
conda_channels: tuple[str, ...] = (),
):
@@ -990,14 +978,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)
gpu_marker = _normalize_gpu_marker(gpu)
runtime = PythonRuntimeSpec(
kind="python_v2" if gpu_requirement is not None else "python",
kind="python_v2" if gpu_marker is not None else "python",
python_version=python_version
or f"{sys.version_info.major}.{sys.version_info.minor}",
environment=environment_spec,
env=environment,
gpu=gpu_requirement,
gpu=gpu_marker,
)
self._function = function
self._request = FunctionRegistrationRequest(
@@ -1043,7 +1031,7 @@ def udf(
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
python_version: Optional[str] = None,
gpu: Optional[int | str] = None,
gpu: bool = False,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
) -> Callable[[Callable[..., Any]], UdfDefinition]: ...
@@ -1058,7 +1046,7 @@ def udf(
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
python_version: Optional[str] = None,
gpu: Optional[int | str] = None,
gpu: bool = False,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
):
@@ -1091,11 +1079,10 @@ def udf(
Environment variables included in the Function definition.
python_version : str, optional
Remote Python major/minor version. Defaults to the client 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.
gpu : bool, default False
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.
The packaged artifact is a snapshot: the function source plus exactly
the module-level names it references (modules as imports, importable
@@ -1120,11 +1107,11 @@ def udf(
... return value * 2
>>> score(1.5)
3.0
>>> @udf(pip=["cupy-cuda12x"], gpu="H100:8")
>>> @udf(pip=["cupy-cuda12x"], gpu=True)
... def gpu_score(value: int) -> int:
... return value * 2
>>> gpu_score.registration_request.runtime.gpu
'H100:8'
True
"""
def decorate(target: Callable[..., Any]) -> UdfDefinition:
python/python/tests/test_first_class_function_slice2.py
+9 -17
View File
@@ -89,21 +89,14 @@ def test_udf_conda_environment():
udf(name="channels", conda_channels=["conda-forge"])(lambda value: value)
def test_udf_gpu_requirement_uses_gpu_runtime():
@udf(pip=["cupy-cuda12x"], gpu=1)
def test_udf_gpu_marker_uses_gpu_runtime():
@udf(pip=["cupy-cuda12x"], gpu=True)
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"] == "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"
assert request["runtime"]["gpu"] is True
@udf(pip=["pyarrow"])
def cpu_function(value: int) -> int:
@@ -118,8 +111,8 @@ def test_udf_gpu_requirement_uses_gpu_runtime():
def identity(value: int) -> int:
return value
for invalid in [0, -1, 1.5, True, "", "0", "01", " 1"]:
with pytest.raises(ValueError):
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)
base_runtime = {
@@ -127,11 +120,10 @@ def test_udf_gpu_requirement_uses_gpu_runtime():
"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):
runtime = PythonRuntimeSpec.model_validate({**base_runtime, "gpu": True})
assert runtime.gpu is True
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})
rust/lancedb/src/function.rs
+25 -74
View File
@@ -197,23 +197,6 @@ pub struct PythonEnvironmentSpec {
pub image: Option<String>,
}
fn validate_gpu_requirement(gpu: &str) -> std::result::Result<(), String> {
if gpu.is_empty() || gpu.trim() != gpu {
return Err("runtime.gpu must be non-empty and trimmed".to_string());
}
if gpu.bytes().all(|byte| byte.is_ascii_digit()) {
let count = gpu
.parse::<u32>()
.ok()
.filter(|count| *count > 0)
.ok_or_else(|| "a numeric runtime.gpu must be a positive uint32".to_string())?;
if count.to_string() != gpu {
return Err("a numeric runtime.gpu must use canonical decimal form".to_string());
}
}
Ok(())
}
/// Reproducible Python runtime definition understood by Sophon.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
@@ -243,15 +226,13 @@ pub enum PythonRuntimeSpec {
/// image: None,
/// },
/// env: BTreeMap::new(),
/// gpu: "1".to_string(),
/// };
/// assert_eq!(runtime.gpu(), Some("1"));
/// assert!(runtime.requires_gpu());
/// ```
PythonV2 {
python_version: String,
environment: PythonEnvironmentSpec,
env: BTreeMap<String, String>,
gpu: String,
},
/// A runtime kind introduced by a newer server.
///
@@ -298,12 +279,9 @@ impl PythonRuntimeSpec {
}
}
/// GPU requirement understood by the execution platform.
pub fn gpu(&self) -> Option<&str> {
match self {
Self::PythonV2 { gpu, .. } => Some(gpu),
Self::Python { .. } | Self::Unrecognized { .. } => None,
}
/// Whether the runtime requires a GPU selected by the execution platform.
pub fn requires_gpu(&self) -> bool {
matches!(self, Self::PythonV2 { .. })
}
}
@@ -323,7 +301,7 @@ struct PythonRuntimeV2Wire {
environment: PythonEnvironmentSpec,
#[serde(default)]
env: BTreeMap<String, String>,
gpu: String,
gpu: bool,
}
impl<'de> Deserialize<'de> for PythonRuntimeSpec {
@@ -353,12 +331,13 @@ impl<'de> Deserialize<'de> for PythonRuntimeSpec {
"python_v2" => {
let wire: PythonRuntimeV2Wire =
serde_json::from_value(value).map_err(de::Error::custom)?;
validate_gpu_requirement(&wire.gpu).map_err(de::Error::custom)?;
if !wire.gpu {
return Err(de::Error::custom("runtime.gpu must be true"));
}
Ok(Self::PythonV2 {
python_version: wire.python_version,
environment: wire.environment,
env: wire.env,
gpu: wire.gpu,
})
}
_ => Ok(Self::Unrecognized { kind }),
@@ -376,7 +355,7 @@ impl Serialize for PythonRuntimeSpec {
#[serde(skip_serializing_if = "BTreeMap::is_empty")]
env: &'a BTreeMap<String, String>,
#[serde(skip_serializing_if = "Option::is_none")]
gpu: Option<&'a str>,
gpu: Option<bool>,
}
#[derive(Serialize)]
@@ -401,18 +380,14 @@ impl Serialize for PythonRuntimeSpec {
python_version,
environment,
env,
gpu,
} => {
validate_gpu_requirement(gpu).map_err(serde::ser::Error::custom)?;
PythonRuntimeRef {
kind: "python_v2",
python_version,
environment,
env,
gpu: Some(gpu),
}
.serialize(serializer)
} => PythonRuntimeRef {
kind: "python_v2",
python_version,
environment,
env,
gpu: Some(true),
}
.serialize(serializer),
Self::Unrecognized { kind } => UnrecognizedRuntimeRef { kind }.serialize(serializer),
}
}
@@ -420,7 +395,7 @@ impl Serialize for PythonRuntimeSpec {
/// Immutable Function version returned by the Enterprise catalog.
///
/// The GPU requirement is part of this identity. CPU and memory sizing,
/// The GPU execution requirement is part of this identity. CPU and memory sizing,
/// priority, concurrency, and retry policy belong to the execution platform.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FunctionVersion {
@@ -717,52 +692,28 @@ mod conda_environment_tests {
}
#[test]
fn gpu_python_runtime_round_trips_and_validates() {
fn gpu_python_runtime_marker_round_trips_and_validates() {
let runtime: PythonRuntimeSpec = serde_json::from_str(
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":true}"#,
)
.unwrap();
assert_eq!(runtime.kind(), "python_v2");
assert_eq!(runtime.gpu(), Some("1"));
assert!(runtime.requires_gpu());
assert_eq!(
super::canonical_json(&runtime).unwrap(),
r#"{"environment":{"kind":"pip"},"gpu":"1","kind":"python_v2","python_version":"3.12"}"#
r#"{"environment":{"kind":"pip"},"gpu":true,"kind":"python_v2","python_version":"3.12"}"#
);
for requirement in ["H100", "H100:8"] {
let encoded = format!(
r#"{{"kind":"python_v2","python_version":"3.12","environment":{{"kind":"pip"}},"gpu":"{requirement}"}}"#
);
let runtime: PythonRuntimeSpec = serde_json::from_str(&encoded).unwrap();
assert_eq!(runtime.gpu(), Some(requirement));
}
for invalid in [
r#"{"kind":"python","python_version":"3.12","environment":{"kind":"pip"},"gpu":"1"}"#,
r#"{"kind":"python","python_version":"3.12","environment":{"kind":"pip"},"gpu":true}"#,
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":"0"}"#,
r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":"01"}"#,
r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":""}"#,
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"}"#,
] {
assert!(serde_json::from_str::<PythonRuntimeSpec>(invalid).is_err());
}
let invalid_runtime = match runtime {
PythonRuntimeSpec::PythonV2 {
python_version,
environment,
env,
..
} => PythonRuntimeSpec::PythonV2 {
python_version,
environment,
env,
gpu: "0".to_string(),
},
_ => unreachable!(),
};
assert!(super::canonical_json(&invalid_runtime).is_err());
}
#[test]