Develop a Custom Image Builder
How to develop a custom image builder
BaseImageBuilderis 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."""
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.
The build context class.
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
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.
The Docker image repo digest or name.
If you want to create your own custom flavor for an image builder, you can follow the following steps:
- 1.Create a class which inherits from the
BaseImageBuilderclass and implement the abstract
buildmethod. This method should use the given build context and build a Docker image with it. If additionally a container registry is passed to the
buildmethod, the image builder is also responsible for pushing the image there.
- 2.If you need to provide any configuration, create a class which inherits from the
BaseImageBuilderConfigclass add your configuration parameters.
- 3.Bring both of the implementation and the configuration together by inheriting from the
BaseImageBuilderFlavorclass. Make sure that you give a
nameto 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 <IMAGE-BUILDER-FLAVOR-SOURCE-PATH>
BaseImageBuilderabstraction uses the
build_context_classto 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
BuildContextclass to customize the structure of your build context. In your image builder implementation, you can then overwrite the
build_context_classproperty to specify your build context subclass.