Extending Image Builders

To integrate with specialized container registries or custom build systems, you can extend the image building capabilities of this SDK by implementing custom ImageBuilder and ImageChecker plugins.

Implement a Custom Image Builder

To create a custom builder, you must implement the ImageBuilder protocol. This protocol requires a build_image method and an optional get_checkers method to optimize the build process by checking for existing images.

import typing
from typing import Optional, Tuple
from flyte import Image, ImageBuild
from flyte.extend import Architecture, ImageBuilder, ImageChecker

class MyImageChecker(ImageChecker):
    @classmethod
    async def image_exists(
        cls, 
        repository: str, 
        tag: str, 
        arch: Tuple[Architecture, ...] = ("linux/amd64",)
    ) -> Optional[str]:
        # Logic to check if the image exists in your custom registry
        # Return the full image URI if it exists, otherwise None
        print(f"Checking existence for {repository}:{tag}...")
        return None

class MyImageBuilder(ImageBuilder):
    def get_checkers(self) -> Optional[typing.List[typing.Type[ImageChecker]]]:
        """Return the list of checkers used to verify image existence."""
        return [MyImageChecker]

    async def build_image(
        self,
        image: Image,
        dry_run: bool = False,
        wait: bool = True,
        force: bool = False,
    ) -> ImageBuild:
        """
        Logic to build and push the image.
        """
        if dry_run:
            print(f"Dry run: would build {image.uri}")
            return ImageBuild(uri=image.uri, remote_run=None)
            
        print(f"Building image {image.name} for platforms {image.platform}...")
        # Your custom build logic here (e.g., calling a remote CI/CD API)
        
        return ImageBuild(uri=image.uri, remote_run=None)

Register the Plugin

Custom builders are discovered via Python entry points. You must register your class in your pyproject.toml file under the flyte.plugins.image_builders group.

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

The name used as the key (my_builder in this example) is how you will reference the builder in your configuration.

Use the Custom Builder

Once registered, you can configure Flyte to use your custom builder during initialization or via the ImageBuildEngine.

Using flyte.init

The most common way to set the builder is during the init call:

import flyte

flyte.init(image_builder="my_builder")

image = (
    flyte.Image.from_debian_base(python_version=(3, 11))
    .with_pip_packages("pandas")
)

Using ImageBuildEngine directly

You can also specify the builder explicitly when calling the build engine:

from flyte._internal.imagebuild.image_builder import ImageBuildEngine

result = await ImageBuildEngine.build(image, builder="my_builder")
print(f"Built image: {result.uri}")

Troubleshooting and Best Practices

• The 'latest' Tag: Images tagged with latest always skip the existence check and trigger a build. If you want to leverage ImageChecker optimizations, ensure your images use specific tags or rely on the SDK's automatic tagging.
• Checker Fallback: If all registered ImageChecker implementations fail (raise an exception) or return None, the ImageBuildEngine will assume the image does not exist and proceed with the build.
• Authentication: If your ImageChecker needs to communicate with a private registry, ensure it handles authentication internally. The default DockerAPIImageChecker primarily supports Docker Hub because it does not handle complex registry authentication tokens.
• Platform Support: When implementing image_exists, respect the arch parameter. A multi-arch image might exist for linux/amd64 but not for linux/arm64.