Pass meta to containerop and pipeline

sdk/python/kfp/compiler/compiler.py

@@ -516,7 +516,7 @@ def _compile(self, pipeline_func):
     if pipeline_func not in registered_pipeline_functions:
       raise ValueError('Please use a function with @dsl.pipeline decorator.')
 
-    pipeline_name, _ = dsl.Pipeline.get_pipeline_functions()[pipeline_func]
+    pipeline_name = dsl.Pipeline.get_pipeline_functions()[pipeline_func].name
     pipeline_name = K8sHelper.sanitize_k8s_name(pipeline_name)
 
     # Create the arg list with no default values and call pipeline function.

sdk/python/kfp/components/_dsl_bridge.py

@@ -16,7 +16,7 @@
 from typing import Mapping
 from ._structures import ConcatPlaceholder, IfPlaceholder, InputValuePlaceholder, InputPathPlaceholder, IsPresentPlaceholder, OutputPathPlaceholder, TaskSpec
 from ._components import _generate_output_file_name, _default_component_name
-
+from kfp.dsl._metadata import ComponentMeta, ParameterMeta, TypeMeta, _annotation_to_typemeta
 
 def create_container_op_from_task(task_spec: TaskSpec):
     argument_values = task_spec.arguments
@@ -125,12 +125,13 @@ def expand_argument_list(argument_list):
         arguments=expanded_args,
         output_paths=output_paths,
         env=container_spec.env,
+        component_spec=component_spec,
     )
 
 
 _dummy_pipeline=None
 
-def _create_container_op_from_resolved_task(name:str, container_image:str, command=None, arguments=None, output_paths=None, env : Mapping[str, str]=None):
+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
     global _dummy_pipeline
     need_dummy = dsl.Pipeline._default_pipeline is None
@@ -150,6 +151,16 @@ def _create_container_op_from_resolved_task(name:str, container_image:str, comma
 
     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, type=_annotation_to_typemeta(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, type=_annotation_to_typemeta(input.type)))
+
     task = dsl.ContainerOp(
         name=name,
         image=container_image,
@@ -162,6 +173,8 @@ def _create_container_op_from_resolved_task(name:str, container_image:str, comma
         for name, value in env.items():
             task.add_env_variable(k8s_client.V1EnvVar(name=name, value=value))
 
+    task._set_metadata(component_meta)
+
     if need_dummy:
         _dummy_pipeline.__exit__()
 

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

@@ -16,10 +16,10 @@
 from . import _pipeline
 from . import _pipeline_param
 from ._pipeline_param import _extract_pipelineparams
+from ._metadata import ComponentMeta
 import re
 from typing import Dict
 
-
 class ContainerOp(object):
   """Represents an op implemented by a docker container image."""
 
@@ -64,6 +64,7 @@ def __init__(self, name: str, image: str, command: str=None, arguments: str=None
     self.pod_annotations = {}
     self.pod_labels = {}
     self.num_retries = 0
+    self._metadata = ComponentMeta()
 
     self.argument_inputs = _extract_pipelineparams([str(arg) for arg in (command or []) + (arguments or [])])
 
@@ -296,3 +297,12 @@ def set_retry(self, num_retries: int):
 
   def __repr__(self):
       return str({self.__class__.__name__: self.__dict__})
+
+  def _set_metadata(self, metadata):
+    '''_set_metadata passes the containerop the metadata information
+    Args:
+      metadata (ComponentMeta): component metadata
+    '''
+    if not isinstance(metadata, ComponentMeta):
+      raise ValueError('_set_medata is expecting ComponentMeta.')
+    self._metadata = metadata

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 BaseType, _check_valid_type_dict, _str_to_dict, _instance_to_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 ]
+            }
+
+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,7 +33,33 @@ def my_pipeline(a: PipelineParam, b: PipelineParam):
   ```
   """
   def _pipeline(func):
-    Pipeline.add_pipeline(name, description, func)
+    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 PipelineMeta
+    pipeline_meta = PipelineMeta(name=name, description=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])
+      pipeline_meta.inputs.append(ParameterMeta(name=arg, description='', param_type=arg_type, default=arg_default))
+
+    #TODO: add descriptions to the metadata
+    #docstring parser:
+    #  https://github.com/rr-/docstring_parser
+    #  https://github.com/terrencepreilly/darglint/blob/master/darglint/parse.py
+
+    Pipeline.add_pipeline(pipeline_meta, func)
     return func
 
   return _pipeline
@@ -93,9 +120,9 @@ def get_pipeline_functions():
     return Pipeline._pipeline_functions
 
   @staticmethod
-  def add_pipeline(name, description, func):
+  def add_pipeline(pipeline_meta, func):
     """Add a pipeline function (decorated with @pipeline)."""
-    Pipeline._pipeline_functions[func] = (name, description)
+    Pipeline._pipeline_functions[func] = pipeline_meta
 
   def __init__(self, name: str):
     """Create a new instance of Pipeline.
@@ -109,6 +136,7 @@ def __init__(self, name: str):
     self.groups = [_ops_group.OpsGroup('pipeline', name=name)]
     self.group_id = 0
     self.conf = PipelineConf()
+    self._metadata = PipelineMeta()
 
   def __enter__(self):
     if Pipeline._default_pipeline:
@@ -163,4 +191,13 @@ def get_next_group_id(self):
     self.group_id += 1
     return self.group_id
 
+  def _set_metadata(self, metadata):
+    '''_set_metadata passes the containerop the metadata information
+    Args:
+      metadata (ComponentMeta): component metadata
+    '''
+    if not isinstance(metadata, PipelineMeta):
+      raise ValueError('_set_medata is expecting PipelineMeta.')
+    self._metadata = metadata
+
 

sdk/python/kfp/dsl/_python_component.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from ._metadata import ComponentMeta, ParameterMeta, TypeMeta, _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 +49,49 @@ 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)
+    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 fullargspec.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
+
+    container_op = func(*args, **kargs)
+    container_op._set_metadata(component_meta)
+    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)