feat(sdk): support writing TASK_FINAL_STATUS parameter type (#9080)

* write new parameter type * update golden snapshots * add compiler tests
kubeflow · May 12, 2023 · c01288d · c01288d
1 parent f35d50c
commit c01288d
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 14 deletions.
diff --git a/sdk/python/kfp/compiler/compiler_test.py b/sdk/python/kfp/compiler/compiler_test.py
@@ -647,6 +647,28 @@ def test_use_task_final_status_in_non_exit_op_yaml(self):
             def my_pipeline(text: bool):
                 print_op()
 
+    def test_task_final_status_parameter_type_is_used(self):
+        # previously compiled to STRUCT type, so checking that this is updated
+
+        @dsl.component
+        def identity(string: str) -> str:
+            return string
+
+        @dsl.component
+        def exit_comp(status: dsl.PipelineTaskFinalStatus):
+            print(status)
+
+        @dsl.pipeline
+        def my_pipeline():
+            exit_task = exit_comp()
+            with dsl.ExitHandler(exit_task=exit_task):
+                identity(string='hi')
+
+        self.assertEqual(
+            my_pipeline.pipeline_spec.components['comp-exit-comp']
+            .input_definitions.parameters['status'].parameter_type,
+            pipeline_spec_pb2.ParameterType.TASK_FINAL_STATUS)
+
     def test_compile_parallel_for_with_valid_parallelism(self):
 
         @dsl.component

diff --git a/sdk/python/kfp/compiler/pipeline_spec_builder.py b/sdk/python/kfp/compiler/pipeline_spec_builder.py
@@ -359,7 +359,7 @@ def _build_component_spec_from_component_spec_structure(
         # Special handling for PipelineTaskFinalStatus first.
         if type_utils.is_task_final_status_type(input_spec.type):
             component_spec.input_definitions.parameters[
-                input_name].parameter_type = pipeline_spec_pb2.ParameterType.STRUCT
+                input_name].parameter_type = pipeline_spec_pb2.ParameterType.TASK_FINAL_STATUS
             component_spec.input_definitions.parameters[
                 input_name].is_optional = True
             if input_spec.description:

diff --git a/sdk/python/test_data/pipelines/pipeline_as_exit_task.yaml b/sdk/python/test_data/pipelines/pipeline_as_exit_task.yaml
@@ -56,7 +56,7 @@ components:
       parameters:
         status:
           isOptional: true
-          parameterType: STRUCT
+          parameterType: TASK_FINAL_STATUS
   comp-exit-handler-1:
     dag:
       tasks:
@@ -129,7 +129,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.16'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.17'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -156,7 +156,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.16'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.17'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -183,7 +183,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.16'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.17'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -210,7 +210,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.16'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.17'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -262,4 +262,4 @@ root:
         isOptional: true
         parameterType: STRING
 schemaVersion: 2.1.0
-sdkVersion: kfp-2.0.0-beta.16
+sdkVersion: kfp-2.0.0-beta.17
diff --git a/sdk/python/test_data/pipelines/pipeline_with_task_final_status.yaml b/sdk/python/test_data/pipelines/pipeline_with_task_final_status.yaml
@@ -39,7 +39,7 @@ components:
       parameters:
         status:
           isOptional: true
-          parameterType: STRUCT
+          parameterType: TASK_FINAL_STATUS
         user_input:
           parameterType: STRING
   comp-fail-op:
@@ -68,7 +68,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.16'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.17'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -99,7 +99,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.16'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.17'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -126,7 +126,7 @@ deploymentSpec:
         - -c
         - "\nif ! [ -x \"$(command -v pip)\" ]; then\n    python3 -m ensurepip ||\
           \ python3 -m ensurepip --user || apt-get install python3-pip\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1\
-          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.16'\
+          \ python3 -m pip install --quiet     --no-warn-script-location 'kfp==2.0.0-beta.17'\
           \ && \"$0\" \"$@\"\n"
         - sh
         - -ec
@@ -180,4 +180,4 @@ root:
         isOptional: true
         parameterType: STRING
 schemaVersion: 2.1.0
-sdkVersion: kfp-2.0.0-beta.16
+sdkVersion: kfp-2.0.0-beta.17
diff --git a/sdk/python/test_data/pipelines/pipeline_with_task_final_status_yaml.yaml b/sdk/python/test_data/pipelines/pipeline_with_task_final_status_yaml.yaml
@@ -27,7 +27,7 @@ components:
       parameters:
         status:
           isOptional: true
-          parameterType: STRUCT
+          parameterType: TASK_FINAL_STATUS
         user_input:
           parameterType: STRING
   comp-print-op:
@@ -92,4 +92,4 @@ root:
         isOptional: true
         parameterType: STRING
 schemaVersion: 2.1.0
-sdkVersion: kfp-2.0.0-beta.16
+sdkVersion: kfp-2.0.0-beta.17