def test_handle_constructing_graph_component(self):
task1 = TaskSpec(component_ref=ComponentReference(name='comp 1'), arguments={'in1 1': 11})
task2 = TaskSpec(component_ref=ComponentReference(name='comp 2'), arguments={'in2 1': 21, 'in2 2': TaskOutputArgument.construct(task_id='task 1', output_name='out1 1')})
task3 = TaskSpec(component_ref=ComponentReference(name='comp 3'), arguments={'in3 1': TaskOutputArgument.construct(task_id='task 2', output_name='out2 1'), 'in3 2': GraphInputReference(input_name='graph in 1').as_argument()})
graph_component1 = ComponentSpec(
inputs=[
InputSpec(name='graph in 1'),
InputSpec(name='graph in 2'),
],
outputs=[
OutputSpec(name='graph out 1'),
OutputSpec(name='graph out 2'),
],
implementation=GraphImplementation(graph=GraphSpec(
tasks={
'task 1': task1,
'task 2': task2,
'task 3': task3,
},
output_values={
'graph out 1': TaskOutputArgument.construct(task_id='task 3', output_name='out3 1'),
'graph out 2': TaskOutputArgument.construct(task_id='task 1', output_name='out1 2'),
}
))
)
def test_decorator_metadata(self):
"""Test @pipeline decorator with metadata."""
@pipeline(name='p1', description='description1')
def my_pipeline1(a: {'Schema':
{
'file_type': 'csv'
}} = 'good',
b: Integer() = 12):
pass
golden_meta = ComponentSpec(name='p1',
description='description1',
inputs=[])
golden_meta.inputs.append(
InputSpec(name='a',
type={'Schema': {
'file_type': 'csv'
}},
default='good'))
golden_meta.inputs.append(
InputSpec(name='b',
type={
'Integer': {
'openapi_schema_validator': {
"type": "integer"
}
}
},
default=12))
pipeline_meta = _extract_pipeline_metadata(my_pipeline1)
self.assertEqual(pipeline_meta, golden_meta)
def test_component_metadata(self):
"""Test component decorator metadata."""
class MockContainerOp:
def _set_metadata(self, component_meta):
self._metadata = component_meta
@component
def componentA(
a: {'ArtifactA': {
'file_type': 'csv'
}},
b: Integer() = 12,
c: {'ArtifactB': {
'path_type': 'file',
'file_type': 'tsv'
}} = 'gs://hello/world'
) -> {
'model': Integer()
}:
return MockContainerOp()
containerOp = componentA(1, 2, c=3)
golden_meta = ComponentSpec(name='ComponentA', inputs=[], outputs=[])
golden_meta.inputs.append(
InputSpec(name='a', type={'ArtifactA': {
'file_type': 'csv'
}}))
golden_meta.inputs.append(
InputSpec(name='b',
type={
'Integer': {
'openapi_schema_validator': {
"type": "integer"
}
}
},
default="12",
optional=True))
golden_meta.inputs.append(
InputSpec(
name='c',
type={'ArtifactB': {
'path_type': 'file',
'file_type': 'tsv'
}},
default='gs://hello/world',
optional=True))
golden_meta.outputs.append(
OutputSpec(name='model',
type={
'Integer': {
'openapi_schema_validator': {
"type": "integer"
}
}
}))
self.assertEqual(containerOp._metadata, golden_meta)
def test_extract_component_interface(self):
from typing import NamedTuple
def my_func( # noqa: F722
required_param,
int_param: int = 42,
float_param: float = 3.14,
str_param: str = 'string',
bool_param: bool = True,
none_param=None,
custom_type_param: 'Custom type' = None,
) -> NamedTuple(
'DummyName',
[
#('required_param',), # All typing.NamedTuple fields must have types
('int_param', int),
('float_param', float),
('str_param', str),
('bool_param', bool),
#('custom_type_param', 'Custom type'), #SyntaxError: Forward reference must be an expression -- got 'Custom type'
('custom_type_param', 'CustomType'),
]):
'''Function docstring'''
pass
component_spec = comp._python_op._extract_component_interface(my_func)
from kfp.components._structures import InputSpec, OutputSpec
self.assertEqual(
component_spec.inputs,
[
InputSpec(name='required_param'),
InputSpec(name='int_param',
type='Integer',
default='42',
optional=True),
InputSpec(name='float_param',
type='Float',
default='3.14',
optional=True),
InputSpec(name='str_param',
type='String',
default='string',
optional=True),
InputSpec(name='bool_param',
type='Boolean',
default='True',
optional=True),
InputSpec(name='none_param',
optional=True), # No default='None'
InputSpec(name='custom_type_param',
type='Custom type',
optional=True),
])
self.assertEqual(
component_spec.outputs,
[
OutputSpec(name='int_param', type='Integer'),
OutputSpec(name='float_param', type='Float'),
OutputSpec(name='str_param', type='String'),
OutputSpec(name='bool_param', type='Boolean'),
#OutputSpec(name='custom_type_param', type='Custom type', default='None'),
OutputSpec(name='custom_type_param', type='CustomType'),
])
self.maxDiff = None
self.assertDictEqual(
component_spec.to_dict(),
{
'name':
'My func',
'description':
'Function docstring\n',
'inputs': [
{
'name': 'required_param'
},
{
'name': 'int_param',
'type': 'Integer',
'default': '42',
'optional': True
},
{
'name': 'float_param',
'type': 'Float',
'default': '3.14',
'optional': True
},
{
'name': 'str_param',
'type': 'String',
'default': 'string',
'optional': True
},
{
'name': 'bool_param',
'type': 'Boolean',
'default': 'True',
'optional': True
},
{
'name': 'none_param',
'optional': True
}, # No default='None'
{
'name': 'custom_type_param',
'type': 'Custom type',
'optional': True
},
],
'outputs': [
{
'name': 'int_param',
'type': 'Integer'
},
{
'name': 'float_param',
'type': 'Float'
},
{
'name': 'str_param',
'type': 'String'
},
{
'name': 'bool_param',
'type': 'Boolean'
},
{
'name': 'custom_type_param',
'type': 'CustomType'
},
]
})
def test_to_dict(self):
component_meta = ComponentSpec(
name='foobar',
description='foobar example',
inputs=[
InputSpec(name='input1',
description='input1 desc',
type={
'GCSPath': {
'bucket_type': 'directory',
'file_type': 'csv'
}
},
default='default1'),
InputSpec(name='input2',
description='input2 desc',
type={
'TFModel': {
'input_data': 'tensor',
'version': '1.8.0'
}
},
default='default2'),
InputSpec(name='input3',
description='input3 desc',
type='Integer',
default='default3'),
],
outputs=[
OutputSpec(
name='output1',
description='output1 desc',
type={'Schema': {
'file_type': 'tsv'
}},
)
])
golden_meta = {
'name':
'foobar',
'description':
'foobar example',
'inputs': [{
'name': 'input1',
'description': 'input1 desc',
'type': {
'GCSPath': {
'bucket_type': 'directory',
'file_type': 'csv'
}
},
'default': 'default1'
}, {
'name': 'input2',
'description': 'input2 desc',
'type': {
'TFModel': {
'input_data': 'tensor',
'version': '1.8.0'
}
},
'default': 'default2'
}, {
'name': 'input3',
'description': 'input3 desc',
'type': 'Integer',
'default': 'default3'
}],
'outputs': [{
'name': 'output1',
'description': 'output1 desc',
'type': {
'Schema': {
'file_type': 'tsv'
}
},
}]
}
self.assertEqual(component_meta.to_dict(), golden_meta)