Better typing for docker build options (#4262)

Author: schustmi

src/zenml/config/docker_settings.py

@@ -55,20 +55,129 @@ class PythonPackageInstaller(Enum):
     UV = "uv"
 
 
+# (docker_sdk_argument, attribute_name)
+BUILD_OPTION_CONVERSIONS = [
+    ("buildargs", "build_args"),
+    ("cachefrom", "cache_from"),
+    ("nocache", "no_cache"),
+    ("shmsize", "shm_size"),
+]
+
+
+class DockerBuildOptions(BaseModel):
+    """Docker build options.
+
+    This class only specifies a subset of the options which require explicit
+    conversion as they require different names in the Docker CLI and Python SDK.
+    However, you can still specify any other options as extra model attributes
+    which will be passed unmodified to the build method of the image builder.
+    """
+
+    pull: Optional[bool] = None
+    rm: Optional[bool] = None
+    no_cache: Optional[bool] = None
+    shm_size: Optional[int] = None
+    labels: Optional[Dict[str, Any]] = None
+    build_args: Optional[Dict[str, Any]] = None
+    cache_from: Optional[List[str]] = None
+
+    model_config = ConfigDict(extra="allow")
+
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _migrate_sdk_arguments(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Migrate Docker SDK arguments to attributes.
+
+        Args:
+            data: The model data.
+
+        Returns:
+            The migrated data.
+        """
+        for sdk_argument, attribute_name in BUILD_OPTION_CONVERSIONS:
+            if sdk_argument in data and attribute_name not in data:
+                data[attribute_name] = data.pop(sdk_argument)
+
+        return data
+
+    def to_docker_cli_options(self) -> List[str]:
+        """Convert the build options to a list of Docker CLI options.
+
+        https://docs.docker.com/reference/cli/docker/buildx/build/#options
+
+        Returns:
+            A list of Docker CLI options.
+        """
+        options = []
+        if self.pull:
+            options.append("--pull")
+        if self.rm:
+            options.append("--rm")
+        if self.no_cache:
+            options.append("--no-cache")
+        if self.shm_size:
+            options.extend(["--shm-size", str(self.shm_size)])
+        if self.labels:
+            for key, value in self.labels.items():
+                options.extend(["--label", f"{key}={value}"])
+        if self.build_args:
+            for key, value in self.build_args.items():
+                options.extend(["--build-arg", f"{key}={value}"])
+        if self.cache_from:
+            for value in self.cache_from:
+                options.extend(["--cache-from", value])
+
+        if self.model_extra:
+            for key, value in self.model_extra.items():
+                option = f"--{key.replace('_', '-')}"
+                if isinstance(value, Dict):
+                    for key, value in value.items():
+                        options.extend([option, f"{key}={value}"])
+                elif isinstance(value, List):
+                    for val in value:
+                        options.extend([option, str(val)])
+                elif value in (True, None):
+                    options.extend([option])
+                elif value is False:
+                    pass
+                else:
+                    options.extend([option, str(value)])
+
+        return options
+
+    def to_docker_python_sdk_options(self) -> Dict[str, Any]:
+        """Get the build options as a dictionary of Docker Python SDK options.
+
+        https://docker-py.readthedocs.io/en/stable/images.html#docker.models.images.ImageCollection.build
+
+        Returns:
+            A dictionary of Docker Python SDK options.
+        """
+        options = self.model_dump(exclude_unset=True)
+        for sdk_argument, attribute_name in BUILD_OPTION_CONVERSIONS:
+            if attribute_name in options:
+                options[sdk_argument] = options.pop(attribute_name)
+
+        return options
+
+
 class DockerBuildConfig(BaseModel):
     """Configuration for a Docker build.
 
     Attributes:
-        build_options: Additional options that will be passed unmodified to the
-            Docker build call when building an image. You can use this to for
-            example specify build args or a target stage. See
-            https://docker-py.readthedocs.io/en/stable/images.html#docker.models.images.ImageCollection.build
-            for a full list of available options.
+        build_options: Additional options that will be passed when building an
+            image. Depending on the image builder that is used, different
+            options are available.
+            For image builders that use the Docker CLI:
+            - https://docs.docker.com/reference/cli/docker/buildx/build/#options
+            For image builders that use the Docker Python SDK:
+            - https://docker-py.readthedocs.io/en/stable/images.html#docker.models.images.ImageCollection.build
         dockerignore: Path to a dockerignore file to use when building the
             Docker image.
     """
 
-    build_options: Dict[str, Any] = {}
+    build_options: Optional[DockerBuildOptions] = None
     dockerignore: Optional[str] = None
 
 

src/zenml/image_builders/base_image_builder.py

@@ -17,7 +17,7 @@
 import os
 import tempfile
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type, cast
+from typing import TYPE_CHECKING, Optional, Type, cast
 
 from zenml.client import Client
 from zenml.enums import StackComponentType
@@ -28,6 +28,7 @@
 from zenml.utils.archivable import ArchiveType
 
 if TYPE_CHECKING:
+    from zenml.config.docker_settings import DockerBuildOptions
     from zenml.container_registries import BaseContainerRegistry
     from zenml.image_builders import BuildContext
 
@@ -79,7 +80,7 @@ def build(
         self,
         image_name: str,
         build_context: "BuildContext",
-        docker_build_options: Dict[str, Any],
+        docker_build_options: Optional["DockerBuildOptions"] = None,
         container_registry: Optional["BaseContainerRegistry"] = None,
     ) -> str:
         """Builds a Docker image.

src/zenml/image_builders/local_image_builder.py

@@ -16,7 +16,7 @@
 import shutil
 import subprocess
 import tempfile
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type, cast
+from typing import TYPE_CHECKING, Optional, Type, cast
 
 from pydantic import Field
 
@@ -28,6 +28,7 @@
 from zenml.utils import docker_utils
 
 if TYPE_CHECKING:
+    from zenml.config.docker_settings import DockerBuildOptions
     from zenml.container_registries import BaseContainerRegistry
     from zenml.image_builders import BuildContext
 
@@ -102,7 +103,7 @@ def build(
         self,
         image_name: str,
         build_context: "BuildContext",
-        docker_build_options: Optional[Dict[str, Any]] = None,
+        docker_build_options: Optional["DockerBuildOptions"] = None,
         container_registry: Optional["BaseContainerRegistry"] = None,
     ) -> str:
         """Builds and optionally pushes an image using the local Docker client.
@@ -141,7 +142,7 @@ def _build_with_python_sdk(
         self,
         image_name: str,
         build_context: "BuildContext",
-        docker_build_options: Optional[Dict[str, Any]] = None,
+        docker_build_options: Optional["DockerBuildOptions"] = None,
         container_registry: Optional["BaseContainerRegistry"] = None,
     ) -> None:
         """Builds an image using the Python Docker SDK.
@@ -159,6 +160,12 @@ def _build_with_python_sdk(
         else:
             docker_client = docker_utils._try_get_docker_client_from_env()
 
+        build_options = (
+            docker_build_options.to_docker_python_sdk_options()
+            if docker_build_options
+            else {}
+        )
+
         with tempfile.TemporaryFile(mode="w+b") as f:
             build_context.write_archive(f)
 
@@ -167,15 +174,15 @@ def _build_with_python_sdk(
                 fileobj=f,
                 custom_context=True,
                 tag=image_name,
-                **(docker_build_options or {}),
+                **build_options,
             )
         docker_utils._process_stream(output_stream)
 
     def _build_with_subprocess_call(
         self,
         image_name: str,
         build_context: "BuildContext",
-        docker_build_options: Optional[Dict[str, Any]] = None,
+        docker_build_options: Optional["DockerBuildOptions"] = None,
     ) -> None:
         """Builds an image using a subprocess `docker build` call.
 
@@ -194,31 +201,8 @@ def _build_with_subprocess_call(
             # file to stdin
             command = ["docker", "build", "-", "-t", image_name]
 
-            docker_build_options = docker_build_options or {}
-            if docker_build_options.pop("pull", False):
-                command.append("--pull")
-
-            build_args = docker_build_options.pop("buildargs", None) or {}
-            for key, value in build_args.items():
-                command.extend(["--build-arg", f"{key}={value}"])
-
-            labels = docker_build_options.pop("labels", None) or {}
-            for key, value in labels.items():
-                command.extend(["--label", f"{key}={value}"])
-
-            cache_from = docker_build_options.pop("cache_from", None) or []
-            for value in cache_from:
-                command.extend(["--cache-from", value])
-
-            for key, value in docker_build_options.items():
-                if isinstance(value, list):
-                    for val in value:
-                        command.extend([f"--{key}", str(val)])
-                elif isinstance(value, bool):
-                    if value:
-                        command.append(f"--{key}")
-                else:
-                    command.extend([f"--{key}", str(value)])
+            if docker_build_options:
+                command.extend(docker_build_options.to_docker_cli_options())
 
             process = subprocess.Popen(command, stdin=f)
 

src/zenml/integrations/aws/image_builders/aws_image_builder.py

@@ -31,6 +31,7 @@
 from zenml.utils.archivable import ArchiveType
 
 if TYPE_CHECKING:
+    from zenml.config.docker_settings import DockerBuildOptions
     from zenml.container_registries import BaseContainerRegistry
     from zenml.image_builders import BuildContext
     from zenml.stack import Stack
@@ -126,7 +127,7 @@ def build(
         self,
         image_name: str,
         build_context: "BuildContext",
-        docker_build_options: Dict[str, Any],
+        docker_build_options: Optional["DockerBuildOptions"] = None,
         container_registry: Optional["BaseContainerRegistry"] = None,
     ) -> str:
         """Builds and pushes a Docker image.
@@ -189,17 +190,11 @@ def build(
                 f"aws ecr get-login-password --region {self.code_build_client._client_config.region_name} | docker login --username AWS --password-stdin {container_registry.config.uri}",
             ]
 
-        # Convert the docker_build_options dictionary to a list of strings
         docker_build_args = ""
-        for key, value in docker_build_options.items():
-            option = f"--{key}"
-            if isinstance(value, list):
-                for val in value:
-                    docker_build_args += f"{option} {val} "
-            elif value is not None and not isinstance(value, bool):
-                docker_build_args += f"{option} {value} "
-            elif value is not False:
-                docker_build_args += f"{option} "
+        if docker_build_options:
+            docker_build_args = " ".join(
+                docker_build_options.to_docker_cli_options()
+            )
 
         pre_build_commands_str = "\n".join(
             [f"            - {command}" for command in pre_build_commands]

src/zenml/integrations/gcp/image_builders/gcp_image_builder.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Google Cloud Builder image builder implementation."""
 
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, cast
+from typing import TYPE_CHECKING, Optional, Tuple, cast
 from urllib.parse import urlparse
 
 from google.cloud.devtools import cloudbuild_v1
@@ -29,6 +29,7 @@
 from zenml.stack import StackValidator
 
 if TYPE_CHECKING:
+    from zenml.config.docker_settings import DockerBuildOptions
     from zenml.container_registries import BaseContainerRegistry
     from zenml.image_builders import BuildContext
     from zenml.stack import Stack
@@ -98,7 +99,7 @@ def build(
         self,
         image_name: str,
         build_context: "BuildContext",
-        docker_build_options: Dict[str, Any],
+        docker_build_options: Optional["DockerBuildOptions"] = None,
         container_registry: Optional["BaseContainerRegistry"] = None,
     ) -> str:
         """Builds and pushes a Docker image.
@@ -141,7 +142,7 @@ def _configure_cloud_build(
         self,
         image_name: str,
         cloud_build_context: str,
-        build_options: Dict[str, Any],
+        build_options: Optional["DockerBuildOptions"] = None,
     ) -> cloudbuild_v1.Build:
         """Configures the build to be run to generate the Docker image.
 
@@ -171,17 +172,9 @@ def _configure_cloud_build(
             cloud_builder_network_option,
         )
 
-        # Convert the docker_build_options dictionary to a list of strings
-        docker_build_args = []
-        for key, value in build_options.items():
-            option = f"--{key}"
-            if isinstance(value, list):
-                for val in value:
-                    docker_build_args.extend([option, val])
-            elif value is not None and not isinstance(value, bool):
-                docker_build_args.extend([option, value])
-            elif value is not False:
-                docker_build_args.extend([option])
+        docker_build_args = (
+            build_options.to_docker_cli_options() if build_options else []
+        )
 
         return cloudbuild_v1.Build(
             source=cloudbuild_v1.Source(