feat: implement force push functionality in ImageSpec (flyteorg#2234)

Signed-off-by: jason.lai <[email protected]>
Signed-off-by: Kevin Su <[email protected]>
Co-authored-by: Kevin Su <[email protected]>

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

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()