Add Echo task

flytekit/core/task.py

@@ -12,7 +12,7 @@
 from flytekit.core import launch_plan as _annotated_launchplan
 from flytekit.core import workflow as _annotated_workflow
 from flytekit.core.base_task import PythonTask, TaskMetadata, TaskResolverMixin
-from flytekit.core.interface import transform_function_to_interface
+from flytekit.core.interface import Interface, output_name_generator, transform_function_to_interface
 from flytekit.core.pod_template import PodTemplate
 from flytekit.core.python_function_task import PythonFunctionTask
 from flytekit.core.reference_entity import ReferenceEntity, TaskReference
@@ -416,3 +416,30 @@ def wrapper(fn) -> ReferenceTask:
         return ReferenceTask(project, domain, name, version, interface.inputs, interface.outputs)
 
     return wrapper
+
+
+class Echo(PythonTask):
+    """
+    A task that simply echoes the inputs back to the user.
+    The task's inputs and outputs interface are the same.
+    FlytePropeller won't create a pod for this task, it will simply pass the inputs to the outputs.
+    https://github.com/flyteorg/flyte/blob/master/flyteplugins/go/tasks/plugins/testing/echo.go
+    """
+
+    _TASK_TYPE = "echo"
+
+    def __init__(self, name: str, inputs: Optional[Dict[str, Type]] = None, **kwargs):
+        outputs = dict(zip(list(output_name_generator(len(inputs.values()))), inputs.values())) if inputs else None
+        super().__init__(
+            task_type=self._TASK_TYPE,
+            name=name,
+            interface=Interface(inputs=inputs, outputs=outputs),
+            **kwargs,
+        )
+
+    def execute(self, **kwargs) -> Any:
+        values = list(kwargs.values())
+        if len(values) == 1:
+            return values[0]
+        else:
+            return tuple(values)

tests/flytekit/unit/core/test_conditions.py

@@ -9,6 +9,7 @@
 from flytekit import task, workflow
 from flytekit.configuration import Image, ImageConfig, SerializationSettings
 from flytekit.core.condition import conditional
+from flytekit.core.task import Echo
 from flytekit.models.core.workflow import Node
 from flytekit.tools.translator import get_serializable
 
@@ -495,3 +496,41 @@ def multiplier_2(my_input: float) -> float:
 
     res = multiplier_2(my_input=10.0)
     assert res == 20
+
+
+def test_echo_in_condition():
+    echo1 = Echo(name="echo", inputs={"a": float})
+
+    @task()
+    def t1(radius: float) -> typing.Optional[float]:
+        return 2 * 3.14 * radius
+
+    @workflow
+    def wf1(radius: float) -> typing.Optional[float]:
+        return (
+            conditional("shape_properties_with_multiple_branches")
+            .if_((radius >= 0.1) & (radius < 1.0))
+            .then(t1(radius=radius))
+            .else_()
+            .then(echo1(a=radius))
+        )
+
+    assert wf1(radius=1.8) == 1.8
+
+    echo2 = Echo(name="echo", inputs={"a": float, "b": float})
+
+    @task()
+    def t2(radius: float) -> (float, float):
+        return 2 * 3.14 * radius, 2 * 3.14 * radius
+
+    @workflow
+    def wf2(radius1: float, radius2: float) -> (float, float):
+        return (
+            conditional("shape_properties_with_multiple_branches")
+            .if_((radius1 >= 0.1) & (radius1 < 1.0))
+            .then(t2(radius=radius2))
+            .else_()
+            .then(echo2(a=radius1, b=radius2))
+        )
+
+    assert wf2(radius1=1.8, radius2=1.8) == (1.8, 1.8)