Feature gate for FlyteMissingReturnValueException (flyteorg#2623)

Signed-off-by: Kevin Su <[email protected]>
Signed-off-by: mao3267 <[email protected]>

flytekit/core/interface.py

@@ -3,6 +3,7 @@
 import collections
 import copy
 import inspect
+import sys
 import typing
 from collections import OrderedDict
 from typing import Any, Dict, Generator, List, Optional, Tuple, Type, TypeVar, Union, cast
@@ -381,10 +382,12 @@ def transform_function_to_interface(fn: typing.Callable, docstring: Optional[Doc
     return_annotation = type_hints.get("return", None)
 
     ctx = FlyteContextManager.current_context()
-    # Only check if the task/workflow has a return statement at compile time locally.
     if (
         ctx.execution_state
+        # Only check if the task/workflow has a return statement at compile time locally.
         and ctx.execution_state.mode is None
+        # inspect module does not work correctly with Python <3.10.10. https://github.com/flyteorg/flyte/issues/5608
+        and sys.version_info >= (3, 10, 10)
         and return_annotation
         and type(None) not in get_args(return_annotation)
         and return_annotation is not type(None)

tests/flytekit/unit/core/test_workflows.py

@@ -246,6 +246,29 @@ def one_output_wf() -> int:  # type: ignore
         one_output_wf()
 
 
+def test_custom_wrapper():
+    def our_task(
+            _task_function: typing.Optional[typing.Callable] = None,
+            **kwargs,
+    ):
+        def wrapped(_func: typing.Callable):
+            return task(_task_function=_func)
+
+        if _task_function:
+            return wrapped(_task_function)
+        else:
+            return wrapped
+
+    @our_task(
+        foo={
+            "bar1": lambda x: print(x),
+            "bar2": lambda x: print(x),
+        },
+    )
+    def missing_func_body() -> str:
+        return "foo"
+
+
 def test_wf_no_output():
     @task
     def t1(a: int) -> int: