add ImageSpecBuilder.should_build method to make building logic custo…

…mizeable by builder plugin

Signed-off-by: Niels Bantilan <[email protected]>

flytekit/image_spec/image_spec.py

@@ -216,6 +216,19 @@ def build_image(self, image_spec: ImageSpec) -> Optional[str]:
         """
         raise NotImplementedError("This method is not implemented in the base class.")
 
+    @abstractmethod
+    def should_build(self, image_spec: ImageSpec) -> bool:
+        """
+        Whether or not the builder should build the ImageSpec.
+
+        Args:
+            image_spec: image spec of the task.
+
+        Returns:
+            True if the image should be built, otherwise it returns False.
+        """
+        raise NotImplementedError("This method is not implemented in the base class.")
+
 
 class ImageBuildEngine:
     """
@@ -247,20 +260,20 @@ def build(cls, image_spec: ImageSpec):
             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])
+            builder: ImageSpecBuilder = max(cls._REGISTRY, key=lambda name: cls._REGISTRY[name][1])
         else:
-            builder = image_spec.builder
+            builder: ImageSpecBuilder = image_spec.builder
 
         img_name = image_spec.image_name()
-        if image_spec.exist():
+        if builder.should_build(image_spec):
             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)
+                msg = f"Image {img_name} found. but overwriting existing image."
             else:
-                click.secho(f"Image {img_name} found. Skip building.", fg="blue")
-        else:
-            click.secho(f"Image {img_name} not found. building...", fg="blue")
+                msg = f"Image {img_name} not found. building..."
+            click.secho(msg, fg="blue")
             cls._build_image(builder, image_spec, img_name)
+        else:
+            click.secho(f"Image {img_name} found. Skip building.", fg="blue")
 
     @classmethod
     def _build_image(cls, builder, image_spec, img_name):

plugins/flytekit-envd/flytekitplugins/envd/image_builder.py

@@ -36,6 +36,9 @@ def build_image(self, image_spec: ImageSpec):
         envd_context_switch(image_spec.registry)
         execute_command(build_command)
 
+    def should_build(self, image_spec: ImageSpec) -> bool:
+        return image_spec.exist()
+
 
 def envd_context_switch(registry: str):
     if registry == FLYTE_LOCAL_REGISTRY: