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:
Wyatt Alt
2026-10-29 15:28:36 +08:00
committed by GitHub
parent 079711d4ad
commit 41c1a2a8b1
3 changed files with 73 additions and 3 deletions
+23 -3
View File
@@ -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:
+26
View File
@@ -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(