Add user-defined builder image in image spec #3207

arbaobao · 2025-03-24T09:23:24Z

Tracking issue

Why are the changes needed?

According to the discussion in #3180, It seems more reasonable to add this configuration in image_spec rather than config.yaml.

What changes were proposed in this pull request?

Add an argument in image_spec to let users specify uv_image and micromamba_image in their own registry.

How was this patch tested?

This patch can be tested by add builder_config into image_spec.

flytekit_hash = "21652e067d474724877cb6440660e4a4271564f9"
flytekit = f"git+https://github.com/arbaobao/flytekit.git@{flytekit_hash}"
image = ImageSpec(
    packages=[flytekit, "pandas", "pyarrow"],
    apt_packages=["git"],
    registry="localhost:30000",
    builder_config={"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 task using the image_spec above, and you can find out that the uv_image and micromamba_image are both using the image you specified.

Screenshots

Check all the applicable boxes

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

Related PRs

#3180

Docs link

Summary by Bito

This PR introduces configurable builder images via the builder_config parameter, replacing hardcoded references with defaults in the image_spec module. The changes support more flexible configuration, allowing custom registry images for uv and micromamba images as discussed in issue #3180. It also updates the Dockerfile template, enhances GeoPandas encoder, and improves Slurm connectors.

Unit tests added: True

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

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

…r_image_in_image_spec

flyte-bot · 2025-03-24T09:23:34Z

Code Review Agent Run #7a262c

Actionable Suggestions - 1

flytekit/image_spec/default_builder.py - 1
- Consider adding validation for builder_config keys · Line 420-421

Review Details

Files reviewed - 4 · Commit Range: da3716a..21652e0
- flytekit/image_spec/default_builder.py
- flytekit/image_spec/image_spec.py
- pydoclint-errors-baseline.txt
- tests/flytekit/unit/core/image_spec/test_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

Bito Usage Guide

Commands

Type the following command in the pull request comment and save the comment.

/review - Manually triggers a full AI review.

Refer to the documentation for additional commands.

Configuration

This repository uses code_review_bito You can customize the agent settings here or contact your Bito workspace admin at [email protected].

Documentation & Help

AI Code Review powered by

flyte-bot · 2025-03-24T09:28:24Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
New Feature - User-Defined Builder Images Integration	- `default_builder.py` - Replaced hard-coded image references with environment variable placeholders and added conditional handling using builder_config. - `image_spec.py` - Introduced a new builder_config attribute to allow users to specify custom builder images and refined attribute documentation.
Documentation - Documentation Cleanup and Enhancement	- `pydoclint-errors-baseline.txt` - Removed obsolete pydoclint errors by updating class attribute documentation to match actual attributes.
Testing - Unit Tests for Builder Configuration	- `test_default_builder.py` - Added tests to verify that custom uv_image from builder_config is applied and defaults are maintained as expected.

flyte-bot · 2025-03-24T09:28:26Z

flytekit/image_spec/default_builder.py

+        UV_IMAGE=image_spec.builder_config["uv_image"],
+        MICROMAMBA_IMAGE=image_spec.builder_config["micromamba_image"],


Consider adding validation for builder_config keys

The PR adds two new variables to the DOCKER_FILE_TEMPLATE.substitute() call: UV_IMAGE and MICROMAMBA_IMAGE. These variables are used in the Dockerfile template on lines 96-97 to specify the base images for the multi-stage build. Consider adding validation to ensure these keys exist in image_spec.builder_config to prevent potential KeyError exceptions.

Code suggestion

Check the AI-generated fix before applying

Suggested change

UV_IMAGE=image_spec.builder_config["uv_image"],

MICROMAMBA_IMAGE=image_spec.builder_config["micromamba_image"],

UV_IMAGE=image_spec.builder_config.get(

"uv_image", "ghcr.io/astral-sh/uv:0.5.1"),

MICROMAMBA_IMAGE=image_spec.builder_config.get(

"micromamba_image", "mambaorg/micromamba:2.0.3-debian12-slim"),

Code Review Run #7a262c

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

Yes, avoid them

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

codecov · 2025-03-24T10:17:20Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 77.84%. Comparing base (5503ee5) to head (51d4fc4).
Report is 2 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #3207      +/-   ##
==========================================
- Coverage   81.95%   77.84%   -4.11%     
==========================================
  Files         346      214     -132     
  Lines       27852    22356    -5496     
  Branches     2920     2921       +1     
==========================================
- Hits        22826    17404    -5422     
+ Misses       4191     4103      -88     
- Partials      835      849      +14

☔ 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-24T10:54:03Z

Code Review Agent Run #ac0eb1

Actionable Suggestions - 0

Review Details

Files reviewed - 2 · Commit Range: 21652e0..4286947
- flytekit/image_spec/default_builder.py
- flytekit/image_spec/image_spec.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

Bito Usage Guide

Commands

Type the following command in the pull request comment and save the comment.

/review - Manually triggers a full AI review.

Refer to the documentation for additional commands.

Configuration

This repository uses code_review_bito You can customize the agent settings here or contact your Bito workspace admin at [email protected].

Documentation & Help

AI Code Review powered by

thomasjpfan · 2025-03-24T16:51:12Z

flytekit/image_spec/image_spec.py

+        if self.builder_config is None:
+            self.builder_config = {
+                "uv_image": "ghcr.io/astral-sh/uv:0.5.1",
+                "micromamba_image": "mambaorg/micromamba:2.0.3-debian12-slim",
+            }


Since this is specific to the image builder, I think the default should be set in flytekit/image_spec/default_builder.py itself.

For example, this configuration does not mean anything for the envd image builder)

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

flyte-bot · 2025-03-25T17:34:03Z

Code Review Agent Run #fb7615

Actionable Suggestions - 0

Review Details

Files reviewed - 2 · Commit Range: 4286947..738f88e
- flytekit/image_spec/default_builder.py
- flytekit/image_spec/image_spec.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

Bito Usage Guide

Commands

Type the following command in the pull request comment and save the comment.

/review - Manually triggers a full AI review.

Refer to the documentation for additional commands.

Configuration

This repository uses code_review_bito You can customize the agent settings here or contact your Bito workspace admin at [email protected].

Documentation & Help

AI Code Review powered by

thomasjpfan · 2025-03-25T19:22:20Z

flytekit/image_spec/default_builder.py

+    uv_image = "ghcr.io/astral-sh/uv:0.5.1"
+    micromamba_image = "mambaorg/micromamba:2.0.3-debian12-slim"
+    if image_spec.builder_config is not None:
+        if image_spec.builder_config.get("uv_image") is not None:
+            uv_image = image_spec.builder_config.get("uv_image")
+
+        if image_spec.builder_config.get("micromamba_image") is not None:
+            micromamba_image = image_spec.builder_config.get("micromamba_image")


Minor nit:

Suggested change

uv_image = "ghcr.io/astral-sh/uv:0.5.1"

micromamba_image = "mambaorg/micromamba:2.0.3-debian12-slim"

if image_spec.builder_config is not None:

if image_spec.builder_config.get("uv_image") is not None:

uv_image = image_spec.builder_config.get("uv_image")

if image_spec.builder_config.get("micromamba_image") is not None:

micromamba_image = image_spec.builder_config.get("micromamba_image")

if image_spec.builder_config is not None:

uv_image = image_spec.builder_config.get("uv_image", DEFAULT_UV_IMAGE)

micromamba_image = image_spec.builder_config.get("micromamba_image", DEFAULT_MICROMAMBA_IMAGE)

where the defaults are constants at the module level:

DEFAULT_UV_IMAGE = "ghcr.io/astral-sh/uv:0.5.1" DEFAULT_MICROMAMBA_IMAGE = "mambaorg/micromamba:2.0.3-debian12-slim"

Thanks for the advice! It is more readable.

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

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

flyte-bot · 2025-03-27T08:24:06Z

Code Review Agent Run #66e237

Actionable Suggestions - 0

Review Details

Files reviewed - 6 · Commit Range: 738f88e..7dbf061
- flytekit/image_spec/default_builder.py
- plugins/flytekit-geopandas/flytekitplugins/geopandas/gdf_transformers.py
- plugins/flytekit-geopandas/tests/test_geopandas_plugin.py
- plugins/flytekit-slurm/flytekitplugins/slurm/function/connector.py
- plugins/flytekit-slurm/flytekitplugins/slurm/script/connector.py
- tests/flytekit/unit/core/image_spec/test_default_builder.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful

Bito Usage Guide

Commands

Type the following command in the pull request comment and save the comment.

/review - Manually triggers a full AI review.

Refer to the documentation for additional commands.

Configuration

This repository uses code_review_bito You can customize the agent settings here or contact your Bito workspace admin at [email protected].

Documentation & Help

AI Code Review powered by

flytekit/image_spec/image_spec.py

One more comment

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

…/flytekit into add_builder_image_in_image_spec

flyte-bot · 2025-04-01T16:01:56Z

Code Review Agent Run #24a266

Actionable Suggestions - 0

Review Details

Files reviewed - 1 · Commit Range: 7dbf061..c6d824c
- flytekit/image_spec/image_spec.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful

Bito Usage Guide

Commands

Type the following command in the pull request comment and save the comment.

/review - Manually triggers a full AI review.

Refer to the documentation for additional commands.

Configuration

This repository uses code_review_bito You can customize the agent settings here or contact your Bito workspace admin at [email protected].

Documentation & Help

AI Code Review powered by

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

flyte-bot · 2025-04-23T17:19:56Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - Bito Automatic Review Skipped - Branch Excluded

Bito didn't auto-review because the source or target branch is excluded from automatic reviews.
To trigger review, type /review in the comment and save.
You can change the settings here, or contact the agent instance creator at [email protected].

arbaobao added 6 commits March 13, 2025 16:51

add builder_config

da3716a

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

add unit test

0c00e4e

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

test

c07af06

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

add docstring

0101fed

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

fix pydoclint error

ffd59ce

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

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

21652e0

…r_image_in_image_spec

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

flyte-bot reviewed Mar 24, 2025

View reviewed changes

arbaobao added 2 commits March 24, 2025 17:48

fix error

9467e32

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

fix unit test error

4286947

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

thomasjpfan reviewed Mar 24, 2025

View reviewed changes

arbaobao added 3 commits March 26, 2025 00:30

move the default builder to default_builder

29da1fd

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

error

c53eb13

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

adjust logic

738f88e

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

thomasjpfan reviewed Mar 25, 2025

View reviewed changes

arbaobao added 3 commits March 27, 2025 15:20

improve logic and test more accurate

30f6b4f

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

lint

51d4fc4

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

Merge branch 'flyteorg:master' into add_builder_image_in_image_spec

7dbf061

thomasjpfan previously approved these changes Mar 31, 2025

View reviewed changes

thomasjpfan reviewed Mar 31, 2025

View reviewed changes

flytekit/image_spec/image_spec.py Outdated Show resolved Hide resolved

arbaobao added 2 commits April 1, 2025 22:56

type hint

9611f07

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

Merge branch 'add_builder_image_in_image_spec' of github.com:arbaobao…

c6d824c

…/flytekit into add_builder_image_in_image_spec

conflicts

e454d7a

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

arbaobao requested a review from davidmirror-ops as a code owner April 23, 2025 16:10

pingsutw approved these changes Apr 23, 2025

View reviewed changes

pingsutw merged commit 1bc8302 into flyteorg:master Apr 23, 2025
114 of 115 checks passed

This was referenced Apr 23, 2025

[BUG] The lockfile at uv.lock needs to be updated, but --locked was provided. To update the lockfile, run uv lock. flyteorg/flyte#6392

Closed

Allowing developer set uv version when building task image #3234

Closed

		UV_IMAGE=image_spec.builder_config["uv_image"],
		MICROMAMBA_IMAGE=image_spec.builder_config["micromamba_image"],

-        UV_IMAGE=image_spec.builder_config["uv_image"],
-        MICROMAMBA_IMAGE=image_spec.builder_config["micromamba_image"],
+        UV_IMAGE=image_spec.builder_config.get(
+            "uv_image", "ghcr.io/astral-sh/uv:0.5.1"),
+        MICROMAMBA_IMAGE=image_spec.builder_config.get(
+            "micromamba_image", "mambaorg/micromamba:2.0.3-debian12-slim"),

Add user-defined builder image in image spec #3207

Add user-defined builder image in image spec #3207

Uh oh!

Conversation

arbaobao commented Mar 24, 2025 • edited by flyte-bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

Uh oh!

flyte-bot commented Mar 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run #7a262c

Bito Usage Guide

Uh oh!

flyte-bot commented Mar 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changelist by Bito

Uh oh!

flyte-bot Mar 24, 2025

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Mar 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

flyte-bot commented Mar 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run #ac0eb1

Bito Usage Guide

Uh oh!

thomasjpfan Mar 24, 2025

Choose a reason for hiding this comment

Uh oh!

flyte-bot commented Mar 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run #fb7615

Bito Usage Guide

Uh oh!

thomasjpfan Mar 25, 2025

Choose a reason for hiding this comment

Uh oh!

arbaobao Mar 27, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

flyte-bot commented Mar 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run #66e237

Bito Usage Guide

Uh oh!

Uh oh!

flyte-bot commented Apr 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run #24a266

Uh oh!

Uh oh!

flyte-bot commented Apr 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run Status

Uh oh!

Uh oh!

arbaobao commented Mar 24, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Mar 24, 2025 •

edited

Loading

flyte-bot commented Mar 24, 2025 •

edited

Loading

codecov bot commented Mar 24, 2025 •

edited

Loading

flyte-bot commented Mar 24, 2025 •

edited

Loading

flyte-bot commented Mar 25, 2025 •

edited

Loading

flyte-bot commented Mar 27, 2025 •

edited

Loading

flyte-bot commented Apr 1, 2025 •

edited

Loading

flyte-bot commented Apr 23, 2025 •

edited

Loading