Finish dynamic yielding of launch plans (#92)

Author: wild-endeavor

flytekit/__init__.py

@@ -1,4 +1,4 @@
 from __future__ import absolute_import
 import flytekit.plugins
 
-__version__ = '0.6.2'
+__version__ = '0.7.0b1'

flytekit/common/component_nodes.py

@@ -65,9 +65,24 @@ def __init__(self, sdk_workflow=None, sdk_launch_plan=None):
         :param flytekit.common.workflow.SdkWorkflow sdk_workflow:
         :param flytekit.common.launch_plan.SdkLaunchPlan sdk_launch_plan:
         """
+        if sdk_workflow and sdk_launch_plan:
+            raise _system_exceptions.FlyteSystemException("SdkWorkflowNode cannot be called with both a workflow and "
+                                                          "a launchplan specified, please pick one. WF: {} LP: {}",
+                                                          sdk_workflow, sdk_launch_plan)
+
         self._sdk_workflow = sdk_workflow
         self._sdk_launch_plan = sdk_launch_plan
-        super(SdkWorkflowNode, self).__init__()
+        sdk_wf_id = sdk_workflow.id if sdk_workflow else None
+        sdk_lp_id = sdk_launch_plan.id if sdk_launch_plan else None
+        super(SdkWorkflowNode, self).__init__(launchplan_ref=sdk_lp_id, sub_workflow_ref=sdk_wf_id)
+
+    def __repr__(self):
+        """
+        :rtype: Text
+        """
+        if self.sdk_workflow is not None:
+            return "SdkWorkflowNode with workflow: {}".format(self.sdk_workflow)
+        return "SdkWorkflowNode with launch plan: {}".format(self.sdk_launch_plan)
 
     @property
     def launchplan_ref(self):

flytekit/common/launch_plan.py

@@ -27,6 +27,7 @@ class SdkLaunchPlan(
 ):
     def __init__(self, *args, **kwargs):
         super(SdkLaunchPlan, self).__init__(*args, **kwargs)
+        # Set all the attributes we expect this class to have
         self._id = None
 
         # The interface is not set explicitly unless fetched in an engine context
@@ -79,26 +80,6 @@ def fetch(cls, project, domain, name, version=None):
         sdk_lp._interface = lp_wf.interface
         return sdk_lp
 
-    @_exception_scopes.system_entry_point
-    def register(self, project, domain, name, version):
-        """
-        :param Text project:
-        :param Text domain:
-        :param Text name:
-        :param Text version:
-        """
-        self.validate()
-        id_to_register = _identifier.Identifier(
-            _identifier_model.ResourceType.LAUNCH_PLAN,
-            project,
-            domain,
-            name,
-            version
-        )
-        _engine_loader.get_engine().get_launch_plan(self).register(id_to_register)
-        self._id = id_to_register
-        return _six.text_type(self.id)
-
     @property
     def id(self):
         """
@@ -225,12 +206,43 @@ def execute_with_literals(self, project, domain, literal_inputs, name=None, noti
 
     @_exception_scopes.system_entry_point
     def __call__(self, *args, **input_map):
-        raise _user_exceptions.FlyteAssertion(
-            "TODO: Implement adding of remote launch plans to workflows. Current workaround is to add remote "
-            "workflows directly."
+        """
+        :param list[T] args: Do not specify.  Kwargs only are supported for this function.
+        :param dict[Text,T] input_map: Map of inputs.  Can be statically defined or OutputReference links.
+        :rtype: flytekit.common.nodes.SdkNode
+        """
+        if len(args) > 0:
+            raise _user_exceptions.FlyteAssertion(
+                "When adding a launchplan as a node in a workflow, all inputs must be specified with kwargs only.  We "
+                "detected {} positional args.".format(self, len(args))
+            )
+
+        # Take the default values from the launch plan
+        default_inputs = {
+            k: v.sdk_default
+            for k, v in _six.iteritems(self.default_inputs.parameters) if not v.required
+        }
+        default_inputs.update(input_map)
+
+        bindings, upstream_nodes = self.interface.create_bindings_for_inputs(default_inputs)
+
+        return _nodes.SdkNode(
+            id=None,
+            metadata=_workflow_models.NodeMetadata("", _datetime.timedelta(), _literal_models.RetryStrategy(0)),
+            bindings=sorted(bindings, key=lambda b: b.var),
+            upstream_nodes=upstream_nodes,
+            sdk_launch_plan=self
         )
 
+    def __repr__(self):
+        """
+        :rtype: Text
+        """
+        return "SdkLaunchPlan(ID: {} Interface: {} WF ID: {})".format(self.id, self.interface, self.workflow_id)
+
 
+# The difference between this and the SdkLaunchPlan class is that this runnable class is supposed to only be used for
+# launch plans loaded alongside the current Python interpreter.
 class SdkRunnableLaunchPlan(
     _hash_mixin.HashOnReferenceMixin,
     SdkLaunchPlan,
@@ -272,14 +284,14 @@ def __init__(
         if role:
             auth = _launch_plan_models.Auth(assumable_iam_role=role)
 
+        # The constructor for SdkLaunchPlan sets the id to None anyways so we don't bother passing in an ID. The ID
+        # should be set in one of three places,
+        #   1) When the object is registered (in the code above)
+        #   2) By the dynamic task code after this runnable object has already been __call__'ed. The SdkNode produced
+        #      maintains a link to this object and will set the ID according to the configuration variables present.
+        #   3) When SdkLaunchPlan.fetch() is run
         super(SdkRunnableLaunchPlan, self).__init__(
-            _identifier.Identifier(
-                _identifier_model.ResourceType.WORKFLOW,
-                _internal_config.PROJECT.get(),
-                _internal_config.DOMAIN.get(),
-                _uuid.uuid4().hex,
-                _internal_config.VERSION.get()
-            ),
+            None,
             _launch_plan_models.LaunchPlanMetadata(
                 schedule=schedule or _schedule_model.Schedule(''),
                 notifications=notifications or []
@@ -303,6 +315,26 @@ def __init__(
         self._upstream_entities = {sdk_workflow}
         self._sdk_workflow = sdk_workflow
 
+    @_exception_scopes.system_entry_point
+    def register(self, project, domain, name, version):
+        """
+        :param Text project:
+        :param Text domain:
+        :param Text name:
+        :param Text version:
+        """
+        self.validate()
+        id_to_register = _identifier.Identifier(
+            _identifier_model.ResourceType.LAUNCH_PLAN,
+            project,
+            domain,
+            name,
+            version
+        )
+        _engine_loader.get_engine().get_launch_plan(self).register(id_to_register)
+        self._id = id_to_register
+        return _six.text_type(self.id)
+
     @classmethod
     def from_flyte_idl(cls, _):
         raise _user_exceptions.FlyteAssertion(
@@ -356,33 +388,8 @@ def workflow_id(self):
         """
         return self._sdk_workflow.id
 
-    @_exception_scopes.system_entry_point
-    def __call__(self, *args, **input_map):
+    def __repr__(self):
         """
-        :param list[T] args: Do not specify.  Kwargs only are supported for this function.
-        :param dict[Text,T] input_map: Map of inputs.  Can be statically defined or OutputReference links.
-        :rtype: flytekit.common.nodes.SdkNode
+        :rtype: Text
         """
-        if len(args) > 0:
-            raise _user_exceptions.FlyteAssertion(
-                "When adding a launchplan as a node in a workflow, all inputs must be specified with kwargs only.  We "
-                "detected {} positional args.".format(self, len(args))
-            )
-
-        # Take the default values from the launch plan
-        default_inputs = {
-            k: v.sdk_default
-            for k, v in _six.iteritems(self.default_inputs.parameters) if not v.required
-        }
-        default_inputs.update(input_map)
-
-        bindings, upstream_nodes = self.interface.create_bindings_for_inputs(default_inputs)
-
-        # TODO: Remove DEADBEEF
-        return _nodes.SdkNode(
-            id=None,
-            metadata=_workflow_models.NodeMetadata("DEADBEEF", _datetime.timedelta(), _literal_models.RetryStrategy(0)),
-            bindings=sorted(bindings, key=lambda b: b.var),
-            upstream_nodes=upstream_nodes,
-            sdk_launch_plan=self
-        )
+        return "SdkRunnableLaunchPlan(ID: {} Interface: {} WF ID: {})".format(self.id, self.interface, self.workflow_id)

flytekit/common/mixins/registerable.py

@@ -2,7 +2,12 @@
 import abc as _abc
 import inspect as _inspect
 import six as _six
+import importlib as _importlib
+import logging as _logging
+
 from flytekit.common import sdk_bases as _sdk_bases
+from flytekit.common.exceptions import system as _system_exceptions
+from flytekit.common import utils as _utils
 
 
 class _InstanceTracker(_sdk_bases.ExtendedSdkType):
@@ -33,6 +38,10 @@ def __call__(cls, *args, **kwargs):
 
 class RegisterableEntity(_six.with_metaclass(_InstanceTracker, object)):
 
+    def __init__(self, *args, **kwargs):
+        self._platform_valid_name = None
+        super(RegisterableEntity, self).__init__(*args, **kwargs)
+
     @_abc.abstractmethod
     def register(self, project, domain, name, version):
         """
@@ -73,3 +82,67 @@ def instantiated_in(self):
         :rtype: Optional[Text]
         """
         return self._instantiated_in
+
+    @property
+    def has_valid_name(self):
+        """
+        :rtype: bool
+        """
+        return self._platform_valid_name is not None and self._platform_valid_name != ""
+
+    @property
+    def platform_valid_name(self):
+        """
+        :rtype: Text
+        """
+        return self._platform_valid_name
+
+    def auto_assign_name(self):
+        """
+        This function is a bit of trickster Python code that goes hand in hand with the _InstanceTracker metaclass
+        defined above. Thanks @matthewphsmith for this bit of ingenuity.
+
+        For instance, if a user has code that looks like this:
+
+            from some.other.module import wf
+            my_launch_plan = wf.create_launch_plan()
+
+            @dynamic_task
+            def sample_task(wf_params):
+                yield my_launch_plan()
+
+        This code means that we should have a launch plan with a name ending in "my_launch_plan", since that is the
+        name of the variable that the created launch plan gets assigned to. That is also the name that the launch plan
+        would be registered with.
+
+        However, when the create_launch_plan() function runs, the Python interpreter has no idea where the created
+        object will be assigned to. It has no idea that the output of the create_launch_plan call is to be paired up
+        with a variable named "my_launch_plan". This function basically does this after the fact. Leveraging the
+        _instantiated_in field provided by the _InstanceTracker class above, this code will re-import the
+        module (ie Python file) that the object is in. Since it's already loaded, it's just retrieved from memory.
+        It then scans all objects in the module, and when an object match is found, it knows it's found the right
+        variable name.
+
+        Just to drive the point home, this function is mostly needed for Launch Plans. Assuming that user code has:
+
+            @python_task
+            def some_task()
+
+        When Flytekit calls the module loader and loads the task, the name of the task is the name of the function
+        itself.  It's known at time of creation. In contrast, when
+
+            xyz = SomeWorflow.create_launch_plan()
+
+        is called, the name of the launch plan isn't known until after creation, it's not "SomeWorkflow", it's "xyz"
+        """
+        _logging.debug("Running name auto assign")
+        m = _importlib.import_module(self.instantiated_in)
+
+        for k in dir(m):
+            if getattr(m, k) == self:
+                self._platform_valid_name = _utils.fqdn(m.__name__, k, entity_type=self.resource_type)
+                _logging.debug("Auto-assigning name to {}".format(self._platform_valid_name))
+                return
+
+        _logging.error("Could not auto-assign name")
+        raise _system_exceptions.FlyteSystemException("Error looking for object while auto-assigning name.")

flytekit/common/nodes.py

@@ -292,7 +292,7 @@ def __repr__(self):
         """
         :rtype: Text
         """
-        return "Node({})".format(self._executable_sdk_object)
+        return "Node(ID: {} Executable: {})".format(self.id, self._executable_sdk_object)
 
 
 class SdkNodeExecution(