diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index dfbd678fb6..572df5f881 100644 --- a/flytekit/core/python_auto_container.py +++ b/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,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: diff --git a/flytekit/core/task.py b/flytekit/core/task.py index 6451e742c5..db58baafb7 100644 --- a/flytekit/core/task.py +++ b/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, diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index fda7758010..32081cf98b 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -382,7 +382,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 diff --git a/tests/flytekit/integration/jupyter/test_notebook_run.py b/tests/flytekit/integration/jupyter/test_notebook_run.py index e72cd44328..df1077841d 100644 --- a/tests/flytekit/integration/jupyter/test_notebook_run.py +++ 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!" diff --git a/tests/flytekit/unit/cli/pyflyte/image_spec_wf.py b/tests/flytekit/unit/cli/pyflyte/image_spec_wf.py index 9d5d74ff1b..29fe53d81c 100644 --- a/tests/flytekit/unit/cli/pyflyte/image_spec_wf.py +++ b/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" diff --git a/tests/flytekit/unit/core/test_python_function_task.py b/tests/flytekit/unit/core/test_python_function_task.py index e775cb922a..82b2b35133 100644 --- a/tests/flytekit/unit/core/test_python_function_task.py +++ b/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"}, diff --git a/tests/flytekit/unit/core/test_serialization.py b/tests/flytekit/unit/core/test_serialization.py index 82b69859bd..8a50ab563b 100644 --- a/tests/flytekit/unit/core/test_serialization.py +++ b/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 diff --git a/tests/flytekit/unit/core/test_task.py b/tests/flytekit/unit/core/test_task.py index b02f7cb08b..fbae686a3f 100644 --- a/tests/flytekit/unit/core/test_task.py +++ b/tests/flytekit/unit/core/test_task.py @@ -1,11 +1,16 @@ +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 @@ -13,3 +18,70 @@ def t1(a: int, b: int) -> int: 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." diff --git a/tests/flytekit/unit/remote/test_remote.py b/tests/flytekit/unit/remote/test_remote.py index 9911cad02f..c8b86834b7 100644 --- a/tests/flytekit/unit/remote/test_remote.py +++ b/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}!"