ImageConfig should contain the image from sandbox.config (#1106)

* ImageConfig should contain the image from sandbox.config Signed-off-by: Kevin Su <[email protected]> * Updated Signed-off-by: Kevin Su <[email protected]> * lint Signed-off-by: Kevin Su <[email protected]> * pr into #1106 (#1113) Signed-off-by: Yee Hing Tong <[email protected]> Co-authored-by: Yee Hing Tong <[email protected]>
flyteorg · Jul 29, 2022 · 9d31342 · 9d31342
1 parent 490597a
commit 9d31342
Show file tree

Hide file tree

Showing 8 changed files with 88 additions and 25 deletions.
diff --git a/flytekit/clis/sdk_in_container/run.py b/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
@@ -15,7 +15,7 @@
 from typing_extensions import get_args
 
 from flytekit import BlobType, Literal, Scalar
-from flytekit.clis.sdk_in_container.constants import CTX_DOMAIN, CTX_PROJECT
+from flytekit.clis.sdk_in_container.constants import CTX_CONFIG_FILE, CTX_DOMAIN, CTX_PROJECT
 from flytekit.clis.sdk_in_container.helpers import FLYTE_REMOTE_INSTANCE_KEY, get_and_save_remote_with_click_context
 from flytekit.configuration import ImageConfig
 from flytekit.configuration.default_images import DefaultImages
@@ -512,12 +512,35 @@ def _run(*args, **kwargs):
             return
 
         remote = ctx.obj[FLYTE_REMOTE_INSTANCE_KEY]
+        config_file = ctx.obj.get(CTX_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,
             project=project,
             domain=domain,
-            image_config=run_level_params.get("image_config", None),
+            image_config=image_config,
             destination_dir=run_level_params.get("destination_dir"),
         )
 

diff --git a/flytekit/configuration/__init__.py b/flytekit/configuration/__init__.py
@@ -230,25 +230,25 @@ 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:
         """
-        if config_file is None and img_name is None:
-            raise ValueError("Either an image or a config with a default image should be provided")
-
         default_img = Image.look_up_image_info("default", img_name) if img_name else None
 
-        other_images = []
-        if config_file:
-            config_file = get_config_file(config_file)
-            other_images = [
-                Image.look_up_image_info(k, tag=v, optional_tag=True)
-                for k, v in _internal.Images.get_specified_images(config_file).items()
-            ]
+        config_file = get_config_file(config_file)
+        other_images = [
+            Image.look_up_image_info(k, tag=v, optional_tag=True)
+            for k, v in _internal.Images.get_specified_images(config_file).items()
+        ]
         return cls.create_from(default_img, other_images)
 
     @classmethod

diff --git a/flytekit/configuration/internal.py b/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
@@ -19,16 +19,20 @@ def get_specified_images(cfg: ConfigFile) -> typing.Dict[str, str]:
         :returns a dictionary of name: image<fqn+version> Version is optional
         """
         images: typing.Dict[str, str] = {}
-        if cfg is None or not cfg.legacy_config:
+        if cfg is None:
             return images
-        try:
-            image_names = cfg.legacy_config.options("images")
-        except configparser.NoSectionError:
-            image_names = None
-        if image_names:
-            for i in image_names:
-                images[str(i)] = cfg.legacy_config.get("images", i)
-        return images
+
+        if cfg.legacy_config:
+            try:
+                image_names = cfg.legacy_config.options("images")
+            except configparser.NoSectionError:
+                return {}
+            if image_names:
+                for i in image_names:
+                    images[str(i)] = cfg.legacy_config.get("images", i)
+            return images
+        if cfg.yaml_config:
+            return cfg.yaml_config.get("images", images)
 
 
 class Deck(object):

diff --git a/tests/flytekit/unit/configuration/configs/no_images.yaml b/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
diff --git a/tests/flytekit/unit/configuration/configs/sample.yaml b/tests/flytekit/unit/configuration/configs/sample.yaml
@@ -11,3 +11,6 @@ storage:
     access-key: minio
     endpoint: http://localhost:30084
     secret-key: miniostorage
+images:
+  xyz: docker.io/xyz:latest
+  abc: docker.io/abc
diff --git a/tests/flytekit/unit/configuration/test_image_config.py b/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

diff --git a/tests/flytekit/unit/configuration/test_internal.py b/tests/flytekit/unit/configuration/test_internal.py
@@ -11,6 +11,10 @@ def test_load_images():
     imgs = Images.get_specified_images(cfg)
     assert imgs == {"abc": "docker.io/abc", "xyz": "docker.io/xyz:latest"}
 
+    cfg = get_config_file(os.path.join(os.path.dirname(os.path.realpath(__file__)), "configs/sample.yaml"))
+    imgs = Images.get_specified_images(cfg)
+    assert imgs == {"abc": "docker.io/abc", "xyz": "docker.io/xyz:latest"}
+
 
 def test_no_images():
     cfg = get_config_file(os.path.join(os.path.dirname(os.path.realpath(__file__)), "configs/good.config"))

diff --git a/tests/flytekit/unit/configuration/test_yaml_file.py b/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.