diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index ac6bf384e..641d75332 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -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: diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index ec452d6b3..1ec8d3b06 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -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}) diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index e81cd057b..4b31a4376 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -197,23 +197,6 @@ pub struct PythonEnvironmentSpec { pub image: Option, } -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. #[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, - 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, - 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, #[serde(skip_serializing_if = "Option::is_none")] - gpu: Option<&'a str>, + gpu: Option, } #[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::(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]