From 1e69da67819d5ead468016f10111d55d07c3d2dd Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
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 <wild-endeavor@users.noreply.github.com>
---
 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 <wild-endeavor@users.noreply.github.com>
Date: Fri, 11 Feb 2022 11:53:25 -0800
Subject: [PATCH 2/5] nit

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 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 <wild-endeavor@users.noreply.github.com>
Date: Fri, 11 Feb 2022 13:29:20 -0800
Subject: [PATCH 3/5] remove duplicate

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 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 <wild-endeavor@users.noreply.github.com>
Date: Fri, 11 Feb 2022 13:32:02 -0800
Subject: [PATCH 4/5] use wrapper type

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 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 <wild-endeavor@users.noreply.github.com>
Date: Fri, 11 Feb 2022 15:09:19 -0800
Subject: [PATCH 5/5] wip

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 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"
 
