def test_create_component_spec_composes_correctly(self):
image_uri = "my-image"
image_tag = "my-tag"
file_path = "fake-path"
expected = ComponentSpec(
name="Dummy component",
description="Dummy description",
inputs=self.DUMMY_IO_ARGS.inputs,
outputs=self.DUMMY_IO_ARGS.outputs,
implementation=ContainerImplementation(container=ContainerSpec(
image="my-image:my-tag",
command=["python3"],
args=[
"fake-path",
"--input1",
InputValuePlaceholder(input_name="input1"),
"--input2",
InputValuePlaceholder(input_name="input2"),
"--output1_output_path",
OutputPathPlaceholder(output_name="output1"),
"--output2_output_path",
OutputPathPlaceholder(output_name="output2"),
],
)),
)
with patch(
"common.component_compiler.SageMakerComponentCompiler._create_io_from_component_spec",
MagicMock(return_value=self.DUMMY_IO_ARGS),
):
response = SageMakerComponentCompiler._create_component_spec(
DummyComponent, file_path, image_uri, image_tag)
self.assertEqual(expected, response)
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',
optional=True))
golden_meta.inputs.append(
InputSpec(name='b',
type={
'Integer': {
'openapi_schema_validator': {
"type": "integer"
}
}
},
default="12",
optional=True))
pipeline_meta = _extract_pipeline_metadata(my_pipeline1)
self.assertEqual(pipeline_meta, golden_meta)
def _create_component_spec(
component_def: Type[SageMakerComponent],
component_file_path: str,
component_image_uri: str,
component_image_tag: str,
) -> ComponentSpec:
"""Create a component YAML specification object based on a component.
Args:
component_def: The type of the SageMaker component.
component_file_path: The path to the component definition file.
component_image_uri: Compiled image URI.
component_image_tag: Compiled image tag.
Returns:
ComponentSpec: A `component.yaml` specification object.
"""
io_args = SageMakerComponentCompiler._create_io_from_component_spec(
component_def.COMPONENT_SPEC)
return ComponentSpec(
name=component_def.COMPONENT_NAME,
description=component_def.COMPONENT_DESCRIPTION,
inputs=io_args.inputs,
outputs=io_args.outputs,
implementation=ContainerImplementation(container=ContainerSpec(
image=f"{component_image_uri}:{component_image_tag}",
command=["python3"],
args=[
component_file_path,
] + io_args.args, # type: ignore
)),
)
def test_component_metadata_standard_type_annotation(self):
"""Test component decorator metadata."""
class MockContainerOp:
def _set_metadata(self, component_meta):
self._metadata = component_meta
@component
def componentA(a: float, b: List[int], c: Optional[str] = None) -> None:
return MockContainerOp()
containerOp = componentA('str_value', '[1,2,3]')
golden_meta = ComponentSpec(name='ComponentA', inputs=[], outputs=None)
golden_meta.inputs.append(InputSpec(name='a', type='Float'))
golden_meta.inputs.append(
InputSpec(
name='b',
type='typing.List[int]' if sys.version_info >=
(3, 7) else 'List'))
golden_meta.inputs.append(
InputSpec(name='c', type='String', default=None, optional=True))
self.assertEqual(containerOp._metadata, 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_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 component_yaml_generator(**kwargs):
input_specs = []
input_args = []
input_kwargs = {}
serialized_args = {INIT_KEY: {}, METHOD_KEY: {}}
init_kwargs = {}
method_kwargs = {}
for key, value in kwargs.items():
if key in init_arg_names:
prefix_key = INIT_KEY
init_kwargs[key] = value
signature = init_signature
else:
prefix_key = METHOD_KEY
method_kwargs[key] = value
signature = method_signature
# no need to add this argument because it's optional
# this param is validated against the signature because
# of init_kwargs, method_kwargs
if value is None:
continue
param_type = signature.parameters[key].annotation
param_type = resolve_annotation(param_type)
serializer = get_serializer(param_type)
if serializer:
param_type = str
value = serializer(value)
# TODO remove PipelineParam check when Metadata Importer component available
# if we serialize we need to include the argument as input
# perhaps, another option is to embed in yaml as json serialized list
component_param_name = component_param_name_to_mb_sdk_param_name.get(
key, key
)
if isinstance(value,
kfp.dsl._pipeline_param.PipelineParam) or serializer:
if is_mb_sdk_resource_noun_type(param_type):
metadata_type = map_resource_to_metadata_type(param_type)[1]
component_param_type = metadata_type
else:
component_param_type = 'String'
input_specs.append(
InputSpec(
name=key,
type=component_param_type,
)
)
input_args.append(f'--{prefix_key}.{component_param_name}')
if is_mb_sdk_resource_noun_type(param_type):
input_args.append(InputUriPlaceholder(input_name=key))
else:
input_args.append(InputValuePlaceholder(input_name=key))
input_kwargs[key] = value
else:
# Serialized arguments must always be strings
value = str(value)
serialized_args[prefix_key][component_param_name] = value
# validate parameters
if should_serialize_init:
init_signature.bind(**init_kwargs)
method_signature.bind(**method_kwargs)
component_spec = ComponentSpec(
name=f'{cls_name}-{method_name}',
inputs=input_specs,
outputs=output_specs,
implementation=ContainerImplementation(
container=ContainerSpec(
image=DEFAULT_CONTAINER_IMAGE,
command=[
'python3',
'-m',
'google_cloud_pipeline_components.aiplatform.remote_runner',
'--cls_name',
cls_name,
'--method_name',
method_name,
],
args=make_args(serialized_args) + output_args + input_args,
)
)
)
component_path = tempfile.mktemp()
component_spec.save(component_path)
return components.load_component_from_file(component_path)(
**input_kwargs
)
class ComponentCompilerTestCase(unittest.TestCase):
# These should always match the dummy spec
DUMMY_IO_ARGS = IOArgs(
inputs=[
InputSpec(
name="input1",
description="The first input.",
type="String",
default="input1-default",
),
InputSpec(name="input2",
description="The second input.",
type="Integer"),
],
outputs=[
OutputSpec(name="output1", description="The first output."),
OutputSpec(name="output2", description="The second output."),
],
args=[
"--input1",
InputValuePlaceholder(input_name="input1"),
"--input2",
InputValuePlaceholder(input_name="input2"),
"--output1_output_path",
OutputPathPlaceholder(output_name="output1"),
"--output2_output_path",
OutputPathPlaceholder(output_name="output2"),
],
)
DUMMY_COMPONENT_SPEC = ComponentSpec(
name="Dummy component",
description="Dummy description",
inputs=DUMMY_IO_ARGS.inputs,
outputs=DUMMY_IO_ARGS.outputs,
implementation=ContainerImplementation(container=ContainerSpec(
image="my-image:my-tag",
command=["python3"],
args=[
"fake-path",
"--input1",
InputValuePlaceholder(input_name="input1"),
"--input2",
InputValuePlaceholder(input_name="input2"),
"--output1_output_path",
OutputPathPlaceholder(output_name="output1"),
"--output2_output_path",
OutputPathPlaceholder(output_name="output2"),
],
)),
)
EXTRA_IO_ARGS = IOArgs(
inputs=[
InputSpec(name="inputStr", description="str", type="String"),
InputSpec(name="inputInt", description="int", type="Integer"),
InputSpec(name="inputBool", description="bool", type="Bool"),
InputSpec(name="inputDict", description="dict", type="JsonObject"),
InputSpec(name="inputList", description="list", type="JsonArray"),
InputSpec(
name="inputOptional",
description="optional",
type="String",
default="default-string",
),
InputSpec(
name="inputOptionalNoDefault",
description="optional",
type="String",
default="",
),
],
outputs=[],
args=[
"--inputStr",
InputValuePlaceholder(input_name="inputStr"),
"--inputInt",
InputValuePlaceholder(input_name="inputInt"),
"--inputBool",
InputValuePlaceholder(input_name="inputBool"),
"--inputDict",
InputValuePlaceholder(input_name="inputDict"),
"--inputList",
InputValuePlaceholder(input_name="inputList"),
"--inputOptional",
InputValuePlaceholder(input_name="inputOptional"),
"--inputOptionalNoDefault",
InputValuePlaceholder(input_name="inputOptionalNoDefault"),
],
)
@classmethod
def setUpClass(cls):
cls.compiler = SageMakerComponentCompiler()
def test_create_io_from_component_spec(self):
response = SageMakerComponentCompiler._create_io_from_component_spec(
DummySpec) # type: ignore
self.assertEqual(self.DUMMY_IO_ARGS, response)
def test_create_io_from_component_spec_extra_types(self):
response = SageMakerComponentCompiler._create_io_from_component_spec(
ExtraSpec) # type: ignore
self.assertEqual(self.EXTRA_IO_ARGS, response)
def test_create_component_spec_composes_correctly(self):
image_uri = "my-image"
image_tag = "my-tag"
file_path = "fake-path"
expected = ComponentSpec(
name="Dummy component",
description="Dummy description",
inputs=self.DUMMY_IO_ARGS.inputs,
outputs=self.DUMMY_IO_ARGS.outputs,
implementation=ContainerImplementation(container=ContainerSpec(
image="my-image:my-tag",
command=["python3"],
args=[
"fake-path",
"--input1",
InputValuePlaceholder(input_name="input1"),
"--input2",
InputValuePlaceholder(input_name="input2"),
"--output1_output_path",
OutputPathPlaceholder(output_name="output1"),
"--output2_output_path",
OutputPathPlaceholder(output_name="output2"),
],
)),
)
with patch(
"common.component_compiler.SageMakerComponentCompiler._create_io_from_component_spec",
MagicMock(return_value=self.DUMMY_IO_ARGS),
):
response = SageMakerComponentCompiler._create_component_spec(
DummyComponent, file_path, image_uri, image_tag)
self.assertEqual(expected, response)
def test_write_component(self):
DummyComponent.save = MagicMock()
SageMakerComponentCompiler._write_component(DummyComponent,
"/tmp/fake-path")
DummyComponent.save.assert_called_once_with("/tmp/fake-path")
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)