feat: implement force push functionality in ImageSpec

flytekit/image_spec/image_spec.py

@@ -44,6 +44,7 @@ class ImageSpec:
         pip_index: Specify the custom pip index url
         registry_config: Specify the path to a JSON registry config file
         commands: Command to run during the building process
+        is_force_push: forcefully pushed to the registry
     """
 
     name: str = "flytekit"
@@ -64,9 +65,11 @@ class ImageSpec:
     pip_index: Optional[str] = None
     registry_config: Optional[str] = None
     commands: Optional[List[str]] = None
+    is_force_push: Optional[bool] = None
 
     def __post_init__(self):
         self.name = self.name.lower()
+        self.is_force_push = False  # False by default
         if self.registry:
             self.registry = self.registry.lower()
 
@@ -181,6 +184,15 @@ def with_apt_packages(self, apt_packages: Union[str, List[str]]) -> "ImageSpec":
 
         return new_image_spec
 
+    def force_push(self, is_force_push: Optional[bool] = True) -> "ImageSpec":
+        """
+        Builder that returns a new image spec with force push enabled.
+        """
+        new_image_spec = copy.deepcopy(self)
+        new_image_spec.is_force_push = is_force_push
+
+        return new_image_spec
+
 
 class ImageSpecBuilder:
     @abstractmethod
@@ -222,10 +234,13 @@ 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():
+        if (img_name in cls._BUILT_IMAGES or image_spec.exist()) and not image_spec.is_force_push:
             click.secho(f"Image {img_name} found. Skip building.", fg="blue")
         else:
-            click.secho(f"Image {img_name} not found. Building...", fg="blue")
+            if image_spec.is_force_push:
+                click.secho(f"Image {img_name} found. but overwriting existing image.", 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":

tests/flytekit/unit/cli/pyflyte/test_run.py

@@ -297,9 +297,9 @@ def test_list_default_arguments(wf_path):
 )
 
 ic_result_4 = ImageConfig(
-    default_image=Image(name="default", fqn="flytekit", tag="DgQMqIi61py4I4P5iOeS0Q.."),
+    default_image=Image(name="default", fqn="flytekit", tag="guqLm2fbPHgwk0EiXmXI9g.."),
     images=[
-        Image(name="default", fqn="flytekit", tag="DgQMqIi61py4I4P5iOeS0Q.."),
+        Image(name="default", fqn="flytekit", tag="guqLm2fbPHgwk0EiXmXI9g.."),
         Image(name="xyz", fqn="docker.io/xyz", tag="latest"),
         Image(name="abc", fqn="docker.io/abc", tag=None),
     ],

tests/flytekit/unit/core/image_spec/test_image_spec.py

@@ -30,6 +30,7 @@ def test_image_spec(mock_image_spec_builder):
     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 +48,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 image_spec.image_name() == f"flytekit:{tag}"