feat: make container task return python type

-e
Signed-off-by: machichima <nary12321@gmail.com>

Author: machichima

flytekit/core/array_node_map_task.py

@@ -17,7 +17,6 @@
 from flytekit.core.context_manager import ExecutionState, FlyteContext, FlyteContextManager
 from flytekit.core.interface import transform_interface_to_list_interface
 from flytekit.core.launch_plan import LaunchPlan
-from flytekit.core.promise import Promise, create_native_named_tuple, create_task_output
 from flytekit.core.python_function_task import PythonFunctionTask, PythonInstanceTask
 from flytekit.core.task import ReferenceTask
 from flytekit.core.type_engine import TypeEngine
@@ -380,13 +379,6 @@ def _raw_execute(self, **kwargs) -> Any:
                     single_instance_inputs[k] = kwargs[k]
             try:
                 o = self._run_task.execute(**single_instance_inputs)
-                # For running container task in local execution, it will return
-                # the LiteralMap. We need to convert it to native type here.
-                if isinstance(o, _literal_models.LiteralMap):
-                    vals = [Promise(var, o.literals[var]) for var in o.literals.keys()]
-                    result = create_task_output(vals, self.python_interface)
-                    ctx = FlyteContextManager.current_context()
-                    o = create_native_named_tuple(ctx, result, self._run_task.python_interface)
                 if outputs_expected:
                     outputs.append(o)
             except Exception as exc:

flytekit/core/container_task.py

@@ -14,7 +14,6 @@
 from flytekit.image_spec.image_spec import ImageSpec
 from flytekit.loggers import logger
 from flytekit.models import task as _task_model
-from flytekit.models.literals import LiteralMap
 from flytekit.models.security import Secret, SecurityContext
 
 _PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
@@ -254,14 +253,12 @@ def _get_output_dict(self, output_directory: str) -> Dict[str, Any]:
                     output_dict[k] = self._convert_output_val_to_correct_type(output_val, output_type)
         return output_dict
 
-    def execute(self, **kwargs) -> LiteralMap:
+    def execute(self, **kwargs) -> Any:
         try:
             import docker
         except ImportError:
             raise ImportError(DOCKER_IMPORT_ERROR_MESSAGE)
 
-        from flytekit.core.type_engine import TypeEngine
-
         ctx = FlyteContext.current_context()
 
         # Normalize the input and output directories
@@ -289,8 +286,12 @@ def execute(self, **kwargs) -> LiteralMap:
         container.wait()
 
         output_dict = self._get_output_dict(output_directory)
-        outputs_literal_map = TypeEngine.dict_to_literal_map(ctx, output_dict)
-        return outputs_literal_map
+        if len(output_dict) == 0:
+            return None
+        elif len(output_dict) == 1:
+            return list(output_dict.values())[0]
+        elif len(output_dict) > 1:
+            return tuple(output_dict.values())
 
     def get_container(self, settings: SerializationSettings) -> _task_model.Container:
         # if pod_template is specified, return None here but in get_k8s_pod, return pod_template merged with container