feat(functions): support GPU resource requirements

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

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,

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",

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:

rust/lancedb/src/function.rs

@@ -197,6 +197,13 @@ pub struct PythonEnvironmentSpec {
     pub image: Option<String>,
 }
 
+/// 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<String, String>,
     },
+    /// The resource-aware Sophon-managed Python runtime.
+    PythonV2 {
+        python_version: String,
+        environment: PythonEnvironmentSpec,
+        env: BTreeMap<String, String>,
+        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<String, String>> {
         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<PythonEnvironmentSpec>,
     #[serde(default)]
     env: BTreeMap<String, String>,
+    #[serde(default)]
+    resources: Option<FunctionResourceRequirements>,
 }
 
 impl<'de> Deserialize<'de> for PythonRuntimeSpec {
     fn deserialize<D: Deserializer<'de>>(deserializer: D) -> std::result::Result<Self, D::Error> {
         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<String, String>,
+            #[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::<PythonRuntimeSpec>(invalid).is_err());
+        }
+    }
 }