Skip to main content

Custom Image Builders

When you use a private container registry or a specialized build tool like Kaniko, the default Docker-based image builder in flyte-sdk may fail to authenticate or utilize your infrastructure's caching. The ImageBuildEngine architecture allows you to swap the default "local" or "remote" builders with custom implementations that handle specific registry protocols and build environments.

The Image Build Engine Orchestrator​

The ImageBuildEngine (located in src/flyte/_internal/imagebuild/image_builder.py) is the central hub for managing the containerization lifecycle of an ImageSpec. It does not perform builds itself; instead, it orchestrates the process by:

  1. Checking for existence: It uses ImageChecker implementations to determine if an image with the specific tag and architecture already exists in the target registry.
  2. Selecting a builder: It resolves which ImageBuilder to use based on configuration or explicit parameters.
  3. Executing the build: It triggers the build_image method of the selected builder.

The engine is designed to be efficient. For example, in ImageBuildEngine.image_exists, it skips existence checks for unmodified default Flyte images (where image._is_cloned is false) because these are guaranteed to exist as part of Flyte releases.

Implementing a Custom Image Builder​

To create a custom builder, you must implement the ImageBuilder protocol defined in flyte.extend. This protocol requires two methods:

  • build_image: The core logic for building and pushing the image. It must return an ImageBuild object containing the final image URI.
  • get_checkers: An optional method that returns a list of ImageChecker classes compatible with the builder's target registries.

The ImageBuilder Protocol​

from typing import List, Optional, Type
from flyte import Image
from flyte.extend import ImageBuilder, ImageChecker, ImageBuild

class MyCustomBuilder(ImageBuilder):
    def get_checkers(self) -> Optional[List[Type[ImageChecker]]]:
        # Return custom checkers for your specific registry
        return [MyRegistryChecker]

    async def build_image(
        self, 
        image: Image, 
        dry_run: bool, 
        wait: bool = True, 
        force: bool = False
    ) -> ImageBuild:
        # Custom logic to build image (e.g., calling a CI/CD API)
        print(f"Building {image.name}...")
        return ImageBuild(uri=image.uri, remote_run=None)

Optimizing with ImageCheckers​

Redundant builds are a significant bottleneck in CI/CD pipelines. flyte-sdk uses the ImageChecker protocol to verify if a build can be skipped. If any checker returns a URI, the ImageBuildEngine assumes the image is ready and skips the build step.

The ImageChecker Protocol​

from typing import Optional, Tuple
from flyte.extend import ImageChecker, Architecture

class MyRegistryChecker(ImageChecker):
    @classmethod
    async def image_exists(
        cls, 
        repository: str, 
        tag: str, 
        arch: Tuple[Architecture, ...] = ("linux/amd64",)
    ) -> Optional[str]:
        # Logic to query your registry's API
        # Return the full image URI if it exists, else None
        return f"{repository}:{tag}"

The ImageBuildEngine handles checker failures gracefully. If all registered checkers for a builder raise an exception or fail to find the image, the engine logs a warning and assumes the image exists to avoid blocking the workflow. This behavior ensures that transient registry issues do not necessarily stop a deployment, though it may lead to "Image Not Found" errors later during task execution.

Registration and Configuration​

Custom builders are discovered via Python entry points. To register your builder, add it to your pyproject.toml or setup.py under the flyte.plugins.image_builders group:

[project.entry-points."flyte.plugins.image_builders"]
my_builder = "my_package.builders:MyCustomBuilder"

Once registered, you can configure flyte-sdk to use your builder globally during initialization:

import flyte

flyte.init(image_builder="my_builder")

Alternatively, you can specify the builder name when calling the build engine directly or via environment variables.

Design Tradeoffs and Behavior​

The implementation in src/flyte/_internal/imagebuild/image_builder.py includes several specific design choices:

  • The "latest" Tag Exception: In ImageBuildEngine.image_exists, if an image is tagged as latest, the existence check is always skipped. This ensures that developers iterating with the latest tag always trigger a fresh build of their code.
  • Architecture Awareness: The ImageChecker.image_exists method accepts a tuple of Architecture values (e.g., linux/amd64, linux/arm64). This allows checkers to verify multi-arch manifests rather than just single-platform images.
  • Builder Fallback: If no builder is specified and no default is configured, the engine defaults to the local builder, which uses the local Docker or Podman daemon.