feat(sdk): supporting bring your own container for arbitrary input and outputs

samples/v2/two_step_pipeline_containerized_test.py

@@ -54,31 +54,26 @@ def verify(t: unittest.TestCase, run: kfp_server_api.ApiRun,
             'outputs': {
                 'artifacts': [{
                     'name': 'output_gcs',
-                    'type': 'system.Artifact',
-                    # 'type': 'system.Dataset'
+                    'type': 'system.Dataset'
                 }],
             },
             'type': 'system.ContainerExecution',
             'state': Execution.State.COMPLETE,
-        },
-        component1_dict)
+        }, component1_dict)
 
     t.assertEqual(
         {
             'name': 'component2',
             'inputs': {
                 'artifacts': [{
                     'name': 'input_gcs',
-                    # TODO(chesu): compiled pipeline spec incorrectly sets importer artifact type to system.Artifact, but in the pipeline, it should be system.Dataset.
-                    'type': 'system.Artifact',
-                    # 'type': 'system.Dataset'
+                    'type': 'system.Dataset'
                 }],
             },
             'outputs': {},
             'type': 'system.ContainerExecution',
             'state': Execution.State.COMPLETE,
-        },
-        component2_dict)
+        }, component2_dict)
 
 
 if __name__ == '__main__':

sdk/python/kfp/compiler/compiler_test.py

@@ -874,12 +874,12 @@ def container_with_artifact_output(
                 ])
 
         with tempfile.TemporaryDirectory() as tempdir:
-            output_json = os.path.join(tempdir, 'component.yaml')
+            output_yaml = os.path.join(tempdir, 'component.yaml')
             compiler.Compiler().compile(
                 pipeline_func=container_with_artifact_output,
-                package_path=output_json,
+                package_path=output_yaml,
                 pipeline_name='container-with-artifact-output')
-            with open(output_json, 'r') as f:
+            with open(output_yaml, 'r') as f:
                 pipeline_spec = yaml.safe_load(f)
         self.assertEqual(
             pipeline_spec['components']['comp-container-with-artifact-output']

sdk/python/kfp/components/component_factory.py

@@ -62,9 +62,7 @@ def __getattr__(
                .InputPathPlaceholder, placeholders.OutputUriPlaceholder,
                placeholders.OutputPathPlaceholder]:
         if _name not in ['uri', 'path']:
-            raise AttributeError(
-                'Accessing artifact attribute other than uri or path is not supported.'
-            )
+            raise AttributeError('Cannot access artifact attribute "{_name}".')
         if self._io_type == 'input':
             if _name == 'uri':
                 return placeholders.InputUriPlaceholder(self._var_name)
@@ -279,8 +277,8 @@ def extract_component_interface(
             inputs[io_name] = input_spec
 
     #Analyzing the return type annotations.
+    return_ann = signature.return_annotation
     if not containerized:
-        return_ann = signature.return_annotation
         if hasattr(return_ann, '_fields'):  #NamedTuple
             # Getting field type annotations.
             # __annotations__ does not exist in python 3.5 and earlier
@@ -321,6 +319,10 @@ def extract_component_interface(
                 signature.return_annotation)
             output_spec = structures.OutputSpec(type=type_struct)
             outputs[output_name] = output_spec
+    elif return_ann != inspect.Parameter.empty and return_ann != structures.ContainerSpec:
+        raise TypeError(
+            'Return annotation should be either ContainerSpec or ignored for container components.'
+        )
 
     # Component name and description are derived from the function's name and
     # docstring.  The name can be overridden by setting setting func.__name__