From a417e46bface04b813f54458b76f7181f4b7bdb7 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sun, 30 Aug 2026 23:10:08 +0800 Subject: [PATCH] feat(functions): support GPU resource requirements (#4085) Functions can describe their Python environment today, but cannot declare accelerator requirements. That prevents Sophon from scheduling computed-column UDF refreshes onto GPU workers from the immutable Function definition. Add `num_gpus` to Python `@udf` through a typed `FunctionResourceRequirements` value and represent resource-aware definitions with the `python_v2` runtime discriminator. CPU Functions retain their existing `python` encoding and canonical identity. The new discriminator is intentional for mixed-version safety: deployments that do not understand execution resources reject the runtime instead of accepting a new field and silently running the Function on CPU. Required resources are part of Function version identity; priority, concurrency, and retry policy remain Job concerns. The actual resource scheduling remains owned by Sophon. --- python/python/lancedb/functions.py | 60 +++++- .../tests/test_first_class_function_slice2.py | 54 +++++- rust/lancedb/src/function.rs | 177 +++++++++++++++--- 3 files changed, 260 insertions(+), 31 deletions(-) diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 3e1b4f2d6..9be63a558 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -54,6 +54,18 @@ _UInt32 = conint(strict=True, ge=0, le=2**32 - 1) _UInt64 = conint(strict=True, ge=0, le=2**64 - 1) +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_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): def _immutable(self, *args, **kwargs): raise TypeError("remote canonical values are immutable") @@ -239,6 +251,23 @@ class PythonRuntimeSpec(_RemoteValue): python_version: Optional[str] = None environment: Optional[PythonEnvironmentSpec] = None env: Optional[Mapping[str, str]] = None + gpu: Optional[bool] = 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_marker(cls, value): + if value is None: + return None + return _validate_gpu_wire_marker(value) @model_validator(mode="after") def _validate_runtime_kind(self): @@ -247,18 +276,28 @@ class PythonRuntimeSpec(_RemoteValue): raise ValueError("python runtime requires python_version") if self.environment is None: raise ValueError("python runtime requires environment") + 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.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, "gpu", None) return self class FunctionVersion(_RemoteValue): """An exact immutable Function version returned by Enterprise. - Scheduling resources, priority, concurrency, and retry policy belong to - the submitting Job and are not part of this identity. + The GPU execution requirement is part of this identity. CPU and memory sizing, + priority, concurrency, and retry policy belong to the execution platform. """ name: str @@ -996,6 +1035,7 @@ class UdfDefinition: pip: tuple[str, ...], env: Mapping[str, str], python_version: Optional[str], + gpu: bool = False, conda: tuple[str, ...] = (), conda_channels: tuple[str, ...] = (), ): @@ -1024,12 +1064,14 @@ class UdfDefinition: signature = _infer_signature(function, input_schema, output_schema) source = _package_source(function) digest = f"sha256:{hashlib.sha256(source).hexdigest()}" + gpu_marker = _normalize_gpu_marker(gpu) runtime = PythonRuntimeSpec( - kind="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_marker, ) self._function = function self._request = FunctionRegistrationRequest( @@ -1075,6 +1117,7 @@ def udf( pip: tuple[str, ...] | list[str] = (), env: Optional[Mapping[str, str]] = None, python_version: Optional[str] = None, + gpu: bool = False, conda: tuple[str, ...] | list[str] = (), conda_channels: tuple[str, ...] | list[str] = (), ) -> Callable[[Callable[..., Any]], UdfDefinition]: ... @@ -1089,6 +1132,7 @@ def udf( pip: tuple[str, ...] | list[str] = (), env: Optional[Mapping[str, str]] = None, python_version: Optional[str] = None, + gpu: bool = False, conda: tuple[str, ...] | list[str] = (), conda_channels: tuple[str, ...] | list[str] = (), ): @@ -1121,6 +1165,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 : 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 @@ -1145,6 +1193,11 @@ def udf( ... return value * 2 >>> score(1.5) 3.0 + >>> @udf(pip=["cupy-cuda12x"], gpu=True) + ... def gpu_score(value: int) -> int: + ... return value * 2 + >>> gpu_score.registration_request.runtime.gpu + True """ def decorate(target: Callable[..., Any]) -> UdfDefinition: @@ -1156,6 +1209,7 @@ def udf( pip=tuple(pip), env={} if env is None else env, python_version=python_version, + gpu=gpu, conda=tuple(conda), conda_channels=tuple(conda_channels), ) diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index ee14043e4..cf1542b55 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,6 +89,58 @@ def test_udf_conda_environment(): udf(name="channels", conda_channels=["conda-forge"])(lambda value: value) +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"] is True + + @udf(pip=["pyarrow"]) + def cpu_function(value: int) -> int: + return value + + cpu_runtime = json.loads(cpu_function.registration_request.to_canonical_json())[ + "runtime" + ] + assert cpu_runtime["kind"] == "python" + assert "gpu" not in cpu_runtime + + def identity(value: int) -> int: + return value + + 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 = { + "kind": "python_v2", + "python_version": "3.12", + "environment": {"kind": "pip"}, + } + 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}) + + +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(): @udf def word_norm(body: str) -> float: diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 5366d984e..4b31a4376 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -207,6 +207,33 @@ pub enum PythonRuntimeSpec { environment: PythonEnvironmentSpec, env: BTreeMap, }, + /// 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(), + /// }; + /// assert!(runtime.requires_gpu()); + /// ``` + PythonV2 { + python_version: String, + environment: PythonEnvironmentSpec, + env: BTreeMap, + }, /// A runtime kind introduced by a newer server. /// /// Unknown payload fields are intentionally not retained because the @@ -219,22 +246,27 @@ impl PythonRuntimeSpec { pub fn kind(&self) -> &str { match self { Self::Python { .. } => "python", + Self::PythonV2 { .. } => "python_v2", Self::Unrecognized { kind } => kind, } } - /// The Python version for the V1 runtime, or `None` for an unknown kind. + /// The Python version for a known Python runtime, or `None` for an unknown kind. pub fn python_version(&self) -> Option<&str> { match self { - Self::Python { python_version, .. } => Some(python_version), + Self::Python { python_version, .. } | Self::PythonV2 { python_version, .. } => { + Some(python_version) + } Self::Unrecognized { .. } => None, } } - /// The Python environment for the V1 runtime, or `None` for an unknown kind. + /// The Python environment for a known Python runtime, or `None` for an unknown kind. pub fn environment(&self) -> Option<&PythonEnvironmentSpec> { match self { - Self::Python { environment, .. } => Some(environment), + Self::Python { environment, .. } | Self::PythonV2 { environment, .. } => { + Some(environment) + } Self::Unrecognized { .. } => None, } } @@ -242,38 +274,73 @@ impl PythonRuntimeSpec { /// Environment variables, or `None` for an unknown kind. pub fn env(&self) -> Option<&BTreeMap> { match self { - Self::Python { env, .. } => Some(env), + Self::Python { env, .. } | Self::PythonV2 { env, .. } => Some(env), Self::Unrecognized { .. } => None, } } + + /// Whether the runtime requires a GPU selected by the execution platform. + pub fn requires_gpu(&self) -> bool { + matches!(self, Self::PythonV2 { .. }) + } } #[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)] + gpu: Option, +} + +#[derive(Deserialize)] +struct PythonRuntimeV2Wire { + python_version: String, + environment: PythonEnvironmentSpec, + #[serde(default)] + env: BTreeMap, + gpu: bool, } impl<'de> Deserialize<'de> for PythonRuntimeSpec { fn deserialize>(deserializer: D) -> std::result::Result { - let wire = PythonRuntimeWire::deserialize(deserializer)?; - if wire.kind == "python" { - Ok(Self::Python { - python_version: wire - .python_version - .ok_or_else(|| de::Error::missing_field("python_version"))?, - environment: wire - .environment - .ok_or_else(|| de::Error::missing_field("environment"))?, - env: wire.env, - }) - } else { - Ok(Self::Unrecognized { kind: wire.kind }) + 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" => { + 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 gpu requires kind='python_v2'", + )); + } + Ok(Self::Python { + python_version: wire.python_version, + environment: wire.environment, + env: wire.env, + }) + } + "python_v2" => { + let wire: PythonRuntimeV2Wire = + serde_json::from_value(value).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, + }) + } + _ => Ok(Self::Unrecognized { kind }), } } } @@ -287,6 +354,8 @@ impl Serialize for PythonRuntimeSpec { environment: &'a PythonEnvironmentSpec, #[serde(skip_serializing_if = "BTreeMap::is_empty")] env: &'a BTreeMap, + #[serde(skip_serializing_if = "Option::is_none")] + gpu: Option, } #[derive(Serialize)] @@ -304,6 +373,19 @@ impl Serialize for PythonRuntimeSpec { python_version, environment, env, + gpu: None, + } + .serialize(serializer), + Self::PythonV2 { + python_version, + environment, + env, + } => PythonRuntimeRef { + kind: "python_v2", + python_version, + environment, + env, + gpu: Some(true), } .serialize(serializer), Self::Unrecognized { kind } => UnrecognizedRuntimeRef { kind }.serialize(serializer), @@ -313,8 +395,8 @@ impl Serialize for PythonRuntimeSpec { /// Immutable Function version returned by the Enterprise catalog. /// -/// Scheduling resources, priority, concurrency, and retry policy belong to -/// the submitting Job and are not part of this identity. +/// 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 { name: String, @@ -589,7 +671,7 @@ impl_json!(RefreshColumnResult); #[cfg(test)] mod conda_environment_tests { - use super::PythonEnvironmentSpec; + use super::{PythonEnvironmentSpec, PythonRuntimeSpec}; #[test] fn conda_channels_round_trip_and_pip_stays_bare() { @@ -608,4 +690,45 @@ mod conda_environment_tests { serde_json::from_str(r#"{"kind":"pip","packages":["numpy"]}"#).unwrap(); assert!(!serde_json::to_string(&pip).unwrap().contains("channels")); } + + #[test] + 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":true}"#, + ) + .unwrap(); + assert_eq!(runtime.kind(), "python_v2"); + assert!(runtime.requires_gpu()); + assert_eq!( + super::canonical_json(&runtime).unwrap(), + r#"{"environment":{"kind":"pip"},"gpu":true,"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_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"}"#, + ] { + assert!(serde_json::from_str::(invalid).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"}"# + ); + } + } }