From 020afea626c21378fe673ad93e8073e5427e3aa0 Mon Sep 17 00:00:00 2001 From: Jason Lai Date: Thu, 14 Mar 2024 04:31:15 +0800 Subject: [PATCH] feat: implement force push functionality in ImageSpec (#2234) Signed-off-by: jason.lai Signed-off-by: Kevin Su Co-authored-by: Kevin Su --- flytekit/image_spec/image_spec.py | 56 ++++++++++++------- .../unit/core/image_spec/test_image_spec.py | 14 +++++ 2 files changed, 50 insertions(+), 20 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 9f84d28dd4d..5a1517f90a0 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -18,6 +18,7 @@ DOCKER_HUB = "docker.io" _F_IMG_ID = "_F_IMG_ID" +FLYTE_FORCE_PUSH_IMAGE_SPEC = "FLYTE_FORCE_PUSH_IMAGE_SPEC" @dataclass @@ -68,6 +69,7 @@ class ImageSpec: def __post_init__(self): self.name = self.name.lower() + self._is_force_push = os.environ.get(FLYTE_FORCE_PUSH_IMAGE_SPEC, False) # False by default if self.registry: self.registry = self.registry.lower() @@ -182,6 +184,15 @@ def with_apt_packages(self, apt_packages: Union[str, List[str]]) -> "ImageSpec": return new_image_spec + def force_push(self) -> "ImageSpec": + """ + Builder that returns a new image spec with force push enabled. + """ + new_image_spec = copy.deepcopy(self) + new_image_spec._is_force_push = True + + return new_image_spec + class ImageSpecBuilder: @abstractmethod @@ -204,7 +215,6 @@ class ImageBuildEngine: """ _REGISTRY: typing.Dict[str, Tuple[ImageSpecBuilder, int]] = {} - _BUILT_IMAGES: typing.Set[str] = set() # _IMAGE_NAME_TO_REAL_NAME is used to keep track of the fully qualified image name # returned by the image builder. This allows ImageSpec to map from `image_spc.image_name()` # to the real qualified name. @@ -223,26 +233,32 @@ def build(cls, image_spec: ImageSpec) -> str: builder = image_spec.builder img_name = image_spec.image_name() - if img_name in cls._BUILT_IMAGES or image_spec.exist(): - click.secho(f"Image {img_name} found. Skip building.", fg="blue") + if image_spec.exist(): + if image_spec._is_force_push: + click.secho(f"Image {img_name} found. but overwriting existing image.", fg="blue") + cls._build_image(builder, image_spec, img_name) + else: + click.secho(f"Image {img_name} found. Skip building.", fg="blue") else: - click.secho(f"Image {img_name} not found. Building...", fg="blue") - if builder not in cls._REGISTRY: - raise Exception(f"Builder {builder} is not registered.") - if builder == "envd": - envd_version = metadata.version("envd") - # flytekit v1.10.2+ copies the workflow code to the WorkDir specified in the Dockerfile. However, envd<0.3.39 - # overwrites the WorkDir when building the image, resulting in a permission issue when flytekit downloads the file. - if Version(envd_version) < Version("0.3.39"): - raise FlyteAssertion( - f"envd version {envd_version} is not compatible with flytekit>v1.10.2." - f" Please upgrade envd to v0.3.39+." - ) - - fully_qualified_image_name = cls._REGISTRY[builder][0].build_image(image_spec) - if fully_qualified_image_name is not None: - cls._IMAGE_NAME_TO_REAL_NAME[img_name] = fully_qualified_image_name - cls._BUILT_IMAGES.add(img_name) + click.secho(f"Image {img_name} not found. building...", fg="blue") + cls._build_image(builder, image_spec, img_name) + + @classmethod + def _build_image(cls, builder, image_spec, img_name): + if builder not in cls._REGISTRY: + raise Exception(f"Builder {builder} is not registered.") + if builder == "envd": + envd_version = metadata.version("envd") + # flytekit v1.10.2+ copies the workflow code to the WorkDir specified in the Dockerfile. However, envd<0.3.39 + # overwrites the WorkDir when building the image, resulting in a permission issue when flytekit downloads the file. + if Version(envd_version) < Version("0.3.39"): + raise FlyteAssertion( + f"envd version {envd_version} is not compatible with flytekit>v1.10.2." + f" Please upgrade envd to v0.3.39+." + ) + fully_qualified_image_name = cls._REGISTRY[builder][0].build_image(image_spec) + if fully_qualified_image_name is not None: + cls._IMAGE_NAME_TO_REAL_NAME[img_name] = fully_qualified_image_name @lru_cache diff --git a/tests/flytekit/unit/core/image_spec/test_image_spec.py b/tests/flytekit/unit/core/image_spec/test_image_spec.py index 9ec1adc77c2..f1b0bbe9755 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -26,10 +26,12 @@ def test_image_spec(mock_image_spec_builder): requirements=REQUIREMENT_FILE, registry_config=REGISTRY_CONFIG_FILE, ) + assert image_spec._is_force_push is False image_spec = image_spec.with_commands("echo hello") image_spec = image_spec.with_packages("numpy") image_spec = image_spec.with_apt_packages("wget") + image_spec = image_spec.force_push() assert image_spec.python_version == "3.8" assert image_spec.base_image == "cr.flyte.org/flyteorg/flytekit:py3.8-latest" @@ -47,6 +49,7 @@ def test_image_spec(mock_image_spec_builder): assert image_spec.pip_index is None assert image_spec.is_container() is True assert image_spec.commands == ["echo hello"] + assert image_spec._is_force_push is True tag = calculate_hash_from_image_spec(image_spec) assert "=" != tag[-1] @@ -98,3 +101,14 @@ def test_image_spec_engine_priority(): assert image_spec.image_name() == new_image_name del ImageBuildEngine._REGISTRY["build_10"] del ImageBuildEngine._REGISTRY["build_default"] + + +def test_build_existing_image_with_force_push(): + image_spec = Mock() + image_spec.exist.return_value = True + image_spec._is_force_push = True + + ImageBuildEngine._build_image = Mock() + + ImageBuildEngine.build(image_spec) + ImageBuildEngine._build_image.assert_called_once()