mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-30 15:26:52 +00:00
feat(python): declare Debian system packages on a Function
Adds `apt` to `@udf` and `system_packages` to the Rust `PythonEnvironmentSpec`. The server adds the packages to the Function's image before installing its Python environment, for shared libraries that pip and conda cannot provide, such as FFmpeg's for torchcodec. The field is omitted when empty, so existing Function identities are unchanged.
This commit is contained in:
@@ -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),
|
||||
)
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -224,6 +224,10 @@ pub struct PythonEnvironmentSpec {
|
||||
pub modules: Vec<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub image: Option<String>,
|
||||
/// 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<String>,
|
||||
}
|
||||
|
||||
/// 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(
|
||||
|
||||
Reference in New Issue
Block a user