Repository: flyteorg/flytekit

Issue:
use getattr to access task_def.disable_deck in entrypoint
Fixes https://github.com/flyteorg/flyte/issues/3841

# TL;DR

https://github.com/flyteorg/flytekit/pull/1693 introduced a bug in task templates like `SQlite3Task`. This PR fixes an issue in the execution entrypoint that assumes the all task objects have the `disable_deck` attribute.

## Type
 - [X] Bug Fix
 - [ ] Feature
 - [ ] Plugin

## Are all requirements met?

 - [X] Code completed
 - [X] Smoke tested
 - [ ] Unit tests added
 - [ ] Code documentation added
 - [ ] Any pending items have an associated Issue

## Complete description

Use the `getattr` function to get `disable_deck`, and return `False` if not present.

## Tracking Issue
[https://github.com/flyteorg/flyte/issues/<number>](https://github.com/flyteorg/flyte/issues/3841)



Discussion:

Resolving patch:
diff --git a/flytekit/core/python_function_task.py b/flytekit/core/python_function_task.py
--- a/flytekit/core/python_function_task.py
+++ b/flytekit/core/python_function_task.py
@@ -26,7 +26,7 @@
 from flytekit.core.docstring import Docstring
 from flytekit.core.interface import transform_signature_to_interface
 from flytekit.core.python_auto_container import PythonAutoContainerTask, default_task_resolver
-from flytekit.core.tracker import isnested, istestfunction
+from flytekit.core.tracker import is_functools_wrapped_module_level, isnested, istestfunction
 from flytekit.core.workflow import (
     PythonFunctionWorkflow,
     WorkflowFailurePolicy,
@@ -130,12 +130,18 @@ def __init__(
         if self._task_resolver is default_task_resolver:
             # The default task resolver can't handle nested functions
             # TODO: Consider moving this to a can_handle function or something inside the resolver itself.
-            if not istestfunction(func=task_function) and isnested(func=task_function):
+            if (
+                not istestfunction(func=task_function)
+                and isnested(func=task_function)
+                and not is_functools_wrapped_module_level(task_function)
+            ):
                 raise ValueError(
                     "TaskFunction cannot be a nested/inner or local function. "
                     "It should be accessible at a module level for Flyte to execute it. Test modules with "
-                    "names beginning with `test_` are allowed to have nested tasks. If you want to create your own tasks"
-                    "use the TaskResolverMixin"
+                    "names beginning with `test_` are allowed to have nested tasks. "
+                    "If you're decorating your task function with custom decorators, use functools.wraps "
+                    "or functools.update_wrapper on the function wrapper. "
+                    "Alternatively if you want to create your own tasks with custom behavior use the TaskResolverMixin"
                 )
         self._task_function = task_function
         self._execution_mode = execution_mode
diff --git a/flytekit/core/tracker.py b/flytekit/core/tracker.py
--- a/flytekit/core/tracker.py
+++ b/flytekit/core/tracker.py
@@ -109,6 +109,48 @@ def foo_inner():
     return func.__code__.co_flags & inspect.CO_NESTED != 0
 
 
+def is_functools_wrapped_module_level(func: Callable) -> bool:
+    """Returns true if the function is a functools.wraps-updated function that is defined in the module-level scope.
+
+    .. code:: python
+
+        import functools
+
+        def decorator(fn):
+            @functools.wraps(fn)
+            def wrapper(*args, **kwargs):
+                return fn(*arks, **kwargs)
+
+            return wrapper
+
+        @decorator
+        def foo():
+            ...
+
+        def define_inner_wrapped_fn():
+
+            @decorator
+            def foo_inner(*args, **kwargs):
+                return fn(*arks, **kwargs)
+
+            return foo_inner
+
+        bar = define_inner_wrapped_fn()
+
+        is_functools_wrapped_module_level(foo)  # True
+        is_functools_wrapped_module_level(bar)  # False
+
+    In this case, applying this function to ``foo`` returns true because ``foo`` was defined in the module-level scope.
+    Applying this function to ``bar`` returns false because it's being assigned to ``foo_inner``, which is a
+    functools-wrapped function but is actually defined in the local scope of ``define_inner_wrapped_fn``.
+
+    This works because functools.wraps updates the __name__ and __qualname__ attributes of the wrapper to match the
+    wrapped function. Since ``define_inner_wrapped_fn`` doesn't update the __qualname__ of ``foo_inner``, the inner
+    function's __qualname__ won't match its __name__.
+    """
+    return hasattr(func, "__wrapped__") and func.__name__ == func.__qualname__
+
+
 def istestfunction(func) -> bool:
     """
     Returns true if the function is defined in a test module. A test module has to have `test_` as the prefix.


Test patch:
diff --git a/tests/flytekit/unit/core/functools/nested_function.py b/tests/flytekit/unit/core/functools/nested_function.py
new file mode 100644
--- /dev/null
+++ b/tests/flytekit/unit/core/functools/nested_function.py
@@ -0,0 +1,35 @@
+"""Script used for testing local execution of nested functions that don't use functools.wraps."""
+
+import os
+
+from flytekit import task, workflow
+
+
+def task_decorator(fn):
+    def wrapper(x: int) -> int:
+        print("running task_decorator")
+        return fn(x=x)
+
+    return wrapper
+
+
+def foo():
+    @task
+    @task_decorator
+    def my_task(x: int) -> int:
+        print("running my_task")
+        return x + 1
+
+    return my_task
+
+
+my_task = foo()
+
+
+@workflow
+def my_workflow(x: int) -> int:
+    return my_task(x=x)
+
+
+if __name__ == "__main__":
+    print(my_workflow(x=int(os.getenv("SCRIPT_INPUT"))))
diff --git a/tests/flytekit/unit/core/functools/nested_wrapped_function.py b/tests/flytekit/unit/core/functools/nested_wrapped_function.py
new file mode 100644
--- /dev/null
+++ b/tests/flytekit/unit/core/functools/nested_wrapped_function.py
@@ -0,0 +1,36 @@
+"""Script used for testing local execution of functool.wraps-wrapped tasks"""
+
+from functools import wraps
+
+from flytekit import task, workflow
+
+
+def task_decorator(fn):
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        return fn(*args, **kwargs)
+
+    return wrapper
+
+
+def foo():
+    @task
+    @task_decorator
+    def my_task(x: int) -> int:
+        assert x > 0, f"my_task failed with input: {x}"
+        print("running my_task")
+        return x + 1
+
+    return my_task
+
+
+my_task = foo()
+
+
+@workflow
+def my_workflow(x: int) -> int:
+    return my_task(x=x)
+
+
+if __name__ == "__main__":
+    print(my_workflow(x=11))
diff --git a/tests/flytekit/unit/core/functools/simple_decorator.py b/tests/flytekit/unit/core/functools/simple_decorator.py
new file mode 100644
--- /dev/null
+++ b/tests/flytekit/unit/core/functools/simple_decorator.py
@@ -0,0 +1,40 @@
+"""Script used for testing local execution of functool.wraps-wrapped tasks"""
+
+import os
+from functools import wraps
+
+from flytekit import task, workflow
+
+
+def task_decorator(fn):
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        print(f"before running {fn.__name__}")
+        try:
+            print(f"try running {fn.__name__}")
+            out = fn(*args, **kwargs)
+        except Exception as exc:
+            print(f"error running {fn.__name__}: {exc}")
+        finally:
+            print(f"finally after running {fn.__name__}")
+        print(f"after running {fn.__name__}")
+        return out
+
+    return wrapper
+
+
+@task
+@task_decorator
+def my_task(x: int) -> int:
+    assert x > 0, f"my_task failed with input: {x}"
+    print("running my_task")
+    return x + 1
+
+
+@workflow
+def my_workflow(x: int) -> int:
+    return my_task(x=x)
+
+
+if __name__ == "__main__":
+    print(my_workflow(x=int(os.getenv("SCRIPT_INPUT"))))
diff --git a/tests/flytekit/unit/core/functools/stacked_decorators.py b/tests/flytekit/unit/core/functools/stacked_decorators.py
new file mode 100644
--- /dev/null
+++ b/tests/flytekit/unit/core/functools/stacked_decorators.py
@@ -0,0 +1,51 @@
+"""Script used for testing local execution of functool.wraps-wrapped tasks for stacked decorators"""
+
+import os
+from functools import wraps
+
+from flytekit import task, workflow
+
+
+def task_decorator_1(fn):
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        print("running task_decorator_1")
+        return fn(*args, **kwargs)
+
+    return wrapper
+
+
+def task_decorator_2(fn):
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        print("running task_decorator_2")
+        return fn(*args, **kwargs)
+
+    return wrapper
+
+
+def task_decorator_3(fn):
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        print("running task_decorator_3")
+        return fn(*args, **kwargs)
+
+    return wrapper
+
+
+@task
+@task_decorator_1
+@task_decorator_2
+@task_decorator_3
+def my_task(x: int) -> int:
+    print("running my_task")
+    return x + 1
+
+
+@workflow
+def my_workflow(x: int) -> int:
+    return my_task(x=x)
+
+
+if __name__ == "__main__":
+    print(my_workflow(x=int(os.getenv("SCRIPT_INPUT"))))
diff --git a/tests/flytekit/unit/core/functools/test_decorators.py b/tests/flytekit/unit/core/functools/test_decorators.py
new file mode 100644
--- /dev/null
+++ b/tests/flytekit/unit/core/functools/test_decorators.py
@@ -0,0 +1,89 @@
+"""Test local execution of files that use functools to decorate tasks and workflows."""
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+test_module_dir = Path(os.path.dirname(__file__))
+
+
+def test_wrapped_tasks_happy_path(capfd):
+    subprocess.run(
+        [sys.executable, str(test_module_dir / "simple_decorator.py")],
+        env={"SCRIPT_INPUT": "10"},
+        text=True,
+    )
+    out = capfd.readouterr().out
+
+    assert out.strip().split("\n") == [
+        "before running my_task",
+        "try running my_task",
+        "running my_task",
+        "finally after running my_task",
+        "after running my_task",
+        "11",
+    ]
+
+
+def test_wrapped_tasks_error(capfd):
+    subprocess.run(
+        [sys.executable, str(test_module_dir / "simple_decorator.py")],
+        env={"SCRIPT_INPUT": "0"},
+        text=True,
+    )
+    out = capfd.readouterr().out
+
+    assert out.strip().split("\n") == [
+        "before running my_task",
+        "try running my_task",
+        "error running my_task: my_task failed with input: 0",
+        "finally after running my_task",
+        "after running my_task",
+    ]
+
+
+def test_stacked_wrapped_tasks(capfd):
+    subprocess.run(
+        [sys.executable, str(test_module_dir / "stacked_decorators.py")],
+        env={"SCRIPT_INPUT": "10"},
+        text=True,
+    )
+    out = capfd.readouterr().out
+    assert out.strip().split("\n") == [
+        "running task_decorator_1",
+        "running task_decorator_2",
+        "running task_decorator_3",
+        "running my_task",
+        "11",
+    ]
+
+
+def test_unwrapped_task():
+    completed_process = subprocess.run(
+        [sys.executable, str(test_module_dir / "unwrapped_decorator.py")],
+        env={"SCRIPT_INPUT": "10"},
+        text=True,
+        capture_output=True,
+    )
+    error = completed_process.stderr
+    error_str = error.strip().split("\n")[-1]
+    assert error_str == (
+        "ValueError: Type <class 'inspect._empty'> not supported currently in Flytekit. "
+        "Please register a new transformer"
+    )
+
+
+@pytest.mark.parametrize("script", ["nested_function.py", "nested_wrapped_function.py"])
+def test_nested_function(script):
+    completed_process = subprocess.run(
+        [sys.executable, str(test_module_dir / script)],
+        env={"SCRIPT_INPUT": "10"},
+        text=True,
+        capture_output=True,
+    )
+    error = completed_process.stderr
+    error_str = error.strip().split("\n")[-1]
+    assert error_str.startswith("ValueError: TaskFunction cannot be a nested/inner or local function.")
diff --git a/tests/flytekit/unit/core/functools/unwrapped_decorator.py b/tests/flytekit/unit/core/functools/unwrapped_decorator.py
new file mode 100644
--- /dev/null
+++ b/tests/flytekit/unit/core/functools/unwrapped_decorator.py
@@ -0,0 +1,29 @@
+"""Script used for testing local execution of non-functools.wraps-wrapped tasks"""
+
+import os
+
+from flytekit import task, workflow
+
+
+def task_decorator(fn):
+    def wrapper(*args, **kwargs):
+        print("running task_decorator")
+        return fn(*args, **kwargs)
+
+    return wrapper
+
+
+@task
+@task_decorator
+def my_task(x: int) -> int:
+    print("running my_task")
+    return x + 1
+
+
+@workflow
+def my_workflow(x: int) -> int:
+    return my_task(x=x)
+
+
+if __name__ == "__main__":
+    print(my_workflow(x=int(os.getenv("SCRIPT_INPUT"))))