From f2e628c3a70ec5bda15eba5397490293ced23515 Mon Sep 17 00:00:00 2001 From: Connor McCarthy Date: Fri, 4 Nov 2022 11:20:16 -0700 Subject: [PATCH] fix(sdk): fix nested placeholders and block illegal IfPresent form in Concat (#8414) * support single element IfPresentPlaceholders that contain a primitive placeholder * block list then and else_ inside concat * support single element IfPresent in v1 * clean up typing * skip loading batch_predict_job yaml --- components/test_load_all_components.sh | 3 + sdk/python/kfp/components/pipeline_task.py | 19 +- sdk/python/kfp/components/placeholders.py | 70 +++++-- .../kfp/components/placeholders_test.py | 171 ++++++++++++++++-- sdk/python/kfp/components/structures.py | 18 +- sdk/python/kfp/components/structures_test.py | 30 +-- 6 files changed, 257 insertions(+), 54 deletions(-) diff --git a/components/test_load_all_components.sh b/components/test_load_all_components.sh index 00f44282f0f..a79d53698dd 100755 --- a/components/test_load_all_components.sh +++ b/components/test_load_all_components.sh @@ -35,6 +35,9 @@ SKIP_COMPONENT_FILES = [ "./contrib/XGBoost/Cross_validation_for_regression/from_CSV/component.yaml", "./contrib/XGBoost/Train_regression_and_calculate_metrics/from_CSV/component.yaml", "./contrib/XGBoost/Train_and_cross-validate_regression/from_CSV/component.yaml", + # TODO: This component uses invalid placeholders. Updated when migrating GCPC to v2. + "./google-cloud/google_cloud_pipeline_components/aiplatform/batch_predict_job/component.yaml", + "./google-cloud/google_cloud_pipeline_components/v1/batch_predict_job/component.yaml" ] for component_file in sys.stdin: diff --git a/sdk/python/kfp/components/pipeline_task.py b/sdk/python/kfp/components/pipeline_task.py index e1a710a20a6..c0e1a1ba5fe 100644 --- a/sdk/python/kfp/components/pipeline_task.py +++ b/sdk/python/kfp/components/pipeline_task.py @@ -474,7 +474,7 @@ def my_pipeline(): def check_primitive_placeholder_is_used_for_correct_io_type( inputs_dict: Dict[str, structures.InputSpec], outputs_dict: Dict[str, structures.OutputSpec], - arg: Union[placeholders.Placeholder, Any], + arg: Union[placeholders.CommandLineElement, Any], ): """Validates input/output placeholders refer to an input/output with an appropriate type for the placeholder. This should only apply to components @@ -514,7 +514,22 @@ def check_primitive_placeholder_is_used_for_correct_io_type( f'"{outputs_dict[output_name].type}" cannot be paired with ' f'{arg.__class__.__name__}.') elif isinstance(arg, placeholders.IfPresentPlaceholder): - for arg in itertools.chain(arg.then or [], arg.else_ or []): + all_normalized_args: List[placeholders.CommandLineElement] = [] + if arg.then is None: + pass + elif isinstance(arg.then, list): + all_normalized_args.extend(arg.then) + else: + all_normalized_args.append(arg.then) + + if arg.else_ is None: + pass + elif isinstance(arg.else_, list): + all_normalized_args.extend(arg.else_) + else: + all_normalized_args.append(arg.else_) + + for arg in all_normalized_args: check_primitive_placeholder_is_used_for_correct_io_type( inputs_dict, outputs_dict, arg) elif isinstance(arg, placeholders.ConcatPlaceholder): diff --git a/sdk/python/kfp/components/placeholders.py b/sdk/python/kfp/components/placeholders.py index eec779e8c83..05aaded6fb7 100644 --- a/sdk/python/kfp/components/placeholders.py +++ b/sdk/python/kfp/components/placeholders.py @@ -141,6 +141,9 @@ def container_with_concat_placeholder(text1: str, text2: Output[Dataset], """ def __init__(self, items: List['CommandLineElement']) -> None: + for item in items: + if isinstance(item, IfPresentPlaceholder): + item._validate_then_and_else_are_only_single_element() self.items = items def _to_dict(self) -> Dict[str, Any]: @@ -197,6 +200,36 @@ def __init__( self.then = then self.else_ = else_ + def _validate_then_and_else_are_only_single_element(self) -> None: + """Rercursively validate that then and else contain only a single + element. + + This method should only be called by a ConcatPlaceholder, which + cannot have an IfPresentPlaceholder with a list in either 'then' + or 'else_'. + """ + + # the illegal state + if isinstance(self.then, list) or isinstance(self.else_, list): + raise ValueError( + f'Cannot use {IfPresentPlaceholder.__name__} within {ConcatPlaceholder.__name__} when `then` and `else_` arguments to {IfPresentPlaceholder.__name__} are lists. Please use a single element for `then` and `else_` only.' + ) + + # check that there is no illegal state found recursively + if isinstance(self.then, ConcatPlaceholder): + for item in self.then.items: + if isinstance(item, IfPresentPlaceholder): + item._validate_then_and_else_are_only_single_element() + elif isinstance(self.then, IfPresentPlaceholder): + self.then._validate_then_and_else_are_only_single_element() + + if isinstance(self.else_, ConcatPlaceholder): + for item in self.else_.items: + if isinstance(item, IfPresentPlaceholder): + item._validate_then_and_else_are_only_single_element() + elif isinstance(self.else_, IfPresentPlaceholder): + self.else_._validate_then_and_else_are_only_single_element() + def _to_dict(self) -> Dict[str, Any]: struct = { 'IfPresent': { @@ -231,11 +264,7 @@ def _to_string(self) -> str: OutputPathPlaceholder, OutputUriPlaceholder, OutputMetadataPlaceholder) -CommandLineElement = Union[str, IfPresentPlaceholder, ConcatPlaceholder, - InputValuePlaceholder, InputPathPlaceholder, - InputUriPlaceholder, InputMetadataPlaceholder, - OutputParameterPlaceholder, OutputPathPlaceholder, - OutputUriPlaceholder, OutputMetadataPlaceholder] +CommandLineElement = Union[str, Placeholder] def convert_command_line_element_to_string( @@ -330,18 +359,31 @@ def maybe_convert_v1_yaml_placeholder_to_v2_placeholder( elif 'if' in arg: if_ = arg['if'] input_name = utils.sanitize_input_name(if_['cond']['isPresent']) - then_ = if_['then'] - else_ = if_.get('else', []) - return IfPresentPlaceholder( - input_name=input_name, - then=[ + then = if_['then'] + else_ = if_.get('else') + + if isinstance(then, list): + then = [ maybe_convert_v1_yaml_placeholder_to_v2_placeholder( - val, component_dict=component_dict) for val in then_ - ], - else_=[ + val, component_dict=component_dict) for val in then + ] + else: + then = maybe_convert_v1_yaml_placeholder_to_v2_placeholder( + then, component_dict=component_dict) + + if else_ is None: + pass + elif isinstance(else_, list): + else_ = [ maybe_convert_v1_yaml_placeholder_to_v2_placeholder( val, component_dict=component_dict) for val in else_ - ]) + ] + else: + maybe_convert_v1_yaml_placeholder_to_v2_placeholder( + else_, component_dict=component_dict) + + return IfPresentPlaceholder( + input_name=input_name, then=then, else_=else_) elif 'concat' in arg: diff --git a/sdk/python/kfp/components/placeholders_test.py b/sdk/python/kfp/components/placeholders_test.py index 2ab5502c439..74d614663be 100644 --- a/sdk/python/kfp/components/placeholders_test.py +++ b/sdk/python/kfp/components/placeholders_test.py @@ -12,9 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. """Contains tests for kfp.components.placeholders.""" +import os +import tempfile from typing import Any from absl.testing import parameterized +from kfp import compiler +from kfp import dsl from kfp.components import placeholders @@ -89,7 +93,7 @@ def test(self): "{{$.outputs.artifacts['output1'].metadata}}") -class TestIfPresentPlaceholderStructure(parameterized.TestCase): +class TestIfPresentPlaceholder(parameterized.TestCase): @parameterized.parameters([ (placeholders.IfPresentPlaceholder( @@ -139,6 +143,56 @@ def test_with_primitive_placeholders( placeholder: str): self.assertEqual(placeholder_obj._to_string(), placeholder) + def test_if_present_with_single_element_simple_can_be_compiled(self): + + @dsl.container_component + def container_component(a: str): + return dsl.ContainerSpec( + image='alpine', + command=[ + placeholders.IfPresentPlaceholder( + input_name='a', then='b', else_='c') + ]) + + with tempfile.TemporaryDirectory() as tempdir: + output_yaml = os.path.join(tempdir, 'component.yaml') + compiler.Compiler().compile( + pipeline_func=container_component, package_path=output_yaml) + + def test_if_present_with_single_element_parameter_reference_can_be_compiled( + self): + + @dsl.container_component + def container_component(a: str): + return dsl.ContainerSpec( + image='alpine', + command=[ + placeholders.IfPresentPlaceholder( + input_name='a', then=a, else_='c') + ]) + + with tempfile.TemporaryDirectory() as tempdir: + output_yaml = os.path.join(tempdir, 'component.yaml') + compiler.Compiler().compile( + pipeline_func=container_component, package_path=output_yaml) + + def test_if_present_with_single_element_artifact_reference_can_be_compiled( + self): + + @dsl.container_component + def container_component(a: dsl.Input[dsl.Artifact]): + return dsl.ContainerSpec( + image='alpine', + command=[ + placeholders.IfPresentPlaceholder( + input_name='a', then=a.path, else_='c') + ]) + + with tempfile.TemporaryDirectory() as tempdir: + output_yaml = os.path.join(tempdir, 'component.yaml') + compiler.Compiler().compile( + pipeline_func=container_component, package_path=output_yaml) + class TestConcatPlaceholder(parameterized.TestCase): @@ -182,16 +236,6 @@ class TestContainerPlaceholdersTogether(parameterized.TestCase): ]), """{"Concat": ["a", {"Concat": ["b", {"Concat": ["c", "{{$.inputs.parameters['input2']}}"]}]}]}""" ), - (placeholders.ConcatPlaceholder([ - 'a', - placeholders.ConcatPlaceholder([ - 'b', - placeholders.IfPresentPlaceholder( - input_name='output1', then=['then'], else_=['something']) - ]) - ]), - '{"Concat": ["a", {"Concat": ["b", {"IfPresent": {"InputName": "output1", "Then": ["then"], "Else": ["something"]}}]}]}' - ), (placeholders.ConcatPlaceholder([ 'a', placeholders.ConcatPlaceholder([ @@ -227,10 +271,111 @@ class TestContainerPlaceholdersTogether(parameterized.TestCase): """{"Concat": ["a", {"IfPresent": {"InputName": "output1", "Then": {"Concat": ["--", "flag"]}, "Else": "{{$.inputs.artifacts['input2'].path}}"}}, "c"]}""" ), ]) - def test(self, placeholder_obj: placeholders.IfPresentPlaceholder, - placeholder: str): + def test_valid(self, placeholder_obj: placeholders.IfPresentPlaceholder, + placeholder: str): self.assertEqual(placeholder_obj._to_string(), placeholder) + def test_only_single_element_ifpresent_inside_concat_outer(self): + with self.assertRaisesRegex( + ValueError, + f'Please use a single element for `then` and `else_` only\.'): + placeholders.ConcatPlaceholder([ + 'b', + placeholders.IfPresentPlaceholder( + input_name='output1', then=['then'], else_=['something']) + ]) + + def test_only_single_element_ifpresent_inside_concat_recursive(self): + with self.assertRaisesRegex( + ValueError, + f'Please use a single element for `then` and `else_` only\.'): + placeholders.ConcatPlaceholder([ + 'a', + placeholders.ConcatPlaceholder([ + 'b', + placeholders.IfPresentPlaceholder( + input_name='output1', + then=['then'], + else_=['something']) + ]) + ]) + + with self.assertRaisesRegex( + ValueError, + f'Please use a single element for `then` and `else_` only\.'): + placeholders.ConcatPlaceholder([ + 'a', + placeholders.ConcatPlaceholder([ + 'b', + placeholders.IfPresentPlaceholder( + input_name='output1', + then=placeholders.ConcatPlaceholder([ + placeholders.IfPresentPlaceholder( + input_name='a', then=['b']) + ]), + else_='something') + ]) + ]) + + def test_only_single_element_in_nested_ifpresent_inside_concat(self): + with self.assertRaisesRegex( + ValueError, + f'Please use a single element for `then` and `else_` only\.'): + dsl.ConcatPlaceholder([ + 'my-prefix-', + dsl.IfPresentPlaceholder( + input_name='input1', + then=[ + dsl.IfPresentPlaceholder( + input_name='input1', + then=dsl.ConcatPlaceholder(['infix-', 'value'])) + ]) + ]) + + def test_recursive_nested_placeholder_validation_does_not_exit_when_first_valid_then_is_found( + self): + with self.assertRaisesRegex( + ValueError, + f'Please use a single element for `then` and `else_` only\.'): + dsl.ConcatPlaceholder([ + 'my-prefix-', + dsl.IfPresentPlaceholder( + input_name='input1', + then=dsl.IfPresentPlaceholder( + input_name='input1', + then=[dsl.ConcatPlaceholder(['infix-', 'value'])])) + ]) + + def test_only_single_element_in_nested_ifpresent_inside_concat_with_outer_ifpresent( + self): + with self.assertRaisesRegex( + ValueError, + f'Please use a single element for `then` and `else_` only\.'): + dsl.IfPresentPlaceholder( + input_name='input_1', + then=dsl.ConcatPlaceholder([ + 'my-prefix-', + dsl.IfPresentPlaceholder( + input_name='input1', + then=dsl.IfPresentPlaceholder( + input_name='input1', + then=[dsl.ConcatPlaceholder(['infix-', 'value'])])) + ])) + + def test_valid_then_but_invalid_else(self): + with self.assertRaisesRegex( + ValueError, + f'Please use a single element for `then` and `else_` only\.'): + dsl.ConcatPlaceholder([ + 'my-prefix-', + dsl.IfPresentPlaceholder( + input_name='input1', + then=dsl.IfPresentPlaceholder( + input_name='input1', + then='single-element', + else_=['one', 'two'])) + ]) + class TestConvertCommandLineElementToStringOrStruct(parameterized.TestCase): diff --git a/sdk/python/kfp/components/structures.py b/sdk/python/kfp/components/structures.py index f0b3b6df182..c3748818f48 100644 --- a/sdk/python/kfp/components/structures.py +++ b/sdk/python/kfp/components/structures.py @@ -505,7 +505,23 @@ def check_placeholder_references_valid_io_name( raise ValueError( f'Argument "{arg}" references nonexistant input: "{arg.input_name}".' ) - for arg in itertools.chain(arg.then or [], arg.else_ or []): + + all_normalized_args: List[placeholders.CommandLineElement] = [] + if arg.then is None: + pass + elif isinstance(arg.then, list): + all_normalized_args.extend(arg.then) + else: + all_normalized_args.append(arg.then) + + if arg.else_ is None: + pass + elif isinstance(arg.else_, list): + all_normalized_args.extend(arg.else_) + else: + all_normalized_args.append(arg.else_) + + for arg in all_normalized_args: check_placeholder_references_valid_io_name(inputs_dict, outputs_dict, arg) elif isinstance(arg, placeholders.ConcatPlaceholder): diff --git a/sdk/python/kfp/components/structures_test.py b/sdk/python/kfp/components/structures_test.py index 66bcb97e41a..b30564f6b28 100644 --- a/sdk/python/kfp/components/structures_test.py +++ b/sdk/python/kfp/components/structures_test.py @@ -103,15 +103,8 @@ - if: cond: isPresent: input_prefix - else: - - --arg2 - - default - - concat: - - --arg1 - - {inputValue: input_prefix} - then: - - --arg1 - - {inputValue: input_prefix} + then: {inputValue: input_prefix} + else: default image: alpine inputs: - {name: input_prefix, optional: false, type: String} @@ -127,21 +120,10 @@ '--arg1', placeholders.IfPresentPlaceholder( input_name='input_prefix', - then=[ - '--arg1', - placeholders.InputValuePlaceholder( - input_name='input_prefix'), - ], - else_=[ - '--arg2', - 'default', - placeholders.ConcatPlaceholder(items=[ - '--arg1', - placeholders.InputValuePlaceholder( - input_name='input_prefix'), - ]), - ]), - ]) + then=placeholders.InputValuePlaceholder( + input_name='input_prefix'), + else_='default'), + ]), ])), inputs={'input_prefix': structures.InputSpec(type='String')}, )