diff --git a/docker/base.Dockerfile b/docker/base.Dockerfile index a20f72241de..415d3597543 100644 --- a/docker/base.Dockerfile +++ b/docker/base.Dockerfile @@ -105,6 +105,8 @@ ENV \ PYTHONFAULTHANDLER=1 \ # Use a random seed for random number generators PYTHONHASHSEED=random \ + # Set environment variable to point to the active virtual env + VIRTUAL_ENV=$VIRTUAL_ENV \ # Signal to ZenML that it is running in a container ZENML_CONTAINER=1 @@ -136,6 +138,8 @@ ENV \ PYTHONFAULTHANDLER=1 \ # Use a random seed for random number generators PYTHONHASHSEED=random \ + # Set environment variable to point to the active virtual env + VIRTUAL_ENV=$VIRTUAL_ENV \ # Signal to ZenML that it is running in a container ZENML_CONTAINER=1 \ # Set the ZenML global configuration path diff --git a/docker/zenml-dev.Dockerfile b/docker/zenml-dev.Dockerfile index 9ac6f56e9bd..f485a5c2d35 100644 --- a/docker/zenml-dev.Dockerfile +++ b/docker/zenml-dev.Dockerfile @@ -85,6 +85,8 @@ ENV \ PYTHONFAULTHANDLER=1 \ # Use a random seed for random number generators PYTHONHASHSEED=random \ + # Set environment variable to point to the active virtual env + VIRTUAL_ENV=$VIRTUAL_ENV \ # Signal to ZenML that it is running in a container ZENML_CONTAINER=1 \ # Set ZenML debug mode to true diff --git a/docker/zenml-server-dev.Dockerfile b/docker/zenml-server-dev.Dockerfile index 80b6f650b97..1696a9539bd 100644 --- a/docker/zenml-server-dev.Dockerfile +++ b/docker/zenml-server-dev.Dockerfile @@ -78,6 +78,8 @@ ENV \ PYTHONFAULTHANDLER=1 \ # Use a random seed for random number generators PYTHONHASHSEED=random \ + # Set environment variable to point to the active virtual env + VIRTUAL_ENV=$VIRTUAL_ENV \ # Set the ZenML global configuration path ZENML_CONFIG_PATH=/zenml/.zenconfig \ # Signal to ZenML that it is running in a container diff --git a/docs/book/user-guide/advanced-guide/infrastructure-management/containerize-your-pipeline.md b/docs/book/user-guide/advanced-guide/infrastructure-management/containerize-your-pipeline.md index 09859f812ff..3ba1f301ef6 100644 --- a/docs/book/user-guide/advanced-guide/infrastructure-management/containerize-your-pipeline.md +++ b/docs/book/user-guide/advanced-guide/infrastructure-management/containerize-your-pipeline.md @@ -261,6 +261,16 @@ Depending on the options specified in your Docker settings, ZenML installs the r * The packages installed in your local Python environment * The packages specified via the `requirements` attribute (step level overwrites pipeline level) * The packages specified via the `required_integrations` and potentially stack requirements +* You can specify additional arguments for the installer used to install your Python packages as follows: +```python +# This will result in a `pip install --timeout=1000 ...` call when installing packages in the +# Docker image +docker_settings = DockerSettings(python_package_installer_args={"timeout": 1000}) + +@pipeline(settings={"docker": docker_settings}) +def my_pipeline(...): + ... +``` * **Experimental**: If you want to use [`uv`](https://github.com/astral-sh/uv) for faster resolving and installation of your Python packages, you can use by it as follows: diff --git a/src/zenml/config/docker_settings.py b/src/zenml/config/docker_settings.py index 0210987ac27..f6ca295d337 100644 --- a/src/zenml/config/docker_settings.py +++ b/src/zenml/config/docker_settings.py @@ -128,6 +128,8 @@ class DockerSettings(BaseSettings): therefore **not** include any registry. python_package_installer: The package installer to use for python packages. + python_package_installer_args: Arguments to pass to the python package + installer. replicate_local_python_environment: If not `None`, ZenML will use the specified method to generate a requirements file that replicates the packages installed in the currently running python environment. @@ -185,6 +187,7 @@ class DockerSettings(BaseSettings): python_package_installer: PythonPackageInstaller = ( PythonPackageInstaller.PIP ) + python_package_installer_args: Dict[str, Any] = {} replicate_local_python_environment: Optional[ Union[List[str], PythonEnvironmentExportMethod] ] = None diff --git a/src/zenml/utils/pipeline_docker_image_builder.py b/src/zenml/utils/pipeline_docker_image_builder.py index d4eeaa80b8f..f462149a1e6 100644 --- a/src/zenml/utils/pipeline_docker_image_builder.py +++ b/src/zenml/utils/pipeline_docker_image_builder.py @@ -20,6 +20,7 @@ from collections import defaultdict from typing import ( TYPE_CHECKING, + Any, DefaultDict, Dict, List, @@ -65,6 +66,14 @@ f"py{sys.version_info.major}.{sys.version_info.minor}" ) +PIP_DEFAULT_ARGS = { + "no-cache-dir": None, + "default-timeout": 60, +} +UV_DEFAULT_ARGS = { + "no-cache-dir": None, +} + class PipelineDockerImageBuilder: """Builds Docker images to run a ZenML pipeline.""" @@ -636,22 +645,32 @@ def _generate_zenml_pipeline_dockerfile( docker_settings.python_package_installer == PythonPackageInstaller.PIP ): - install_command = "pip install --default-timeout=60" + install_command = "pip install" + default_installer_args: Dict[str, Any] = PIP_DEFAULT_ARGS elif ( docker_settings.python_package_installer == PythonPackageInstaller.UV ): lines.append("RUN pip install uv") - install_command = "uv pip install --system" + install_command = "uv pip install" + default_installer_args = UV_DEFAULT_ARGS else: raise ValueError("Unsupported python package installer.") + installer_args = { + **default_installer_args, + **docker_settings.python_package_installer_args, + } + installer_args_string = " ".join( + f"--{key}" if value is None else f"--{key}={value}" + for key, value in installer_args.items() + ) for file, _, options in requirements_files: lines.append(f"COPY {file} .") option_string = " ".join(options) lines.append( - f"RUN {install_command} --no-cache-dir " + f"RUN {install_command} {installer_args_string}" f"{option_string} -r {file}" ) diff --git a/tests/unit/utils/test_pipeline_docker_image_builder.py b/tests/unit/utils/test_pipeline_docker_image_builder.py index da2b498f845..091c3c3f9ac 100644 --- a/tests/unit/utils/test_pipeline_docker_image_builder.py +++ b/tests/unit/utils/test_pipeline_docker_image_builder.py @@ -144,3 +144,29 @@ def test_build_skipping(): download_files=False, ) assert image_digest + + +def test_python_package_installer_args(): + """Tests that the python package installer args get passed correctly.""" + docker_settings = DockerSettings( + python_package_installer_args={ + "default-timeout": 99, + "other-arg": "value", + "option": None, + } + ) + + requirements_files = [("requirements.txt", "numpy", [])] + generated_dockerfile = ( + PipelineDockerImageBuilder._generate_zenml_pipeline_dockerfile( + "image:tag", + docker_settings, + download_files=False, + requirements_files=requirements_files, + ) + ) + + assert ( + "RUN pip install --no-cache-dir --default-timeout=99 --other-arg=value --option" + in generated_dockerfile + )