[BUG] Flytekit's "error message for missing return" broken in python 3.9

[fellhorn]

Describe the bug

CPython versions <3.10.10 have a bug in the inspect module, which causes missing source lines for certain decorated functions.

E.g.

import inspect
from flytekit import task

from typing import Callable, Optional


def our_task(
    _task_function: Optional[Callable] = None,
    **kwargs,
):
    def wrapped(_func: 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"


print(inspect.getsourcelines(missing_func_body)[0])

prints

['@our_task(\n', '    foo={\n', '        "bar1": lambda x: print(x),\n', '        "bar2": lambda x: print(x),\n', '    },\n', ')\n']

and misses the return statement expected by flytekit since 1.13.1 and introduced in flyteorg/flytekit#2551. This breaks registration as a FlyteMissingReturnValueException is raised.

Expected behavior

As flytekit still supports python 3.9, I would expect registration to work with such wrapped functions.

Potentially we could have a feature gate to only run it for python >=3.11

Additional context to reproduce

No response

Screenshots

No response

Are you sure this issue hasn't been raised already?

• Yes

Have you read the Code of Conduct?

• Yes

[runllm]

The issue you're describing seems to be related to a known bug in CPython versions <3.10.10, which affects the inspect module and causes it to miss certain source lines for decorated functions. This bug impacts Flytekit's error handling for missing return statements, introduced in Flytekit PR #2551.

To address this, you could consider implementing a feature gate to only run the return statement check for Python versions >=3.11. This would ensure compatibility with older Python versions still supported by Flytekit.

Would you like more details on how to implement this feature gate or any other specific aspect of the issue?

Sources:

• Flytekit PR #2551
• CPython issue #102647

---

This is an AI-generated response and your feedback is appreciated! Please leave a 👍 if this is helpful and 👎 if it is not.

Mention @RunLLM to ask follow-up questions. See runllm.com for more!

[fellhorn]

@pingsutw What do you think?

[pingsutw]

@fellhorn, sorry for the breaking change. I just opened a PR to add a feature gate. mind taking a look? flyteorg/flytekit#2623