Skip to content

Add image builder user defined #3180

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 7 commits into
base: master
Choose a base branch
from
Open

Add image builder user defined #3180

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
44c7cc0
add image builder comment
arbaobao Feb 25, 2025
baea4ee
Merge branch 'master' of github.com:flyteorg/flytekit into add_image_…
arbaobao Feb 27, 2025
e10a65c
typo
arbaobao Mar 7, 2025
8d95724
lint
arbaobao Mar 7, 2025
4a6394e
make ImageBuilder independent
arbaobao Mar 7, 2025
d6c0901
remove redundant code
arbaobao Mar 7, 2025
a84654b
rewrite for the elasticity
arbaobao Mar 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions flytekit/configuration/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -639,6 +639,45 @@
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)

Check warning on line 659 in flytekit/configuration/__init__.py

View check run for this annotation

Codecov / codecov/patch

flytekit/configuration/__init__.py#L654-L659 

Added lines #L654 - L659 were not covered by tests


@dataclass(init=True, repr=True, eq=True, frozen=True)
class ImageBuilderConfig(object):
"""
Any image builder specific configuration.
"""

default: DefaultConfig = DefaultConfig()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Function call in dataclass default

The dataclass field default is using a function call DefaultConfig() as a default value, which can lead to unexpected behavior. Consider using field(default_factory=DefaultConfig) instead.

Code suggestion
Check the AI-generated fix before applying
Suggested change
default: DefaultConfig = DefaultConfig()
default: DefaultConfig = field(default_factory=DefaultConfig)

Code Review Run #df86ab

Should Bito avoid suggestions like this for future reviews? (Manage Rules)

  • Yes, avoid them

"""
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(

Check warning on line 676 in flytekit/configuration/__init__.py

View check run for this annotation

Codecov / codecov/patch

flytekit/configuration/__init__.py#L675-L676 

Added lines #L675 - L676 were not covered by tests
default=DefaultConfig.auto(config_file),
)


@dataclass(init=True, repr=True, eq=True, frozen=True)
class DataConfig(object):
"""
Expand Down
5 changes: 5 additions & 0 deletions flytekit/configuration/internal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
12 changes: 10 additions & 2 deletions flytekit/image_spec/default_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +96 to +97
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider defining default values for template variables

Consider defining the default values for UV_IMAGE and MICROMAMBA_IMAGE variables to maintain backward compatibility. Currently, these variables are used in the Dockerfile template but their default values are not defined, which might cause issues if they're not provided when the template is rendered.

Code suggestion
Check the AI-generated fix before applying
Suggested change
FROM $UV_IMAGE as uv
FROM $MICROMAMBA_IMAGE as micromamba
# Default to the previously hardcoded versions if not specified
FROM ${UV_IMAGE:-ghcr.io/astral-sh/uv:0.5.1} as uv
FROM ${MICROMAMBA_IMAGE:-mambaorg/micromamba:2.0.3-debian12-slim} as micromamba

Code Review Run #4675f8

Should Bito avoid suggestions like this for future reviews? (Manage Rules)

  • Yes, avoid them


FROM $BASE_IMAGE

Expand Down Expand Up @@ -313,6 +313,12 @@
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

Check warning on line 317 in flytekit/image_spec/default_builder.py

View check run for this annotation

Codecov / codecov/patch

flytekit/image_spec/default_builder.py#L316-L317 

Added lines #L316 - L317 were not covered by tests
image_builder = ImageBuilderConfig().auto()

Check warning on line 319 in flytekit/image_spec/default_builder.py

View check run for this annotation

Codecov / codecov/patch

flytekit/image_spec/default_builder.py#L319 

Added line #L319 was not covered by tests
image_builder_uv = image_builder.default.uv_image
image_builder_micromamba_image = image_builder.default.micromamba_image

Check warning on line 321 in flytekit/image_spec/default_builder.py

View check run for this annotation

Codecov / codecov/patch

flytekit/image_spec/default_builder.py#L321 

Added line #L321 was not covered by tests

if image_spec.cuda is not None or image_spec.cudnn is not None:
msg = (
Expand Down Expand Up @@ -416,6 +422,8 @@
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"
Expand Down
Loading