def test_decorator_metadata(self): """Test @pipeline decorator with metadata.""" @pipeline(name='p1', description='description1') def my_pipeline1(a: {'Schema': { 'file_type': 'csv' }} = 'good', b: Integer() = 12): pass golden_meta = PipelineMeta(name='p1', description='description1') golden_meta.inputs.append( ParameterMeta(name='a', description='', param_type=TypeMeta(name='Schema', properties={'file_type': 'csv'}), default='good')) golden_meta.inputs.append( ParameterMeta(name='b', description='', param_type=TypeMeta(name='Integer'), default=12)) pipeline_meta = Pipeline.get_pipeline_functions()[my_pipeline1] self.assertEqual(pipeline_meta, golden_meta)
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)
def test_decorator_metadata(self): """Test @pipeline decorator with metadata.""" @pipeline(name='p1', description='description1') def my_pipeline1(a: {'Schema': { 'file_type': 'csv' }} = 'good', b: Integer() = 12): pass golden_meta = PipelineMeta(name='p1', description='description1') golden_meta.inputs.append( ParameterMeta(name='a', description='', param_type={'Schema': { 'file_type': 'csv' }}, default='good')) golden_meta.inputs.append( ParameterMeta(name='b', description='', param_type={ 'Integer': { 'openapi_schema_validator': { "type": "integer" } } }, default=12)) pipeline_meta = _extract_pipeline_metadata(my_pipeline1) self.assertEqual(pipeline_meta, golden_meta)
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
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)