feat(sdk): support display_name and description in @dsl.pipeline deco…

…rator

sdk/python/kfp/compiler/compiler_test.py

@@ -282,20 +282,55 @@ def my_pipeline():
 
     def test_set_pipeline_root_through_pipeline_decorator(self):
 
-        with tempfile.TemporaryDirectory() as tmpdir:
+        @dsl.pipeline(name='test-pipeline', pipeline_root='gs://path')
+        def my_pipeline():
+            VALID_PRODUCER_COMPONENT_SAMPLE(input_param='input')
 
-            @dsl.pipeline(name='test-pipeline', pipeline_root='gs://path')
-            def my_pipeline():
-                VALID_PRODUCER_COMPONENT_SAMPLE(input_param='input')
+        self.assertEqual(my_pipeline.pipeline_spec.default_pipeline_root,
+                         'gs://path')
 
-            target_json_file = os.path.join(tmpdir, 'result.yaml')
-            compiler.Compiler().compile(
-                pipeline_func=my_pipeline, package_path=target_json_file)
+    def test_set_display_name_through_pipeline_decorator(self):
 
-            self.assertTrue(os.path.exists(target_json_file))
-            with open(target_json_file) as f:
-                pipeline_spec = yaml.safe_load(f)
-            self.assertEqual('gs://path', pipeline_spec['defaultPipelineRoot'])
+        @dsl.pipeline(display_name='my display name')
+        def my_pipeline():
+            VALID_PRODUCER_COMPONENT_SAMPLE(input_param='input')
+
+        self.assertEqual(my_pipeline.pipeline_spec.pipeline_info.display_name,
+                         'my display name')
+
+    def test_set_description_through_pipeline_decorator(self):
+
+        @dsl.pipeline(description='Prefer me.')
+        def my_pipeline():
+            """Don't prefer me"""
+            VALID_PRODUCER_COMPONENT_SAMPLE(input_param='input')
+
+        self.assertEqual(my_pipeline.pipeline_spec.pipeline_info.description,
+                         'Prefer me.')
+
+    def test_set_description_through_pipeline_docstring_short(self):
+
+        @dsl.pipeline
+        def my_pipeline():
+            """Docstring-specified description."""
+            VALID_PRODUCER_COMPONENT_SAMPLE(input_param='input')
+
+        self.assertEqual(my_pipeline.pipeline_spec.pipeline_info.description,
+                         'Docstring-specified description.')
+
+    def test_set_description_through_pipeline_docstring_long(self):
+
+        @dsl.pipeline
+        def my_pipeline():
+            """Docstring-specified description.
+
+            More information about this pipeline."""
+            VALID_PRODUCER_COMPONENT_SAMPLE(input_param='input')
+
+        self.assertEqual(
+            my_pipeline.pipeline_spec.pipeline_info.description,
+            'Docstring-specified description.\nMore information about this pipeline.'
+        )
 
     def test_passing_string_parameter_to_artifact_should_error(self):
 
@@ -1916,14 +1951,15 @@ def my_pipeline():
                 pipeline_func=my_pipeline, package_path=package_path)
 
 
+@dsl.component
+def identity(string: str, model: bool) -> str:
+    return string
+
+
 class TestYamlComments(unittest.TestCase):
 
     def test_comments_include_inputs_and_outputs_and_pipeline_name(self):
 
-        @dsl.component
-        def identity(string: str, model: bool) -> str:
-            return string
-
         @dsl.pipeline()
         def my_pipeline(sample_input1: bool = True,
                         sample_input2: str = 'string') -> str:
@@ -1957,57 +1993,124 @@ def my_pipeline(sample_input1: bool = True,
 
         self.assertIn(outputs_string, yaml_content)
 
-    def test_comments_include_definition(self):
-
-        @dsl.component
-        def identity(string: str, model: bool) -> str:
-            return string
+    def test_no_description(self):
 
         @dsl.pipeline()
-        def pipeline_with_no_definition(sample_input1: bool = True,
-                                        sample_input2: str = 'string') -> str:
+        def pipeline_with_no_description(sample_input1: bool = True,
+                                         sample_input2: str = 'string') -> str:
             op1 = identity(string=sample_input2, model=sample_input1)
             result = op1.output
             return result
 
         with tempfile.TemporaryDirectory() as tmpdir:
             pipeline_spec_path = os.path.join(tmpdir, 'output.yaml')
             compiler.Compiler().compile(
-                pipeline_func=pipeline_with_no_definition,
+                pipeline_func=pipeline_with_no_description,
                 package_path=pipeline_spec_path)
             with open(pipeline_spec_path, 'r+') as f:
                 yaml_content = f.read()
 
-            description_string = '# Description:'
+            # load and recompile to ensure idempotent description
+            loaded_pipeline = components.load_component_from_file(
+                pipeline_spec_path)
+
+            compiler.Compiler().compile(
+                pipeline_func=loaded_pipeline, package_path=pipeline_spec_path)
+
+            with open(pipeline_spec_path, 'r+') as f:
+                reloaded_yaml_content = f.read()
+
+        comment_description = '# Description:'
+        self.assertNotIn(comment_description, yaml_content)
+        self.assertNotIn(comment_description, reloaded_yaml_content)
+        proto_description = ''
+        self.assertEqual(
+            pipeline_with_no_description.pipeline_spec.pipeline_info
+            .description, proto_description)
+        self.assertEqual(
+            loaded_pipeline.pipeline_spec.pipeline_info.description,
+            proto_description)
 
-        self.assertNotIn(description_string, yaml_content)
+    def test_description_from_docstring(self):
 
         @dsl.pipeline()
-        def pipeline_with_definition(sample_input1: bool = True,
-                                     sample_input2: str = 'string') -> str:
-            """This is a definition of this pipeline."""
+        def pipeline_with_description(sample_input1: bool = True,
+                                      sample_input2: str = 'string') -> str:
+            """This is a description of this pipeline."""
             op1 = identity(string=sample_input2, model=sample_input1)
             result = op1.output
             return result
 
         with tempfile.TemporaryDirectory() as tmpdir:
             pipeline_spec_path = os.path.join(tmpdir, 'output.yaml')
             compiler.Compiler().compile(
-                pipeline_func=pipeline_with_definition,
+                pipeline_func=pipeline_with_description,
                 package_path=pipeline_spec_path)
+            with open(pipeline_spec_path, 'r+') as f:
+                yaml_content = f.read()
+
+            # load and recompile to ensure idempotent description
+            loaded_pipeline = components.load_component_from_file(
+                pipeline_spec_path)
 
+            compiler.Compiler().compile(
+                pipeline_func=loaded_pipeline, package_path=pipeline_spec_path)
+
+            with open(pipeline_spec_path, 'r+') as f:
+                reloaded_yaml_content = f.read()
+
+        comment_description = '# Description: This is a description of this pipeline.'
+        self.assertIn(comment_description, yaml_content)
+        self.assertIn(comment_description, reloaded_yaml_content)
+        proto_description = 'This is a description of this pipeline.'
+        self.assertEqual(
+            pipeline_with_description.pipeline_spec.pipeline_info.description,
+            proto_description)
+        self.assertEqual(
+            loaded_pipeline.pipeline_spec.pipeline_info.description,
+            proto_description)
+
+    def test_description_from_decorator(self):
+
+        @dsl.pipeline(description='Prefer this description.')
+        def pipeline_with_description(sample_input1: bool = True,
+                                      sample_input2: str = 'string') -> str:
+            """Don't prefer this description."""
+            op1 = identity(string=sample_input2, model=sample_input1)
+            result = op1.output
+            return result
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            pipeline_spec_path = os.path.join(tmpdir, 'output.yaml')
+            compiler.Compiler().compile(
+                pipeline_func=pipeline_with_description,
+                package_path=pipeline_spec_path)
             with open(pipeline_spec_path, 'r+') as f:
                 yaml_content = f.read()
 
-            description_string = '# Description:'
+            # load and recompile to ensure idempotent description
+            loaded_pipeline = components.load_component_from_file(
+                pipeline_spec_path)
 
-        self.assertIn(description_string, yaml_content)
+            compiler.Compiler().compile(
+                pipeline_func=loaded_pipeline, package_path=pipeline_spec_path)
 
-    def test_comments_on_pipeline_with_no_inputs_or_outputs(self):
+            with open(pipeline_spec_path, 'r+') as f:
+                reloaded_yaml_content = f.read()
 
-        @dsl.component
-        def identity(string: str, model: bool) -> str:
-            return string
+        comment_description = '# Description: Prefer this description.'
+        self.assertIn(comment_description, yaml_content)
+        self.assertIn(loaded_pipeline.pipeline_spec.pipeline_info.description,
+                      reloaded_yaml_content)
+        proto_description = 'Prefer this description.'
+        self.assertEqual(
+            pipeline_with_description.pipeline_spec.pipeline_info.description,
+            proto_description)
+        self.assertEqual(
+            loaded_pipeline.pipeline_spec.pipeline_info.description,
+            proto_description)
+
+    def test_comments_on_pipeline_with_no_inputs_or_outputs(self):
 
         @dsl.pipeline()
         def pipeline_with_no_inputs() -> str:
@@ -2048,10 +2151,6 @@ def pipeline_with_no_outputs(sample_input1: bool = True,
 
     def test_comments_follow_pattern(self):
 
-        @dsl.component
-        def identity(string: str, model: bool) -> str:
-            return string
-
         @dsl.pipeline()
         def my_pipeline(sample_input1: bool = True,
                         sample_input2: str = 'string') -> str:
@@ -2184,10 +2283,6 @@ def my_container_component(text: str, output_path: OutputPath(str)):
 
     def test_comments_idempotency(self):
 
-        @dsl.component
-        def identity(string: str, model: bool) -> str:
-            return string
-
         @dsl.pipeline()
         def my_pipeline(sample_input1: bool = True,
                         sample_input2: str = 'string') -> str:
@@ -2227,10 +2322,6 @@ def my_pipeline(sample_input1: bool = True,
 
     def test_comment_with_multiline_docstring(self):
 
-        @dsl.component
-        def identity(string: str, model: bool) -> str:
-            return string
-
         @dsl.pipeline()
         def pipeline_with_multiline_definition(
                 sample_input1: bool = True,
@@ -2290,10 +2381,6 @@ def pipeline_with_multiline_definition(
 
     def test_idempotency_on_comment_with_multiline_docstring(self):
 
-        @dsl.component
-        def identity(string: str, model: bool) -> str:
-            return string
-
         @dsl.pipeline()
         def my_pipeline(sample_input1: bool = True,
                         sample_input2: str = 'string') -> str:

sdk/python/kfp/components/component_factory.py

@@ -163,14 +163,18 @@ def _maybe_make_unique(name: str, names: List[str]):
 
 
 def extract_component_interface(
-        func: Callable,
-        containerized: bool = False) -> structures.ComponentSpec:
+    func: Callable,
+    containerized: bool = False,
+    description: Optional[str] = None,
+    name: Optional[str] = None,
+) -> structures.ComponentSpec:
     single_output_name_const = 'Output'
 
     signature = inspect.signature(func)
     parameters = list(signature.parameters.values())
 
-    parsed_docstring = docstring_parser.parse(inspect.getdoc(func))
+    original_docstring = inspect.getdoc(func)
+    parsed_docstring = docstring_parser.parse(original_docstring)
 
     inputs = {}
     outputs = {}
@@ -341,33 +345,21 @@ def extract_component_interface(
             'Return annotation should be either ContainerSpec or omitted 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__
-    # attribute (of the legacy func._component_human_name attribute).  The
-    # description can be overridden by setting the func.__doc__ attribute (or
-    # the legacy func._component_description attribute).
-    component_name = getattr(
-        func, '_component_human_name',
-        _python_function_name_to_component_name(func.__name__))
-
-    short_description = parsed_docstring.short_description
-    long_description = parsed_docstring.long_description
-    docstring_description = short_description + '\n' + long_description if long_description else short_description
+    component_name = name or _python_function_name_to_component_name(
+        func.__name__)
 
-    description = getattr(func, '_component_description', docstring_description)
-
-    if description:
-        description = description.strip()
+    description = get_pipeline_description(
+        decorator_description=description,
+        docstring=parsed_docstring,
+    )
 
-    component_spec = structures.ComponentSpec(
+    return structures.ComponentSpec(
         name=component_name,
         description=description,
-        inputs=inputs if inputs else None,
-        outputs=outputs if outputs else None,
-        # Dummy implementation to bypass model validation.
+        inputs=inputs or None,
+        outputs=outputs or None,
         implementation=structures.Implementation(),
     )
-    return component_spec
 
 
 def _get_command_and_args_for_lightweight_component(
@@ -562,20 +554,43 @@ def create_container_component_from_func(
 
 
 def create_graph_component_from_func(
-        func: Callable) -> graph_component.GraphComponent:
+    func: Callable,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    display_name: Optional[str] = None,
+) -> graph_component.GraphComponent:
     """Implementation for the @pipeline decorator.
 
     The decorator is defined under pipeline_context.py. See the
     decorator for the canonical documentation for this function.
     """
 
-    component_spec = extract_component_interface(func)
-    component_name = getattr(
-        func, '_component_human_name',
-        _python_function_name_to_component_name(func.__name__))
-
+    component_spec = extract_component_interface(
+        func,
+        description=description,
+        name=name,
+    )
     return graph_component.GraphComponent(
         component_spec=component_spec,
         pipeline_func=func,
-        name=component_name,
+        display_name=display_name,
     )
+
+
+def get_pipeline_description(
+    decorator_description: Union[str, None],
+    docstring: docstring_parser.Docstring,
+) -> Union[str, None]:
+    """Obtains the correct pipeline description from the pipeline decorator's
+    description argument and the parsed docstring.
+
+    Gives precedence to the decorator argument.
+    """
+    if decorator_description:
+        return decorator_description
+
+    short_description = docstring.short_description
+    long_description = docstring.long_description
+    docstring_description = short_description + '\n' + long_description if (
+        short_description and long_description) else short_description
+    return docstring_description.strip() if docstring_description else None