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 _create_io_from_component_spec(
spec: Type[SageMakerComponentSpec]) -> IOArgs:
"""Parses the set of inputs and outputs from a component spec into the
YAML spec form.
Args:
spec: A component specification definition.
Returns:
IOArgs: The IO arguments object filled with the fields from the
component spec definition.
"""
inputs = []
outputs = []
args = []
# Iterate through all inputs adding them to the argument list
for key, _input in spec.INPUTS.__dict__.items():
# We know all of these values are validators as we have validated the spec
input_validator: SageMakerComponentInputValidator = cast(
SageMakerComponentInputValidator, _input)
# Map from argsparser to KFP component
input_spec = InputSpec(
name=key,
description=input_validator.description,
type=SageMakerComponentCompiler.KFP_TYPE_FROM_ARGS.get(
input_validator.input_type, "String"),
)
# Add optional fields
if input_validator.default is not None:
input_spec.__dict__["default"] = str(input_validator.default)
elif not input_validator.required:
# If not required and has no default, add empty string
input_spec.__dict__["default"] = ""
inputs.append(input_spec)
# Add arguments to input list
args.append(f"--{key}")
args.append(InputValuePlaceholder(input_name=key))
for key, _output in spec.OUTPUTS.__dict__.items():
output_validator: SageMakerComponentOutputValidator = cast(
SageMakerComponentOutputValidator, _output)
outputs.append(
OutputSpec(name=key, description=output_validator.description))
# Add arguments to input list
args.append(
f"--{key}{SageMakerComponentSpec.OUTPUT_ARGUMENT_SUFFIX}")
args.append(OutputPathPlaceholder(output_name=key))
return IOArgs(inputs=inputs, outputs=outputs, args=args)
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")