Beispiel #1
0
    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 = ComponentMeta(name='componentA', description='')
        golden_meta.inputs.append(
            ParameterMeta(name='a',
                          description='',
                          param_type=TypeMeta(name='ArtifactA',
                                              properties={'file_type':
                                                          'csv'})))
        golden_meta.inputs.append(
            ParameterMeta(name='b',
                          description='',
                          param_type=TypeMeta(name='Integer',
                                              properties={
                                                  'openapi_schema_validator': {
                                                      "type": "integer"
                                                  }
                                              }),
                          default=12))
        golden_meta.inputs.append(
            ParameterMeta(name='c',
                          description='',
                          param_type=TypeMeta(name='ArtifactB',
                                              properties={
                                                  'path_type': 'file',
                                                  'file_type': 'tsv'
                                              }),
                          default='gs://hello/world'))
        golden_meta.outputs.append(
            ParameterMeta(name='model',
                          description='',
                          param_type=TypeMeta(name='Integer',
                                              properties={
                                                  'openapi_schema_validator': {
                                                      "type": "integer"
                                                  }
                                              })))

        self.assertEqual(containerOp._metadata, golden_meta)
Beispiel #2
0
def _create_container_op_from_resolved_task(name: str,
                                            container_image: str,
                                            command=None,
                                            arguments=None,
                                            output_paths=None,
                                            env: Mapping[str, str] = None,
                                            component_spec=None):
    from .. import dsl

    #Renaming outputs to conform with ContainerOp/Argo
    from ._naming import _sanitize_python_function_name, generate_unique_name_conversion_table
    output_names = (output_paths or {}).keys()
    output_name_to_kubernetes = generate_unique_name_conversion_table(
        output_names, _sanitize_python_function_name)
    output_paths_for_container_op = {
        output_name_to_kubernetes[name]: path
        for name, path in output_paths.items()
    }

    # Construct the ComponentMeta
    component_meta = ComponentMeta(name=component_spec.name,
                                   description=component_spec.description)
    # Inputs
    if component_spec.inputs is not None:
        for input in component_spec.inputs:
            component_meta.inputs.append(
                ParameterMeta(name=input.name,
                              description=input.description,
                              param_type=input.type,
                              default=input.default))
    if component_spec.outputs is not None:
        for output in component_spec.outputs:
            component_meta.outputs.append(
                ParameterMeta(name=output.name,
                              description=output.description,
                              param_type=output.type))

    task = dsl.ContainerOp(
        name=name,
        image=container_image,
        command=command,
        arguments=arguments,
        file_outputs=output_paths_for_container_op,
    )

    task._set_metadata(component_meta)

    if env:
        from kubernetes import client as k8s_client
        for name, value in env.items():
            task.container.add_env_variable(
                k8s_client.V1EnvVar(name=name, value=value))

    if component_spec.metadata:
        for key, value in (component_spec.metadata.annotations or {}).items():
            task.add_pod_annotation(key, value)
        for key, value in (component_spec.metadata.labels or {}).items():
            task.add_pod_label(key, value)

    return task
Beispiel #3
0
 def test_to_dict(self):
     component_meta = ComponentMeta(
         name='foobar',
         description='foobar example',
         inputs=[
             ParameterMeta(name='input1',
                           description='input1 desc',
                           param_type={
                               'GCSPath': {
                                   'bucket_type': 'directory',
                                   'file_type': 'csv'
                               }
                           },
                           default='default1'),
             ParameterMeta(name='input2',
                           description='input2 desc',
                           param_type={
                               'TFModel': {
                                   'input_data': 'tensor',
                                   'version': '1.8.0'
                               }
                           },
                           default='default2'),
             ParameterMeta(name='input3',
                           description='input3 desc',
                           param_type='Integer',
                           default='default3'),
         ],
         outputs=[
             ParameterMeta(name='output1',
                           description='output1 desc',
                           param_type={'Schema': {
                               'file_type': 'tsv'
                           }},
                           default='default_output1')
         ])
     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'
                 }
             },
             'default': 'default_output1'
         }]
     }
     self.assertEqual(component_meta.to_dict(), golden_meta)