Develop a Custom Image Builder

How to develop a custom image builder

Base Abstraction

The BaseImageBuilder is the abstract base class that needs to be subclassed in order to create a custom component that can be used to build Docker images. As image builders can come in many shapes and forms, the base class exposes a deliberately basic and generic interface:

from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, Type, cast

from zenml.container_registries import BaseContainerRegistry
from zenml.enums import StackComponentType
from zenml.image_builders import BuildContext
from zenml.stack import Flavor, StackComponent
from zenml.stack.stack_component import StackComponentConfig

class BaseImageBuilder(StackComponent, ABC):
    """Base class for all ZenML image builders."""

    @property
    def build_context_class(self) -> Type["BuildContext"]:
        """Build context class to use.

        The default build context class creates a build context that works
        for the Docker daemon. Override this method if your image builder
        requires a custom context.

        Returns:
            The build context class.
        """
        return BuildContext

    @abstractmethod
    def build(
        self,
        image_name: str,
        build_context: "BuildContext",
        docker_build_options: Dict[str, Any],
        container_registry: Optional["BaseContainerRegistry"] = None,
    ) -> str:
        """Builds a Docker image.

        If a container registry is passed, the image will be pushed to that
        registry.

        Args:
            image_name: Name of the image to build.
            build_context: The build context to use for the image.
            docker_build_options: Docker build options.
            container_registry: Optional container registry to push to.

        Returns:
            The Docker image repo digest or name.
        """

This is a slimmed-down version of the base implementation which aims to highlight the abstraction layer. In order to see the full implementation and get the complete docstrings, please check the source code on GitHub.

Build your own custom image builder

If you want to create your own custom flavor for an image builder, you can follow the following steps:

Create a class which inherits from the BaseImageBuilder class and implement the abstract build method. This method should use the given build context and build a Docker image with it. If additionally a container registry is passed to the build method, the image builder is also responsible for pushing the image there.

If you need to provide any configuration, create a class which inherits from the BaseImageBuilderConfig class add your configuration parameters.

Bring both of the implementation and the configuration together by inheriting from the BaseImageBuilderFlavor class. Make sure that you give a name to the flavor through its abstract property.

Once you are done with the implementation, you can register it through the CLI as:

zenml image-builder flavor register 

It is important to draw attention to when and how these base abstractions are coming into play in a ZenML workflow.

• The CustomImageBuilderFlavor class is imported and utilized upon the creation of the custom flavor through the CLI.

• The CustomImageBuilderConfig class is imported when someone tries to register/update a stack component with this custom flavor. Especially, during the registration process of the stack component, the config will be used to validate the values given by the user. As Config object are inherently pydantic objects, you can also add your own custom validators here.

• The CustomImageBuilder only comes into play when the component is ultimately in use.

The design behind this interaction lets us separate the configuration of the flavor from its implementation. This way we can register flavors and components even when the major dependencies behind their implementation are not installed in our local setting (assuming the CustomImageBuilderFlavor and the CustomImageBuilderConfig are implemented in a different module/path than the actual CustomImageBuilder).

Using a custom build context

The BaseImageBuilder abstraction uses the build_context_class to provide a class that should be used as the build context. In case your custom image builder requires a different build context than the default Docker build context, you can subclass the BuildContext class to customize the structure of your build context. In your image builder implementation, you can then overwrite the build_context_class property to specify your build context subclass.