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)