Add python conf to the metadata

sdk/python/kfp/dsl/__init__.py

@@ -18,3 +18,4 @@
 from ._container_op import ContainerOp
 from ._ops_group import OpsGroup, ExitHandler, Condition
 from ._python_component import python_component
+from ._metadata import ComponentMeta, ParameterMeta, TypeMeta

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,126 @@
+# 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())
+
+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
+
+  def __eq__(self, other):
+    return self.__dict__ == other.__dict__
+
+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 ]
+            }
+
+def _annotation_to_typemeta(annotation):
+  '''_annotation_to_type_meta converts an annotation to an instance of TypeMeta
+  Args:
+    annotation(BaseType/str/dict): input/output annotations
+  Returns:
+    TypeMeta
+    '''
+  if isinstance(annotation, BaseType):
+    arg_type = TypeMeta.from_dict(_instance_to_dict(annotation))
+  elif isinstance(annotation, str):
+    arg_type = TypeMeta.from_dict(_str_to_dict(annotation))
+  elif isinstance(annotation, dict):
+    if not _check_valid_type_dict(annotation):
+      raise ValueError('Annotation ' + str(annotation) + ' is not a valid type dictionary.')
+    arg_type = TypeMeta.from_dict(annotation)
+  else:
+    return TypeMeta()
+  return arg_type

sdk/python/kfp/dsl/_pipeline.py

@@ -14,6 +14,7 @@
 
 
 from . import _container_op
+from ._metadata import  PipelineMeta, ParameterMeta, TypeMeta, _annotation_to_typemeta
 from . import _ops_group
 import sys
 
@@ -32,6 +33,27 @@ def my_pipeline(a: PipelineParam, b: PipelineParam):
   ```
   """
   def _pipeline(func):
+    import inspect
+    fullargspec = inspect.getfullargspec(func)
+    args = fullargspec.args
+    annotations = fullargspec.annotations
+
+    # Construct the PipelineMeta
+    pipeline_meta = PipelineMeta(name=func.__name__, description='')
+    # Inputs
+    for arg in args:
+      arg_type = TypeMeta()
+      if arg in annotations:
+        arg_type = _annotation_to_typemeta(annotations[arg])
+      pipeline_meta.inputs.append(ParameterMeta(name=arg, description='', param_type=arg_type))
+
+    #TODO: add descriptions to the metadata
+    #docstring parser:
+    #  https://github.com/rr-/docstring_parser
+    #  https://github.com/terrencepreilly/darglint/blob/master/darglint/parse.py
+    print(pipeline_meta.serialize())
+    #TODO: parse the metadata to the Pipeline.
+
     Pipeline.add_pipeline(name, description, func)
     return func
 

sdk/python/kfp/dsl/_python_component.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from kfp.dsl import ComponentMeta, ParameterMeta, TypeMeta
+from ._metadata import _annotation_to_typemeta
+
 def python_component(name, description=None, base_image=None, target_component_file: str = None):
   """Decorator for Python component functions.
   This decorator adds the metadata to the function object itself.
@@ -47,3 +50,51 @@ def _python_component(func):
     return func
 
   return _python_component
+
+def component(func):
+  """Decorator for component functions that use ContainerOp.
+  This is useful to enable type checking in the DSL compiler
+
+  Usage:
+  ```python
+  @dsl.component
+  def foobar(model: TFModel(), step: MLStep()):
+    return dsl.ContainerOp()
+  """
+  def _component(*args, **kargs):
+    import inspect
+    fullargspec = inspect.getfullargspec(func)
+    args = fullargspec.args
+    annotations = fullargspec.annotations
+
+    # defaults
+    arg_defaults = {}
+    if fullargspec.defaults:
+      for arg, default in zip(reversed(fullargspec.args), reversed(fullargspec.defaults)):
+        arg_defaults[arg] = default
+
+    # Construct the ComponentMeta
+    component_meta = ComponentMeta(name=func.__name__, description='')
+    # Inputs
+    for arg in args:
+      arg_type = TypeMeta()
+      arg_default = arg_defaults[arg] if arg in arg_defaults else ''
+      if arg in annotations:
+        arg_type = _annotation_to_typemeta(annotations[arg])
+      component_meta.inputs.append(ParameterMeta(name=arg, description='', param_type=arg_type, default=arg_default))
+    # Outputs
+    for output in annotations['return']:
+      arg_type = _annotation_to_typemeta(annotations['return'][output])
+      component_meta.outputs.append(ParameterMeta(name=output, description='', param_type=arg_type))
+
+    #TODO: add descriptions to the metadata
+    #docstring parser:
+    #  https://github.com/rr-/docstring_parser
+    #  https://github.com/terrencepreilly/darglint/blob/master/darglint/parse.py
+
+    print(component_meta.serialize())
+    #TODO: parse the metadata to the ContainerOp.
+    container_op = func(*args, **kargs)
+    return container_op
+
+  return _component

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,8 @@
 import container_op_tests
 import ops_group_tests
 import type_tests
+import python_component_tests
+import metadata_tests
 
 if __name__ == '__main__':
   suite = unittest.TestSuite()
@@ -29,6 +31,8 @@
   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(python_component_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,103 @@
+# 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 import ComponentMeta, ParameterMeta, TypeMeta
+import unittest
+
+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)
+
+  def test_type_meta_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)