Improve error message for ImageSpec (flyteorg#2498)

Signed-off-by: Kevin Su <[email protected]>
Signed-off-by: bugra.gedik <[email protected]>

flytekit/image_spec/image_spec.py

@@ -104,9 +104,11 @@ def is_container(self) -> bool:
             return os.environ.get(_F_IMG_ID) == self.image_name()
         return True
 
-    def exist(self) -> bool:
+    def exist(self) -> Optional[bool]:
         """
         Check if the image exists in the registry.
+        Return True if the image exists in the registry, False otherwise.
+        Return None if failed to check if the image exists due to the permission issue or other reasons.
         """
         import docker
         from docker.errors import APIError, ImageNotFound
@@ -121,7 +123,9 @@ def exist(self) -> bool:
         except APIError as e:
             if e.response.status_code == 404:
                 return False
-            return True
+
+            click.secho(f"Failed to check if the image exists with error:\n {e}", fg="red")
+            return None
         except ImageNotFound:
             return False
         except Exception as e:
@@ -153,9 +157,8 @@ def exist(self) -> bool:
                     f"    pip install --upgrade docker"
                 )
 
-            click.secho(f"Failed to check if the image exists with error : {e}", fg="red")
-            click.secho("Flytekit assumes that the image already exists.", fg="blue")
-            return True
+            click.secho(f"Failed to check if the image exists with error:\n {e}", fg="red")
+            return None
 
     def __hash__(self):
         return hash(asdict(self).__str__())
@@ -240,14 +243,17 @@ def should_build(self, image_spec: ImageSpec) -> bool:
             True if the image should be built, otherwise it returns False.
         """
         img_name = image_spec.image_name()
-        if not image_spec.exist():
+        exist = image_spec.exist()
+        if exist is False:
             click.secho(f"Image {img_name} not found. building...", fg="blue")
             return True
-        if image_spec._is_force_push:
-            click.secho(f"Image {img_name} found but overwriting existing image.", fg="blue")
-            return True
-
-        click.secho(f"Image {img_name} found. Skip building.", fg="blue")
+        elif exist is True:
+            if image_spec._is_force_push:
+                click.secho(f"Overwriting existing image {img_name}.", fg="blue")
+                return True
+            click.secho(f"Image {img_name} found. Skip building.", fg="blue")
+        else:
+            click.secho(f"Flytekit assumes the image {img_name} already exists.", fg="blue")
         return False
 
 

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

@@ -2,10 +2,13 @@
 import pathlib
 import shutil
 import subprocess
+from dataclasses import asdict
 from importlib import metadata
 
 import click
 from packaging.version import Version
+from rich import print
+from rich.pretty import Pretty
 
 from flytekit.configuration import DefaultImages
 from flytekit.core import context_manager
@@ -35,7 +38,16 @@ def build_image(self, image_spec: ImageSpec):
         else:
             build_command += f" --tag {image_spec.image_name()}"
         envd_context_switch(image_spec.registry)
-        execute_command(build_command)
+        try:
+            execute_command(build_command)
+        except Exception as e:
+            click.secho("❌ Failed to build image spec:", fg="red")
+            print(
+                Pretty(
+                    asdict(image_spec, dict_factory=lambda x: {k: v for (k, v) in x if v is not None}), indent_size=2
+                )
+            )
+            raise e from None
 
 
 def envd_context_switch(registry: str):
@@ -72,7 +84,7 @@ def execute_command(command: str):
 
     if p.returncode != 0:
         _, stderr = p.communicate()
-        raise Exception(f"failed to run command {command} with error {stderr}")
+        raise RuntimeError(f"failed to run command {command} with error:\n {stderr.decode()}")
 
     return result
 

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

@@ -66,7 +66,7 @@ def test_image_spec(mock_image_spec_builder):
 
     assert "dummy" in ImageBuildEngine._REGISTRY
     assert calculate_hash_from_image_spec(image_spec) == tag
-    assert image_spec.exist() is True
+    assert image_spec.exist() is None
 
     # Remove the dummy builder, and build the image again
     # The image has already been built, so it shouldn't fail.