Add user-defined builder image in image spec (#3207)

Signed-off-by: Nelson Chen <[email protected]>

Author: arbaobao

flytekit/image_spec/default_builder.py

@@ -93,8 +93,8 @@
 
 DOCKER_FILE_TEMPLATE = Template("""\
 #syntax=docker/dockerfile:1.5
-FROM ghcr.io/astral-sh/uv:0.5.1 as uv
-FROM mambaorg/micromamba:2.0.3-debian12-slim as micromamba
+FROM $UV_IMAGE as uv
+FROM $MICROMAMBA_IMAGE as micromamba
 
 FROM $BASE_IMAGE
 
@@ -141,6 +141,9 @@
     echo "export PATH=$$PATH" >> $$HOME/.profile
 """)
 
+DEFAULT_UV_IMAGE = "ghcr.io/astral-sh/uv:0.5.1"
+DEFAULT_MICROMAMBA_IMAGE = "mambaorg/micromamba:2.0.3-debian12-slim"
+
 
 def get_flytekit_for_pypi():
     """Get flytekit version on PyPI."""
@@ -405,6 +408,12 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path):
     else:
         extra_copy_cmds = ""
 
+    uv_image = DEFAULT_UV_IMAGE
+    micromamba_image = DEFAULT_MICROMAMBA_IMAGE
+    if image_spec.builder_config is not None:
+        uv_image = image_spec.builder_config.get("uv_image", uv_image)
+        micromamba_image = image_spec.builder_config.get("micromamba_image", micromamba_image)
+
     docker_content = DOCKER_FILE_TEMPLATE.substitute(
         UV_PYTHON_INSTALL_COMMAND=uv_python_install_command,
         APT_INSTALL_COMMAND=apt_install_command,
@@ -417,6 +426,8 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path):
         ENTRYPOINT=entrypoint,
         RUN_COMMANDS=run_commands,
         EXTRA_COPY_CMDS=extra_copy_cmds,
+        UV_IMAGE=uv_image,
+        MICROMAMBA_IMAGE=micromamba_image,
     )
 
     dockerfile_path = tmp_dir / "Dockerfile"
@@ -452,6 +463,7 @@ class DefaultImageBuilder(ImageSpecBuilder):
         # "registry_config",
         "commands",
         "copy",
+        "builder_config",
     }
 
     def build_image(self, image_spec: ImageSpec) -> str:

flytekit/image_spec/image_spec.py

@@ -74,6 +74,7 @@ class ImageSpec:
                   use `builder="noop"`: `ImageSpec(base_image="ghcr.io/name/my-custom-image", builder="noop").with_runtime_packages(["numpy"])`.
         builder_options (Optional[Dict[str, Any]]): Additional options for the builder. This is a dictionary that will be passed to the builder.
             The options are builder-specific and may not be supported by all builders.
+        builder_config (Optional[typing.Dict[str, typing.Any]]): Custom builder images configuration, such as uv and micromamba images.
     """
 
     name: str = "flytekit"
@@ -104,6 +105,7 @@ class ImageSpec:
     python_exec: Optional[str] = None
     runtime_packages: Optional[List[str]] = None
     builder_options: Optional[Dict[str, Any]] = None
+    builder_config: Optional[typing.Dict[str, typing.Any]] = None
 
     def __post_init__(self):
         self.name = self.name.lower()

tests/flytekit/unit/core/image_spec/test_default_builder.py

@@ -43,6 +43,7 @@ def test_create_docker_context(tmp_path):
             pip_secret_mounts=[(".gitconfig", '/etc/gitconfig'), ("secret_src_2", "secret_dst_2")],
             source_copy_mode=CopyFileDetection.ALL,
             copy=[tmp_file.relative_to(Path.cwd()).as_posix()],
+            builder_config={"uv_image": "uv_image_local"},
         )
 
         create_docker_context(image_spec, docker_context_path)
@@ -60,6 +61,8 @@ def test_create_docker_context(tmp_path):
     assert "--mount=type=secret,id=secret_0,target=/etc/gitconfig" in dockerfile_content
     assert "--mount=type=secret,id=secret_1,target=secret_dst_2" in dockerfile_content
     assert "COPY --chown=flytekit ./src /root" in dockerfile_content
+    assert "uv_image_local" in dockerfile_content
+    assert "mambaorg/micromamba:2.0.3-debian12-slim" in dockerfile_content
 
     run_match = re.search(r"RUN.+mkdir my_dir", dockerfile_content)
     assert run_match