Develop a Custom Image Builder

Learning how to develop a custom image builder.

Before diving into the specifics of this component type, it is beneficial to familiarize yourself with our general guide to writing custom component flavors in ZenML. This guide provides an essential understanding of ZenML's component flavor concepts.

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.

"""