diff --git a/flytekit/configuration/__init__.py b/flytekit/configuration/__init__.py index bee6feb8be..061a450318 100644 --- a/flytekit/configuration/__init__.py +++ b/flytekit/configuration/__init__.py @@ -639,6 +639,45 @@ def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> GCSConfig: return AzureBlobStorageConfig(**kwargs) +@dataclass(init=True, repr=True, eq=True, frozen=True) +class DefaultConfig(object): + """ + Any image builder specific configuration. + """ + + name: str = "" + uv_image: str = "ghcr.io/astral-sh/uv:0.5.1" + micromamba_image: str = "mambaorg/micromamba:2.0.3-debian12-slim" + + @classmethod + def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> DefaultConfig: + config_file = get_config_file(config_file) + kwargs = {} + kwargs = set_if_exists(kwargs, "name", _internal.Local.IMAGE_BUILDER_NAME.read(config_file)) + kwargs = set_if_exists(kwargs, "uv_image", _internal.Local.UV_IMAGE.read(config_file)) + kwargs = set_if_exists(kwargs, "micromamba_image", _internal.Local.MICROMAMBA_IMAGE.read(config_file)) + return DefaultConfig(**kwargs) + + +@dataclass(init=True, repr=True, eq=True, frozen=True) +class ImageBuilderConfig(object): + """ + Any image builder specific configuration. + """ + + default: DefaultConfig = DefaultConfig() + """ + We can add envd in the future. + """ + + @classmethod + def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> ImageBuilderConfig: + config_file = get_config_file(config_file) + return ImageBuilderConfig( + default=DefaultConfig.auto(config_file), + ) + + @dataclass(init=True, repr=True, eq=True, frozen=True) class DataConfig(object): """ diff --git a/flytekit/configuration/internal.py b/flytekit/configuration/internal.py index c93e65e635..a33cc9cbae 100644 --- a/flytekit/configuration/internal.py +++ b/flytekit/configuration/internal.py @@ -75,6 +75,11 @@ class Local(object): SECTION = "local" CACHE_ENABLED = ConfigEntry(LegacyConfigEntry(SECTION, "cache_enabled", bool)) CACHE_OVERWRITE = ConfigEntry(LegacyConfigEntry(SECTION, "cache_overwrite", bool)) + IMAGE_BUILDER_NAME = ConfigEntry(LegacyConfigEntry(SECTION, "name"), YamlConfigEntry("image_builder.name")) + UV_IMAGE = ConfigEntry(LegacyConfigEntry(SECTION, "uv_image"), YamlConfigEntry("image_builder.uv_image")) + MICROMAMBA_IMAGE = ConfigEntry( + LegacyConfigEntry(SECTION, "micromamba_image"), YamlConfigEntry("image_builder.micromamba_image") + ) class Credentials(object): diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index f341648283..dcb8a0a5f9 100644 --- a/flytekit/image_spec/default_builder.py +++ b/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 @@ -313,6 +313,12 @@ def prepare_python_executable(image_spec: ImageSpec) -> _PythonInstallTemplate: def create_docker_context(image_spec: ImageSpec, tmp_dir: Path): """Populate tmp_dir with Dockerfile as specified by the `image_spec`.""" base_image = image_spec.base_image or "debian:bookworm-slim" + from flytekit.configuration import ImageBuilderConfig + + image_builder = ImageBuilderConfig().auto() + + image_builder_uv = image_builder.default.uv_image + image_builder_micromamba_image = image_builder.default.micromamba_image if image_spec.cuda is not None or image_spec.cudnn is not None: msg = ( @@ -416,6 +422,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=image_builder_uv, + MICROMAMBA_IMAGE=image_builder_micromamba_image, ) dockerfile_path = tmp_dir / "Dockerfile"