def test_hydrate_registration_parameters__workflow_nothing_set(): workflow = _workflow_pb2.WorkflowSpec( template=_core_workflow_pb2.WorkflowTemplate( id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, name="name", ), nodes=[ _core_workflow_pb2.Node( id="foo", task_node=_core_workflow_pb2.TaskNode( reference_id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.TASK, name="task1")), ) ], )) identifier, entity = hydrate_registration_parameters( _identifier_pb2.WORKFLOW, "project", "domain", "12345", workflow) assert (identifier == _identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project", domain="domain", name="name", version="12345", ) == entity.template.id) assert len(workflow.template.nodes) == 1 assert workflow.template.nodes[ 0].task_node.reference_id == _identifier_pb2.Identifier( resource_type=_identifier_pb2.TASK, project="project", domain="domain", name="task1", version="12345", )
def to_flyte_idl(self): """ :rtype: flyteidl.core.workflow_pb2.WorkflowTemplate """ return _core_workflow.WorkflowTemplate( id=self.id.to_flyte_idl(), metadata=self.metadata.to_flyte_idl(), interface=self.interface.to_flyte_idl(), nodes=[n.to_flyte_idl() for n in self.nodes], outputs=[o.to_flyte_idl() for o in self.outputs], failure_node=self.failure_node.to_flyte_idl() if self.failure_node is not None else None )
def test_hydrate_registration_parameters__subworkflows(): workflow_template = _core_workflow_pb2.WorkflowTemplate() workflow_template.id.CopyFrom( _identifier_pb2.Identifier(resource_type=_identifier_pb2.WORKFLOW, name="workflow")) sub_workflow_template = _core_workflow_pb2.WorkflowTemplate() sub_workflow_template.id.CopyFrom( _identifier_pb2.Identifier(resource_type=_identifier_pb2.WORKFLOW, name="subworkflow")) sub_workflow_template.nodes.append( _core_workflow_pb2.Node( id="task_node", task_node=_core_workflow_pb2.TaskNode( reference_id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.TASK)), )) workflow_spec = _workflow_pb2.WorkflowSpec(template=workflow_template) workflow_spec.sub_workflows.append(sub_workflow_template) identifier, entity = hydrate_registration_parameters( _identifier_pb2.WORKFLOW, "project", "domain", "12345", workflow_spec) assert (identifier == _identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project", domain="domain", name="workflow", version="12345", ) == entity.template.id) assert entity.sub_workflows[0].id == _identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project", domain="domain", name="subworkflow", version="12345", )
def get_workflow_template(): """ This function retrieves a TasKTemplate object from the pb file in the resources directory. It was created by reading from Flyte Admin, the following workflow, after registration. from __future__ import absolute_import from flytekit.common.types.primitives import Integer from flytekit.sdk.tasks import ( python_task, inputs, outputs, ) from flytekit.sdk.types import Types from flytekit.sdk.workflow import workflow_class, Input, Output @inputs(a=Types.Integer) @outputs(b=Types.Integer, c=Types.Integer) @python_task() def demo_task_for_promote(wf_params, a, b, c): b.set(a + 1) c.set(a + 2) @workflow_class() class OneTaskWFForPromote(object): wf_input = Input(Types.Integer, required=True) my_task_node = demo_task_for_promote(a=wf_input) wf_output_b = Output(my_task_node.outputs.b, sdk_type=Integer) wf_output_c = Output(my_task_node.outputs.c, sdk_type=Integer) :rtype: flytekit.models.core.workflow.WorkflowTemplate """ workflow_template_pb = _workflow_pb2.WorkflowTemplate() # So that tests that use this work when run from any directory basepath = _path.dirname(__file__) filepath = _path.abspath( _path.join(basepath, "resources/protos", "OneTaskWFForPromote.pb")) with open(filepath, "rb") as fh: workflow_template_pb.ParseFromString(fh.read()) wt = _workflow_model.WorkflowTemplate.from_flyte_idl(workflow_template_pb) return wt
def test_hydrate_registration_parameters__workflow_already_set(): workflow = _workflow_pb2.WorkflowSpec( template=_core_workflow_pb2.WorkflowTemplate( id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project2", domain="domain2", name="name", version="abc", ), )) identifier, entity = hydrate_registration_parameters( _identifier_pb2.WORKFLOW, "project", "domain", "12345", workflow) assert (identifier == _identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project2", domain="domain2", name="name", version="abc", ) == entity.template.id)
def test_hydrate_workflow_template__branch_node(): workflow_template = _core_workflow_pb2.WorkflowTemplate() branch_node = _core_workflow_pb2.Node( id="branch_node", branch_node=_core_workflow_pb2.BranchNode( if_else=_core_workflow_pb2.IfElseBlock( case=_core_workflow_pb2.IfBlock( then_node=_core_workflow_pb2.Node( task_node=_core_workflow_pb2.TaskNode( reference_id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.TASK, name="if_case"), ), )), else_node=_core_workflow_pb2.Node( task_node=_core_workflow_pb2.TaskNode( reference_id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.TASK, name="else_node"), ), ), ), ), )
def test_hydrate_workflow_template(): workflow_template = _core_workflow_pb2.WorkflowTemplate() workflow_template.nodes.append( _core_workflow_pb2.Node( id="task_node", task_node=_core_workflow_pb2.TaskNode( reference_id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.TASK)), )) workflow_template.nodes.append( _core_workflow_pb2.Node( id="launchplan_ref", workflow_node=_core_workflow_pb2.WorkflowNode( launchplan_ref=_identifier_pb2.Identifier( resource_type=_identifier_pb2.LAUNCH_PLAN, project="project2", )), )) workflow_template.nodes.append( _core_workflow_pb2.Node( id="sub_workflow_ref", workflow_node=_core_workflow_pb2.WorkflowNode( sub_workflow_ref=_identifier_pb2.Identifier( resource_type=_identifier_pb2.WORKFLOW, project="project2", domain="domain2", )), )) workflow_template.nodes.append( _core_workflow_pb2.Node( id="unchanged", task_node=_core_workflow_pb2.TaskNode( reference_id=_identifier_pb2.Identifier( resource_type=_identifier_pb2.TASK, project="project2", domain="domain2", version="abc")), )) hydrated_workflow_template = _hydrate_workflow_template_nodes( "project", "domain", "12345", workflow_template) assert len(hydrated_workflow_template.nodes) == 4 task_node_identifier = hydrated_workflow_template.nodes[ 0].task_node.reference_id assert task_node_identifier.project == "project" assert task_node_identifier.domain == "domain" assert task_node_identifier.version == "12345" launchplan_ref_identifier = hydrated_workflow_template.nodes[ 1].workflow_node.launchplan_ref assert launchplan_ref_identifier.project == "project2" assert launchplan_ref_identifier.domain == "domain" assert launchplan_ref_identifier.version == "12345" sub_workflow_ref_identifier = hydrated_workflow_template.nodes[ 2].workflow_node.sub_workflow_ref assert sub_workflow_ref_identifier.project == "project2" assert sub_workflow_ref_identifier.domain == "domain2" assert sub_workflow_ref_identifier.version == "12345" unchanged_identifier = hydrated_workflow_template.nodes[ 3].task_node.reference_id assert unchanged_identifier.project == "project2" assert unchanged_identifier.domain == "domain2" assert unchanged_identifier.version == "abc"