feat(sdk): support compiling components with `dsl.PipelineTaskFinalStatus` type (#9082)

sdk/RELEASE.md

@@ -1,6 +1,7 @@
 # Current Version (in development)
 
 ## Features
+* Support compiling primitive components with `dsl.PipelineTaskFinalStatus` input [\#9080](https://github.com/kubeflow/pipelines/pull/9080), [\#9082](https://github.com/kubeflow/pipelines/pull/9082)
 
 ## Breaking changes
 

sdk/python/kfp/compiler/pipeline_spec_builder.py

@@ -320,11 +320,13 @@ def build_component_spec_for_exit_task(
     Returns:
         A ComponentSpec object for the exit task.
     """
-    return build_component_spec_for_task(task=task, is_exit_task=True)
+    return build_component_spec_for_task(
+        task=task, is_compiled_component=False, is_exit_task=True)
 
 
 def build_component_spec_for_task(
     task: pipeline_task.PipelineTask,
+    is_compiled_component: bool,
     is_exit_task: bool = False,
 ) -> pipeline_spec_pb2.ComponentSpec:
     """Builds ComponentSpec for a pipeline task.
@@ -338,7 +340,7 @@ def build_component_spec_for_task(
     """
     for input_name, input_spec in (task.component_spec.inputs or {}).items():
         if not is_exit_task and type_utils.is_task_final_status_type(
-                input_spec.type):
+                input_spec.type) and not is_compiled_component:
             raise ValueError(
                 f'PipelineTaskFinalStatus can only be used in an exit task. Parameter {input_name} of a non exit task has type PipelineTaskFinalStatus.'
             )
@@ -1077,6 +1079,7 @@ def build_spec_by_group(
     group_name_to_parent_groups: Mapping[str, List[tasks_group.TasksGroup]],
     name_to_for_loop_group: Mapping[str, tasks_group.ParallelFor],
     platform_spec: pipeline_spec_pb2.PlatformSpec,
+    is_compiled_component: bool,
 ) -> None:
     """Generates IR spec given a TasksGroup.
 
@@ -1131,16 +1134,16 @@ def build_spec_by_group(
         is_parent_component_root = (group_component_spec == pipeline_spec.root)
 
         if isinstance(subgroup, pipeline_task.PipelineTask):
-
             subgroup_task_spec = build_task_spec_for_task(
                 task=subgroup,
                 parent_component_inputs=group_component_spec.input_definitions,
                 tasks_in_current_dag=tasks_in_current_dag,
             )
             task_name_to_task_spec[subgroup.name] = subgroup_task_spec
-
             subgroup_component_spec = build_component_spec_for_task(
-                task=subgroup)
+                task=subgroup,
+                is_compiled_component=is_compiled_component,
+            )
             task_name_to_component_spec[subgroup.name] = subgroup_component_spec
 
             if subgroup_component_spec.executor_label:
@@ -1710,6 +1713,7 @@ def create_pipeline_spec(
             group_name_to_parent_groups=group_name_to_parent_groups,
             name_to_for_loop_group=name_to_for_loop_group,
             platform_spec=platform_spec,
+            is_compiled_component=False,
         )
 
     build_exit_handler_groups_recursively(
@@ -1813,14 +1817,6 @@ def write_pipeline_spec_to_file(
 
 def extract_comments_from_pipeline_spec(pipeline_spec: dict,
                                         pipeline_description: str) -> str:
-    map_parameter_types = {
-        'NUMBER_INTEGER': int.__name__,
-        'NUMBER_DOUBLE': float.__name__,
-        'STRING': str.__name__,
-        'BOOLEAN': bool.__name__,
-        'LIST': list.__name__,
-        'STRUCT': dict.__name__
-    }
 
     map_headings = {
         'inputDefinitions': '# Inputs:',
@@ -1840,8 +1836,9 @@ def extract_comments_from_pipeline_spec(pipeline_spec: dict,
                     'parameters', {}).items():
                 data = {}
                 data['name'] = parameter_name
-                data['parameterType'] = map_parameter_types[
-                    parameter_body['parameterType']]
+                data[
+                    'parameterType'] = type_utils.IR_TYPE_TO_COMMENT_TYPE_STRING[
+                        parameter_body['parameterType']]
                 if 'defaultValue' in signature['parameters'][parameter_name]:
                     data['defaultValue'] = signature['parameters'][
                         parameter_name]['defaultValue']

sdk/python/kfp/components/structures.py

@@ -1000,6 +1000,7 @@ class ComponentSpec:
             name_to_for_loop_group={},  # no for loop in single-component pipeline
             platform_spec=pipeline_spec_pb2.PlatformSpec(
             ),  # no PlatformSpec single-component pipeline
+            is_compiled_component=True,
         )
 
         return pipeline_spec

sdk/python/kfp/components/types/type_utils.py

@@ -442,6 +442,17 @@ IR_TYPE_TO_IN_MEMORY_SPEC_TYPE = {
     'LIST': 'List',
     'STRUCT': 'Dict',
     'BOOLEAN': 'Boolean',
+    'TASK_FINAL_STATUS': task_final_status.PipelineTaskFinalStatus.__name__,
+}
+
+IR_TYPE_TO_COMMENT_TYPE_STRING = {
+    'STRING': str.__name__,
+    'NUMBER_INTEGER': int.__name__,
+    'NUMBER_DOUBLE': float.__name__,
+    'LIST': list.__name__,
+    'STRUCT': dict.__name__,
+    'BOOLEAN': bool.__name__,
+    'TASK_FINAL_STATUS': task_final_status.PipelineTaskFinalStatus.__name__,
 }
 
 IN_MEMORY_SPEC_TYPE_TO_IR_TYPE = {

sdk/python/test_data/components/component_with_task_final_status.py

@@ -0,0 +1,25 @@
+# Copyright 2023 The Kubeflow Authors
+#
+# 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 kfp import dsl
+
+
+@dsl.component
+def exit_comp(status: dsl.PipelineTaskFinalStatus):
+    print(status)
+
+
+if __name__ == '__main__':
+    from kfp import compiler
+    compiler.Compiler().compile(
+        pipeline_func=exit_comp, package_path=__file__.replace('.py', '.yaml'))

sdk/python/test_data/components/component_with_task_final_status.yaml

@@ -0,0 +1,64 @@
+# PIPELINE DEFINITION
+# Name: exit-comp
+# Inputs:
+#    status: PipelineTaskFinalStatus
+components:
+  comp-exit-comp:
+    executorLabel: exec-exit-comp
+    inputDefinitions:
+      parameters:
+        status:
+          isOptional: true
+          parameterType: TASK_FINAL_STATUS
+deploymentSpec:
+  executors:
+    exec-exit-comp:
+      container:
+        args:
+        - --executor_input
+        - '{{$}}'
+        - --function_to_execute
+        - exit_comp
+        command:
+        - sh
+        - -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.17'\
+          \ && \"$0\" \"$@\"\n"
+        - sh
+        - -ec
+        - 'program_path=$(mktemp -d)
+
+          printf "%s" "$0" > "$program_path/ephemeral_component.py"
+
+          python3 -m kfp.components.executor_main                         --component_module_path                         "$program_path/ephemeral_component.py"                         "$@"
+
+          '
+        - "\nimport kfp\nfrom kfp import dsl\nfrom kfp.dsl import *\nfrom typing import\
+          \ *\n\ndef exit_comp(status: dsl.PipelineTaskFinalStatus):\n    print(status)\n\
+          \n"
+        image: python:3.7
+pipelineInfo:
+  name: exit-comp
+root:
+  dag:
+    tasks:
+      exit-comp:
+        cachingOptions:
+          enableCache: true
+        componentRef:
+          name: comp-exit-comp
+        inputs:
+          parameters:
+            status:
+              componentInputParameter: status
+        taskInfo:
+          name: exit-comp
+  inputDefinitions:
+    parameters:
+      status:
+        isOptional: true
+        parameterType: TASK_FINAL_STATUS
+schemaVersion: 2.1.0
+sdkVersion: kfp-2.0.0-beta.17

sdk/python/test_data/test_data_config.yaml

@@ -230,6 +230,8 @@ components:
       execute: false
     - module: component_with_metadata_fields
       name: dataset_joiner
+    - module: component_with_task_final_status
+      name: exit_comp
       execute: false
 v1_components:
   test_data_dir: sdk/python/test_data/v1_component_yaml