Esempio n. 1
0
    def to_v1_component_spec(self) -> v1_components.ComponentSpec:
        """Convert to v1 ComponentSpec.

    Needed until downstream accept new ComponentSpec."""
        if isinstance(self.implementation, DagSpec):
            raise NotImplementedError

        def _transform_arg(arg: Union[str, BasePlaceholder]) -> Any:
            if isinstance(arg, str):
                return arg
            elif isinstance(arg, InputValuePlaceholder):
                return v1_components.InputValuePlaceholder(arg.name)
            elif isinstance(arg, InputPathPlaceholder):
                return v1_components.InputPathPlaceholder(arg.name)
            elif isinstance(arg, InputUriPlaceholder):
                return v1_components.InputUriPlaceholder(arg.name)
            elif isinstance(arg, OutputPathPlaceholder):
                return v1_components.OutputPathPlaceholder(arg.name)
            elif isinstance(arg, OutputUriPlaceholder):
                return v1_components.OutputUriPlaceholder(arg.name)
            else:
                # TODO(chensun): transform additional placeholders: if, concat, etc.?
                raise ValueError(
                    f'Unexpected command/argument type: "{arg}" of type "{type(arg)}".'
                )

        return v1_components.ComponentSpec(
            name=self.name,
            inputs=[
                v1_components.InputSpec(
                    name=input_spec.name,
                    type=_data_passing.get_canonical_type_struct_for_type(
                        input_spec.type),
                    default=input_spec.default,
                ) for input_spec in self.input_specs
            ],
            outputs=[
                v1_components.OutputSpec(
                    name=output_spec.name,
                    type=_data_passing.get_canonical_type_struct_for_type(
                        output_spec.type),
                ) for output_spec in self.output_specs
            ],
            implementation=v1_components.ContainerImplementation(
                container=v1_components.ContainerSpec(
                    image=self.implementation.image,
                    command=[
                        _transform_arg(cmd)
                        for cmd in self.implementation.commands or []
                    ],
                    args=[
                        _transform_arg(arg)
                        for arg in self.implementation.arguments or []
                    ],
                    env={
                        name: _transform_arg(value)
                        for name, value in self.implementation.env or {}
                    },
                )),
        )
    def test_build_component_spec_from_structure(self):
        structure_component_spec = structures.ComponentSpec(
            name='component1',
            description='component1 desc',
            inputs=[
                structures.InputSpec(name='input1',
                                     description='input1 desc',
                                     type='Dataset'),
                structures.InputSpec(name='input2',
                                     description='input2 desc',
                                     type='String'),
                structures.InputSpec(name='input3',
                                     description='input3 desc',
                                     type='Integer'),
            ],
            outputs=[
                structures.OutputSpec(name='output1',
                                      description='output1 desc',
                                      type='Model')
            ])
        expected_dict = {
            'inputDefinitions': {
                'artifacts': {
                    'input1': {
                        'artifactType': {
                            'instanceSchema':
                            'properties:\ntitle: kfp.Dataset\ntype: object\n'
                        }
                    }
                },
                'parameters': {
                    'input2': {
                        'type': 'STRING'
                    },
                    'input3': {
                        'type': 'INT'
                    }
                }
            },
            'outputDefinitions': {
                'artifacts': {
                    'output1': {
                        'artifactType': {
                            'instanceSchema':
                            'properties:\ntitle: kfp.Model\ntype: object\n'
                        }
                    }
                }
            },
            'executorLabel': 'exec-component1'
        }
        expected_spec = pipeline_spec_pb2.ComponentSpec()
        json_format.ParseDict(expected_dict, expected_spec)

        component_spec = (
            dsl_component_spec.build_component_spec_from_structure(
                structure_component_spec))

        self.assertEqual(expected_spec, component_spec)
  def test_build_component_spec_from_structure(self):
    structure_component_spec = structures.ComponentSpec(
        name='component1',
        description='component1 desc',
        inputs=[
            structures.InputSpec(
                name='input1', description='input1 desc', type='Dataset'),
            structures.InputSpec(
                name='input2', description='input2 desc', type='String'),
            structures.InputSpec(
                name='input3', description='input3 desc', type='Integer'),
            structures.InputSpec(
                name='input4', description='optional inputs', optional=True),
        ],
        outputs=[
            structures.OutputSpec(
                name='output1', description='output1 desc', type='Model')
        ])
    expected_dict = {
        'inputDefinitions': {
            'artifacts': {
                'input1': {
                    'artifactType': {
                        'schemaTitle': 'system.Dataset'
                    }
                }
            },
            'parameters': {
                'input2': {
                    'type': 'STRING'
                },
                'input3': {
                    'type': 'INT'
                }
            }
        },
        'outputDefinitions': {
            'artifacts': {
                'output1': {
                    'artifactType': {
                        'schemaTitle': 'system.Model'
                    }
                }
            }
        },
        'executorLabel': 'exec-component1'
    }
    expected_spec = pipeline_spec_pb2.ComponentSpec()
    json_format.ParseDict(expected_dict, expected_spec)

    component_spec = (
        dsl_component_spec.build_component_spec_from_structure(
            component_spec=structure_component_spec,
            executor_label='exec-component1',
            actual_inputs=['input1', 'input2', 'input3'],
        ))

    self.assertEqual(expected_spec, component_spec)