diff --git a/flytekit/core/base_task.py b/flytekit/core/base_task.py index 6430aa9eac..2a314e1aab 100644 --- a/flytekit/core/base_task.py +++ b/flytekit/core/base_task.py @@ -75,7 +75,10 @@ FlyteNonRecoverableSystemException, FlyteUploadDataException, ) -from flytekit.exceptions.user import FlyteUserRuntimeException +from flytekit.exceptions.user import ( + FlyteUserRuntimeException, + FlyteValueException, +) from flytekit.loggers import logger from flytekit.models import dynamic_job as _dynamic_job from flytekit.models import interface as _interface_models @@ -770,6 +773,15 @@ def dispatch_execute( ): return native_outputs + if isinstance(native_outputs, VoidPromise): + return _literal_models.LiteralMap(literals={}) + + if native_outputs is not None and len(list(self._outputs_interface.keys())) == 0: + raise FlyteValueException( + native_outputs, + f"Interface has {len(self.python_interface.outputs)} outputs.", + ) + try: with timeit("dispatch execute"): literals_map, native_outputs_as_map = run_sync( diff --git a/plugins/flytekit-flyteinteractive/tests/test_flyteinteractive_jupyter.py b/plugins/flytekit-flyteinteractive/tests/test_flyteinteractive_jupyter.py index 8be6477600..cf69b5f4bd 100644 --- a/plugins/flytekit-flyteinteractive/tests/test_flyteinteractive_jupyter.py +++ b/plugins/flytekit-flyteinteractive/tests/test_flyteinteractive_jupyter.py @@ -27,7 +27,8 @@ def jupyter_patches(): with mock.patch("multiprocessing.Process") as mock_process, mock.patch( "flytekitplugins.flyteinteractive.jupyter_lib.decorator.write_example_notebook" ) as mock_write_example_notebook, mock.patch( - "flytekitplugins.flyteinteractive.jupyter_lib.decorator.exit_handler" + "flytekitplugins.flyteinteractive.jupyter_lib.decorator.exit_handler", + return_value=None, ) as mock_exit_handler: yield (mock_process, mock_write_example_notebook, mock_exit_handler) diff --git a/plugins/flytekit-flyteinteractive/tests/test_flyteinteractive_vscode.py b/plugins/flytekit-flyteinteractive/tests/test_flyteinteractive_vscode.py index 7996a0da41..4b1336f4db 100644 --- a/plugins/flytekit-flyteinteractive/tests/test_flyteinteractive_vscode.py +++ b/plugins/flytekit-flyteinteractive/tests/test_flyteinteractive_vscode.py @@ -55,7 +55,8 @@ def vscode_patches(): with mock.patch("multiprocessing.Process") as mock_process, mock.patch( "flytekit.interactive.vscode_lib.decorator.prepare_interactive_python" ) as mock_prepare_interactive_python, mock.patch( - "flytekit.interactive.vscode_lib.decorator.exit_handler" + "flytekit.interactive.vscode_lib.decorator.exit_handler", + return_value=None, ) as mock_exit_handler, mock.patch( "flytekit.interactive.vscode_lib.decorator.download_vscode" ) as mock_download_vscode, mock.patch("signal.signal") as mock_signal, mock.patch( diff --git a/plugins/flytekit-kf-pytorch/tests/test_elastic_task.py b/plugins/flytekit-kf-pytorch/tests/test_elastic_task.py index faadc1019f..a85b9c569a 100644 --- a/plugins/flytekit-kf-pytorch/tests/test_elastic_task.py +++ b/plugins/flytekit-kf-pytorch/tests/test_elastic_task.py @@ -95,7 +95,7 @@ def test_execution_params(start_method: str, target_exec_id: str, monkeypatch_ex monkeypatch.setenv("FLYTE_INTERNAL_EXECUTION_ID", target_exec_id) @task(task_config=Elastic(nnodes=1, nproc_per_node=1, start_method=start_method)) - def test_task(n: int): + def test_task(n: int) -> int: ctx = flytekit.current_context() assert ctx.execution_id.name == target_exec_id diff --git a/plugins/flytekit-pandera/tests/test_plugin.py b/plugins/flytekit-pandera/tests/test_plugin.py index 9c15864176..90edbb2c39 100644 --- a/plugins/flytekit-pandera/tests/test_plugin.py +++ b/plugins/flytekit-pandera/tests/test_plugin.py @@ -108,7 +108,7 @@ class Schema(pandera.DataFrameModel): col2: float @task - def fn_input(df: typing.Annotated[DataFrame[Schema], config]): + def fn_input(df: typing.Annotated[DataFrame[Schema], config]) -> typing.Annotated[DataFrame[Schema], config]: return df @task diff --git a/tests/flytekit/integration/remote/workflows/basic/generic_idl_flytetypes.py b/tests/flytekit/integration/remote/workflows/basic/generic_idl_flytetypes.py index c157432f21..d679042dc1 100644 --- a/tests/flytekit/integration/remote/workflows/basic/generic_idl_flytetypes.py +++ b/tests/flytekit/integration/remote/workflows/basic/generic_idl_flytetypes.py @@ -23,7 +23,7 @@ def t1(dc: DC = DC(ff=FlyteFile(os.path.realpath(__file__)), uri="tests/flytekit/integration/remote/workflows/basic/data/df.parquet", file_format="parquet"), fd=FlyteDirectory("tests/flytekit/integration/remote/workflows/basic/data/") - )): + )) -> DC: with open(dc.ff, "r") as f: print("File Content: ", f.read()) diff --git a/tests/flytekit/integration/remote/workflows/basic/msgpack_idl_flytetypes.py b/tests/flytekit/integration/remote/workflows/basic/msgpack_idl_flytetypes.py index c157432f21..d679042dc1 100644 --- a/tests/flytekit/integration/remote/workflows/basic/msgpack_idl_flytetypes.py +++ b/tests/flytekit/integration/remote/workflows/basic/msgpack_idl_flytetypes.py @@ -23,7 +23,7 @@ def t1(dc: DC = DC(ff=FlyteFile(os.path.realpath(__file__)), uri="tests/flytekit/integration/remote/workflows/basic/data/df.parquet", file_format="parquet"), fd=FlyteDirectory("tests/flytekit/integration/remote/workflows/basic/data/") - )): + )) -> DC: with open(dc.ff, "r") as f: print("File Content: ", f.read()) diff --git a/tests/flytekit/unit/core/test_task_return.py b/tests/flytekit/unit/core/test_task_return.py new file mode 100644 index 0000000000..0e0df5901b --- /dev/null +++ b/tests/flytekit/unit/core/test_task_return.py @@ -0,0 +1,33 @@ +import pytest +from flytekit import task +from flytekit.exceptions.user import FlyteValueException +from typing import Optional + + +def test_task_return(): + @task + def foo(a: int) -> int: + return a + 1 + + assert foo(1) == 2 + + +def test_task_optional_return(): + @task + def foo(return_none: bool) -> Optional[int]: + return None if return_none else 1 + + assert foo(True) is None + assert foo(False) == 1 + + +def test_task_no_return(): + @task + def foo(a: int): + return a + 1 + + with pytest.raises( + FlyteValueException, + match="Interface has 0 outputs.", + ): + foo(1) diff --git a/tests/flytekit/unit/interactive/test_flyteinteractive_vscode.py b/tests/flytekit/unit/interactive/test_flyteinteractive_vscode.py index 792e030018..fe4935beda 100644 --- a/tests/flytekit/unit/interactive/test_flyteinteractive_vscode.py +++ b/tests/flytekit/unit/interactive/test_flyteinteractive_vscode.py @@ -47,7 +47,8 @@ def vscode_patches(): with mock.patch("multiprocessing.Process") as mock_process, mock.patch( "flytekit.interactive.vscode_lib.decorator.prepare_interactive_python" ) as mock_prepare_interactive_python, mock.patch( - "flytekit.interactive.vscode_lib.decorator.exit_handler" + "flytekit.interactive.vscode_lib.decorator.exit_handler", + return_value=None, ) as mock_exit_handler, mock.patch( "flytekit.interactive.vscode_lib.decorator.download_vscode" ) as mock_download_vscode, mock.patch("signal.signal") as mock_signal, mock.patch(