Add support cache for dynamic spec

flytekit/core/base_task.py

@@ -140,6 +140,7 @@ class TaskMetadata(object):
     pod_template_name: Optional[str] = None
     generates_deck: bool = False
     is_eager: bool = False
+    execution_mode: "PythonFunctionTask.ExecutionBehavior" = None
 
     def __post_init__(self):
         if self.timeout:
@@ -155,6 +156,10 @@ def __post_init__(self):
             raise ValueError(
                 f"Cache ignore input vars are specified ``cache_ignore_input_vars={self.cache_ignore_input_vars}`` but ``cache`` is not enabled."
             )
+        if self.execution_mode is None:
+            from flytekit.core.python_function_task import PythonFunctionTask
+
+            self.execution_mode = PythonFunctionTask.ExecutionBehavior.DEFAULT
 
     @property
     def retry_strategy(self) -> _literal_models.RetryStrategy:
@@ -179,6 +184,7 @@ def to_taskmetadata_model(self) -> _task_model.TaskMetadata:
             cache_serializable=self.cache_serialize,
             generates_deck=self.generates_deck,
             pod_template_name=self.pod_template_name,
+            execution_mode=self.execution_mode,
             cache_ignore_input_vars=self.cache_ignore_input_vars,
             is_eager=self.is_eager,
         )

flytekit/core/python_function_task.py

@@ -28,6 +28,7 @@
 from enum import Enum
 from typing import Any, Callable, Iterable, List, Optional, Tuple, TypeVar, Union, cast
 
+from flyteidl.core import tasks_pb2
 from flytekit.configuration import ImageConfig, ImageSpec, SerializationSettings
 from flytekit.core import launch_plan as _annotated_launch_plan
 from flytekit.core.base_task import Task, TaskMetadata, TaskResolverMixin
@@ -127,6 +128,21 @@ class ExecutionBehavior(Enum):
         DYNAMIC = 2
         EAGER = 3
 
+        def to_flyte_idl(self) -> tasks_pb2.TaskMetadata.ExecutionMode:
+            """convert ExecutionBehavior into flyteidl type"""
+            return {
+                self.DEFAULT: tasks_pb2.TaskMetadata.ExecutionMode.DEFAULT,
+                self.DYNAMIC: tasks_pb2.TaskMetadata.ExecutionMode.DYNAMIC,
+            }[self]
+
+        @classmethod
+        def from_flyte_idl(cls, execution_mode: tasks_pb2.TaskMetadata.ExecutionMode):
+            """convert flyteidl type into ExecutionBehavior"""
+            return {
+                tasks_pb2.TaskMetadata.ExecutionMode.DEFAULT: cls.DEFAULT,
+                tasks_pb2.TaskMetadata.ExecutionMode.DYNAMIC: cls.DYNAMIC,
+            }[execution_mode]
+
     def __init__(
         self,
         task_config: T,

flytekit/core/task.py

@@ -404,6 +404,7 @@ def wrapper(fn: Callable[P, FuncOut]) -> PythonFunctionTask[T]:
             interruptible=interruptible,
             deprecated=deprecated,
             timeout=timeout,
+            execution_mode=execution_mode,
         )
 
         if inspect.iscoroutinefunction(fn):

flytekit/models/task.py

@@ -183,6 +183,7 @@ def __init__(
         cache_serializable,
         pod_template_name,
         cache_ignore_input_vars,
+        execution_mode,
         is_eager: bool = False,
         generates_deck: bool = False,
     ):
@@ -221,6 +222,13 @@ def __init__(
         self._cache_ignore_input_vars = cache_ignore_input_vars
         self._is_eager = is_eager
         self._generates_deck = generates_deck
+        self._execution_mode = execution_mode
+
+    def __post_init__(self):
+        if self.execution_mode is None:
+            from flytekit.core.python_function_task import PythonFunctionTask
+
+            self.execution_mode = PythonFunctionTask.ExecutionBehavior.DEFAULT
 
     @property
     def is_eager(self):
@@ -318,6 +326,13 @@ def cache_ignore_input_vars(self):
         """
         return self._cache_ignore_input_vars
 
+    @property
+    def execution_mode(self):
+        """
+        The execution mode is either default or dynamic
+        """
+        return self._execution_mode
+
     def to_flyte_idl(self):
         """
         :rtype: flyteidl.admin.task_pb2.TaskMetadata
@@ -334,6 +349,7 @@ def to_flyte_idl(self):
             cache_ignore_input_vars=self.cache_ignore_input_vars,
             is_eager=self.is_eager,
             generates_deck=BoolValue(value=self.generates_deck),
+            mode=self.execution_mode.to_flyte_idl(),
         )
         if self.timeout:
             tm.timeout.FromTimedelta(self.timeout)
@@ -345,6 +361,8 @@ def from_flyte_idl(cls, pb2_object: _core_task.TaskMetadata):
         :param flyteidl.core.task_pb2.TaskMetadata pb2_object:
         :rtype: TaskMetadata
         """
+        from flytekit.core.python_function_task import PythonFunctionTask
+
         return cls(
             discoverable=pb2_object.discoverable,
             runtime=RuntimeMetadata.from_flyte_idl(pb2_object.runtime),
@@ -358,6 +376,7 @@ def from_flyte_idl(cls, pb2_object: _core_task.TaskMetadata):
             cache_ignore_input_vars=pb2_object.cache_ignore_input_vars,
             is_eager=pb2_object.is_eager,
             generates_deck=pb2_object.generates_deck.value if pb2_object.HasField("generates_deck") else False,
+            execution_mode=PythonFunctionTask.ExecutionBehavior.from_flyte_idl(pb2_object.mode),
         )
 
 

tests/flytekit/common/parameterizers.py

@@ -7,6 +7,7 @@
 from flytekit.models import interface, literals, security, task, types
 from flytekit.models.core import identifier
 from flytekit.models.core import types as _core_types
+from flytekit.core.python_function_task import PythonFunctionTask
 
 LIST_OF_SCALAR_LITERAL_TYPES = [
     types.LiteralType(simple=types.SimpleType.BINARY),
@@ -125,8 +126,9 @@
         cache_serializable,
         pod_template_name,
         cache_ignore_input_vars,
+        execution_mode,
     )
-    for discoverable, runtime_metadata, timeout, retry_strategy, interruptible, discovery_version, deprecated, cache_serializable, pod_template_name, cache_ignore_input_vars in product(
+    for discoverable, runtime_metadata, timeout, retry_strategy, interruptible, discovery_version, deprecated, cache_serializable, pod_template_name, cache_ignore_input_vars, execution_mode in product(
         [True, False],
         LIST_OF_RUNTIME_METADATA,
         [timedelta(days=i) for i in range(3)],
@@ -137,6 +139,7 @@
         [True, False],
         ["A", "B"],
         [()],
+        [PythonFunctionTask.ExecutionBehavior.DEFAULT]
     )
 ]
 

tests/flytekit/unit/core/test_python_function_task.py

@@ -11,6 +11,7 @@
 from flytekit.core.tracker import isnested, istestfunction
 from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec
 from flytekit.tools.translator import get_serializable_task
+from flyteidl.core import tasks_pb2
 from tests.flytekit.unit.core import tasks
 
 
@@ -284,3 +285,18 @@ def foo3(x: int, y: str) -> int:
         return x
 
     assert foo3.python_interface.inputs_with_defaults == {"x": (int, None), "y": (str, None)}
+
+def test_execution_behavior_conversion():
+    """Test conversion between ExecutionBehavior and flyteidl ExecutionMode"""
+
+    # Test conversion for DEFAULT mode
+    assert (PythonFunctionTask.ExecutionBehavior.DEFAULT.to_flyte_idl() == 
+            tasks_pb2.TaskMetadata.ExecutionMode.DEFAULT)
+    assert (PythonFunctionTask.ExecutionBehavior.from_flyte_idl(tasks_pb2.TaskMetadata.ExecutionMode.DEFAULT) == 
+            PythonFunctionTask.ExecutionBehavior.DEFAULT)
+
+    # Test conversion for DYNAMIC mode
+    assert (PythonFunctionTask.ExecutionBehavior.DYNAMIC.to_flyte_idl() == 
+            tasks_pb2.TaskMetadata.ExecutionMode.DYNAMIC)
+    assert (PythonFunctionTask.ExecutionBehavior.from_flyte_idl(tasks_pb2.TaskMetadata.ExecutionMode.DYNAMIC) == 
+            PythonFunctionTask.ExecutionBehavior.DYNAMIC)

tests/flytekit/unit/models/test_tasks.py

@@ -12,6 +12,7 @@
 from flytekit.extras.accelerators import T4
 from flytekit.models import literals, task, types
 from flytekit.models.core import identifier
+from flytekit.core.python_function_task import PythonFunctionTask
 from tests.flytekit.common import parameterizers
 
 
@@ -74,7 +75,7 @@ def test_task_metadata():
         "This is deprecated!",
         True,
         "A",
-        (),
+        ()
     )
 
     assert obj.discoverable is True
@@ -88,6 +89,7 @@ def test_task_metadata():
     assert obj.discovery_version == "0.1.1b0"
     assert obj.pod_template_name == "A"
     assert obj == task.TaskMetadata.from_flyte_idl(obj.to_flyte_idl())
+    assert obj.mode == PythonFunctionTask.ExecutionBehavior.DEFAULT
 
 
 @pytest.mark.parametrize(