diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 5f359f4b7..636c1bf90 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -68,6 +68,8 @@ 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 21ffc8860..28658bd85 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -27,6 +27,7 @@ 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 8a19a9d37..2d53fe69b 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -51,6 +51,7 @@ 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) @@ -228,6 +229,12 @@ 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. @@ -239,6 +246,7 @@ class PythonRuntimeSpec(_RemoteValue): python_version: Optional[str] = None environment: Optional[PythonEnvironmentSpec] = None env: Optional[Mapping[str, str]] = None + resources: Optional[FunctionResourceRequirements] = None @model_validator(mode="after") def _validate_runtime_kind(self): @@ -247,18 +255,30 @@ 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'" + ) + 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") else: object.__setattr__(self, "python_version", None) object.__setattr__(self, "environment", None) object.__setattr__(self, "env", None) + object.__setattr__(self, "resources", 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. + Required execution resources are part of this identity. Priority, + concurrency, and retry policy belong to the submitting Job. """ name: str @@ -910,12 +930,18 @@ class UdfDefinition: pip: tuple[str, ...], env: Mapping[str, str], python_version: Optional[str], + num_gpus: Optional[int], 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: @@ -939,11 +965,12 @@ class UdfDefinition: source = _package_source(function) digest = f"sha256:{hashlib.sha256(source).hexdigest()}" runtime = PythonRuntimeSpec( - kind="python", + kind="python_v2" if num_gpus 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, ) self._function = function self._request = FunctionRegistrationRequest( @@ -989,6 +1016,7 @@ def udf( pip: tuple[str, ...] | list[str] = (), env: Optional[Mapping[str, str]] = None, python_version: Optional[str] = None, + num_gpus: Optional[int] = None, conda: tuple[str, ...] | list[str] = (), conda_channels: tuple[str, ...] | list[str] = (), ) -> Callable[[Callable[..., Any]], UdfDefinition]: ... @@ -1003,6 +1031,7 @@ def udf( pip: tuple[str, ...] | list[str] = (), env: Optional[Mapping[str, str]] = None, python_version: Optional[str] = None, + num_gpus: Optional[int] = None, conda: tuple[str, ...] | list[str] = (), conda_channels: tuple[str, ...] | list[str] = (), ): @@ -1035,6 +1064,10 @@ 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. The packaged artifact is a snapshot: the function source plus exactly the module-level names it references (modules as imports, importable @@ -1070,6 +1103,7 @@ def udf( pip=tuple(pip), env={} if env is None else env, python_version=python_version, + num_gpus=num_gpus, conda=tuple(conda), conda_channels=tuple(conda_channels), ) @@ -1089,6 +1123,7 @@ __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 7ce6b6b91..6074000a6 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -89,6 +89,30 @@ 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 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} + + @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 "resources" not in cpu_runtime + + for invalid in [0, -1, 1.5, True]: + with pytest.raises(ValueError): + udf(name="invalid_gpu", num_gpus=invalid)(lambda value: value) + + 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..2b862ba9c 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -197,6 +197,13 @@ 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, +} + /// Reproducible Python runtime definition understood by Sophon. #[derive(Debug, Clone, PartialEq, Eq)] #[non_exhaustive] @@ -207,6 +214,13 @@ pub enum PythonRuntimeSpec { environment: PythonEnvironmentSpec, env: BTreeMap, }, + /// The resource-aware Sophon-managed Python runtime. + PythonV2 { + python_version: String, + environment: PythonEnvironmentSpec, + env: BTreeMap, + resources: FunctionResourceRequirements, + }, /// A runtime kind introduced by a newer server. /// /// Unknown payload fields are intentionally not retained because the @@ -219,22 +233,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,10 +261,18 @@ 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, } } + + /// Immutable execution resources, or `None` for the V1 or an unknown runtime. + pub fn resources(&self) -> Option<&FunctionResourceRequirements> { + match self { + Self::PythonV2 { resources, .. } => Some(resources), + Self::Python { .. } | Self::Unrecognized { .. } => None, + } + } } #[derive(Deserialize)] @@ -257,23 +284,53 @@ struct PythonRuntimeWire { environment: Option, #[serde(default)] env: BTreeMap, + #[serde(default)] + resources: Option, } 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 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() { + "python" => { + if wire.resources.is_some() { + return Err(de::Error::custom( + "python runtime with resources requires kind='python_v2'", + )); + } + Ok(Self::Python { + python_version: python_version()?, + environment: 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", + )); + } + Ok(Self::PythonV2 { + python_version: python_version()?, + environment: environment()?, + env: wire.env, + resources, + }) + } + _ => Ok(Self::Unrecognized { kind: wire.kind }), } } } @@ -287,6 +344,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")] + resources: Option<&'a FunctionResourceRequirements>, } #[derive(Serialize)] @@ -304,6 +363,20 @@ impl Serialize for PythonRuntimeSpec { python_version, environment, env, + resources: None, + } + .serialize(serializer), + Self::PythonV2 { + python_version, + environment, + env, + resources, + } => PythonRuntimeRef { + kind: "python_v2", + python_version, + environment, + env, + resources: Some(resources), } .serialize(serializer), Self::Unrecognized { kind } => UnrecognizedRuntimeRef { kind }.serialize(serializer), @@ -313,8 +386,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. +/// Required execution resources are part of this identity. Priority, +/// concurrency, and retry policy belong to the submitting Job. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct FunctionVersion { name: String, @@ -589,7 +662,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 +681,26 @@ 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 resource_aware_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}}"#, + ) + .unwrap(); + assert_eq!(runtime.kind(), "python_v2"); + assert_eq!(runtime.resources().unwrap().num_gpus, 1); + assert_eq!( + super::canonical_json(&runtime).unwrap(), + r#"{"environment":{"kind":"pip"},"kind":"python_v2","python_version":"3.12","resources":{"num_gpus":1}}"# + ); + + for invalid in [ + r#"{"kind":"python","python_version":"3.12","environment":{"kind":"pip"},"resources":{"num_gpus":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}}"#, + ] { + assert!(serde_json::from_str::(invalid).is_err()); + } + } }