From 1e69da67819d5ead468016f10111d55d07c3d2dd Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 11 Feb 2022 11:52:09 -0800 Subject: [PATCH 1/5] adding changes to task and launch plan Signed-off-by: Yee Hing Tong --- flytekit/remote/launch_plan.py | 47 ++++++++++-- flytekit/remote/task.py | 37 +++++++++- flytekit/remote/workflow.py | 1 - tests/flytekit/unit/remote/test_calling.py | 86 ++++++++++++++++++++++ 4 files changed, 162 insertions(+), 9 deletions(-) create mode 100644 tests/flytekit/unit/remote/test_calling.py diff --git a/flytekit/remote/launch_plan.py b/flytekit/remote/launch_plan.py index b5dddea03f..865b891b5a 100644 --- a/flytekit/remote/launch_plan.py +++ b/flytekit/remote/launch_plan.py @@ -1,7 +1,9 @@ from typing import Optional from flytekit.core.interface import Interface +from flytekit.core.launch_plan import ReferenceLaunchPlan from flytekit.core.type_engine import TypeEngine +from flytekit.loggers import remote_logger as logger from flytekit.models import interface as _interface_models from flytekit.models import launch_plan as _launch_plan_models from flytekit.models.core import identifier as id_models @@ -18,8 +20,45 @@ def __init__(self, id, *args, **kwargs): # The interface is not set explicitly unless fetched in an engine context self._interface = None - self._python_interface = None + self._reference_entity = None + + @property + def interface(self) -> Optional[_interface_models.TypedInterface]: + return self._interface + + def __call__(self, *args, **kwargs): + if self.reference_entity is None: + logger.warning( + f"FlyteLaunchPlan {self} is not callable, most likely because flytekit could not " + f"guess the python interface. The workflow calling this launch plan may not behave correctly." + ) + return + return self.reference_entity(*args, **kwargs) + + # TODO: Refactor behind mixin + @property + def reference_entity(self) -> Optional[ReferenceLaunchPlan]: + if self._reference_entity is None: + if self.guessed_python_interface is None: + try: + self.guessed_python_interface = Interface( + TypeEngine.guess_python_types(self.interface.inputs), + TypeEngine.guess_python_types(self.interface.outputs), + ) + except Exception as e: + logger.warning(f"Error backing out interface {e}, Flyte interface {self.interface}") + return None + + self._reference_entity = ReferenceLaunchPlan( + self.id.project, + self.id.domain, + self.id.name, + self.id.version, + inputs=self.guessed_python_interface.inputs, + outputs=self.guessed_python_interface.outputs, + ) + return self._reference_entity @classmethod def promote_from_model( @@ -37,12 +76,6 @@ def promote_from_model( raw_output_data_config=model.raw_output_data_config, ) - if lp.interface is not None: - lp.guessed_python_interface = Interface( - inputs=TypeEngine.guess_python_types(lp.interface.inputs), - outputs=TypeEngine.guess_python_types(lp.interface.outputs), - ) - return lp @property diff --git a/flytekit/remote/task.py b/flytekit/remote/task.py index 1ff99549d6..34b4f1d9c2 100644 --- a/flytekit/remote/task.py +++ b/flytekit/remote/task.py @@ -2,8 +2,9 @@ from flytekit.core import hash as _hash_mixin from flytekit.core.interface import Interface +from flytekit.core.task import ReferenceTask from flytekit.core.type_engine import TypeEngine -from flytekit.loggers import logger +from flytekit.loggers import remote_logger as logger from flytekit.models import task as _task_model from flytekit.models.core import identifier as _identifier_model from flytekit.remote import interface as _interfaces @@ -24,6 +25,40 @@ def __init__(self, id, type, metadata, interface, custom, container=None, task_t config=config, ) self._python_interface = None + self._reference_entity = None + + def __call__(self, *args, **kwargs): + if self.reference_entity is None: + logger.warning( + f"FlyteTask {self} is not callable, most likely because flytekit could not " + f"guess the python interface. The workflow calling this task may not behave correctly" + ) + return + return self.reference_entity(*args, **kwargs) + + # TODO: Refactor behind mixin + @property + def reference_entity(self) -> Optional[ReferenceTask]: + if self._reference_entity is None: + if self.guessed_python_interface is None: + try: + self.guessed_python_interface = Interface( + TypeEngine.guess_python_types(self.interface.inputs), + TypeEngine.guess_python_types(self.interface.outputs), + ) + except Exception as e: + logger.warning(f"Error backing out interface {e}, Flyte interface {self.interface}") + return None + + self._reference_entity = ReferenceTask( + self.id.project, + self.id.domain, + self.id.name, + self.id.version, + inputs=self.guessed_python_interface.inputs, + outputs=self.guessed_python_interface.outputs, + ) + return self._reference_entity @property def interface(self) -> _interfaces.TypedInterface: diff --git a/flytekit/remote/workflow.py b/flytekit/remote/workflow.py index 010a703187..14ef6c91bf 100644 --- a/flytekit/remote/workflow.py +++ b/flytekit/remote/workflow.py @@ -57,7 +57,6 @@ def __init__( self._tasks = tasks self._launch_plans = launch_plans self._compiled_closure = compiled_closure - self._node_map = None @property diff --git a/tests/flytekit/unit/remote/test_calling.py b/tests/flytekit/unit/remote/test_calling.py new file mode 100644 index 0000000000..5c3f13c73a --- /dev/null +++ b/tests/flytekit/unit/remote/test_calling.py @@ -0,0 +1,86 @@ +import typing +from collections import OrderedDict + +import pytest + +from flytekit.core import context_manager +from flytekit.core.context_manager import Image, ImageConfig +from flytekit.core.launch_plan import LaunchPlan +from flytekit.core.reference_entity import ReferenceSpec +from flytekit.core.task import task +from flytekit.core.workflow import workflow +from flytekit.remote import FlyteTask, FlyteLaunchPlan +from flytekit.tools.translator import gather_dependent_entities, get_serializable + +default_img = Image(name="default", fqn="test", tag="tag") +serialization_settings = context_manager.SerializationSettings( + project="project", + domain="domain", + version="version", + env=None, + image_config=ImageConfig(default_image=default_img, images=[default_img]), +) + + +@task +def t1(a: int) -> int: + return a + 2 + + +@task +def t2(a: int, b: str) -> str: + return b + str(a) + + +@workflow +def sub_wf(a: int, b: str) -> (int, str): + x = t1(a=a) + d = t2(a=x, b=b) + return x, d + + +serialized = OrderedDict() +t1_spec = get_serializable(serialized, serialization_settings, t1) +ft = FlyteTask.promote_from_model(t1_spec.template) + + + + +def test_fetched_task(): + @workflow + def wf(a: int) -> int: + return ft(a=a) + + # Should not work unless mocked out. + with pytest.raises(Exception, match="cannot be run locally"): + wf(a=3) + + # Should have one reference entity + serialized = OrderedDict() + get_serializable(serialized, serialization_settings, wf) + vals = [v for v in serialized.values()] + refs = [f for f in filter(lambda x: isinstance(x, ReferenceSpec), vals)] + assert len(refs) == 1 + + +def test_calling_lp(): + sub_wf_lp = LaunchPlan.get_or_create(sub_wf) + serialized = OrderedDict() + lp_model = get_serializable(serialized, serialization_settings, sub_wf_lp) + task_templates, wf_specs, lp_specs = gather_dependent_entities(serialized) + for wf_id, spec in wf_specs.items(): + break + + remote_lp = FlyteLaunchPlan.promote_from_model(lp_model.id, lp_model.spec) + # To pretend that we've fetched this launch plan from Admin, also fill in the Flyte interface, which isn't + # part of the IDL object but is something FlyteRemote does + remote_lp._interface = spec.template.interface + serialized = OrderedDict() + + @workflow + def wf2(a: int) -> typing.Tuple[int, str]: + return remote_lp(a=a, b="hello") + + wf_spec = get_serializable(serialized, serialization_settings, wf2) + print(wf_spec.template.nodes[0].workflow_node.launchplan_ref) + assert wf_spec.template.nodes[0].workflow_node.launchplan_ref == lp_model.id From 84fa52074163b8482ea3196963807eae823f048d Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 11 Feb 2022 11:53:25 -0800 Subject: [PATCH 2/5] nit Signed-off-by: Yee Hing Tong --- tests/flytekit/unit/remote/test_calling.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/flytekit/unit/remote/test_calling.py b/tests/flytekit/unit/remote/test_calling.py index 5c3f13c73a..d07c5f749f 100644 --- a/tests/flytekit/unit/remote/test_calling.py +++ b/tests/flytekit/unit/remote/test_calling.py @@ -9,7 +9,7 @@ from flytekit.core.reference_entity import ReferenceSpec from flytekit.core.task import task from flytekit.core.workflow import workflow -from flytekit.remote import FlyteTask, FlyteLaunchPlan +from flytekit.remote import FlyteLaunchPlan, FlyteTask from flytekit.tools.translator import gather_dependent_entities, get_serializable default_img = Image(name="default", fqn="test", tag="tag") @@ -44,8 +44,6 @@ def sub_wf(a: int, b: str) -> (int, str): ft = FlyteTask.promote_from_model(t1_spec.template) - - def test_fetched_task(): @workflow def wf(a: int) -> int: From 8b51e7392085c57b8b2a314dfb2d8ff5bacc69b8 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 11 Feb 2022 13:29:20 -0800 Subject: [PATCH 3/5] remove duplicate Signed-off-by: Yee Hing Tong --- flytekit/remote/launch_plan.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/flytekit/remote/launch_plan.py b/flytekit/remote/launch_plan.py index 865b891b5a..3bf845fcab 100644 --- a/flytekit/remote/launch_plan.py +++ b/flytekit/remote/launch_plan.py @@ -23,10 +23,6 @@ def __init__(self, id, *args, **kwargs): self._python_interface = None self._reference_entity = None - @property - def interface(self) -> Optional[_interface_models.TypedInterface]: - return self._interface - def __call__(self, *args, **kwargs): if self.reference_entity is None: logger.warning( @@ -98,7 +94,7 @@ def workflow_id(self) -> id_models.Identifier: return self._workflow_id @property - def interface(self) -> _interface.TypedInterface: + def interface(self) -> Optional[_interface.TypedInterface]: """ The interface is not technically part of the admin.LaunchPlanSpec in the IDL, however the workflow ID is, and from the workflow ID, fetch will fill in the interface. This is nice because then you can __call__ the= From 7f3fde3e616d8ac5e6349a0072410125fced6ffc Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 11 Feb 2022 13:32:02 -0800 Subject: [PATCH 4/5] use wrapper type Signed-off-by: Yee Hing Tong --- tests/flytekit/unit/remote/test_calling.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/flytekit/unit/remote/test_calling.py b/tests/flytekit/unit/remote/test_calling.py index d07c5f749f..97a1a001dc 100644 --- a/tests/flytekit/unit/remote/test_calling.py +++ b/tests/flytekit/unit/remote/test_calling.py @@ -10,6 +10,7 @@ from flytekit.core.task import task from flytekit.core.workflow import workflow from flytekit.remote import FlyteLaunchPlan, FlyteTask +from flytekit.remote.interface import TypedInterface from flytekit.tools.translator import gather_dependent_entities, get_serializable default_img = Image(name="default", fqn="test", tag="tag") @@ -72,7 +73,7 @@ def test_calling_lp(): remote_lp = FlyteLaunchPlan.promote_from_model(lp_model.id, lp_model.spec) # To pretend that we've fetched this launch plan from Admin, also fill in the Flyte interface, which isn't # part of the IDL object but is something FlyteRemote does - remote_lp._interface = spec.template.interface + remote_lp._interface = TypedInterface.promote_from_model(spec.template.interface) serialized = OrderedDict() @workflow From 790259e74713b450015623f57dc700fb3e989875 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 11 Feb 2022 15:09:19 -0800 Subject: [PATCH 5/5] wip Signed-off-by: Yee Hing Tong --- plugins/flytekit-greatexpectations/setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/flytekit-greatexpectations/setup.py b/plugins/flytekit-greatexpectations/setup.py index 3541415664..ac7d9481ef 100644 --- a/plugins/flytekit-greatexpectations/setup.py +++ b/plugins/flytekit-greatexpectations/setup.py @@ -4,7 +4,7 @@ microlib_name = f"flytekitplugins-{PLUGIN_NAME}" -plugin_requires = ["flytekit>=0.21.0,<1.0.0", "great-expectations>=0.13.30", "sqlalchemy>=1.4.23"] +plugin_requires = ["flytekit>=0.21.0,<1.0.0", "great-expectations>=0.13.30,<0.14.6", "sqlalchemy>=1.4.23"] __version__ = "0.0.0+develop"