diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index ae40b4e3f..2125542af 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -283,6 +283,7 @@ class PythonEnvironmentSpec(_RemoteValue): path: Optional[str] = None modules: tuple[str, ...] = () image: Optional[str] = None + system_packages: tuple[str, ...] = () class PythonRuntimeSpec(_RemoteValue): @@ -1854,6 +1855,7 @@ class UdfDefinition: gpu: bool = False, conda: tuple[str, ...] = (), conda_channels: tuple[str, ...] = (), + apt: tuple[str, ...] = (), code: tuple[types.ModuleType, ...] = (), ): is_class = inspect.isclass(function) @@ -1869,14 +1871,23 @@ class UdfDefinition: if conda_channels and not conda: raise ValueError("conda_channels requires conda packages") packages = tuple(sorted(set(conda if conda else pip))) - if any(not package or package != package.strip() for package in packages): + system_packages = tuple(sorted(set(apt))) + if any( + not package or package != package.strip() + for package in packages + system_packages + ): raise ValueError("package requirements must be non-empty and trimmed") if conda: environment_spec = PythonEnvironmentSpec( - kind="conda", packages=packages, channels=tuple(conda_channels) + kind="conda", + packages=packages, + channels=tuple(conda_channels), + system_packages=system_packages, ) else: - environment_spec = PythonEnvironmentSpec(kind="pip", packages=packages) + environment_spec = PythonEnvironmentSpec( + kind="pip", packages=packages, system_packages=system_packages + ) environment = dict(env) if any( not isinstance(key, str) or not isinstance(value, str) @@ -2045,6 +2056,7 @@ def udf( gpu: bool = False, conda: tuple[str, ...] | list[str] = (), conda_channels: tuple[str, ...] | list[str] = (), + apt: tuple[str, ...] | list[str] = (), code: Sequence[types.ModuleType] = (), ) -> Callable[[Callable[..., Any]], UdfDefinition]: ... @@ -2061,6 +2073,7 @@ def udf( gpu: bool = False, conda: tuple[str, ...] | list[str] = (), conda_channels: tuple[str, ...] | list[str] = (), + apt: tuple[str, ...] | list[str] = (), code: Sequence[types.ModuleType] = (), ): """Prepare a Python function or class for remote Function registration. @@ -2150,6 +2163,12 @@ def udf( Conda packages for the remote environment, instead of ``pip``. conda_channels : sequence of str, optional Conda channels in priority order; requires ``conda``. + apt : sequence of str, optional + Debian packages (``"name"`` or ``"name=version"``) added to the remote + image before the Python environment, for shared libraries pip and + conda cannot provide, such as FFmpeg's for ``torchcodec``. Packages + are unpacked without running their maintainer scripts. Works with + ``pip`` or ``conda``; requires image execution. env : mapping of str to str, optional Environment variables included in the Function definition. Not for credentials -- these are ordinary configuration, stored with the @@ -2219,6 +2238,7 @@ def udf( gpu=gpu, conda=tuple(conda), conda_channels=tuple(conda_channels), + apt=tuple(apt), code=tuple(code), ) diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 0e8aca4e9..562a90eb4 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -494,6 +494,30 @@ def test_udf_conda_environment(): udf(name="channels", conda_channels=["conda-forge"])(lambda value: value) +def test_udf_apt_system_packages(): + @udf(pip=["torchcodec"], apt=["libavutil58", "libavcodec60", "libavutil58"]) + def decode(value: float) -> float: + return value + + request = json.loads(decode.registration_request.to_canonical_json()) + assert request["runtime"]["environment"] == { + "kind": "pip", + "packages": ["torchcodec"], + "system_packages": ["libavcodec60", "libavutil58"], + } + + @udf(conda=["numpy"], apt=["libsndfile1"]) + def conda_decode(value: float) -> float: + return value + + request = json.loads(conda_decode.registration_request.to_canonical_json()) + assert request["runtime"]["environment"]["system_packages"] == ["libsndfile1"] + pip_request = json.loads(normalize_score.registration_request.to_canonical_json()) + assert "system_packages" not in pip_request["runtime"]["environment"] + with pytest.raises(ValueError, match="trimmed"): + udf(name="padded", apt=[" libsndfile1"])(lambda value: value) + + def test_udf_gpu_marker_uses_gpu_runtime(): @udf(pip=["cupy-cuda12x"], gpu=True) def double_on_gpu(value: int) -> int: diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index db93f0071..128a7efdf 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -224,6 +224,10 @@ pub struct PythonEnvironmentSpec { pub modules: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] pub image: Option, + /// Debian packages (`name` or `name=version`) added to the Function's + /// image before its Python environment is installed. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub system_packages: Vec, } /// Reproducible Python runtime definition understood by Sophon. @@ -253,6 +257,7 @@ pub enum PythonRuntimeSpec { /// path: None, /// modules: Vec::new(), /// image: None, + /// system_packages: Vec::new(), /// }, /// env: BTreeMap::new(), /// }; @@ -938,6 +943,27 @@ mod conda_environment_tests { assert!(!serde_json::to_string(&pip).unwrap().contains("channels")); } + #[test] + fn system_packages_round_trip_and_stay_absent_when_empty() { + let apt: PythonEnvironmentSpec = serde_json::from_str( + r#"{"kind":"pip","packages":["torchcodec"],"system_packages":["libavutil58"]}"#, + ) + .unwrap(); + assert_eq!(apt.system_packages, ["libavutil58"]); + assert!( + serde_json::to_string(&apt) + .unwrap() + .contains(r#""system_packages":["libavutil58"]"#) + ); + let pip: PythonEnvironmentSpec = + serde_json::from_str(r#"{"kind":"pip","packages":["numpy"]}"#).unwrap(); + assert!( + !serde_json::to_string(&pip) + .unwrap() + .contains("system_packages") + ); + } + #[test] fn gpu_python_runtime_marker_round_trips_and_validates() { let runtime: PythonRuntimeSpec = serde_json::from_str(