Rename container_image to image for improved UX

flytekit/core/python_auto_container.py

@@ -2,6 +2,7 @@
 
 import importlib
 import re
+import warnings
 from abc import ABC
 from dataclasses import dataclass
 from typing import Callable, Dict, List, Optional, TypeVar, Union
@@ -42,6 +43,7 @@ def __init__(
         name: str,
         task_config: T,
         task_type="python-task",
+        image: Optional[Union[str, ImageSpec]] = None,
         container_image: Optional[Union[str, ImageSpec]] = None,
         requests: Optional[Resources] = None,
         limits: Optional[Resources] = None,
@@ -57,7 +59,8 @@ def __init__(
         :param name: unique name for the task, usually the function's module and name.
         :param task_config: Configuration object for Task. Should be a unique type for that specific Task.
         :param task_type: String task type to be associated with this Task
-        :param container_image: String FQN for the image.
+        :param image: String FQN for the image.
+        :param container_image: Deprecated, please use `image` instead.
         :param requests: custom resource request settings.
         :param limits: custom resource limit settings.
         :param environment: Environment variables you want the task to have when run.
@@ -90,7 +93,23 @@ def __init__(
         kwargs["metadata"] = kwargs["metadata"] if "metadata" in kwargs else TaskMetadata()
         kwargs["metadata"].pod_template_name = pod_template_name
 
-        self._container_image = container_image
+        # Rename the `container_image` parameter to `image` for improved user experience.
+        # Currently, both `image` and `container_image` are supported to maintain backward compatibility.
+        # For more details, please refer to https://github.com/flyteorg/flyte/issues/6140.
+        if image is not None and container_image is not None:
+            raise ValueError(
+                "Cannot specify both image and container_image. "
+                "Please use image because container_image is deprecated and will be removed in the future."
+            )
+        elif container_image is not None:
+            warnings.warn(
+                "container_image is deprecated and will be removed in the future. Please use image instead.",
+                DeprecationWarning,
+            )
+            self._image = container_image
+        else:
+            self._image = image
+
         # TODO(katrogan): Implement resource overrides
         self._resources = ResourceSpec(
             requests=requests if requests else Resources(), limits=limits if limits else Resources()
@@ -133,9 +152,23 @@ def __init__(
     def task_resolver(self) -> TaskResolverMixin:
         return self._task_resolver
 
+    @property
+    def image(self) -> Optional[Union[str, ImageSpec]]:
+        return self._image
+
     @property
     def container_image(self) -> Optional[Union[str, ImageSpec]]:
-        return self._container_image
+        """Deprecated, please use `image` instead."""
+        return self._image
+
+    @property
+    def _container_image(self) -> Optional[Union[str, ImageSpec]]:
+        """Deprecated, please use `image` instead."""
+        return self._image
+
+    @_container_image.setter
+    def _container_image(self, image: Optional[Union[str, ImageSpec]]):
+        self._image = image
 
     @property
     def resources(self) -> ResourceSpec:

flytekit/core/task.py

@@ -104,6 +104,7 @@ def task(
     interruptible: Optional[bool] = ...,
     deprecated: str = ...,
     timeout: Union[datetime.timedelta, int] = ...,
+    image: Optional[Union[str, ImageSpec]] = ...,
     container_image: Optional[Union[str, ImageSpec]] = ...,
     environment: Optional[Dict[str, str]] = ...,
     requests: Optional[Resources] = ...,
@@ -143,6 +144,7 @@ def task(
     interruptible: Optional[bool] = ...,
     deprecated: str = ...,
     timeout: Union[datetime.timedelta, int] = ...,
+    image: Optional[Union[str, ImageSpec]] = ...,
     container_image: Optional[Union[str, ImageSpec]] = ...,
     environment: Optional[Dict[str, str]] = ...,
     requests: Optional[Resources] = ...,
@@ -181,6 +183,7 @@ def task(
     interruptible: Optional[bool] = None,
     deprecated: str = "",
     timeout: Union[datetime.timedelta, int] = 0,
+    image: Optional[Union[str, ImageSpec]] = None,
     container_image: Optional[Union[str, ImageSpec]] = None,
     environment: Optional[Dict[str, str]] = None,
     requests: Optional[Resources] = None,
@@ -272,7 +275,7 @@ def my_task(x: int, y: typing.Dict[str, str]) -> str:
                        indicates that the task is active and not deprecated
     :param timeout: the max amount of time for which one execution of this task should be executed for. The execution
                     will be terminated if the runtime exceeds the given timeout (approximately).
-    :param container_image: By default the configured FLYTE_INTERNAL_IMAGE is used for every task. This directive can be
+    :param image: By default the configured FLYTE_INTERNAL_IMAGE is used for every task. This directive can be
                 used to provide an alternate image for a specific task. This is useful for the cases in which images
                 bloat because of various dependencies and a dependency is only required for this or a set of tasks,
                 and they vary from the default.
@@ -282,15 +285,16 @@ def my_task(x: int, y: typing.Dict[str, str]) -> str:
                     # Use default image name `fqn` and alter the tag to `tag-{{default.tag}}` tag of the default image
                     # with a prefix. In this case, it is assumed that the image like
                     # flytecookbook:tag-gitsha is published alongwith the default of flytecookbook:gitsha
-                    @task(container_image='{{.images.default.fqn}}:tag-{{images.default.tag}}')
+                    @task(image='{{.images.default.fqn}}:tag-{{images.default.tag}}')
                     def foo():
                         ...
 
                     # Refer to configurations to configure fqns for other images besides default. In this case it will
                     # lookup for an image named xyz
-                    @task(container_image='{{.images.xyz.fqn}}:{{images.default.tag}}')
+                    @task(image='{{.images.xyz.fqn}}:{{images.default.tag}}')
                     def foo2():
                         ...
+    :param container_image: Deprecated, please use `image` instead.
     :param environment: Environment variables that should be added for this tasks execution
     :param requests: Specify compute resource requests for your task. For Pod-plugin tasks, these values will apply only
       to the primary container.
@@ -374,6 +378,7 @@ def wrapper(fn: Callable[P, Any]) -> PythonFunctionTask[T]:
             task_config,
             decorated_fn,
             metadata=_metadata,
+            image=image,
             container_image=container_image,
             environment=environment,
             requests=requests,
@@ -584,7 +589,6 @@ async def eager_workflow(x: int) -> int:
             async def eager_workflow(x: int) -> int:
                 ...
     """
-
     if _fn is None:
         return partial(
             eager,

flytekit/core/type_engine.py

@@ -380,7 +380,7 @@ class DC:
                     a: Union[int, bool, str, float]
                     b: Union[int, bool, str, float]
 
-                @task(container_image=custom_image)
+                @task(image=custom_image)
                 def add(a: Union[int, bool, str, float], b: Union[int, bool, str, float]) -> Union[int, bool, str, float]:
                     return a + b
 

tests/flytekit/integration/jupyter/test_notebook_run.py

@@ -66,11 +66,11 @@ def execute_code_in_kernel(kc: BlockingKernelClient, code: str):
     interactive_mode_enabled=True,
 )
 
-@task(container_image="{IMAGE}")
+@task(image="{IMAGE}")
 def hello(name: str) -> str:
     return f"Hello {{name}}"
 
-@task(container_image="{IMAGE}")
+@task(image="{IMAGE}")
 def world(pre: str) -> str:
     return f"{{pre}}, Welcome to the world!"
 

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

@@ -4,12 +4,12 @@
 image_spec = ImageSpec(packages=["numpy", "pandas"], apt_packages=["git"], registry="", builder="test")
 
 
-@task(container_image=image_spec)
+@task(image=image_spec)
 def t2() -> str:
     return "flyte"
 
 
-@task(container_image=image_spec)
+@task(image=image_spec)
 def t1() -> str:
     return "flyte"
 

tests/flytekit/unit/core/test_array_node_map_task.py

@@ -356,7 +356,7 @@ def my_mappable_task(a: int) -> typing.Optional[str]:
 
 
 @task(
-    container_image="original-image",
+    image="original-image",
     timeout=timedelta(seconds=10),
     interruptible=False,
     retries=10,

tests/flytekit/unit/core/test_python_function_task.py

@@ -176,7 +176,7 @@ def foo_missing_cache(i: str):
 
 def test_pod_template():
     @task(
-        container_image="repo/image:0.0.0",
+        image="repo/image:0.0.0",
         pod_template=PodTemplate(
             primary_container_name="primary",
             labels={"lKeyA": "lValA"},

tests/flytekit/unit/core/test_serialization.py

@@ -267,23 +267,23 @@ def test_bad_configuration():
 def test_serialization_images(mock_image_spec_builder):
     ImageBuildEngine.register("test", mock_image_spec_builder)
 
-    @task(container_image="{{.image.xyz.fqn}}:{{.image.xyz.version}}")
+    @task(image="{{.image.xyz.fqn}}:{{.image.xyz.version}}")
     def t1(a: int) -> int:
         return a
 
-    @task(container_image="{{.image.abc.fqn}}:{{.image.xyz.version}}")
+    @task(image="{{.image.abc.fqn}}:{{.image.xyz.version}}")
     def t2():
         pass
 
-    @task(container_image="docker.io/org/myimage:latest")
+    @task(image="docker.io/org/myimage:latest")
     def t4():
         pass
 
-    @task(container_image="docker.io/org/myimage:{{.image.xyz.version}}")
+    @task(image="docker.io/org/myimage:{{.image.xyz.version}}")
     def t5(a: int) -> int:
         return a
 
-    @task(container_image="{{.image.xyz_123.fqn}}:{{.image.xyz_123.version}}")
+    @task(image="{{.image.xyz_123.fqn}}:{{.image.xyz_123.version}}")
     def t6(a: int) -> int:
         return a
 
@@ -294,7 +294,7 @@ def t6(a: int) -> int:
         builder="test",
     )
 
-    @task(container_image=image_spec)
+    @task(image=image_spec)
     def t7(a: int) -> int:
         return a
 

tests/flytekit/unit/core/test_task.py

@@ -1,15 +1,87 @@
+import os
 import pytest
 
+from flytekit import task, dynamic, eager
 from flytekit.core.task import decorate_function
 from flytekit.core.utils import str2bool
 from flytekit.interactive import vscode
 from flytekit.interactive.constants import FLYTE_ENABLE_VSCODE_KEY
 
 
+IMAGE = os.environ.get("FLYTEKIT_IMAGE", "localhost:30000/flytekit:dev")
+
+
 def test_decorate_python_task(monkeypatch: pytest.MonkeyPatch):
     def t1(a: int, b: int) -> int:
         return a + b
 
     assert decorate_function(t1) is t1
     monkeypatch.setenv(FLYTE_ENABLE_VSCODE_KEY, str2bool("True"))
     assert isinstance(decorate_function(t1), vscode)
+
+
+def test_image():
+    # Define expected warning and error messages
+    WARN_MSG = "container_image is deprecated and will be removed in the future. Please use image instead."
+    ERR_MSG = (
+        "Cannot specify both image and container_image. "
+        "Please use image because container_image is deprecated and will be removed in the future."
+    )
+
+    # Plain tasks
+    @task(image=IMAGE)
+    def t1() -> str:
+        return "Use image in @task."
+    assert t1._container_image == IMAGE, f"_container_image of t1 should match the user-specified {IMAGE}"
+
+    with pytest.warns(DeprecationWarning, match=WARN_MSG):
+        @task(container_image=IMAGE)
+        def t2() -> str:
+            return "Use container_image in @task."
+    assert t2._container_image == IMAGE, f"_container_image of t2 should match the user-specified {IMAGE}"
+
+    with pytest.raises(ValueError, match=ERR_MSG):
+        @task(image=IMAGE, container_image=IMAGE)
+        def t3() -> str:
+            return "Use both image and container_image in @task."
+
+    # Dynamic workflow tasks
+    @dynamic(image=IMAGE)
+    def dy_t1(i: int) -> str:
+        return "Use image in @dynamic."
+    assert dy_t1._container_image == IMAGE, f"_container_image of dy_t1 should match the user-specified {IMAGE}"
+
+    with pytest.warns(DeprecationWarning, match=WARN_MSG):
+        @dynamic(container_image=IMAGE)
+        def dy_t2(i: int) -> str:
+            return "Use container_image in @dynamic."
+    assert dy_t2._container_image == IMAGE, f"_container_image of dy_t2 should match the user-specified {IMAGE}"
+
+    with pytest.raises(ValueError, match=ERR_MSG):
+        @dynamic(image=IMAGE, container_image=IMAGE)
+        def dy_t3(i: int) -> str:
+            return "Use both image and container_image in @dynamic."
+
+
+    # Eager workflow task
+    @eager(image=IMAGE)
+    def eager_t1(dummy: bool) -> str:
+        if dummy:
+            print(f"Eager!")
+        return "Use image in @eager."
+    assert eager_t1._container_image == IMAGE, f"_container_image of eager_t1 should match the user-specified {IMAGE}"
+
+    with pytest.warns(DeprecationWarning, match=WARN_MSG):
+        @eager(container_image=IMAGE)
+        def eager_t2(dummy: bool) -> str:
+            if dummy:
+                print(f"Eager!")
+            return "Use container_image in @eager."
+    assert eager_t2._container_image == IMAGE, f"_container_image of eager_t2 should match the user-specified {IMAGE}"
+
+    with pytest.raises(ValueError, match=ERR_MSG):
+        @eager(image=IMAGE, container_image=IMAGE)
+        def eager_t3(dummy: bool) -> str:
+            if dummy:
+                print(f"Eager!")
+            return "Use both image and container_image in @eager."

tests/flytekit/unit/remote/test_remote.py

@@ -512,7 +512,7 @@ def test_get_image_names(
 
     image_spec = ImageSpec(requirements="requirements.txt", registry="flyteorg")
 
-    @task(container_image=image_spec)
+    @task(image=image_spec)
     def say_hello(name: str) -> str:
         return f"hello {name}!"