From 1443bf703fd7100bffc8e2ca75ab67a6fd9eecc5 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 29 Aug 2026 19:12:45 +0800 Subject: [PATCH] refactor(functions): simplify GPU requirements --- docs/src/python/python.md | 2 - python/python/lancedb/__init__.py | 1 - python/python/lancedb/functions.py | 98 ++++++--- .../tests/test_first_class_function_slice2.py | 50 ++++- rust/lancedb/src/function.rs | 203 ++++++++++++------ 5 files changed, 248 insertions(+), 106 deletions(-) diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 636c1bf90..5f359f4b7 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -68,8 +68,6 @@ listing a storage directory. ::: lancedb.functions.PythonEnvironmentSpec -::: lancedb.functions.FunctionResourceRequirements - ::: lancedb.functions.udf ::: lancedb.functions.UdfDefinition diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 28658bd85..21ffc8860 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -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, diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 2d53fe69b..ac6bf384e 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -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", diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 6074000a6..ec452d6b3 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -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(): diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 2b862ba9c..e81cd057b 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -197,11 +197,21 @@ pub struct PythonEnvironmentSpec { pub image: Option, } -/// Immutable resources required every time a Function version executes. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct FunctionResourceRequirements { - /// Whole NVIDIA GPUs requested from the execution substrate. - pub num_gpus: u32, +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::() + .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. @@ -214,12 +224,34 @@ pub enum PythonRuntimeSpec { environment: PythonEnvironmentSpec, env: BTreeMap, }, - /// The resource-aware Sophon-managed Python runtime. + /// The GPU-enabled Sophon-managed Python runtime. + /// + /// # Examples + /// + /// ``` + /// use std::collections::BTreeMap; + /// use lancedb::function::{PythonEnvironmentSpec, PythonRuntimeSpec}; + /// + /// let runtime = PythonRuntimeSpec::PythonV2 { + /// python_version: "3.12".to_string(), + /// environment: PythonEnvironmentSpec { + /// kind: "pip".to_string(), + /// packages: vec!["cupy-cuda12x".to_string()], + /// channels: Vec::new(), + /// path: None, + /// modules: Vec::new(), + /// image: None, + /// }, + /// env: BTreeMap::new(), + /// gpu: "1".to_string(), + /// }; + /// assert_eq!(runtime.gpu(), Some("1")); + /// ``` PythonV2 { python_version: String, environment: PythonEnvironmentSpec, env: BTreeMap, - resources: FunctionResourceRequirements, + gpu: String, }, /// A runtime kind introduced by a newer server. /// @@ -266,71 +298,70 @@ impl PythonRuntimeSpec { } } - /// Immutable execution resources, or `None` for the V1 or an unknown runtime. - pub fn resources(&self) -> Option<&FunctionResourceRequirements> { + /// GPU requirement understood by the execution platform. + pub fn gpu(&self) -> Option<&str> { match self { - Self::PythonV2 { resources, .. } => Some(resources), + Self::PythonV2 { gpu, .. } => Some(gpu), Self::Python { .. } | Self::Unrecognized { .. } => None, } } } #[derive(Deserialize)] -struct PythonRuntimeWire { - kind: String, - #[serde(default)] - python_version: Option, - #[serde(default)] - environment: Option, +struct PythonRuntimeV1Wire { + python_version: String, + environment: PythonEnvironmentSpec, #[serde(default)] env: BTreeMap, #[serde(default)] - resources: Option, + gpu: Option, +} + +#[derive(Deserialize)] +struct PythonRuntimeV2Wire { + python_version: String, + environment: PythonEnvironmentSpec, + #[serde(default)] + env: BTreeMap, + gpu: String, } impl<'de> Deserialize<'de> for PythonRuntimeSpec { fn deserialize>(deserializer: D) -> std::result::Result { - let wire = PythonRuntimeWire::deserialize(deserializer)?; - let python_version = || { - wire.python_version - .clone() - .ok_or_else(|| de::Error::missing_field("python_version")) - }; - let environment = || { - wire.environment - .clone() - .ok_or_else(|| de::Error::missing_field("environment")) - }; - match wire.kind.as_str() { + let value = Value::deserialize(deserializer)?; + let kind = value + .get("kind") + .ok_or_else(|| de::Error::missing_field("kind"))? + .as_str() + .ok_or_else(|| de::Error::custom("runtime.kind must be a string"))? + .to_string(); + match kind.as_str() { "python" => { - if wire.resources.is_some() { + let wire: PythonRuntimeV1Wire = + serde_json::from_value(value).map_err(de::Error::custom)?; + if wire.gpu.is_some() { return Err(de::Error::custom( - "python runtime with resources requires kind='python_v2'", + "python runtime with gpu requires kind='python_v2'", )); } Ok(Self::Python { - python_version: python_version()?, - environment: environment()?, + python_version: wire.python_version, + environment: wire.environment, env: wire.env, }) } "python_v2" => { - let resources = wire - .resources - .ok_or_else(|| de::Error::missing_field("resources"))?; - if resources.num_gpus == 0 { - return Err(de::Error::custom( - "resources.num_gpus must be greater than zero", - )); - } + let wire: PythonRuntimeV2Wire = + serde_json::from_value(value).map_err(de::Error::custom)?; + validate_gpu_requirement(&wire.gpu).map_err(de::Error::custom)?; Ok(Self::PythonV2 { - python_version: python_version()?, - environment: environment()?, + python_version: wire.python_version, + environment: wire.environment, env: wire.env, - resources, + gpu: wire.gpu, }) } - _ => Ok(Self::Unrecognized { kind: wire.kind }), + _ => Ok(Self::Unrecognized { kind }), } } } @@ -345,7 +376,7 @@ impl Serialize for PythonRuntimeSpec { #[serde(skip_serializing_if = "BTreeMap::is_empty")] env: &'a BTreeMap, #[serde(skip_serializing_if = "Option::is_none")] - resources: Option<&'a FunctionResourceRequirements>, + gpu: Option<&'a str>, } #[derive(Serialize)] @@ -363,22 +394,25 @@ impl Serialize for PythonRuntimeSpec { python_version, environment, env, - resources: None, + gpu: None, } .serialize(serializer), Self::PythonV2 { python_version, environment, env, - resources, - } => PythonRuntimeRef { - kind: "python_v2", - python_version, - environment, - env, - resources: Some(resources), + gpu, + } => { + validate_gpu_requirement(gpu).map_err(serde::ser::Error::custom)?; + PythonRuntimeRef { + kind: "python_v2", + python_version, + environment, + env, + gpu: Some(gpu), + } + .serialize(serializer) } - .serialize(serializer), Self::Unrecognized { kind } => UnrecognizedRuntimeRef { kind }.serialize(serializer), } } @@ -386,8 +420,8 @@ impl Serialize for PythonRuntimeSpec { /// Immutable Function version returned by the Enterprise catalog. /// -/// 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. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct FunctionVersion { name: String, @@ -683,24 +717,67 @@ mod conda_environment_tests { } #[test] - fn resource_aware_python_runtime_round_trips_and_validates() { + fn gpu_python_runtime_round_trips_and_validates() { let runtime: PythonRuntimeSpec = serde_json::from_str( - r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"resources":{"num_gpus":1}}"#, + r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"gpu":"1"}"#, ) .unwrap(); assert_eq!(runtime.kind(), "python_v2"); - assert_eq!(runtime.resources().unwrap().num_gpus, 1); + assert_eq!(runtime.gpu(), Some("1")); assert_eq!( super::canonical_json(&runtime).unwrap(), - r#"{"environment":{"kind":"pip"},"kind":"python_v2","python_version":"3.12","resources":{"num_gpus":1}}"# + r#"{"environment":{"kind":"pip"},"gpu":"1","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"},"resources":{"num_gpus":1}}"#, + r#"{"kind":"python","python_version":"3.12","environment":{"kind":"pip"},"gpu":"1"}"#, r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"}}"#, - r#"{"kind":"python_v2","python_version":"3.12","environment":{"kind":"pip"},"resources":{"num_gpus":0}}"#, + 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":""}"#, ] { assert!(serde_json::from_str::(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] + fn unknown_runtime_discards_payload_before_known_field_validation() { + for encoded in [ + r#"{"kind":"python_v3","gpu":{"model":"H100"}}"#, + r#"{"kind":"python_v3","resources":[]}"#, + r#"{"kind":"python_v3","python_version":3.15,"environment":{"kind":[]}}"#, + ] { + let runtime: PythonRuntimeSpec = serde_json::from_str(encoded).unwrap(); + assert_eq!(runtime.kind(), "python_v3"); + assert_eq!( + super::canonical_json(&runtime).unwrap(), + r#"{"kind":"python_v3"}"# + ); + } } }