From 64730c8ffb4e76cb268817e30d31c6d209452e75 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 18 Mar 2024 01:44:44 -0700 Subject: [PATCH 1/2] feat: Support ImageSpec as base_image in ImageSpec --- flytekit/image_spec/image_spec.py | 8 ++++++- .../flytekit-envd/tests/test_image_spec.py | 21 ++++++++++++++++--- 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 326e8a0e80..2fb00d0ebe 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -60,7 +60,7 @@ class ImageSpec: apt_packages: Optional[List[str]] = None cuda: Optional[str] = None cudnn: Optional[str] = None - base_image: Optional[str] = None + base_image: Optional[Union[str, "ImageSpec"]] = None platform: str = "linux/amd64" pip_index: Optional[str] = None registry_config: Optional[str] = None @@ -226,6 +226,10 @@ def register(cls, builder_type: str, image_spec_builder: ImageSpecBuilder, prior @classmethod @lru_cache def build(cls, image_spec: ImageSpec) -> str: + if isinstance(image_spec.base_image, ImageSpec): + cls.build(image_spec.base_image) + image_spec.base_image = image_spec.base_image.image_name() + if image_spec.builder is None and cls._REGISTRY: builder = max(cls._REGISTRY, key=lambda name: cls._REGISTRY[name][1]) else: @@ -267,6 +271,8 @@ def calculate_hash_from_image_spec(image_spec: ImageSpec): """ # copy the image spec to avoid modifying the original image spec. otherwise, the hash will be different. spec = copy.deepcopy(image_spec) + if isinstance(spec.base_image, ImageSpec): + spec.base_image = spec.base_image.image_name() spec.source_root = hash_directory(image_spec.source_root) if image_spec.source_root else b"" if spec.requirements: spec.requirements = hashlib.sha1(pathlib.Path(spec.requirements).read_bytes()).hexdigest() diff --git a/plugins/flytekit-envd/tests/test_image_spec.py b/plugins/flytekit-envd/tests/test_image_spec.py index d77b2ca89b..4d11c10279 100644 --- a/plugins/flytekit-envd/tests/test_image_spec.py +++ b/plugins/flytekit-envd/tests/test_image_spec.py @@ -4,7 +4,7 @@ import pytest from flytekitplugins.envd.image_builder import EnvdImageSpecBuilder, create_envd_config -from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec +from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec, ImageSpecBuilder @pytest.fixture(scope="module", autouse=True) @@ -21,17 +21,32 @@ def register_envd_higher_priority(): def test_image_spec(): + class MockImageSpecBuilder(ImageSpecBuilder): + def build_image(self, img): + print("Building an image...") + + ImageBuildEngine.register("dummy", MockImageSpecBuilder()) + base_image = ImageSpec( + packages=["numpy"], + python_version="3.8", + registry="", + base_image="cr.flyte.org/flyteorg/flytekit:py3.8-latest", + ) + # Replace the base image name with the default flytekit image name, + # so Envd can find the base image when building imageSpec below + ImageBuildEngine._IMAGE_NAME_TO_REAL_NAME[base_image.image_name()] = "cr.flyte.org/flyteorg/flytekit:py3.8-latest" + image_spec = ImageSpec( packages=["pandas"], apt_packages=["git"], python_version="3.8", - base_image="cr.flyte.org/flyteorg/flytekit:py3.8-latest", + base_image=base_image, pip_index="https://private-pip-index/simple", ) image_spec = image_spec.with_commands("echo hello") - EnvdImageSpecBuilder().build_image(image_spec) + ImageBuildEngine.build(image_spec) config_path = create_envd_config(image_spec) assert image_spec.platform == "linux/amd64" image_name = image_spec.image_name() From edbaf855a4a0ed2879be27151701deb87299cfad Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 18 Mar 2024 02:12:57 -0700 Subject: [PATCH 2/2] nit Signed-off-by: Kevin Su --- plugins/flytekit-envd/tests/test_image_spec.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/plugins/flytekit-envd/tests/test_image_spec.py b/plugins/flytekit-envd/tests/test_image_spec.py index 4d11c10279..c495308f37 100644 --- a/plugins/flytekit-envd/tests/test_image_spec.py +++ b/plugins/flytekit-envd/tests/test_image_spec.py @@ -4,7 +4,7 @@ import pytest from flytekitplugins.envd.image_builder import EnvdImageSpecBuilder, create_envd_config -from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec, ImageSpecBuilder +from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec @pytest.fixture(scope="module", autouse=True) @@ -21,11 +21,6 @@ def register_envd_higher_priority(): def test_image_spec(): - class MockImageSpecBuilder(ImageSpecBuilder): - def build_image(self, img): - print("Building an image...") - - ImageBuildEngine.register("dummy", MockImageSpecBuilder()) base_image = ImageSpec( packages=["numpy"], python_version="3.8",