Beispiel #1
0
    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)
Beispiel #3
0
    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
        )
Beispiel #4
0
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")