feat: add the option to inject binary libraries to install in Dockerf…

…ile (#249)

Signed-off-by: SdgJlbl <[email protected]>

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.

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)

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

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,

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/*