diff --git a/flytekit/clis/flyte_cli/main.py b/flytekit/clis/flyte_cli/main.py index 82300a9c7f..6415e92b15 100644 --- a/flytekit/clis/flyte_cli/main.py +++ b/flytekit/clis/flyte_cli/main.py @@ -314,7 +314,6 @@ def _render_schedule_expr(lp): ) _insecure_option = _click.option(*_INSECURE_FLAGS, is_flag=True, help="Do not use SSL") _urn_option = _click.option("-u", "--urn", required=True, help="The unique identifier for an entity.") - _optional_urn_option = _click.option("-u", "--urn", required=False, help="The unique identifier for an entity.") _host_option = _click.option( diff --git a/flytekit/models/core/workflow.py b/flytekit/models/core/workflow.py index 6338a6695f..e9e68ce485 100644 --- a/flytekit/models/core/workflow.py +++ b/flytekit/models/core/workflow.py @@ -99,7 +99,7 @@ def else_node(self): def error(self): """ An error to throw in case none of the branches were taken. - :rtype: flytekit.models.core.errors.ContainerError + :rtype: flytekit.models.types.Error """ return self._error diff --git a/flytekit/remote/component_nodes.py b/flytekit/remote/component_nodes.py index 877a6d6494..2d6acca1ff 100644 --- a/flytekit/remote/component_nodes.py +++ b/flytekit/remote/component_nodes.py @@ -41,7 +41,7 @@ def promote_from_model( if base_model.reference_id in tasks: task = tasks[base_model.reference_id] - _logging.info(f"Found existing task template for {task.id}, will not retrieve from Admin") + _logging.debug(f"Found existing task template for {task.id}, will not retrieve from Admin") flyte_task = FlyteTask.promote_from_model(task) return cls(flyte_task) @@ -124,3 +124,36 @@ def promote_from_model( raise _system_exceptions.FlyteSystemException( "Bad workflow node model, neither subworkflow nor launchplan specified." ) + + +class FlyteBranchNode(_workflow_model.BranchNode): + def __init__(self, if_else: _workflow_model.IfElseBlock): + super().__init__(if_else) + + @classmethod + def promote_from_model( + cls, + base_model: _workflow_model.BranchNode, + sub_workflows: Dict[id_models.Identifier, _workflow_model.WorkflowTemplate], + node_launch_plans: Dict[id_models.Identifier, _launch_plan_model.LaunchPlanSpec], + tasks: Dict[id_models.Identifier, _task_model.TaskTemplate], + ) -> "FlyteBranchNode": + + from flytekit.remote.nodes import FlyteNode + + block = base_model.if_else + + else_node = None + if block.else_node: + else_node = FlyteNode.promote_from_model(block.else_node, sub_workflows, node_launch_plans, tasks) + + block.case._then_node = FlyteNode.promote_from_model( + block.case.then_node, sub_workflows, node_launch_plans, tasks + ) + + for o in block.other: + o._then_node = FlyteNode.promote_from_model(o.then_node, sub_workflows, node_launch_plans, tasks) + + new_if_else_block = _workflow_model.IfElseBlock(block.case, block.other, else_node, block.error) + + return cls(new_if_else_block) diff --git a/flytekit/remote/launch_plan.py b/flytekit/remote/launch_plan.py index 3bf845fcab..7a3b8d915b 100644 --- a/flytekit/remote/launch_plan.py +++ b/flytekit/remote/launch_plan.py @@ -56,6 +56,13 @@ def reference_entity(self) -> Optional[ReferenceLaunchPlan]: ) return self._reference_entity + # If fetched when creating this object, can store it here. + self._flyte_workflow = None + + @property + def flyte_workflow(self) -> Optional["FlyteWorkflow"]: + return self._flyte_workflow + @classmethod def promote_from_model( cls, id: id_models.Identifier, model: _launch_plan_models.LaunchPlanSpec diff --git a/flytekit/remote/nodes.py b/flytekit/remote/nodes.py index 4131905c5b..e7ee0dc487 100644 --- a/flytekit/remote/nodes.py +++ b/flytekit/remote/nodes.py @@ -27,15 +27,19 @@ def __init__( flyte_task: Optional["FlyteTask"] = None, flyte_workflow: Optional["FlyteWorkflow"] = None, flyte_launch_plan: Optional["FlyteLaunchPlan"] = None, - flyte_branch=None, + flyte_branch_node: Optional["FlyteBranchNode"] = None, ): - non_none_entities = list(filter(None, [flyte_task, flyte_workflow, flyte_launch_plan, flyte_branch])) + # todo: flyte_branch_node is the only non-entity here, feels wrong, it should probably be a Condition + # or the other ones changed. + non_none_entities = list(filter(None, [flyte_task, flyte_workflow, flyte_launch_plan, flyte_branch_node])) if len(non_none_entities) != 1: raise _user_exceptions.FlyteAssertion( "An Flyte node must have one underlying entity specified at once. Received the following " "entities: {}".format(non_none_entities) ) - self._flyte_entity = flyte_task or flyte_workflow or flyte_launch_plan or flyte_branch + # todo: wip - flyte_branch_node is a hack, it should be a Condition, but backing out a Condition object from + # the compiled IfElseBlock is cumbersome, shouldn't do it if we can get away with it. + self._flyte_entity = flyte_task or flyte_workflow or flyte_launch_plan or flyte_branch_node workflow_node = None if flyte_workflow is not None: @@ -46,7 +50,6 @@ def __init__( task_node = None if flyte_task: task_node = _component_nodes.FlyteTaskNode(flyte_task) - branch_node = None super(FlyteNode, self).__init__( id=id, @@ -56,7 +59,7 @@ def __init__( output_aliases=[], task_node=task_node, workflow_node=workflow_node, - branch_node=branch_node, + branch_node=flyte_branch_node, ) self._upstream = upstream_nodes @@ -78,7 +81,7 @@ def promote_from_model( _logging.warning(f"Should not call promote from model on a start node or end node {model}") return None - flyte_task_node, flyte_workflow_node = None, None + flyte_task_node, flyte_workflow_node, flyte_branch_node = None, None, None if model.task_node is not None: flyte_task_node = _component_nodes.FlyteTaskNode.promote_from_model(model.task_node, tasks) elif model.workflow_node is not None: @@ -88,7 +91,10 @@ def promote_from_model( node_launch_plans, tasks, ) - # TODO: Implement branch node https://github.com/flyteorg/flyte/issues/1116 + elif model.branch_node is not None: + flyte_branch_node = _component_nodes.FlyteBranchNode.promote_from_model( + model.branch_node, sub_workflows, node_launch_plans, tasks + ) else: raise _system_exceptions.FlyteSystemException( f"Bad Node model, neither task nor workflow detected, node: {model}" @@ -132,6 +138,14 @@ def promote_from_model( raise _system_exceptions.FlyteSystemException( "Bad FlyteWorkflowNode model, both launch plan and workflow are None" ) + elif flyte_branch_node is not None: + return cls( + id=node_model_id, + upstream_nodes=[], # set downstream, model doesn't contain this information + bindings=model.inputs, + metadata=model.metadata, + flyte_branch_node=flyte_branch_node, + ) raise _system_exceptions.FlyteSystemException("Bad FlyteNode model, both task and workflow nodes are empty") @property diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 93adc04db8..50d29e8e23 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -440,6 +440,7 @@ def fetch_launch_plan( wf_id = flyte_launch_plan.workflow_id workflow = self.fetch_workflow(wf_id.project, wf_id.domain, wf_id.name, wf_id.version) flyte_launch_plan._interface = workflow.interface + flyte_launch_plan._flyte_workflow = workflow flyte_launch_plan.guessed_python_interface = Interface( inputs=TypeEngine.guess_python_types(flyte_launch_plan.interface.inputs), outputs=TypeEngine.guess_python_types(flyte_launch_plan.interface.outputs), @@ -1053,6 +1054,7 @@ def sync_workflow_execution( if execution.spec.launch_plan.resource_type == ResourceType.TASK: # This condition is only true for single-task executions flyte_entity = self.fetch_task(lp_id.project, lp_id.domain, lp_id.name, lp_id.version) + node_interface = flyte_entity.interface if sync_nodes: # Need to construct the mapping. There should've been returned exactly three nodes, a start, # an end, and a task node. @@ -1080,10 +1082,10 @@ def sync_workflow_execution( ) else: # This is the default case, an execution of a normal workflow through a launch plan - wf_id = self.fetch_launch_plan(lp_id.project, lp_id.domain, lp_id.name, lp_id.version).workflow_id - flyte_entity = self.fetch_workflow(wf_id.project, wf_id.domain, wf_id.name, wf_id.version) - execution._flyte_workflow = flyte_entity - node_mapping = flyte_entity._node_map + fetched_lp = self.fetch_launch_plan(lp_id.project, lp_id.domain, lp_id.name, lp_id.version) + node_interface = fetched_lp.flyte_workflow.interface + execution._flyte_workflow = fetched_lp.flyte_workflow + node_mapping = fetched_lp.flyte_workflow._node_map # update node executions (if requested), and inputs/outputs if sync_nodes: @@ -1091,10 +1093,12 @@ def sync_workflow_execution( for n in underlying_node_executions: node_execs[n.id.node_id] = self.sync_node_execution(n, node_mapping) execution._node_executions = node_execs - return self._assign_inputs_and_outputs(execution, execution_data, flyte_entity.interface) + return self._assign_inputs_and_outputs(execution, execution_data, node_interface) def sync_node_execution( - self, execution: FlyteNodeExecution, node_mapping: typing.Dict[str, FlyteNode] + self, + execution: FlyteNodeExecution, + node_mapping: typing.Dict[str, FlyteNode], ) -> FlyteNodeExecution: """ Get data backing a node execution. These FlyteNodeExecution objects should've come from Admin with the model @@ -1201,13 +1205,9 @@ def sync_node_execution( for t in iterate_task_executions(self.client, execution.id) ] execution._interface = dynamic_flyte_wf.interface - else: - # If it does not, then it should be a static subworkflow - if not isinstance(execution._node.flyte_entity, FlyteWorkflow): - remote_logger.error( - f"NE {execution} entity should be a workflow, {type(execution._node)}, {execution._node}" - ) - raise Exception(f"Node entity has type {type(execution._node)}") + + # Handle the case where it's a static subworkflow + elif isinstance(execution._node.flyte_entity, FlyteWorkflow): sub_flyte_workflow = execution._node.flyte_entity sub_node_mapping = {n.id: n for n in sub_flyte_workflow.flyte_nodes} execution._underlying_node_executions = [ @@ -1216,6 +1216,14 @@ def sync_node_execution( ] execution._interface = sub_flyte_workflow.interface + # Handle the case where it's a branch node + elif execution._node.branch_node is not None: + remote_logger.debug("Skipping remote node execution for now") + return execution + else: + remote_logger.error(f"NE {execution} undeterminable, {type(execution._node)}, {execution._node}") + raise Exception(f"Node execution undeterminable, entity has type {type(execution._node)}") + # This is the plain ol' task execution case else: execution._task_executions = [