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)
