diff --git a/changes/249.added b/changes/249.added new file mode 100644 index 00000000..a23447ec --- /dev/null +++ b/changes/249.added @@ -0,0 +1,2 @@ +The `Dependency` object now takes an optional parameter `binary_dependencies` to specify binary packages to be +installed in the computation container. diff --git a/substrafl/dependency/schemas.py b/substrafl/dependency/schemas.py index 9bef2c2d..d2c1378b 100644 --- a/substrafl/dependency/schemas.py +++ b/substrafl/dependency/schemas.py @@ -45,6 +45,7 @@ class Dependency(BaseModel): `setup.py` or `pyproject.toml`). See the documentation of pip wheel for more details. local_code (List[pathlib.Path]): Local relative imports used by your script. All files / folders will be pasted to the level of the running script. + binary_dependencies (List[str]): Binary dependencies to install in the computation container. excluded_paths (List[pathlib.Path]): Local paths excluded from `local_installable_dependencies` / `local_code`. Default to []. excluded_regex (List[pathlib.Path]): Regex used to exclude files from `local_installable_dependencies` / @@ -64,6 +65,7 @@ class Dependency(BaseModel): pypi_dependencies: List[str] = Field(default_factory=list) local_installable_dependencies: List[Path] = Field(default_factory=list) local_code: List[Path] = Field(default_factory=list) + binary_dependencies: List[str] = Field(default_factory=list) excluded_paths: List[Path] = Field(default_factory=list) excluded_regex: List[str] = Field(default_factory=list) force_included_paths: List[Path] = Field(default_factory=list) diff --git a/substrafl/remote/register/register.py b/substrafl/remote/register/register.py index 7b352369..d9bbb0f0 100644 --- a/substrafl/remote/register/register.py +++ b/substrafl/remote/register/register.py @@ -36,6 +36,17 @@ RUN apt-get update -y && pip uninstall -y setuptools """ +_CPU_BASE_IMAGE_WITH_DEPENDENCIES = """ +FROM python:{python_version}-slim + +# update image +RUN apt-get update -y\ + && pip uninstall -y setuptools\ + && apt-get install -y {binary_dependencies}\ + && apt-get clean\ + && rm -rf /var/lib/apt/lists/* +""" + _GPU_BASE_IMAGE = """ FROM nvidia/cuda:12.6.1-runtime-ubuntu24.04 @@ -45,7 +56,7 @@ && apt-get install -y software-properties-common\ && add-apt-repository -y ppa:deadsnakes/ppa\ && apt-get -y upgrade\ - && apt-get install -y python{python_version} python{python_version}-venv python3-pip\ + && apt-get install -y python{python_version} python{python_version}-venv python3-pip {binary_dependencies}\ && apt-get clean\ && rm -rf /var/lib/apt/lists/* @@ -131,12 +142,18 @@ def _check_python_version(python_major_minor: str) -> None: ) -def _get_base_docker_image(python_major_minor: str, use_gpu: bool) -> str: +def _get_base_docker_image( + python_major_minor: str, use_gpu: bool, custom_binary_dependencies: typing.Optional[list] = None +) -> str: """Get the base Docker image for the Dockerfile""" if use_gpu: base_docker_image = _GPU_BASE_IMAGE.format( - python_version=python_major_minor, + python_version=python_major_minor, binary_dependencies=" ".join(custom_binary_dependencies or []) + ) + elif custom_binary_dependencies: + base_docker_image = _CPU_BASE_IMAGE_WITH_DEPENDENCIES.format( + python_version=python_major_minor, binary_dependencies=" ".join(custom_binary_dependencies) ) else: base_docker_image = _CPU_BASE_IMAGE.format( @@ -161,7 +178,11 @@ def _create_dockerfile(install_libraries: bool, dependencies: Dependency, operat _check_python_version(python_major_minor) # Get the base Docker image - base_docker_image = _get_base_docker_image(python_major_minor=python_major_minor, use_gpu=dependencies.use_gpu) + base_docker_image = _get_base_docker_image( + python_major_minor=python_major_minor, + use_gpu=dependencies.use_gpu, + custom_binary_dependencies=dependencies.binary_dependencies, + ) # Build Substrafl, Substra and Substratools, and local dependencies wheels if necessary if install_libraries: # generate the copy wheel command diff --git a/tests/dependency/test_dependency.py b/tests/dependency/test_dependency.py index fe702f6a..a8b30f87 100644 --- a/tests/dependency/test_dependency.py +++ b/tests/dependency/test_dependency.py @@ -294,6 +294,27 @@ def train( train_task = self._register_train_task(function_key, numpy_datasets[0], constant_samples[0], client) client.wait_task(train_task.key, raise_on_failure=True) + @pytest.mark.docker_only + def test_binary_dependencies( + self, + network, + numpy_datasets, + constant_samples, + session_dir, + dummy_algo_class, + ): + """Test that you can install binary dependencies""" + + client = network.clients[0] + algo_deps = Dependency( + binary_dependencies=["gcc"], + editable_mode=True, + ) + function_key = self._register_function(dummy_algo_class(), algo_deps, client, session_dir) + + train_task = self._register_train_task(function_key, numpy_datasets[0], constant_samples[0], client) + client.wait_task(train_task.key, raise_on_failure=True) + @pytest.mark.docker_only def test_force_editable_mode( self, diff --git a/tests/remote/register/test_register.py b/tests/remote/register/test_register.py index d401952c..8ad17e70 100644 --- a/tests/remote/register/test_register.py +++ b/tests/remote/register/test_register.py @@ -64,7 +64,7 @@ def test_get_base_docker_image_gpu(): && apt-get install -y software-properties-common\ && add-apt-repository -y ppa:deadsnakes/ppa\ && apt-get -y upgrade\ - && apt-get install -y python3.11 python3.11-venv python3-pip\ + && apt-get install -y python3.11 python3.11-venv python3-pip \ && apt-get clean\ && rm -rf /var/lib/apt/lists/*