feat(sdk): support `display_name` and `description` in `@dsl.pipeline` decorator (#9153)

* feat(sdk): support display_name and description in @dsl.pipeline decorator

* add release note

* test additional case
This commit is contained in:
Connor McCarthy 2023-04-17 14:15:39 -07:00 committed by GitHub
parent df8565c73e
commit 91abbeaf2f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 221 additions and 93 deletions

View File

@ -1,6 +1,7 @@
# Current Version (in development)
## Features
* Support `display_name` and `description` in `@dsl.pipeline` decorator [\#9153](https://github.com/kubeflow/pipelines/pull/9153)
## Breaking changes

View File

@ -282,20 +282,69 @@ class TestCompilePipeline(parameterized.TestCase):
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')
target_json_file = os.path.join(tmpdir, 'result.yaml')
compiler.Compiler().compile(
pipeline_func=my_pipeline, package_path=target_json_file)
self.assertEqual(my_pipeline.pipeline_spec.default_pipeline_root,
'gs://path')
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'])
def test_set_display_name_through_pipeline_decorator(self):
@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_name_and_display_name_through_pipeline_decorator(self):
@dsl.pipeline(
name='my-pipeline-name',
display_name='my display name',
)
def my_pipeline():
VALID_PRODUCER_COMPONENT_SAMPLE(input_param='input')
self.assertEqual(my_pipeline.pipeline_spec.pipeline_info.name,
'my-pipeline-name')
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 +1965,15 @@ class TestCannotUseAfterCrossDAG(unittest.TestCase):
pipeline_func=my_pipeline, package_path=package_path)
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
class TestYamlComments(unittest.TestCase):
def test_comments_include_inputs_and_outputs_and_pipeline_name(self):
@dsl.pipeline()
def my_pipeline(sample_input1: bool = True,
sample_input2: str = 'string') -> str:
@ -1957,14 +2007,10 @@ class TestYamlComments(unittest.TestCase):
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,
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
@ -1973,19 +2019,38 @@ class TestYamlComments(unittest.TestCase):
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)
self.assertNotIn(description_string, yaml_content)
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)
def test_description_from_docstring(self):
@dsl.pipeline()
def pipeline_with_definition(sample_input1: bool = True,
def pipeline_with_description(sample_input1: bool = True,
sample_input2: str = 'string') -> str:
"""This is a definition of this pipeline."""
"""This is a description of this pipeline."""
op1 = identity(string=sample_input2, model=sample_input1)
result = op1.output
return result
@ -1993,22 +2058,74 @@ class TestYamlComments(unittest.TestCase):
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()
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)
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()
# 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: 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.component
def identity(string: str, model: bool) -> str:
return string
@dsl.pipeline()
def pipeline_with_no_inputs() -> str:
op1 = identity(string='string', model=True)
@ -2048,10 +2165,6 @@ class TestYamlComments(unittest.TestCase):
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 +2297,6 @@ class TestYamlComments(unittest.TestCase):
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 +2336,6 @@ class TestYamlComments(unittest.TestCase):
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 +2395,6 @@ class TestYamlComments(unittest.TestCase):
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:

View File

@ -164,13 +164,17 @@ def _maybe_make_unique(name: str, names: List[str]):
def extract_component_interface(
func: Callable,
containerized: bool = False) -> structures.ComponentSpec:
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__))
component_name = name or _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
description = get_pipeline_description(
decorator_description=description,
docstring=parsed_docstring,
)
description = getattr(func, '_component_description', docstring_description)
if description:
description = description.strip()
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

View File

@ -14,7 +14,7 @@
"""Pipeline as a component (aka graph component)."""
import inspect
from typing import Callable
from typing import Callable, Optional
import uuid
from kfp.compiler import pipeline_spec_builder as builder
@ -36,11 +36,10 @@ class GraphComponent(base_component.BaseComponent):
self,
component_spec: structures.ComponentSpec,
pipeline_func: Callable,
name: str,
display_name: Optional[str] = None,
):
super().__init__(component_spec=component_spec)
self.pipeline_func = pipeline_func
self.name = name
args_list = []
signature = inspect.signature(pipeline_func)
@ -75,6 +74,10 @@ class GraphComponent(base_component.BaseComponent):
pipeline_root = getattr(pipeline_func, 'pipeline_root', None)
if pipeline_root is not None:
pipeline_spec.default_pipeline_root = pipeline_root
if display_name is not None:
pipeline_spec.pipeline_info.display_name = display_name
if component_spec.description is not None:
pipeline_spec.pipeline_info.description = component_spec.description
self.component_spec.implementation.graph = pipeline_spec
self.component_spec.platform_spec = platform_spec

View File

@ -17,7 +17,6 @@ import functools
from typing import Callable, Optional
from kfp.components import component_factory
from kfp.components import graph_component
from kfp.components import pipeline_task
from kfp.components import tasks_group
from kfp.components import utils
@ -27,7 +26,8 @@ def pipeline(func: Optional[Callable] = None,
*,
name: Optional[str] = None,
description: Optional[str] = None,
pipeline_root: Optional[str] = None) -> Callable:
pipeline_root: Optional[str] = None,
display_name: Optional[str] = None) -> Callable:
"""Decorator used to construct a pipeline.
Example
@ -48,6 +48,7 @@ def pipeline(func: Optional[Callable] = None,
description: A human-readable description of the pipeline.
pipeline_root: The root directory from which to read input and output
parameters and artifacts.
display_name: A human-readable name for the pipeline.
"""
if func is None:
return functools.partial(
@ -55,16 +56,18 @@ def pipeline(func: Optional[Callable] = None,
name=name,
description=description,
pipeline_root=pipeline_root,
display_name=display_name,
)
if name:
func._component_human_name = name
if description:
func._component_description = description
if pipeline_root:
func.pipeline_root = pipeline_root
return component_factory.create_graph_component_from_func(func)
return component_factory.create_graph_component_from_func(
func,
name=name,
description=description,
display_name=display_name,
)
class Pipeline:

View File

@ -67,6 +67,7 @@ print_op = components.load_component_from_text("""
@dsl.pipeline(
name='conditional-execution-pipeline',
display_name='Conditional execution pipeline.',
pipeline_root='dummy_root',
description='Shows how to use dsl.Condition().')
def my_pipeline():

View File

@ -302,6 +302,8 @@ deploymentSpec:
- '{{$.inputs.parameters[''msg'']}}'
image: python:alpine3.6
pipelineInfo:
description: Shows how to use dsl.Condition().
displayName: Conditional execution pipeline.
name: conditional-execution-pipeline
root:
dag:

View File

@ -61,6 +61,7 @@ deploymentSpec:
- '{{$.outputs.artifacts[''model''].uri}}'
image: gcr.io/my-project/my-fancy-trainer
pipelineInfo:
description: A linear two-step pipeline with artifact ontology types.
name: two-step-pipeline-with-ontology
root:
dag:

View File

@ -69,6 +69,7 @@ deploymentSpec:
memoryLimit: 15.032385536
memoryRequest: 4.294967296
pipelineInfo:
description: A linear two-step pipeline with resource specification.
name: two-step-pipeline-with-resource-spec
root:
dag: