flytekit icon indicating copy to clipboard operation
flytekit copied to clipboard

Add support imageSpec chaining

Open pingsutw opened this issue 1 year ago • 0 comments

Tracking issue

NA

Why are the changes needed?

Cannot control the order of the docker command

What changes were proposed in this pull request?

.with_packages will return a new ImageSpec. In this new ImageSpec, the base image will be the current image, and packages will be the input of with_packages. Therefore, the new imageSpec will build and push the base image, and another new imageSpec will be built with new packages.

ImageSpec(registry="ghcr.io/flyteorg", packages=["tensorflow"])  # build and push image 1
         .with_pacakges("wget")  # build and push image 2 (base_image is image 1)
         .with_apt_packages("numpy")  -> build image  # build and push image 3 (base image is image 2)

How was this patch tested?

from flytekit import task, workflow, ImageSpec

custom_image = ImageSpec(
    registry="ghcr.io/flyteorg",
    packages=["tensorflow"],
    apt_packages=["wget"],
    commands=["echo 6"]
).with_apt_packages("wget").with_packages("numpy")


@task(container_image=custom_image)
def t1(a: int) -> int:
    return a + 1


@workflow
def wf() -> int:
    return t1(a=3)

Setup process

Screenshots

Check all the applicable boxes

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

Related PRs

NA

Docs link

NA

pingsutw avatar Jun 30 '24 06:06 pingsutw