Add component metadata

sdk/python/kfp/dsl/__init__.py

@@ -17,4 +17,4 @@
 from ._pipeline import Pipeline, pipeline, get_pipeline_conf
 from ._container_op import ContainerOp
 from ._ops_group import OpsGroup, ExitHandler, Condition
-from ._python_component import python_component
+from ._python_component import python_component

sdk/python/kfp/dsl/_container_op.py

@@ -19,7 +19,6 @@
 import re
 from typing import Dict
 
-
 class ContainerOp(object):
   """Represents an op implemented by a docker container image."""
 

sdk/python/kfp/dsl/_metadata.py

@@ -0,0 +1,107 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List
+from abc import ABCMeta, abstractmethod
+from ._types import _check_valid_type_dict
+
+class BaseMeta(object):
+  __metaclass__ = ABCMeta
+  def __init__(self):
+    pass
+
+  @abstractmethod
+  def to_dict(self):
+    pass
+
+  def serialize(self):
+    import yaml
+    return yaml.dump(self.to_dict())
+
+  def __eq__(self, other):
+    return self.__dict__ == other.__dict__
+
+class TypeMeta(BaseMeta):
+  def __init__(self,
+      name: str = '',
+      properties: Dict = None):
+    self.name = name
+    self.properties = {} if properties is None else properties
+
+  def to_dict(self):
+    return {self.name: self.properties}
+
+  @staticmethod
+  def from_dict(json_dict):
+    if not _check_valid_type_dict(json_dict):
+      raise ValueError(json_dict + ' is not a valid type string')
+    type_meta = TypeMeta()
+    type_meta.name, type_meta.properties = list(json_dict.items())[0]
+    return type_meta
+
+class ParameterMeta(BaseMeta):
+  def __init__(self,
+      name: str = '',
+      description: str = '',
+      param_type: TypeMeta = None,
+      default = ''):
+    self.name = name
+    self.description = description
+    self.param_type = TypeMeta() if param_type is None else param_type
+    self.default = default
+
+  def to_dict(self):
+    return {'name': self.name,
+            'description': self.description,
+            'type': self.param_type.to_dict(),
+            'default': self.default}
+
+class ComponentMeta(BaseMeta):
+  def __init__(
+      self,
+      name: str = '',
+      description: str = '',
+      inputs: List[ParameterMeta] = None,
+      outputs: List[ParameterMeta] = None
+  ):
+    self.name = name
+    self.description = description
+    self.inputs = [] if inputs is None else inputs
+    self.outputs = [] if outputs is None else outputs
+
+  def to_dict(self):
+    return {'name': self.name,
+            'description': self.description,
+            'inputs': [ input.to_dict() for input in self.inputs ],
+            'outputs': [ output.to_dict() for output in self.outputs ]
+            }
+
+# Add a pipeline level metadata calss here.
+# If one day we combine the component and pipeline yaml, ComponentMeta and PipelineMeta will become one, too.
+class PipelineMeta(BaseMeta):
+  def __init__(
+      self,
+      name: str = '',
+      description: str = '',
+      inputs: List[ParameterMeta] = None
+  ):
+    self.name = name
+    self.description = description
+    self.inputs = [] if inputs is None else inputs
+
+  def to_dict(self):
+    return {'name': self.name,
+            'description': self.description,
+            'inputs': [ input.to_dict() for input in self.inputs ]
+            }

sdk/python/kfp/dsl/_types.py

@@ -144,7 +144,7 @@ def _str_to_dict(payload):
 	return json_dict
 
 def _check_dict_types(checked_type, expected_type):
-	'''_check_type_types checks the type consistency.
+	'''_check_dict_types checks the type consistency.
 	Args:
   	checked_type (dict): A dict that describes a type from the upstream component output
   	expected_type (dict): A dict that describes a type from the downstream component input

sdk/python/tests/dsl/container_op_tests.py

@@ -16,7 +16,6 @@
 from kfp.dsl import Pipeline, PipelineParam, ContainerOp
 import unittest
 
-
 class TestContainerOp(unittest.TestCase):
 
   def test_basic(self):

sdk/python/tests/dsl/main.py

@@ -21,6 +21,7 @@
 import container_op_tests
 import ops_group_tests
 import type_tests
+import metadata_tests
 
 if __name__ == '__main__':
   suite = unittest.TestSuite()
@@ -29,6 +30,7 @@
   suite.addTests(unittest.defaultTestLoader.loadTestsFromModule(container_op_tests))
   suite.addTests(unittest.defaultTestLoader.loadTestsFromModule(ops_group_tests))
   suite.addTests(unittest.defaultTestLoader.loadTestsFromModule(type_tests))
+  suite.addTests(unittest.defaultTestLoader.loadTestsFromModule(metadata_tests))
   runner = unittest.TextTestRunner()
   if not runner.run(suite).wasSuccessful():
     sys.exit(1)

sdk/python/tests/dsl/metadata_tests.py

@@ -0,0 +1,118 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from kfp.dsl._metadata import ComponentMeta, ParameterMeta, TypeMeta
+import unittest
+
+class TestTypeMeta(unittest.TestCase):
+  def test_from_dict(self):
+    component_dict = {
+        'GCSPath': {
+            'bucket_type': 'directory',
+            'file_type': 'csv'
+        }
+    }
+    golden_type_meta = TypeMeta(name='GCSPath', properties={'bucket_type': 'directory',
+                                                          'file_type': 'csv'})
+    self.assertEqual(TypeMeta.from_dict(component_dict), golden_type_meta)
+
+  def test_eq(self):
+    type_a = TypeMeta(name='GCSPath', properties={'bucket_type': 'directory',
+                                                  'file_type': 'csv'})
+    type_b = TypeMeta(name='GCSPath', properties={'bucket_type': 'directory',
+                                                  'file_type': 'tsv'})
+    type_c = TypeMeta(name='GCSPatha', properties={'bucket_type': 'directory',
+                                                  'file_type': 'csv'})
+    type_d = TypeMeta(name='GCSPath', properties={'bucket_type': 'directory',
+                                                  'file_type': 'csv'})
+    self.assertNotEqual(type_a, type_b)
+    self.assertNotEqual(type_a, type_c)
+    self.assertEqual(type_a, type_d)
+
+
+class TestComponentMeta(unittest.TestCase):
+
+  def test_to_dict(self):
+    component_meta = ComponentMeta(name='foobar',
+                                   description='foobar example',
+                                   inputs=[ParameterMeta(name='input1',
+                                                         description='input1 desc',
+                                                         param_type=TypeMeta(name='GCSPath',
+                                                                             properties={'bucket_type': 'directory',
+                                                                                         'file_type': 'csv'
+                                                                                         }
+                                                                             ),
+                                                         default='default1'
+                                                         ),
+                                           ParameterMeta(name='input2',
+                                                         description='input2 desc',
+                                                         param_type=TypeMeta(name='TFModel',
+                                                                             properties={'input_data': 'tensor',
+                                                                                         'version': '1.8.0'
+                                                                                         }
+                                                                             ),
+                                                         default='default2'
+                                                         ),
+                                           ],
+                                   outputs=[ParameterMeta(name='output1',
+                                                          description='output1 desc',
+                                                          param_type=TypeMeta(name='Schema',
+                                                                              properties={'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'
+            }
+        ],
+        'outputs': [
+            {
+                'name': 'output1',
+                'description': 'output1 desc',
+                'type': {
+                    'Schema': {
+                        'file_type': 'tsv'
+                    }
+                },
+                'default': 'default_output1'
+            }
+        ]
+    }
+    self.assertEqual(component_meta.to_dict(), golden_meta)