pr into #1106 (#1113)

Signed-off-by: Yee Hing Tong <[email protected]>

flytekit/clis/sdk_in_container/run.py

@@ -6,7 +6,7 @@
 import os
 import pathlib
 import typing
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import cast
 
 import click
@@ -514,11 +514,27 @@ def _run(*args, **kwargs):
         remote = ctx.obj[FLYTE_REMOTE_INSTANCE_KEY]
         config_file = ctx.obj.get(CTX_CONFIG_FILE)
 
-        image_config = run_level_params.get("image_config", None)
-        if image_config:
-            image_config.images.extend(ImageConfig.auto(config_file).images)
-        else:
-            image_config = ImageConfig.auto(config_file)
+        # Images come from three places:
+        # * The default flytekit images, which are already supplied by the base run_level_params.
+        # * The images provided by the user on the command line.
+        # * The images provided by the user via the config file, if there is one. (Images on the command line should
+        #   override all).
+        #
+        # However, the run_level_params already contains both the default flytekit images (lowest priority), as well
+        # as the images from the command line (highest priority). So when we read from the config file, we only
+        # want to add in the images that are missing, including the default, if that's also missing.
+        image_config_from_parent_cmd = run_level_params.get("image_config", None)
+        additional_image_names = set([v.name for v in image_config_from_parent_cmd.images])
+        new_additional_images = [v for v in image_config_from_parent_cmd.images]
+        new_default = image_config_from_parent_cmd.default_image
+        if config_file:
+            cfg_ic = ImageConfig.auto(config_file=config_file)
+            new_default = new_default or cfg_ic.default_image
+            for addl in cfg_ic.images:
+                if addl.name not in additional_image_names:
+                    new_additional_images.append(addl)
+
+        image_config = replace(image_config_from_parent_cmd, default_image=new_default, images=new_additional_images)
 
         remote_entity = remote.register_script(
             entity,

flytekit/configuration/__init__.py

@@ -230,9 +230,14 @@ def create_from(
         return ImageConfig(default_image=default_image, images=all_images)
 
     @classmethod
-    def auto(cls, config_file: typing.Union[str, ConfigFile] = None, img_name: Optional[str] = None) -> ImageConfig:
+    def auto(
+        cls, config_file: typing.Union[str, ConfigFile, None] = None, img_name: Optional[str] = None
+    ) -> ImageConfig:
         """
         Reads from config file or from img_name
+        Note that this function does not take into account the flytekit default images (see the Dockerfiles at the
+        base of this repo). To pick those up, see the auto_default_image function..
+
         :param config_file:
         :param img_name:
         :return:

flytekit/configuration/internal.py

@@ -7,7 +7,7 @@
 
 class Images(object):
     @staticmethod
-    def get_specified_images(cfg: ConfigFile) -> typing.Dict[str, str]:
+    def get_specified_images(cfg: typing.Optional[ConfigFile]) -> typing.Dict[str, str]:
         """
         This section should contain options, where the option name is the friendly name of the image and the corresponding
         value is actual FQN of the image. Example of how the section is structured
@@ -32,7 +32,7 @@ def get_specified_images(cfg: ConfigFile) -> typing.Dict[str, str]:
                     images[str(i)] = cfg.legacy_config.get("images", i)
             return images
         if cfg.yaml_config:
-            return cfg.yaml_config.get("images")
+            return cfg.yaml_config.get("images", images)
 
 
 class Deck(object):

tests/flytekit/unit/configuration/configs/no_images.yaml

@@ -0,0 +1,13 @@
+admin:
+  # For GRPC endpoints you might want to use dns:///flyte.myexample.com
+  endpoint: dns:///flyte.mycorp.io
+  authType: Pkce
+  insecure: true
+  clientId: propeller
+  scopes:
+    - all
+storage:
+  connection:
+    access-key: minio
+    endpoint: http://localhost:30084
+    secret-key: miniostorage

tests/flytekit/unit/configuration/test_image_config.py

@@ -1,3 +1,4 @@
+import os
 import sys
 
 import mock
@@ -38,6 +39,13 @@ def test_image_config_auto():
     assert x.images[0].full == f"ghcr.io/flyteorg/flytekit:py{version_str}-latest"
 
 
+def test_image_from_flytectl_config():
+    os.environ["FLYTECTL_CONFIG"] = os.path.join(os.path.dirname(os.path.realpath(__file__)), "configs/sample.yaml")
+    image_config = ImageConfig.auto(config_file=None)
+    assert image_config.images[0].full == "docker.io/xyz:latest"
+    assert image_config.images[1].full == "docker.io/abc:None"
+
+
 @mock.patch("flytekit.configuration.default_images.sys")
 def test_not_version(mock_sys):
     mock_sys.version_info.major.return_value = 2

tests/flytekit/unit/configuration/test_yaml_file.py

@@ -4,7 +4,7 @@
 
 from flytekit.configuration import ConfigEntry, get_config_file
 from flytekit.configuration.file import LegacyConfigEntry, YamlConfigEntry
-from flytekit.configuration.internal import AWS, Credentials, Platform
+from flytekit.configuration.internal import AWS, Credentials, Images, Platform
 
 
 def test_config_entry_file():
@@ -20,6 +20,14 @@ def test_config_entry_file():
     assert c.read(cfg) is None
 
 
+def test_config_entry_file_normal():
+    # Most yaml config files will not have images, make sure that a normal one without an image section doesn't
+    # return None
+    cfg = get_config_file(os.path.join(os.path.dirname(os.path.realpath(__file__)), "configs/no_images.yaml"))
+    images_dict = Images.get_specified_images(cfg)
+    assert images_dict == {}
+
+
 @mock.patch("flytekit.configuration.file.getenv")
 def test_config_entry_file_2(mock_get):
     # Test reading of the environment variable that flytectl asks users to set.