def test_typed_interface(literal_type): typed_interface = interface.TypedInterface( {'a': interface.Variable(literal_type, "description1")}, { 'b': interface.Variable(literal_type, "description2"), 'c': interface.Variable(literal_type, "description3") }) assert typed_interface.inputs['a'].type == literal_type assert typed_interface.outputs['b'].type == literal_type assert typed_interface.outputs['c'].type == literal_type assert typed_interface.inputs['a'].description == "description1" assert typed_interface.outputs['b'].description == "description2" assert typed_interface.outputs['c'].description == "description3" assert len(typed_interface.inputs) == 1 assert len(typed_interface.outputs) == 2 pb = typed_interface.to_flyte_idl() deserialized_typed_interface = interface.TypedInterface.from_flyte_idl(pb) assert typed_interface == deserialized_typed_interface assert deserialized_typed_interface.inputs['a'].type == literal_type assert deserialized_typed_interface.outputs['b'].type == literal_type assert deserialized_typed_interface.outputs['c'].type == literal_type assert deserialized_typed_interface.inputs[ 'a'].description == "description1" assert deserialized_typed_interface.outputs[ 'b'].description == "description2" assert deserialized_typed_interface.outputs[ 'c'].description == "description3" assert len(deserialized_typed_interface.inputs) == 1 assert len(deserialized_typed_interface.outputs) == 2
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_typed_interface(literal_type): typed_interface = interface.TypedInterface( {"a": interface.Variable(literal_type, "description1")}, { "b": interface.Variable(literal_type, "description2"), "c": interface.Variable(literal_type, "description3") }, ) assert typed_interface.inputs["a"].type == literal_type assert typed_interface.outputs["b"].type == literal_type assert typed_interface.outputs["c"].type == literal_type assert typed_interface.inputs["a"].description == "description1" assert typed_interface.outputs["b"].description == "description2" assert typed_interface.outputs["c"].description == "description3" assert len(typed_interface.inputs) == 1 assert len(typed_interface.outputs) == 2 pb = typed_interface.to_flyte_idl() deserialized_typed_interface = interface.TypedInterface.from_flyte_idl(pb) assert typed_interface == deserialized_typed_interface assert deserialized_typed_interface.inputs["a"].type == literal_type assert deserialized_typed_interface.outputs["b"].type == literal_type assert deserialized_typed_interface.outputs["c"].type == literal_type assert deserialized_typed_interface.inputs[ "a"].description == "description1" assert deserialized_typed_interface.outputs[ "b"].description == "description2" assert deserialized_typed_interface.outputs[ "c"].description == "description3" assert len(deserialized_typed_interface.inputs) == 1 assert len(deserialized_typed_interface.outputs) == 2
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_task_template__k8s_pod_target(): int_type = types.LiteralType(types.SimpleType.INTEGER) obj = task.TaskTemplate( identifier.Identifier(identifier.ResourceType.TASK, "project", "domain", "name", "version"), "python", task.TaskMetadata( False, task.RuntimeMetadata(1, "v", "f"), timedelta(days=1), literal_models.RetryStrategy(5), False, "1.0", "deprecated", False, ), interface_models.TypedInterface( # inputs {"a": interface_models.Variable(int_type, "description1")}, # outputs { "b": interface_models.Variable(int_type, "description2"), "c": interface_models.Variable(int_type, "description3"), }, ), { "a": 1, "b": { "c": 2, "d": 3 } }, config={"a": "b"}, k8s_pod=task.K8sPod( metadata=task.K8sObjectMetadata(labels={"label": "foo"}, annotations={"anno": "bar"}), pod_spec={ "str": "val", "int": 1 }, ), ) assert obj.id.resource_type == identifier.ResourceType.TASK assert obj.id.project == "project" assert obj.id.domain == "domain" assert obj.id.name == "name" assert obj.id.version == "version" assert obj.type == "python" assert obj.custom == {"a": 1, "b": {"c": 2, "d": 3}} assert obj.k8s_pod.metadata == task.K8sObjectMetadata( labels={"label": "foo"}, annotations={"anno": "bar"}) assert obj.k8s_pod.pod_spec == {"str": "val", "int": 1} assert text_format.MessageToString( obj.to_flyte_idl()) == text_format.MessageToString( task.TaskTemplate.from_flyte_idl( obj.to_flyte_idl()).to_flyte_idl()) assert obj.config == {"a": "b"}
def test_basic_workflow_promote(mock_task_fetch): # This section defines a sample workflow from a user @_sdk_tasks.inputs(a=_Types.Integer) @_sdk_tasks.outputs(b=_Types.Integer, c=_Types.Integer) @_sdk_tasks.python_task() def demo_task_for_promote(wf_params, a, b, c): b.set(a + 1) c.set(a + 2) @_sdk_workflow.workflow_class() class TestPromoteExampleWf(object): wf_input = _sdk_workflow.Input(_Types.Integer, required=True) my_task_node = demo_task_for_promote(a=wf_input) wf_output_b = _sdk_workflow.Output(my_task_node.outputs.b, sdk_type=_Types.Integer) wf_output_c = _sdk_workflow.Output(my_task_node.outputs.c, sdk_type=_Types.Integer) # This section uses the TaskTemplate stored in Admin to promote back to an Sdk Workflow int_type = _types.LiteralType(_types.SimpleType.INTEGER) task_interface = _interface.TypedInterface( # inputs {'a': _interface.Variable(int_type, "description1")}, # outputs { 'b': _interface.Variable(int_type, "description2"), 'c': _interface.Variable(int_type, "description3") }) # Since the promotion of a workflow requires retrieving the task from Admin, we mock the SdkTask to return task_template = _task_model.TaskTemplate(_identifier.Identifier( _identifier.ResourceType.TASK, "project", "domain", "tests.flytekit.unit.common_tests.test_workflow_promote.demo_task_for_promote", "version"), "python_container", get_sample_task_metadata(), task_interface, custom={}, container=get_sample_container()) sdk_promoted_task = _task.SdkTask.promote_from_model(task_template) mock_task_fetch.return_value = sdk_promoted_task workflow_template = get_workflow_template() promoted_wf = _workflow_common.SdkWorkflow.promote_from_model( workflow_template) assert promoted_wf.interface.inputs[ "wf_input"] == TestPromoteExampleWf.interface.inputs["wf_input"] assert promoted_wf.interface.outputs[ "wf_output_b"] == TestPromoteExampleWf.interface.outputs["wf_output_b"] assert promoted_wf.interface.outputs[ "wf_output_c"] == TestPromoteExampleWf.interface.outputs["wf_output_c"] assert len(promoted_wf.nodes) == 1 assert len(TestPromoteExampleWf.nodes) == 1 assert promoted_wf.nodes[0].inputs[0] == TestPromoteExampleWf.nodes[ 0].inputs[0]
def transform_interface_to_typed_interface( interface: typing.Optional[Interface], ) -> typing.Optional[_interface_models.TypedInterface]: """ Transform the given simple python native interface to FlyteIDL's interface """ if interface is None: return None inputs_map = transform_variable_map(interface.inputs) outputs_map = transform_variable_map(interface.outputs) return _interface_models.TypedInterface(inputs_map, outputs_map)
def __init__( self, hive_job, metadata, ): """ :param _qubole.QuboleHiveJob hive_job: Hive job spec :param TaskMetadata metadata: This contains information needed at runtime to determine behavior such as whether or not outputs are discoverable, timeouts, and retries. """ super(SdkHiveJob, self).__init__( _constants.SdkTaskType.HIVE_JOB, metadata, # Individual hive tasks never take anything, or return anything. They just run a query that's already # got the location set. _interface_model.TypedInterface({}, {}), _MessageToDict(hive_job), )
def transform_interface_to_typed_interface( interface: typing.Optional[Interface], ) -> typing.Optional[_interface_models.TypedInterface]: """ Transform the given simple python native interface to FlyteIDL's interface """ if interface is None: return None if interface.docstring is None: input_descriptions = output_descriptions = {} else: input_descriptions = interface.docstring.input_descriptions output_descriptions = remap_shared_output_descriptions( interface.docstring.output_descriptions, interface.outputs) inputs_map = transform_variable_map(interface.inputs, input_descriptions) outputs_map = transform_variable_map(interface.outputs, output_descriptions) return _interface_models.TypedInterface(inputs_map, outputs_map)
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
types.LiteralType(collection_type=literal_type) for literal_type in LIST_OF_SCALAR_LITERAL_TYPES ] LIST_OF_NESTED_COLLECTION_LITERAL_TYPES = [ types.LiteralType(collection_type=literal_type) for literal_type in LIST_OF_COLLECTION_LITERAL_TYPES ] LIST_OF_ALL_LITERAL_TYPES = \ LIST_OF_SCALAR_LITERAL_TYPES + \ LIST_OF_COLLECTION_LITERAL_TYPES + \ LIST_OF_NESTED_COLLECTION_LITERAL_TYPES LIST_OF_INTERFACES = [ interface.TypedInterface({'a': interface.Variable(t, "description 1")}, {'b': interface.Variable(t, "description 2")}) for t in LIST_OF_ALL_LITERAL_TYPES ] LIST_OF_RESOURCE_ENTRIES = [ task.Resources.ResourceEntry(task.Resources.ResourceName.CPU, "1"), task.Resources.ResourceEntry(task.Resources.ResourceName.GPU, "1"), task.Resources.ResourceEntry(task.Resources.ResourceName.MEMORY, "1G"), task.Resources.ResourceEntry(task.Resources.ResourceName.STORAGE, "1G") ] LIST_OF_RESOURCE_ENTRY_LISTS = [LIST_OF_RESOURCE_ENTRIES] LIST_OF_RESOURCES = [ task.Resources(request, limit) for request, limit in product( LIST_OF_RESOURCE_ENTRY_LISTS, LIST_OF_RESOURCE_ENTRY_LISTS)
types.LiteralType(collection_type=literal_type) for literal_type in LIST_OF_SCALAR_LITERAL_TYPES ] LIST_OF_NESTED_COLLECTION_LITERAL_TYPES = [ types.LiteralType(collection_type=literal_type) for literal_type in LIST_OF_COLLECTION_LITERAL_TYPES ] LIST_OF_ALL_LITERAL_TYPES = (LIST_OF_SCALAR_LITERAL_TYPES + LIST_OF_COLLECTION_LITERAL_TYPES + LIST_OF_NESTED_COLLECTION_LITERAL_TYPES) LIST_OF_INTERFACES = [ interface.TypedInterface( {"a": interface.Variable(t, "description 1")}, {"b": interface.Variable(t, "description 2")}, ) for t in LIST_OF_ALL_LITERAL_TYPES ] LIST_OF_RESOURCE_ENTRIES = [ task.Resources.ResourceEntry(task.Resources.ResourceName.CPU, "1"), task.Resources.ResourceEntry(task.Resources.ResourceName.GPU, "1"), task.Resources.ResourceEntry(task.Resources.ResourceName.MEMORY, "1G"), task.Resources.ResourceEntry(task.Resources.ResourceName.STORAGE, "1G"), ] LIST_OF_RESOURCE_ENTRY_LISTS = [LIST_OF_RESOURCE_ENTRIES] LIST_OF_RESOURCES = [ task.Resources(request, limit) for request, limit in product( LIST_OF_RESOURCE_ENTRY_LISTS, LIST_OF_RESOURCE_ENTRY_LISTS)