Always a default image and streamline image handling serialize vs pac…

…kage (#1610)

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

flytekit/clis/sdk_in_container/serialize.py

@@ -29,7 +29,7 @@ def serialize_all(
     local_source_root: typing.Optional[str] = None,
     folder: typing.Optional[str] = None,
     mode: typing.Optional[SerializationMode] = None,
-    image: typing.Optional[str] = None,
+    image_config: typing.Optional[ImageConfig] = None,
     flytekit_virtualenv_root: typing.Optional[str] = None,
     python_interpreter: typing.Optional[str] = None,
     config_file: typing.Optional[str] = None,
@@ -49,15 +49,15 @@ def serialize_all(
     :param local_source_root: Where to start looking for the code.
     :param folder: Where to write the output protobuf files
     :param mode: Regular vs fast
-    :param image: The fully qualified and versioned default image to use
+    :param image_config: ImageConfig object to use
     :param flytekit_virtualenv_root: The full path of the virtual env in the container.
     """
 
     if not (mode == SerializationMode.DEFAULT or mode == SerializationMode.FAST):
         raise AssertionError(f"Unrecognized serialization mode: {mode}")
 
     serialization_settings = SerializationSettings(
-        image_config=ImageConfig.auto(config_file, img_name=image),
+        image_config=image_config or ImageConfig.auto(config_file),
         fast_serialization_settings=FastSerializationSettings(
             enabled=mode == SerializationMode.FAST,
             # TODO: if we want to move the destination dir as a serialization argument, we should initialize it here
@@ -71,10 +71,18 @@ def serialize_all(
 
 @click.group("serialize", cls=click.RichGroup)
 @click.option(
+    "-i",
     "--image",
+    "image_config",
     required=False,
-    default=lambda: os.environ.get("FLYTE_INTERNAL_IMAGE", ""),
-    help="Text tag, for example ``somedocker.com/myimage:someversion123``",
+    multiple=True,
+    type=click.UNPROCESSED,
+    callback=ImageConfig.validate_image,
+    help="A fully qualified tag for an docker image, for example ``somedocker.com/myimage:someversion123``. This is a "
+    "multi-option and can be of the form ``--image xyz.io/docker:latest"
+    " --image my_image=xyz.io/docker2:latest``. Note, the ``name=image_uri``. The name is optional, if not "
+    "provided the image will be used as the default image. All the names have to be unique, and thus "
+    "there can only be one ``--image`` option with no name.",
 )
 @click.option(
     "--local-source-root",
@@ -99,17 +107,19 @@ def serialize_all(
     "your container installs the flytekit virtualenv outside of the default `/opt/venv`",
 )
 @click.pass_context
-def serialize(ctx, image, local_source_root, in_container_config_path, in_container_virtualenv_root):
+def serialize(
+    ctx, image_config: ImageConfig, local_source_root, in_container_config_path, in_container_virtualenv_root
+):
     """
     This command produces protobufs for tasks and templates.
     For tasks, one pb file is produced for each task, representing one TaskTemplate object.
     For workflows, one pb file is produced for each workflow, representing a WorkflowClosure object. The closure
     object contains the WorkflowTemplate, along with the relevant tasks for that workflow. In lieu of Admin,
     this serialization step will set the URN of the tasks to the fully qualified name of the task function.
     """
-    ctx.obj[CTX_IMAGE] = image
+    ctx.obj[CTX_IMAGE] = image_config
     ctx.obj[CTX_LOCAL_SRC_ROOT] = local_source_root
-    click.echo("Serializing Flyte elements with image {}".format(image))
+    click.echo(f"Serializing Flyte elements with image {image_config}")
 
     if in_container_virtualenv_root:
         ctx.obj[CTX_FLYTEKIT_VIRTUALENV_ROOT] = in_container_virtualenv_root
@@ -141,7 +151,7 @@ def workflows(ctx, folder=None):
         dir,
         folder,
         SerializationMode.DEFAULT,
-        image=ctx.obj[CTX_IMAGE],
+        image_config=ctx.obj[CTX_IMAGE],
         flytekit_virtualenv_root=ctx.obj[CTX_FLYTEKIT_VIRTUALENV_ROOT],
         python_interpreter=ctx.obj[CTX_PYTHON_INTERPRETER],
         config_file=ctx.obj.get(constants.CTX_CONFIG_FILE, None),
@@ -180,7 +190,7 @@ def fast_workflows(ctx, folder=None, deref_symlinks=False):
         dir,
         folder,
         SerializationMode.FAST,
-        image=ctx.obj[CTX_IMAGE],
+        image_config=ctx.obj[CTX_IMAGE],
         flytekit_virtualenv_root=ctx.obj[CTX_FLYTEKIT_VIRTUALENV_ROOT],
         python_interpreter=ctx.obj[CTX_PYTHON_INTERPRETER],
         config_file=ctx.obj.get(constants.CTX_CONFIG_FILE, None),

flytekit/configuration/__init__.py

@@ -282,6 +282,9 @@ def validate_image(_: typing.Any, param: str, values: tuple) -> ImageConfig:
             else:
                 images.append(img)
 
+        if default_image is None:
+            default_image_str = os.environ.get("FLYTE_INTERNAL_IMAGE", 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)
 
     @classmethod

tests/flytekit/unit/cli/pyflyte/test_run.py

@@ -255,17 +255,22 @@ def test_list_default_arguments(wf_path):
 )
 # test that command line args are merged with the file
 ic_result_2 = ImageConfig(
-    default_image=None,
+    default_image=Image(name="default", fqn="cr.flyte.org/flyteorg/flytekit", tag="py3.9-latest"),
     images=[
+        Image(name="default", fqn="cr.flyte.org/flyteorg/flytekit", tag="py3.9-latest"),
         Image(name="asdf", fqn="ghcr.io/asdf/asdf", tag="latest"),
         Image(name="xyz", fqn="docker.io/xyz", tag="latest"),
         Image(name="abc", fqn="docker.io/abc", tag=None),
     ],
 )
 # test that command line args override the file
 ic_result_3 = ImageConfig(
-    default_image=None,
-    images=[Image(name="xyz", fqn="ghcr.io/asdf/asdf", tag="latest"), Image(name="abc", fqn="docker.io/abc", tag=None)],
+    default_image=Image(name="default", fqn="cr.flyte.org/flyteorg/flytekit", tag="py3.9-latest"),
+    images=[
+        Image(name="default", fqn="cr.flyte.org/flyteorg/flytekit", tag="py3.9-latest"),
+        Image(name="xyz", fqn="ghcr.io/asdf/asdf", tag="latest"),
+        Image(name="abc", fqn="docker.io/abc", tag=None),
+    ],
 )
 
 ic_result_4 = ImageConfig(
@@ -280,6 +285,7 @@ def test_list_default_arguments(wf_path):
 IMAGE_SPEC = os.path.join(os.path.dirname(os.path.realpath(__file__)), "imageSpec.yaml")
 
 
+@mock.patch("flytekit.configuration.default_images.DefaultImages.default_image")
 @pytest.mark.parametrize(
     "image_string, leaf_configuration_file_name, final_image_config",
     [
@@ -289,7 +295,9 @@ def test_list_default_arguments(wf_path):
         (IMAGE_SPEC, "sample.yaml", ic_result_4),
     ],
 )
-def test_pyflyte_run_run(image_string, leaf_configuration_file_name, final_image_config):
+def test_pyflyte_run_run(mock_image, image_string, leaf_configuration_file_name, final_image_config):
+    mock_image.return_value = "cr.flyte.org/flyteorg/flytekit:py3.9-latest"
+
     class TestImageSpecBuilder(ImageSpecBuilder):
         def build_image(self, img):
             ...

tests/flytekit/unit/core/test_context_manager.py

@@ -1,6 +1,7 @@
 import os
 from datetime import datetime
 
+import mock
 import py
 import pytest
 
@@ -64,10 +65,12 @@ def test_look_up_image_info():
     assert img.fqn == "localhost:5000/xyz"
 
 
-def test_validate_image():
+@mock.patch("flytekit.configuration.default_images.DefaultImages.default_image")
+def test_validate_image(mock_image):
+    mock_image.return_value = "cr.flyte.org/flyteorg/flytekit:py3.9-latest"
     ic = ImageConfig.validate_image(None, "image", ())
     assert ic
-    assert ic.default_image is None
+    assert ic.default_image == Image(name="default", fqn="cr.flyte.org/flyteorg/flytekit", tag="py3.9-latest")
 
     img1 = "xyz:latest"
     img2 = "docker.io/xyz:latest"