def test_workflow_template(): task = _workflow.TaskNode(reference_id=_generic_id) nm = _get_sample_node_metadata() int_type = _types.LiteralType(_types.SimpleType.INTEGER) wf_metadata = _workflow.WorkflowMetadata() wf_metadata_defaults = _workflow.WorkflowMetadataDefaults() typed_interface = _interface.TypedInterface( {"a": _interface.Variable(int_type, "description1")}, { "b": _interface.Variable(int_type, "description2"), "c": _interface.Variable(int_type, "description3") }, ) wf_node = _workflow.Node( id="some:node:id", metadata=nm, inputs=[], upstream_node_ids=[], output_aliases=[], task_node=task, ) obj = _workflow.WorkflowTemplate( id=_generic_id, metadata=wf_metadata, metadata_defaults=wf_metadata_defaults, interface=typed_interface, nodes=[wf_node], outputs=[], ) obj2 = _workflow.WorkflowTemplate.from_flyte_idl(obj.to_flyte_idl()) assert obj2 == obj
def test_workflow_template_with_queuing_budget(): task = _workflow.TaskNode(reference_id=_generic_id) nm = _get_sample_node_metadata() int_type = _types.LiteralType(_types.SimpleType.INTEGER) wf_metadata = _workflow.WorkflowMetadata(queuing_budget=timedelta( seconds=10)) wf_metadata_defaults = _workflow.WorkflowMetadataDefaults() typed_interface = _interface.TypedInterface( {'a': _interface.Variable(int_type, "description1")}, { 'b': _interface.Variable(int_type, "description2"), 'c': _interface.Variable(int_type, "description3") }) wf_node = _workflow.Node(id='some:node:id', metadata=nm, inputs=[], upstream_node_ids=[], output_aliases=[], task_node=task) obj = _workflow.WorkflowTemplate(id=_generic_id, metadata=wf_metadata, metadata_defaults=wf_metadata_defaults, interface=typed_interface, nodes=[wf_node], outputs=[]) obj2 = _workflow.WorkflowTemplate.from_flyte_idl(obj.to_flyte_idl()) assert obj2 == obj
def test_workflow_closure(): int_type = _types.LiteralType(_types.SimpleType.INTEGER) typed_interface = _interface.TypedInterface( {'a': _interface.Variable(int_type, "description1")}, { 'b': _interface.Variable(int_type, "description2"), 'c': _interface.Variable(int_type, "description3") }) b0 = _literals.Binding( 'a', _literals.BindingData(scalar=_literals.Scalar( primitive=_literals.Primitive(integer=5)))) b1 = _literals.Binding( 'b', _literals.BindingData(promise=_types.OutputReference('my_node', 'b'))) b2 = _literals.Binding( 'b', _literals.BindingData(promise=_types.OutputReference('my_node', 'c'))) node_metadata = _workflow.NodeMetadata(name='node1', timeout=timedelta(seconds=10), retries=_literals.RetryStrategy(0)) task_metadata = _task.TaskMetadata( True, _task.RuntimeMetadata(_task.RuntimeMetadata.RuntimeType.FLYTE_SDK, "1.0.0", "python"), timedelta(days=1), _literals.RetryStrategy(3), "0.1.1b0", "This is deprecated!") cpu_resource = _task.Resources.ResourceEntry( _task.Resources.ResourceName.CPU, "1") resources = _task.Resources(requests=[cpu_resource], limits=[cpu_resource]) task = _task.TaskTemplate( _identifier.Identifier(_identifier.ResourceType.TASK, "project", "domain", "name", "version"), "python", task_metadata, typed_interface, { 'a': 1, 'b': { 'c': 2, 'd': 3 } }, container=_task.Container("my_image", ["this", "is", "a", "cmd"], ["this", "is", "an", "arg"], resources, {}, {})) task_node = _workflow.TaskNode(task.id) node = _workflow.Node(id='my_node', metadata=node_metadata, inputs=[b0], upstream_node_ids=[], output_aliases=[], task_node=task_node) template = _workflow.WorkflowTemplate( id=_identifier.Identifier(_identifier.ResourceType.WORKFLOW, "project", "domain", "name", "version"), metadata=_workflow.WorkflowMetadata(), interface=typed_interface, nodes=[node], outputs=[b1, b2], ) obj = _workflow_closure.WorkflowClosure(workflow=template, tasks=[task]) assert len(obj.tasks) == 1 obj2 = _workflow_closure.WorkflowClosure.from_flyte_idl(obj.to_flyte_idl()) assert obj == obj2
def get_serializable_workflow( entity_mapping: OrderedDict, settings: SerializationSettings, entity: WorkflowBase, ) -> admin_workflow_models.WorkflowSpec: # Get node models upstream_node_models = [ get_serializable(entity_mapping, settings, n) for n in entity.nodes if n.id != _common_constants.GLOBAL_INPUT_NODE_ID ] sub_wfs = [] for n in entity.nodes: if isinstance(n.flyte_entity, WorkflowBase): if isinstance(n.flyte_entity, ReferenceEntity): raise Exception( f"Sorry, reference subworkflows do not work right now, please use the launch plan instead for the " f"subworkflow you're trying to invoke. Node: {n}" ) sub_wf_spec = get_serializable(entity_mapping, settings, n.flyte_entity) if not isinstance(sub_wf_spec, admin_workflow_models.WorkflowSpec): raise Exception( f"Serialized form of a workflow should be an admin.WorkflowSpec but {type(sub_wf_spec)} found instead" ) sub_wfs.append(sub_wf_spec.template) sub_wfs.extend(sub_wf_spec.sub_workflows) if isinstance(n.flyte_entity, BranchNode): if_else: workflow_model.IfElseBlock = n.flyte_entity._ifelse_block # See comment in get_serializable_branch_node also. Again this is a List[Node] even though it's supposed # to be a List[workflow_model.Node] leaf_nodes: List[Node] = filter( # noqa None, [ if_else.case.then_node, *([] if if_else.other is None else [x.then_node for x in if_else.other]), if_else.else_node, ], ) for leaf_node in leaf_nodes: if isinstance(leaf_node.flyte_entity, WorkflowBase): sub_wf_spec = get_serializable(entity_mapping, settings, leaf_node.flyte_entity) sub_wfs.append(sub_wf_spec.template) sub_wfs.extend(sub_wf_spec.sub_workflows) wf_id = _identifier_model.Identifier( resource_type=_identifier_model.ResourceType.WORKFLOW, project=settings.project, domain=settings.domain, name=entity.name, version=settings.version, ) wf_t = workflow_model.WorkflowTemplate( id=wf_id, metadata=entity.workflow_metadata.to_flyte_model(), metadata_defaults=entity.workflow_metadata_defaults.to_flyte_model(), interface=entity.interface, nodes=upstream_node_models, outputs=entity.output_bindings, ) return admin_workflow_models.WorkflowSpec(template=wf_t, sub_workflows=list(set(sub_wfs)))
def get_serializable_workflow( entity_mapping: OrderedDict, settings: SerializationSettings, entity: WorkflowBase, options: Optional[Options] = None, ) -> admin_workflow_models.WorkflowSpec: # TODO: Try to move up following config refactor - https://github.com/flyteorg/flyte/issues/2214 from flytekit.remote.workflow import FlyteWorkflow # Get node models upstream_node_models = [ get_serializable(entity_mapping, settings, n, options) for n in entity.nodes if n.id != _common_constants.GLOBAL_INPUT_NODE_ID ] sub_wfs = [] for n in entity.nodes: if isinstance(n.flyte_entity, WorkflowBase): # We are currently not supporting reference workflows since these will # require a network call to flyteadmin to populate the WorkflowTemplate # object if isinstance(n.flyte_entity, ReferenceWorkflow): raise Exception( "Reference sub-workflows are currently unsupported. Use reference launch plans instead." ) sub_wf_spec = get_serializable(entity_mapping, settings, n.flyte_entity, options) if not isinstance(sub_wf_spec, admin_workflow_models.WorkflowSpec): raise TypeError( f"Unexpected type for serialized form of workflow. Expected {admin_workflow_models.WorkflowSpec}, but got {type(sub_wf_spec)}" ) sub_wfs.append(sub_wf_spec.template) sub_wfs.extend(sub_wf_spec.sub_workflows) if isinstance(n.flyte_entity, FlyteWorkflow): get_serializable(entity_mapping, settings, n.flyte_entity, options) sub_wfs.append(n.flyte_entity) sub_wfs.extend([s for s in n.flyte_entity.sub_workflows.values()]) if isinstance(n.flyte_entity, BranchNode): if_else: workflow_model.IfElseBlock = n.flyte_entity._ifelse_block # See comment in get_serializable_branch_node also. Again this is a List[Node] even though it's supposed # to be a List[workflow_model.Node] leaf_nodes: List[Node] = filter( # noqa None, [ if_else.case.then_node, *([] if if_else.other is None else [x.then_node for x in if_else.other]), if_else.else_node, ], ) for leaf_node in leaf_nodes: if isinstance(leaf_node.flyte_entity, WorkflowBase): sub_wf_spec = get_serializable(entity_mapping, settings, leaf_node.flyte_entity, options) sub_wfs.append(sub_wf_spec.template) sub_wfs.extend(sub_wf_spec.sub_workflows) elif isinstance(leaf_node.flyte_entity, FlyteWorkflow): get_serializable(entity_mapping, settings, leaf_node.flyte_entity, options) sub_wfs.append(leaf_node.flyte_entity) sub_wfs.extend([ s for s in leaf_node.flyte_entity.sub_workflows.values() ]) wf_id = _identifier_model.Identifier( resource_type=_identifier_model.ResourceType.WORKFLOW, project=settings.project, domain=settings.domain, name=entity.name, version=settings.version, ) wf_t = workflow_model.WorkflowTemplate( id=wf_id, metadata=entity.workflow_metadata.to_flyte_model(), metadata_defaults=entity.workflow_metadata_defaults.to_flyte_model(), interface=entity.interface, nodes=upstream_node_models, outputs=entity.output_bindings, ) return admin_workflow_models.WorkflowSpec( template=wf_t, sub_workflows=sorted(set(sub_wfs), key=lambda x: x.short_string()))