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 = ComponentMeta(name='componentA', description='')
golden_meta.inputs.append(
ParameterMeta(name='a',
description='',
param_type=TypeMeta(name='ArtifactA',
properties={'file_type':
'csv'})))
golden_meta.inputs.append(
ParameterMeta(name='b',
description='',
param_type=TypeMeta(name='Integer',
properties={
'openapi_schema_validator': {
"type": "integer"
}
}),
default=12))
golden_meta.inputs.append(
ParameterMeta(name='c',
description='',
param_type=TypeMeta(name='ArtifactB',
properties={
'path_type': 'file',
'file_type': 'tsv'
}),
default='gs://hello/world'))
golden_meta.outputs.append(
ParameterMeta(name='model',
description='',
param_type=TypeMeta(name='Integer',
properties={
'openapi_schema_validator': {
"type": "integer"
}
})))
self.assertEqual(containerOp._metadata, golden_meta)
def _create_container_op_from_resolved_task(name: str,
container_image: str,
command=None,
arguments=None,
output_paths=None,
env: Mapping[str, str] = None,
component_spec=None):
from .. import dsl
#Renaming outputs to conform with ContainerOp/Argo
from ._naming import _sanitize_python_function_name, generate_unique_name_conversion_table
output_names = (output_paths or {}).keys()
output_name_to_kubernetes = generate_unique_name_conversion_table(
output_names, _sanitize_python_function_name)
output_paths_for_container_op = {
output_name_to_kubernetes[name]: path
for name, path in output_paths.items()
}
# Construct the ComponentMeta
component_meta = ComponentMeta(name=component_spec.name,
description=component_spec.description)
# Inputs
if component_spec.inputs is not None:
for input in component_spec.inputs:
component_meta.inputs.append(
ParameterMeta(name=input.name,
description=input.description,
param_type=input.type,
default=input.default))
if component_spec.outputs is not None:
for output in component_spec.outputs:
component_meta.outputs.append(
ParameterMeta(name=output.name,
description=output.description,
param_type=output.type))
task = dsl.ContainerOp(
name=name,
image=container_image,
command=command,
arguments=arguments,
file_outputs=output_paths_for_container_op,
)
task._set_metadata(component_meta)
if env:
from kubernetes import client as k8s_client
for name, value in env.items():
task.container.add_env_variable(
k8s_client.V1EnvVar(name=name, value=value))
if component_spec.metadata:
for key, value in (component_spec.metadata.annotations or {}).items():
task.add_pod_annotation(key, value)
for key, value in (component_spec.metadata.labels or {}).items():
task.add_pod_label(key, value)
return task
def test_to_dict(self):
component_meta = ComponentMeta(
name='foobar',
description='foobar example',
inputs=[
ParameterMeta(name='input1',
description='input1 desc',
param_type={
'GCSPath': {
'bucket_type': 'directory',
'file_type': 'csv'
}
},
default='default1'),
ParameterMeta(name='input2',
description='input2 desc',
param_type={
'TFModel': {
'input_data': 'tensor',
'version': '1.8.0'
}
},
default='default2'),
ParameterMeta(name='input3',
description='input3 desc',
param_type='Integer',
default='default3'),
],
outputs=[
ParameterMeta(name='output1',
description='output1 desc',
param_type={'Schema': {
'file_type': 'tsv'
}},
default='default_output1')
])
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'
}
},
'default': 'default_output1'
}]
}
self.assertEqual(component_meta.to_dict(), golden_meta)