Allow specifying python package installer args (#2727)

* Allow specifying python package installer args

* Add test

* Test docstring

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

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

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

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:
 

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

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}"
             )
 

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
+    )