address comments

kubeflow · Mar 30, 2022 · 9709a1f · 9709a1f
1 parent 8ccb372
commit 9709a1f
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/sdk/python/kfp/v2/compiler/compiler.py b/sdk/python/kfp/v2/compiler/compiler.py
@@ -36,6 +36,7 @@
 from kfp.pipeline_spec import pipeline_spec_pb2
 from kfp.v2.compiler import compiler_utils
 from kfp.v2.components import component_factory
+from kfp.v2.components import task_final_status
 from kfp.v2.components.types import artifact_types, type_utils
 
 _GroupOrOp = Union[dsl.OpsGroup, dsl.BaseOp]
@@ -733,8 +734,8 @@ def _group_to_dag_spec(
                         if type_utils.is_task_final_status_type(
                                 input_spec.type):
                             raise ValueError(
-                                'PipelineTaskFinalStatus can only be used in an exit task.'
-                            )
+                                f'{task_final_status.PipelineTaskFinalStatus.__name__}'
+                                ' can only be used in an exit task.')
 
                 # Task level caching option.
                 subgroup.task_spec.caching_options.enable_cache = subgroup.enable_caching

diff --git a/sdk/python/kfp/v2/components/executor.py b/sdk/python/kfp/v2/components/executor.py
@@ -274,14 +274,14 @@ def execute(self):
             # `Optional[]` to get the actual parameter type.
             v = type_annotations.maybe_strip_optional_from_annotation(v)
 
-            if v == task_final_status.PipelineTaskFinalStatus:
+            if v is task_final_status.PipelineTaskFinalStatus:
                 value = self._get_input_parameter_value(k, v)
                 func_kwargs[k] = task_final_status.PipelineTaskFinalStatus(
                     state=value.get('state'),
                     pipeline_job_resource_name=value.get(
                         'pipelineJobResourceName'),
-                    # pipelineTaskName won't be None once the BE change is
-                    # rolled out
+                    # pipelineTaskName won't be None once the Vertex Pipelines
+                    # BE change is rolled out
                     pipeline_task_name=value.get('pipelineTaskName', None),
                     error_code=value.get('error').get('code', None),
                     error_message=value.get('error').get('message', None),

diff --git a/sdk/python/kfp/v2/components/task_final_status.py b/sdk/python/kfp/v2/components/task_final_status.py
@@ -22,7 +22,7 @@ class PipelineTaskFinalStatus:
     """The final status of a pipeline task.
 
     This is the Python representation of the proto: PipelineTaskFinalStatus
-    https://github.com/kubeflow/pipelines/blob/61d9210c61ff9be780b678ab78a7f7200d50cffa/api/v2alpha1/pipeline_spec.proto#L886
+    https://github.com/kubeflow/pipelines/blob/1c3e2768e6177d5d6e3f4b8eff8fafb9a3b76c1f/api/v2alpha1/pipeline_spec.proto#L886
 
     Attributes:
         state: The final state of the task. The value could be one of