diff --git a/substrafl/dependency/schemas.py b/substrafl/dependency/schemas.py index 9bef2c2d..00ca67a5 100644 --- a/substrafl/dependency/schemas.py +++ b/substrafl/dependency/schemas.py @@ -64,6 +64,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 d562b463..2ff964c8 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:11.8.0-runtime-ubuntu22.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/remote/register/test_register.py b/tests/remote/register/test_register.py index 1c2c43d6..40cdb3fb 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/*