remove check of settings lists in partials (#3187)

* remove check of settings lists in partials

Signed-off-by: Paul Dittamo <pvdittamo@gmail.com>

* update unit test

Signed-off-by: Paul Dittamo <pvdittamo@gmail.com>

* test

Signed-off-by: Paul Dittamo <pvdittamo@gmail.com>

* clean up

Signed-off-by: Paul Dittamo <pvdittamo@gmail.com>

---------

Signed-off-by: Paul Dittamo <pvdittamo@gmail.com>

Author: pvditt

flytekit/core/array_node_map_task.py

@@ -55,10 +55,6 @@ def __init__(
         """
         self._partial = None
         if isinstance(python_function_task, functools.partial):
-            # TODO: We should be able to support partial tasks with lists as inputs
-            for arg in python_function_task.keywords.values():
-                if isinstance(arg, list):
-                    raise ValueError("Cannot use a partial task with lists as inputs")
             self._partial = python_function_task
             actual_task = self._partial.func
         else:

tests/flytekit/unit/core/test_array_node_map_task.py

@@ -290,7 +290,7 @@ def many_outputs(a: int) -> (int, str):
         _ = map_task(many_outputs)
 
 
-def test_parameter_order():
+def test_partials_local_execute():
     @task()
     def task1(a: int, b: float, c: str) -> str:
         return f"{a} - {b} - {c}"
@@ -303,20 +303,40 @@ def task2(b: float, c: str, a: int) -> str:
     def task3(c: str, a: int, b: float) -> str:
         return f"{a} - {b} - {c}"
 
+    @task()
+    def task4(c: list[str], a: list[int], b: list[float]) -> list[str]:
+        return [f"{a[i]} - {b[i]} - {c[i]}" for i in range(len(a))]
+
     param_a = [1, 2, 3]
     param_b = [0.1, 0.2, 0.3]
-    param_c = "c"
+    fixed_param_c = "c"
 
-    m1 = map_task(functools.partial(task1, c=param_c))(a=param_a, b=param_b)
-    m2 = map_task(functools.partial(task2, c=param_c))(a=param_a, b=param_b)
-    m3 = map_task(functools.partial(task3, c=param_c))(a=param_a, b=param_b)
+    m1 = map_task(functools.partial(task1, c=fixed_param_c))(a=param_a, b=param_b)
+    m2 = map_task(functools.partial(task2, c=fixed_param_c))(a=param_a, b=param_b)
+    m3 = map_task(functools.partial(task3, c=fixed_param_c))(a=param_a, b=param_b)
 
-    m4 = ArrayNodeMapTask(task1, bound_inputs_values={"c": param_c})(a=param_a, b=param_b)
-    m5 = ArrayNodeMapTask(task2, bound_inputs_values={"c": param_c})(a=param_a, b=param_b)
-    m6 = ArrayNodeMapTask(task3, bound_inputs_values={"c": param_c})(a=param_a, b=param_b)
+    m4 = ArrayNodeMapTask(task1, bound_inputs_values={"c": fixed_param_c})(a=param_a, b=param_b)
+    m5 = ArrayNodeMapTask(task2, bound_inputs_values={"c": fixed_param_c})(a=param_a, b=param_b)
+    m6 = ArrayNodeMapTask(task3, bound_inputs_values={"c": fixed_param_c})(a=param_a, b=param_b)
 
     assert m1 == m2 == m3 == m4 == m5 == m6 == ["1 - 0.1 - c", "2 - 0.2 - c", "3 - 0.3 - c"]
 
+    list_param_a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+    list_param_b = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]]
+    fixed_list_param_c = ["c", "d", "e"]
+
+    m7 = map_task(functools.partial(task4, c=fixed_list_param_c))(a=list_param_a, b=list_param_b)
+    m8 = ArrayNodeMapTask(task4, bound_inputs_values={"c": fixed_list_param_c})(a=list_param_a, b=list_param_b)
+
+    assert m7 == m8 == [
+        ['1 - 0.1 - c', '2 - 0.2 - d', '3 - 0.3 - e'],
+        ['4 - 0.4 - c', '5 - 0.5 - d', '6 - 0.6 - e'],
+        ['7 - 0.7 - c', '8 - 0.8 - d', '9 - 0.9 - e']
+    ]
+
+    with pytest.raises(ValueError):
+        map_task(functools.partial(task1, c=fixed_list_param_c))(a=param_a, b=param_b)
+
 
 def test_bounded_inputs_vars_order(serialization_settings):
     @task()

tests/flytekit/unit/core/test_partials.py

@@ -142,7 +142,6 @@ def wf_in(a: typing.List[int]):
     "map_task_fn",
     [
         legacy_map_task,
-        array_node_map_task,
     ],
 )
 def test_lists_cannot_be_used_in_partials(map_task_fn):