fix(ImageSpec): Ensure all APIErrors return a value (#2408)

Signed-off-by: Kevin Su <pingsutw@apache.org>
flyteorg · May 13, 2024 · cc3a7a9 · cc3a7a9
1 parent c36f322
commit cc3a7a9
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py
@@ -122,6 +122,7 @@ def exist(self) -> bool:
         except APIError as e:
             if e.response.status_code == 404:
                 return False
+            return True
         except ImageNotFound:
             return False
         except Exception as e:
@@ -136,7 +137,7 @@ def exist(self) -> bool:
                 if response.status_code == 200:
                     return True
 
-                if response.status_code == 404:
+                if response.status_code == 404 and "not found" in str(response.content):
                     return False
 
             click.secho(f"Failed to check if the image exists with error : {e}", fg="red")

diff --git a/tests/flytekit/unit/core/image_spec/test_image_spec.py b/tests/flytekit/unit/core/image_spec/test_image_spec.py
@@ -19,7 +19,7 @@ def test_image_spec(mock_image_spec_builder):
         packages=["pandas"],
         apt_packages=["git"],
         python_version="3.8",
-        registry="",
+        registry="localhost:30001",
         base_image="cr.flyte.org/flyteorg/flytekit:py3.8-latest",
         cuda="11.2.2",
         cudnn="8",
@@ -37,7 +37,7 @@ def test_image_spec(mock_image_spec_builder):
     assert image_spec.base_image == "cr.flyte.org/flyteorg/flytekit:py3.8-latest"
     assert image_spec.packages == ["pandas", "numpy"]
     assert image_spec.apt_packages == ["git", "wget"]
-    assert image_spec.registry == ""
+    assert image_spec.registry == "localhost:30001"
     assert image_spec.requirements == REQUIREMENT_FILE
     assert image_spec.registry_config == REGISTRY_CONFIG_FILE
     assert image_spec.cuda == "11.2.2"
@@ -53,20 +53,20 @@ def test_image_spec(mock_image_spec_builder):
 
     tag = calculate_hash_from_image_spec(image_spec)
     assert "=" != tag[-1]
-    assert image_spec.image_name() == f"flytekit:{tag}"
+    assert image_spec.image_name() == f"localhost:30001/flytekit:{tag}"
     ctx = context_manager.FlyteContext.current_context()
     with context_manager.FlyteContextManager.with_context(
         ctx.with_execution_state(ctx.execution_state.with_params(mode=ExecutionState.Mode.TASK_EXECUTION))
     ):
-        os.environ[_F_IMG_ID] = "flytekit:123"
+        os.environ[_F_IMG_ID] = "localhost:30001/flytekit:123"
         assert image_spec.is_container() is False
 
     ImageBuildEngine.register("dummy", mock_image_spec_builder)
     ImageBuildEngine.build(image_spec)
 
     assert "dummy" in ImageBuildEngine._REGISTRY
     assert calculate_hash_from_image_spec(image_spec) == tag
-    assert image_spec.exist() is False
+    assert image_spec.exist() is True
 
     # Remove the dummy builder, and build the image again
     # The image has already been built, so it shouldn't fail.