Skip to content

Add user-defined builder image in image spec #3207

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.

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

Merged
merged 17 commits into from
Apr 23, 2025

Conversation

arbaobao
Copy link
Contributor

@arbaobao arbaobao commented Mar 24, 2025

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

Screenshot 2025-03-24 at 5 21 55 PM

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]>
Signed-off-by: Nelson Chen <[email protected]>
Signed-off-by: Nelson Chen <[email protected]>
Signed-off-by: Nelson Chen <[email protected]>
Signed-off-by: Nelson Chen <[email protected]>
@flyte-bot
Copy link
Contributor

flyte-bot commented Mar 24, 2025

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 Bito Logo

@flyte-bot
Copy link
Contributor

flyte-bot commented Mar 24, 2025

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.

Comment on lines 420 to 421
UV_IMAGE=image_spec.builder_config["uv_image"],
MICROMAMBA_IMAGE=image_spec.builder_config["micromamba_image"],
Copy link
Contributor

Choose a reason for hiding this comment

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

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]>
Signed-off-by: Nelson Chen <[email protected]>
Copy link

codecov bot commented Mar 24, 2025

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
Copy link
Contributor

flyte-bot commented Mar 24, 2025

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 Bito Logo

Comment on lines 150 to 154
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",
}
Copy link
Contributor

Choose a reason for hiding this comment

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

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]>
Signed-off-by: Nelson Chen <[email protected]>
@flyte-bot
Copy link
Contributor

flyte-bot commented Mar 25, 2025

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 Bito Logo

Comment on lines 408 to 415
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")
Copy link
Contributor

Choose a reason for hiding this comment

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

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"

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Thanks for the advice! It is more readable.

@flyte-bot
Copy link
Contributor

flyte-bot commented Mar 27, 2025

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 Bito Logo

thomasjpfan
thomasjpfan previously approved these changes Mar 31, 2025
@thomasjpfan thomasjpfan dismissed their stale review March 31, 2025 14:06

One more comment

arbaobao added 2 commits April 1, 2025 22:56
Signed-off-by: Nelson Chen <[email protected]>
…/flytekit into add_builder_image_in_image_spec
@flyte-bot
Copy link
Contributor

flyte-bot commented Apr 1, 2025

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 Bito Logo

Signed-off-by: Nelson Chen <[email protected]>
@pingsutw pingsutw merged commit 1bc8302 into flyteorg:master Apr 23, 2025
114 of 115 checks passed
@flyte-bot
Copy link
Contributor

flyte-bot commented Apr 23, 2025

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].

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants