From e18253d427830c2d428cdb351d0e0032feae7622 Mon Sep 17 00:00:00 2001 From: SdgJlbl Date: Tue, 17 Sep 2024 17:17:44 +0200 Subject: [PATCH] feat: add the option to inject binary libraries to install in Dockerfile Signed-off-by: SdgJlbl --- substrafl/dependency/schemas.py | 1 + substrafl/remote/register/register.py | 25 +++++++++++++++++++++---- 2 files changed, 22 insertions(+), 4 deletions(-) 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 f7170421..c951b4c8 100644 --- a/substrafl/remote/register/register.py +++ b/substrafl/remote/register/register.py @@ -36,6 +36,13 @@ RUN apt-get update -y """ +_CPU_BASE_IMAGE_WITH_DEPENDENCIES = """ +FROM python:{python_version}-slim + +# update image +RUN apt-get update -y && 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 +52,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 +138,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 +174,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