Use default image (flyteorg#2223)

Signed-off-by: ddl-rliu <[email protected]>
jasonlai1218 · Mar 4, 2024 · 1bb5067 · 1bb5067
1 parent 4a58a67
commit 1bb5067
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/flytekit/configuration/__init__.py b/flytekit/configuration/__init__.py
@@ -285,7 +285,7 @@ def validate_image(_: typing.Any, param: str, values: tuple) -> ImageConfig:
                 images.append(img)
 
         if default_image is None:
-            default_image_str = os.environ.get("FLYTE_INTERNAL_IMAGE", DefaultImages.default_image())
+            default_image_str = DefaultImages.default_image()
             default_image = Image.look_up_image_info(DEFAULT_IMAGE_NAME, default_image_str, False)
         return ImageConfig.create_from(default_image=default_image, other_images=images)
 

diff --git a/flytekit/configuration/default_images.py b/flytekit/configuration/default_images.py
@@ -1,4 +1,5 @@
 import enum
+import os
 import sys
 import typing
 from contextlib import suppress
@@ -34,7 +35,8 @@ def default_image(cls) -> str:
             if default_image is not None:
                 return default_image
 
-        return cls.find_image_for()
+        default_image_str = os.environ.get("FLYTE_INTERNAL_IMAGE", cls.find_image_for())
+        return default_image_str
 
     @classmethod
     def find_image_for(