Add image builder user defined #3180

arbaobao · 2025-03-07T07:14:38Z

Tracking issue

Why are the changes needed?

In some situations, there isn't internet to build image from pulling uv and micromamba images from the net.

What changes were proposed in this pull request?

I add an ImageBuilderConfig to detect if there is user defined image_builder in config.yaml.

How was this patch tested?

This patch can be tested by adding image_builder to ~/.flyte/config.yaml

image_builder:
  uv_image: arbaobao/flytekit:uv
  micromamba_image: arbaobao/flyte-test-images:micromamba-amd64

if we don't specify the image, the default images are ghcr.io/astral-sh/uv:0.5.1 and mambaorg/micromamba:2.0.3-debian12-slim.

Setup process

Run a flyte sandbox and add image_builder to the config.yaml.

2. if you are using image_spec to run the workflow, it will use the image that you specified to build the task image.

Screenshots

Check all the applicable boxes

I updated the documentation accordingly.
All new and existing tests passed.
All commits are signed-off.

Related PRs

Docs link

Summary by Bito

This PR adds support for user-defined image builder configuration to enhance flexibility in offline or air-gapped environments. It refactors configuration classes to better separate default from user-defined settings, replacing legacy configuration with new classes including ImageBuilderConfig and updated DataConfig. Docker file templates are revised to allow users to override default images through configuration files.

Unit tests added: False

Estimated effort to review (1-5, lower is better): 2

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

…builder_user_defined

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

flyte-bot · 2025-03-07T07:14:55Z

Code Review Agent Run #4675f8

Actionable Suggestions - 2

flytekit/image_spec/default_builder.py - 1
- Consider defining default values for template variables · Line 96-97
flytekit/configuration/__init__.py - 1
- Function call in dataclass default · Line 673-673

Review Details

Files reviewed - 3 · Commit Range: 44c7cc0..e10a65c
- flytekit/configuration/__init__.py
- flytekit/configuration/internal.py
- flytekit/image_spec/default_builder.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2025-03-07T07:19:05Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
New Feature - User-defined Image Builder Configuration	- `__init__.py` - Introduces new dataclass-based configuration classes (DefaultConfig and ImageBuilderConfig) and their auto methods to load user-defined image builder settings. - `internal.py` - Adds new configuration entries (IMAGE_BUILDER_NAME, UV_IMAGE, and MICROMAMBA_IMAGE) to support the image builder feature.
Feature Improvement - Docker Builder Template Parameterization	- `default_builder.py` - Replaces hardcoded image sources with configurable variables in the Docker template and updates the docker context creation to utilize the user-specified image values.

flyte-bot · 2025-03-07T07:19:07Z

flytekit/image_spec/default_builder.py

+FROM $UV_IMAGE as uv
+FROM $MICROMAMBA_IMAGE as micromamba


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

flyte-bot · 2025-03-07T07:19:09Z

flytekit/configuration/__init__.py

@@ -651,6 +670,8 @@ class DataConfig(object):
    gcs: GCSConfig = GCSConfig()
    azure: AzureBlobStorageConfig = AzureBlobStorageConfig()
    generic: GenericPersistenceConfig = GenericPersistenceConfig()
+    image_builder: ImageBuilderConfig = ImageBuilderConfig()


Function call in dataclass default

The default value for image_builder in DataConfig is a function call to ImageBuilderConfig(). Function calls in dataclass defaults can lead to unexpected behavior as they are evaluated only once at definition time.

Code suggestion

Check the AI-generated fix before applying

Suggested change

image_builder: ImageBuilderConfig = ImageBuilderConfig()

image_builder: ImageBuilderConfig = field(default_factory=ImageBuilderConfig)

Code Review Run #4675f8

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

Yes, avoid them

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

pingsutw · 2025-03-07T07:18:36Z

flytekit/configuration/__init__.py

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


Suggested change

Any GCS specific configuration.

Any image builder configuration.

pingsutw · 2025-03-07T07:19:21Z

flytekit/configuration/__init__.py

+    micromamba_image: str = "mambaorg/micromamba:2.0.3-debian12-slim"
+
+    @classmethod
+    def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> GCSConfig:


Suggested change

def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> GCSConfig:

def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> ImageBuilderConfig:

pingsutw · 2025-03-07T07:20:59Z

flytekit/configuration/__init__.py

@@ -651,6 +670,8 @@ class DataConfig(object):
    gcs: GCSConfig = GCSConfig()
    azure: AzureBlobStorageConfig = AzureBlobStorageConfig()
    generic: GenericPersistenceConfig = GenericPersistenceConfig()
+    image_builder: ImageBuilderConfig = ImageBuilderConfig()


Why do we put it in the DataConfig?

I think I can make it independent.

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

codecov · 2025-03-07T08:05:13Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 77.22%. Comparing base (2e12f43) to head (a84654b).
Report is 13 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #3180      +/-   ##
==========================================
- Coverage   83.58%   77.22%   -6.37%     
==========================================
  Files           3      213     +210     
  Lines         195    22275   +22080     
  Branches        0     2901    +2901     
==========================================
+ Hits          163    17201   +17038     
- Misses         32     4227    +4195     
- Partials        0      847     +847

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

flyte-bot · 2025-03-07T08:29:02Z

Code Review Agent Run #df86ab

Actionable Suggestions - 1

flytekit/configuration/__init__.py - 1
- Function call in dataclass default · Line 668-668

Review Details

Files reviewed - 3 · Commit Range: e10a65c..a84654b
- flytekit/configuration/__init__.py
- flytekit/configuration/internal.py
- flytekit/image_spec/default_builder.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2025-03-07T08:32:38Z

flytekit/configuration/__init__.py

+    Any image builder specific configuration.
+    """
+
+    default: DefaultConfig = DefaultConfig()


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

thomasjpfan

Thank you for the PR @arbaobao !

I prefer to have these configuration options in ImageSpec itself (or a way to for ImageSpec to pass in builder specific configuration, like a task_config but for image builders)

Design-wise, when two people running the same image spec, it should produce the same image. If we depend on an external config file for building, then it'll make it harder to debug and reason about.

The options I see are:

builder_config

ImageSpec(builder_config={"uv_image": "..."})

Add another kwarg

ImageSpec(uv_image="...")

pingsutw · 2025-03-08T00:08:38Z

@thomasjpfan I see your point, but if users want to use the UV image from a private registry, they must update every image spec

pingsutw · 2025-03-08T00:10:07Z

or we can probably show the builder config in the error message, will that be helpful for debugging

thomasjpfan · 2025-03-08T03:06:36Z

I see your point, but if users want to use the UV image from a private registry, they must update every image spec

I think it is reasonable for users to change all their ImageSpecs if their uv image is in a private registry. I consider considering the uv image to be on the same level as ImageSpec.base_image.

arbaobao · 2025-03-12T15:31:15Z

@pingsutw @thomasjpfan
According to the discuss above, I think I will implement this config in image_spec, but user will have to add this config at every image_spec(xxx) repeatedly.

Is there anything still need to discuss?

arbaobao added 3 commits February 25, 2025 16:59

add image builder comment

44c7cc0

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

Merge branch 'master' of github.com:flyteorg/flytekit into add_image_…

baea4ee

…builder_user_defined

typo

e10a65c

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

arbaobao requested review from wild-endeavor, kumare3, eapolinario, pingsutw, cosmicBboy, samhita-alla, thomasjpfan and Future-Outlier as code owners March 7, 2025 07:14

flyte-bot reviewed Mar 7, 2025

View reviewed changes

lint

8d95724

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

pingsutw reviewed Mar 7, 2025

View reviewed changes

arbaobao added 3 commits March 7, 2025 15:31

make ImageBuilder independent

4a6394e

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

remove redundant code

d6c0901

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

rewrite for the elasticity

a84654b

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

flyte-bot reviewed Mar 7, 2025

View reviewed changes

thomasjpfan reviewed Mar 7, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add image builder user defined #3180

Add image builder user defined #3180

arbaobao commented Mar 7, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Mar 7, 2025 •

edited

Loading

Code Review Agent Run #4675f8

flyte-bot commented Mar 7, 2025 •

edited

Loading

Changelist by Bito

flyte-bot Mar 7, 2025

flyte-bot Mar 7, 2025

pingsutw Mar 7, 2025

pingsutw Mar 7, 2025

pingsutw Mar 7, 2025

arbaobao Mar 7, 2025

codecov bot commented Mar 7, 2025 •

edited

Loading

flyte-bot commented Mar 7, 2025 •

edited

Loading

Code Review Agent Run #df86ab

flyte-bot Mar 7, 2025

thomasjpfan left a comment •

edited

Loading

pingsutw commented Mar 8, 2025

pingsutw commented Mar 8, 2025

thomasjpfan commented Mar 8, 2025

arbaobao commented Mar 12, 2025

-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

	image_builder: ImageBuilderConfig = ImageBuilderConfig()
	image_builder: ImageBuilderConfig = field(default_factory=ImageBuilderConfig)

	Any GCS specific configuration.
	Any image builder configuration.

	def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> GCSConfig:
	def auto(cls, config_file: typing.Union[str, ConfigFile] = None) -> ImageBuilderConfig:

	default: DefaultConfig = DefaultConfig()
	default: DefaultConfig = field(default_factory=DefaultConfig)

Add image builder user defined #3180

Are you sure you want to change the base?

Add image builder user defined #3180

Conversation

arbaobao commented Mar 7, 2025 • edited by flyte-bot Loading

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

Related PRs

Docs link

Summary by Bito

flyte-bot commented Mar 7, 2025 • edited Loading

Code Review Agent Run #4675f8

flyte-bot commented Mar 7, 2025 • edited Loading

Changelist by Bito

flyte-bot Mar 7, 2025

Choose a reason for hiding this comment

flyte-bot Mar 7, 2025

Choose a reason for hiding this comment

pingsutw Mar 7, 2025

Choose a reason for hiding this comment

pingsutw Mar 7, 2025

Choose a reason for hiding this comment

pingsutw Mar 7, 2025

Choose a reason for hiding this comment

arbaobao Mar 7, 2025

Choose a reason for hiding this comment

codecov bot commented Mar 7, 2025 • edited Loading

Codecov Report

flyte-bot commented Mar 7, 2025 • edited Loading

Code Review Agent Run #df86ab

flyte-bot Mar 7, 2025

Choose a reason for hiding this comment

thomasjpfan left a comment • edited Loading

Choose a reason for hiding this comment

pingsutw commented Mar 8, 2025

pingsutw commented Mar 8, 2025

thomasjpfan commented Mar 8, 2025

arbaobao commented Mar 12, 2025

arbaobao commented Mar 7, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Mar 7, 2025 •

edited

Loading

flyte-bot commented Mar 7, 2025 •

edited

Loading

codecov bot commented Mar 7, 2025 •

edited

Loading

flyte-bot commented Mar 7, 2025 •

edited

Loading

thomasjpfan left a comment •

edited

Loading