Rename container_image to image for improved UX

flytekit/core/node.py

@@ -2,6 +2,7 @@
 
 import datetime
 import typing
+import warnings
 from typing import Any, Dict, List, Optional, Union
 from typing import Literal as L
 
@@ -68,7 +69,7 @@ def __init__(
         self._outputs = None
         self._resources: typing.Optional[_resources_model] = None
         self._extended_resources: typing.Optional[tasks_pb2.ExtendedResources] = None
-        self._container_image: typing.Optional[str] = None
+        self._image: typing.Optional[str] = None
         self._pod_template: typing.Optional[PodTemplate] = None
 
     def runs_before(self, other: Node):
@@ -189,7 +190,7 @@ def with_overrides(
         interruptible: Optional[bool] = None,
         name: Optional[str] = None,
         task_config: Optional[Any] = None,
-        container_image: Optional[str] = None,
+        image: Optional[str] = None,
         accelerator: Optional[BaseAccelerator] = None,
         cache: Optional[bool] = None,
         cache_version: Optional[str] = None,
@@ -236,9 +237,24 @@ def with_overrides(
                 raise ValueError("can't change the type of the task config")
             self.run_entity._task_config = task_config
 
-        if container_image is not None:
-            assert_not_promise(container_image, "container_image")
-            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 kwargs.get("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 kwargs.get("container_image") is not None:
+            warnings.warn(
+                "container_image is deprecated and will be removed in the future. Please use image instead.",
+                DeprecationWarning,
+            )
+            assert_not_promise(kwargs["container_image"], "container_image")
+            self._image = kwargs["container_image"]
+        else:
+            assert_not_promise(image, "image")
+            self._image = image
 
         if accelerator is not None:
             assert_not_promise(accelerator, "accelerator")
@@ -256,6 +272,11 @@ def with_overrides(
 
         return self
 
+    @property
+    def _container_image(self) -> typing.Optional[str]:
+        """Deprecated, please use `_image` instead."""
+        return self._image
+
 
 def _convert_resource_overrides(
     resources: typing.Optional[Resources], resource_name: str

flytekit/core/promise.py

@@ -591,7 +591,7 @@ def with_overrides(
         interruptible: Optional[bool] = None,
         name: Optional[str] = None,
         task_config: Optional[Any] = None,
-        container_image: Optional[str] = None,
+        image: Optional[str] = None,
         accelerator: Optional[BaseAccelerator] = None,
         cache: Optional[bool] = None,
         cache_version: Optional[str] = None,
@@ -611,7 +611,7 @@ def with_overrides(
                 interruptible=interruptible,
                 name=name,
                 task_config=task_config,
-                container_image=container_image,
+                image=image,
                 accelerator=accelerator,
                 cache=cache,
                 cache_version=cache_version,

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
@@ -43,7 +44,7 @@ def __init__(
         name: str,
         task_config: T,
         task_type="python-task",
-        container_image: Optional[Union[str, ImageSpec]] = None,
+        image: Optional[Union[str, ImageSpec]] = None,
         requests: Optional[Resources] = None,
         limits: Optional[Resources] = None,
         environment: Optional[Dict[str, str]] = None,
@@ -59,7 +60,7 @@ 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 or ImageSpec for the image.
         :param requests: custom resource request settings.
         :param limits: custom resource limit settings.
         :param environment: Environment variables you want the task to have when run.
@@ -82,6 +83,7 @@ def __init__(
         :param accelerator: The accelerator to use for this task.
         :param shared_memory: If True, then shared memory will be attached to the container where the size is equal
             to the allocated memory. If str, then the shared memory is set to that size.
+        :param container_image: Deprecated, please use `image` instead.
         """
         sec_ctx = None
         if secret_requests:
@@ -94,7 +96,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 kwargs.get("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 kwargs.get("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 = kwargs["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()
@@ -138,9 +156,28 @@ 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]]):
+        """Deprecated, please use `image` instead.
+
+        This setter is for backward compatibility, so that setting `_container_image`
+        will adjust the new `_image` parameter directly.
+        """
+        self._image = image
 
     @property
     def resources(self) -> ResourceSpec:

flytekit/core/task.py

@@ -103,7 +103,7 @@ def task(
     interruptible: Optional[bool] = ...,
     deprecated: str = ...,
     timeout: Union[datetime.timedelta, int] = ...,
-    container_image: Optional[Union[str, ImageSpec]] = ...,
+    image: Optional[Union[str, ImageSpec]] = ...,
     environment: Optional[Dict[str, str]] = ...,
     requests: Optional[Resources] = ...,
     limits: Optional[Resources] = ...,
@@ -141,7 +141,7 @@ def task(
     interruptible: Optional[bool] = ...,
     deprecated: str = ...,
     timeout: Union[datetime.timedelta, int] = ...,
-    container_image: Optional[Union[str, ImageSpec]] = ...,
+    image: Optional[Union[str, ImageSpec]] = ...,
     environment: Optional[Dict[str, str]] = ...,
     requests: Optional[Resources] = ...,
     limits: Optional[Resources] = ...,
@@ -178,7 +178,7 @@ def task(
     interruptible: Optional[bool] = None,
     deprecated: str = "",
     timeout: Union[datetime.timedelta, int] = 0,
-    container_image: Optional[Union[str, ImageSpec]] = None,
+    image: Optional[Union[str, ImageSpec]] = None,
     environment: Optional[Dict[str, str]] = None,
     requests: Optional[Resources] = None,
     limits: Optional[Resources] = None,
@@ -273,7 +273,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.
@@ -283,13 +283,13 @@ 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 environment: Environment variables that should be added for this tasks execution
@@ -344,6 +344,7 @@ def launch_dynamically():
     :param pickle_untyped: Boolean that indicates if the task allows unspecified data types.
     :param shared_memory: If True, then shared memory will be attached to the container where the size is equal
         to the allocated memory. If int, then the shared memory is set to that size.
+    :param container_image: Deprecated, please use `image` instead.
     """
     # Maintain backwards compatibility with the old cache parameters, while cleaning up the task function definition.
     cache_serialize = kwargs.get("cache_serialize")
@@ -370,7 +371,7 @@ def wrapper(fn: Callable[P, FuncOut]) -> PythonFunctionTask[T]:
             cache_version = cache.get_version(
                 VersionParameters(
                     func=fn,
-                    container_image=container_image,
+                    container_image=image or kwargs.get("container_image"),
                     pod_template=pod_template,
                     pod_template_name=pod_template_name,
                 )
@@ -418,7 +419,7 @@ def wrapper(fn: Callable[P, FuncOut]) -> PythonFunctionTask[T]:
             task_config,
             decorated_fn,
             metadata=_metadata,
-            container_image=container_image,
+            image=image,
             environment=environment,
             requests=requests,
             limits=limits,
@@ -435,6 +436,7 @@ def wrapper(fn: Callable[P, FuncOut]) -> PythonFunctionTask[T]:
             accelerator=accelerator,
             pickle_untyped=pickle_untyped,
             shared_memory=shared_memory,
+            **kwargs,
         )
         update_wrapper(task_instance, decorated_fn)
         return task_instance
@@ -629,7 +631,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

@@ -383,7 +383,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

@@ -349,18 +349,6 @@ def my_wf1() -> typing.List[typing.Optional[int]]:
         assert my_wf1() == [1, None, 3, 4]
 
 
-def test_map_task_override(serialization_settings):
-    @task
-    def my_mappable_task(a: int) -> typing.Optional[str]:
-        return str(a)
-
-    @workflow
-    def wf(x: typing.List[int]):
-        map_task(my_mappable_task)(a=x).with_overrides(container_image="random:image")
-
-    assert wf.nodes[0]._container_image == "random:image"
-
-
 def test_serialization_metadata(serialization_settings):
     @task(interruptible=True)
     def t1(a: int) -> int:

tests/flytekit/unit/core/test_map_task.py

@@ -351,9 +351,14 @@ def my_mappable_task(a: int) -> typing.Optional[str]:
 
     @workflow
     def wf(x: typing.List[int]):
+        map_task(my_mappable_task)(a=x).with_overrides(image="random:image")
         map_task(my_mappable_task)(a=x).with_overrides(container_image="random:image")
 
+    assert wf.nodes[0]._image == "random:image"
     assert wf.nodes[0]._container_image == "random:image"
+    assert wf.nodes[1]._image == "random:image"
+    assert wf.nodes[1]._container_image == "random:image"
+
 
 def test_map_task_pod_template_override(serialization_settings):
     @task

tests/flytekit/unit/core/test_node_creation.py

@@ -466,10 +466,15 @@ def bar():
 
     @workflow
     def wf() -> str:
+        bar().with_overrides(image="hello/world")
         bar().with_overrides(container_image="hello/world")
         return "hi"
 
+    assert wf.nodes[0]._image == "hello/world"
     assert wf.nodes[0]._container_image == "hello/world"
+    assert wf.nodes[1]._image == "hello/world"
+    assert wf.nodes[1]._container_image == "hello/world"
+
 
 def test_pod_template_override():
     @task
@@ -575,3 +580,38 @@ def my_wf(a: str) -> str:
     assert wf_spec.template.nodes[0].metadata.cache_serializable
     assert wf_spec.template.nodes[0].metadata.cacheable
     assert wf_spec.template.nodes[0].metadata.cache_version == "foo"
+
+
+def test_override_image_behavior():
+    # 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."
+    )
+
+    @task
+    def dummy_task():
+        print("hello")
+
+
+    # Specify both image and container_image
+    @workflow
+    def wf1():
+        create_node(dummy_task).with_overrides(image="hello/world", container_image="hello/world:v2")
+    with pytest.raises(ValueError, match=ERR_MSG):
+        wf1()
+
+    # Specify only container_image
+    @workflow
+    def wf2():
+        create_node(dummy_task).with_overrides(container_image="hello/world")
+    with pytest.warns(DeprecationWarning, match=WARN_MSG):
+        wf2()
+    assert wf2.nodes[0]._image == wf2.nodes[0]._container_image == "hello/world"
+
+    # Specify only image
+    @workflow
+    def wf3():
+        create_node(dummy_task).with_overrides(image="hello/world")
+    assert wf3.nodes[0]._image == wf3.nodes[0]._container_image == "hello/world"