Skip to content

Commit

Permalink
Restrict eager task & add unit tests (#2871)
Browse files Browse the repository at this point in the history
Signed-off-by: Mecoli1219 <[email protected]>
  • Loading branch information
Mecoli1219 authored and kumare3 committed Nov 8, 2024
1 parent 6005b34 commit 32d6479
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 0 deletions.
4 changes: 4 additions & 0 deletions flytekit/remote/remote.py
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,10 @@ def _get_pickled_target_dict(
raise FlyteAssertion(
f"Dynamic tasks are not supported in interactive mode. {entity.name} is a dynamic task."
)
if entity.execution_mode == PythonFunctionTask.ExecutionBehavior.EAGER:
raise FlyteAssertion(
f"Eager tasks are not supported in interactive mode. {entity.name} is an eager task."
)

if isinstance(entity, PythonTask):
if isinstance(entity, (PythonAutoContainerTask, ArrayNodeMapTask)):
Expand Down
20 changes: 20 additions & 0 deletions tests/flytekit/unit/remote/test_remote.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from flytekit.core.type_engine import TypeEngine
from flytekit.exceptions import user as user_exceptions
from flytekit.exceptions.user import FlyteEntityNotExistException, FlyteAssertion
from flytekit.experimental.eager_function import eager
from flytekit.models import common as common_models
from flytekit.models import security
from flytekit.models.admin.workflow import Workflow, WorkflowClosure
Expand Down Expand Up @@ -752,3 +753,22 @@ def my_wf(a: int) -> typing.List[str]:

with pytest.raises(FlyteAssertion):
_get_pickled_target_dict(my_wf)

def test_get_pickled_target_dict_with_eager():
@task
def t1(a: int) -> int:
return a + 1

@task
def t2(a: int) -> int:
return a * 2

@eager
async def eager_wf(a: int) -> int:
out = await t1(a=a)
if out < 0:
return -1
return await t2(a=out)

with pytest.raises(FlyteAssertion):
_get_pickled_target_dict(eager_wf)

0 comments on commit 32d6479

Please sign in to comment.