Fail registration if output isn't Optional when using map tasks with min_success_ratio < 1

flytekit/core/map_task.py

@@ -18,6 +18,7 @@
 from flytekit.core.interface import transform_interface_to_list_interface
 from flytekit.core.python_function_task import PythonFunctionTask, PythonInstanceTask
 from flytekit.core.tracker import TrackedInstance
+from flytekit.core.type_engine import UnionTransformer
 from flytekit.core.utils import timeit
 from flytekit.exceptions import scopes as exception_scopes
 from flytekit.models.array_job import ArrayJob
@@ -73,6 +74,13 @@ def __init__(
         if len(actual_task.python_interface.outputs.keys()) > 1:
             raise ValueError("Map tasks only accept python function tasks with 0 or 1 outputs")
 
+        if (
+            min_success_ratio
+            and min_success_ratio != 1
+            and not UnionTransformer.is_optional_type(actual_task.python_interface.outputs["o0"])
+        ):
+            raise ValueError("Map tasks with min_success_ratio < 1 must have an optional output")
+
         self._bound_inputs: typing.Set[str] = set(bound_inputs) if bound_inputs else set()
         if self._partial:
             self._bound_inputs = set(self._partial.keywords.keys())

tests/flytekit/unit/core/test_map_task.py

@@ -26,7 +26,7 @@ def serialization_settings():
 
 
 @task
-def t1(a: int) -> str:
+def t1(a: int) -> typing.Optional[str]:
     b = a + 2
     return str(b)
 
@@ -46,11 +46,11 @@ def t3(a: int, b: str, c: float) -> str:
 def test_map_docs():
     # test_map_task_start
     @task
-    def my_mappable_task(a: int) -> str:
+    def my_mappable_task(a: int) -> typing.Optional[str]:
         return str(a)
 
     @workflow
-    def my_wf(x: typing.List[int]) -> typing.List[str]:
+    def my_wf(x: typing.List[int]) -> typing.List[typing.Optional[str]]:
         return map_task(my_mappable_task, metadata=TaskMetadata(retries=1), concurrency=10, min_success_ratio=0.75,)(
             a=x
         ).with_overrides(cpu="10M")
@@ -262,3 +262,26 @@ def test_map_task_resolver(serialization_settings):
     t = mtr.load_task(loader_args=args)
     assert t.python_interface.inputs == mt.python_interface.inputs
     assert t.python_interface.outputs == mt.python_interface.outputs
+
+
+def test_map_task_min_success_ratio():
+    @task
+    def some_task1(inputs: int) -> int:
+        return inputs
+
+    @workflow
+    def my_wf1() -> typing.List[int]:
+        return map_task(some_task1, min_success_ratio=0.5)(inputs=[1, 2, 3, 4])
+
+    with pytest.raises(ValueError, match="Map tasks with min_success_ratio < 1 must have an optional output"):
+        my_wf1()
+
+    @task
+    def some_task2(inputs: int) -> typing.Optional[int]:
+        return inputs
+
+    @workflow
+    def my_wf2() -> typing.List[typing.Optional[int]]:
+        return map_task(some_task2, min_success_ratio=0.5)(inputs=[1, 2, 3, 4])
+
+    my_wf2()