Add reference entities to FlyteTask and FlyteLaunchPlan

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,41 @@ 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
+
+    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 +72,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
@@ -65,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=

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:

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

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"
 

tests/flytekit/unit/remote/test_calling.py

@@ -0,0 +1,85 @@
+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 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")
+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 = TypedInterface.promote_from_model(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