From 180ac1cf043a5558ea57b1ef751f2bc0e995d5e6 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 18 Jun 2024 14:07:24 -0700 Subject: [PATCH 1/5] Improve error message for ImageSpec Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 5 ++++- .../flytekitplugins/envd/image_builder.py | 13 +++++++++++-- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 05abacdbf9..76555b904d 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -121,6 +121,9 @@ def exist(self) -> bool: except APIError as e: if e.response.status_code == 404: return False + + click.secho(f"Failed to check if the image exists with error:\n {e}", fg="red") + click.secho("Flytekit assumes that the image already exists.", fg="blue") return True except ImageNotFound: return False @@ -153,7 +156,7 @@ 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(f"Failed to check if the image exists with error:\n {e}", fg="red") click.secho("Flytekit assumes that the image already exists.", fg="blue") return True diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py index 344003bafd..0b89754a6f 100644 --- a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py +++ b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py @@ -2,16 +2,20 @@ import pathlib import shutil import subprocess +from dataclasses import asdict from importlib import metadata import click from packaging.version import Version +from rich.console import Console from flytekit.configuration import DefaultImages from flytekit.core import context_manager from flytekit.core.constants import REQUIREMENTS_FILE_NAME from flytekit.image_spec.image_spec import _F_IMG_ID, ImageBuildEngine, ImageSpec, ImageSpecBuilder from flytekit.tools.ignore import DockerIgnore, GitIgnore, IgnoreGroup, StandardIgnore +from rich.pretty import Pretty +from rich import print FLYTE_LOCAL_REGISTRY = "localhost:30000" @@ -34,7 +38,12 @@ 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(f"❌ 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): @@ -68,7 +77,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 Exception(f"failed to run command {command} with error:\n {stderr.decode()}") return result From 25c19b48ccc99744325489935cf6b863d2c43177 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 18 Jun 2024 14:24:17 -0700 Subject: [PATCH 2/5] nit Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 17 ++++++++++------- .../flytekitplugins/envd/image_builder.py | 13 ++++++++----- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 76555b904d..0f446a02ac 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -123,8 +123,8 @@ def exist(self) -> bool: return False click.secho(f"Failed to check if the image exists with error:\n {e}", fg="red") - click.secho("Flytekit assumes that the image already exists.", fg="blue") - return True + click.secho(f"Flytekit assumes the image {self.image_name()} already exists.", fg="blue") + return None except ImageNotFound: return False except Exception as e: @@ -157,8 +157,8 @@ def exist(self) -> bool: ) click.secho(f"Failed to check if the image exists with error:\n {e}", fg="red") - click.secho("Flytekit assumes that the image already exists.", fg="blue") - return True + click.secho(f"Flytekit assumes the image {self.image_name()} already exists.", fg="blue") + return None def __hash__(self): return hash(asdict(self).__str__()) @@ -243,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") + click.secho(f"Overwriting existing image {img_name}.", fg="blue") return True - click.secho(f"Image {img_name} found. Skip building.", fg="blue") + if exist is True: + click.secho(f"Image {img_name} found. Skip building.", fg="blue") return False diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py index 0b89754a6f..82df8e8dee 100644 --- a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py +++ b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py @@ -7,15 +7,14 @@ import click from packaging.version import Version -from rich.console import Console +from rich import print +from rich.pretty import Pretty from flytekit.configuration import DefaultImages from flytekit.core import context_manager from flytekit.core.constants import REQUIREMENTS_FILE_NAME from flytekit.image_spec.image_spec import _F_IMG_ID, ImageBuildEngine, ImageSpec, ImageSpecBuilder from flytekit.tools.ignore import DockerIgnore, GitIgnore, IgnoreGroup, StandardIgnore -from rich.pretty import Pretty -from rich import print FLYTE_LOCAL_REGISTRY = "localhost:30000" @@ -41,8 +40,12 @@ def build_image(self, image_spec: ImageSpec): try: execute_command(build_command) except Exception as e: - click.secho(f"❌ 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)) + 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 From 672e72d65c041af30fc707fb33b172eca761d843 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 18 Jun 2024 14:31:27 -0700 Subject: [PATCH 3/5] update comment Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 0f446a02ac..589cd67807 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -107,6 +107,8 @@ def is_container(self) -> bool: def exist(self) -> 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 From 793199e4554335d222df379f591566273f03d994 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 18 Jun 2024 15:02:51 -0700 Subject: [PATCH 4/5] fix tests Signed-off-by: Kevin Su --- tests/flytekit/unit/core/image_spec/test_image_spec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 011828d4ce..02c734b119 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/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. From e146d88ceee00a8066ef654cceee6565e4c589b8 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 28 Jun 2024 13:46:32 -0700 Subject: [PATCH 5/5] Address comment Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 16 +++++++--------- .../flytekitplugins/envd/image_builder.py | 2 +- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 589cd67807..a17a2e272e 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -104,7 +104,7 @@ 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. @@ -125,7 +125,6 @@ def exist(self) -> bool: return False click.secho(f"Failed to check if the image exists with error:\n {e}", fg="red") - click.secho(f"Flytekit assumes the image {self.image_name()} already exists.", fg="blue") return None except ImageNotFound: return False @@ -159,7 +158,6 @@ def exist(self) -> bool: ) click.secho(f"Failed to check if the image exists with error:\n {e}", fg="red") - click.secho(f"Flytekit assumes the image {self.image_name()} already exists.", fg="blue") return None def __hash__(self): @@ -249,13 +247,13 @@ def should_build(self, image_spec: ImageSpec) -> bool: 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"Overwriting existing image {img_name}.", fg="blue") - return True - - if exist is True: + 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 diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py index 82df8e8dee..13121cce4a 100644 --- a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py +++ b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py @@ -80,7 +80,7 @@ def execute_command(command: str): if p.returncode != 0: _, stderr = p.communicate() - raise Exception(f"failed to run command {command} with error:\n {stderr.decode()}") + raise RuntimeError(f"failed to run command {command} with error:\n {stderr.decode()}") return result