don't transform bound inputs to list for remote entities (#3168)

* don't transform bound inputs to list for remote entities

Signed-off-by: Paul Dittamo <[email protected]>

* clean up

Signed-off-by: Paul Dittamo <[email protected]>

---------

Signed-off-by: Paul Dittamo <[email protected]>

Author: pvditt

flytekit/core/array_node.py

@@ -88,7 +88,7 @@ def __init__(
                 self.target.python_interface, self._bound_inputs, output_as_list_of_optionals
             )
         elif self.target.interface:
-            self._remote_interface = self.target.interface.transform_interface_to_list()
+            self._remote_interface = self.target.interface.transform_interface_to_list(self.bound_inputs)
         else:
             raise ValueError("No interface found for the target entity.")
 

flytekit/core/interface.py

@@ -372,7 +372,7 @@ def transform_interface_to_list_interface(
     Takes a single task interface and interpolates it to an array interface - to allow performing distributed python map
     like functions
     :param interface: Interface to be upgraded to a list interface
-    :param bound_inputs: fixed inputs that should not upgraded to a list and will be maintained as scalars.
+    :param bound_inputs: fixed inputs that should not be updated to a list and will be maintained as is
     """
     map_inputs = transform_types_to_list_of_type(interface.inputs, bound_inputs)
     map_outputs = transform_types_to_list_of_type(interface.outputs, set(), optional_outputs)

flytekit/models/interface.py

@@ -158,13 +158,19 @@ def from_flyte_idl(cls, proto: _interface_pb2.TypedInterface) -> "TypedInterface
             outputs={k: Variable.from_flyte_idl(v) for k, v in proto.outputs.variables.items()},
         )
 
-    def transform_interface_to_list(self) -> "TypedInterface":
+    def transform_interface_to_list(self, bound_inputs: typing.Set[str]) -> "TypedInterface":
         """
         Takes a single task interface and interpolates it to an array interface - to allow performing distributed
         python map like functions
+        :param bound_inputs: fixed inputs that should not be updated to a list and will be maintained as is
         """
         list_interface = _interface_pb2.TypedInterface(
-            inputs=_interface_pb2.VariableMap(variables={k: v.to_flyte_idl_list() for k, v in self.inputs.items()}),
+            inputs=_interface_pb2.VariableMap(
+                variables={
+                    k: (v.to_flyte_idl_list() if k not in bound_inputs else v.to_flyte_idl())
+                    for k, v in self.inputs.items()
+                }
+            ),
             outputs=_interface_pb2.VariableMap(variables={k: v.to_flyte_idl_list() for k, v in self.outputs.items()}),
         )
         return self.from_flyte_idl(list_interface)

flytekit/remote/remote.py

@@ -2639,7 +2639,7 @@ def sync_node_execution(
                 launch_plan = self.fetch_launch_plan(
                     launch_plan_id.project, launch_plan_id.domain, launch_plan_id.name, launch_plan_id.version
                 )
-                task_execution_interface = launch_plan.interface.transform_interface_to_list()
+                task_execution_interface = launch_plan.interface.transform_interface_to_list(bound_inputs=set())
                 execution._task_executions = [
                     self.sync_task_execution(
                         FlyteTaskExecution.promote_from_model(task_execution), task_execution_interface

tests/flytekit/unit/models/test_interface.py

@@ -51,7 +51,7 @@ def test_typed_interface(literal_type):
     assert len(deserialized_typed_interface.inputs) == 1
     assert len(deserialized_typed_interface.outputs) == 2
 
-    deserialized_typed_interface_list = typed_interface.transform_interface_to_list()
+    deserialized_typed_interface_list = typed_interface.transform_interface_to_list(set())
     assert deserialized_typed_interface_list.inputs["a"].type == types.LiteralType(collection_type=literal_type)
     assert deserialized_typed_interface_list.outputs["b"].type == types.LiteralType(collection_type=literal_type)
     assert deserialized_typed_interface_list.outputs["c"].type == types.LiteralType(collection_type=literal_type)
@@ -61,6 +61,16 @@ def test_typed_interface(literal_type):
     assert len(deserialized_typed_interface_list.inputs) == 1
     assert len(deserialized_typed_interface_list.outputs) == 2
 
+    deserialized_typed_interface_list = typed_interface.transform_interface_to_list({"a"})
+    assert deserialized_typed_interface_list.inputs["a"].type == literal_type
+    assert deserialized_typed_interface_list.outputs["b"].type == types.LiteralType(collection_type=literal_type)
+    assert deserialized_typed_interface_list.outputs["c"].type == types.LiteralType(collection_type=literal_type)
+    assert deserialized_typed_interface_list.inputs["a"].description == "description1"
+    assert deserialized_typed_interface_list.outputs["b"].description == "description2"
+    assert deserialized_typed_interface_list.outputs["c"].description == "description3"
+    assert len(deserialized_typed_interface_list.inputs) == 1
+    assert len(deserialized_typed_interface_list.outputs) == 2
+
 
 def test_parameter():
     v = interface.Variable(types.LiteralType(simple=types.SimpleType.BOOLEAN), "asdf asdf asdf")