fix(sdk compiler):Fix v2 compile for exit handler Fixes #5854 (#5899)

* punch the hole to exit handler * add pipeline parameter to test data
kubeflow · Jun 22, 2021 · a1740b4 · a1740b4
1 parent 1983037
commit a1740b4
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 9 deletions.
diff --git a/sdk/python/kfp/v2/compiler/compiler.py b/sdk/python/kfp/v2/compiler/compiler.py
@@ -790,6 +790,21 @@ def _group_to_dag_spec(
             pipeline_spec_pb2.PipelineTaskSpec.TriggerPolicy(
                 condition=condition_string))
 
+      if isinstance(subgroup, dsl.OpsGroup) and subgroup.type == 'exit_handler':
+
+        # "punch the hole", adding inputs needed by its subgroup or tasks.
+        dsl_component_spec.build_component_inputs_spec(
+            component_spec=subgroup_component_spec,
+            pipeline_params=subgroup_params,
+            is_root_component=False,
+        )
+        dsl_component_spec.build_task_inputs_spec(
+            subgroup_task_spec,
+            subgroup_params,
+            tasks_in_current_dag,
+            is_parent_component_root,
+        )
+
       # Generate dependencies section for this task.
       if dependencies.get(subgroup.name, None):
         group_dependencies = list(dependencies[subgroup.name])

diff --git a/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_exit_handler.json b/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_exit_handler.json
@@ -30,11 +30,7 @@
               "inputs": {
                 "parameters": {
                   "message": {
-                    "runtimeValue": {
-                      "constantValue": {
-                        "stringValue": "Hello World!"
-                      }
-                    }
+                    "componentInputParameter": "pipelineparam--message"
                   }
                 }
               },
@@ -43,6 +39,13 @@
               }
             }
           }
+        },
+        "inputDefinitions": {
+          "parameters": {
+            "pipelineparam--message": {
+              "type": "STRING"
+            }
+          }
         }
       },
       "comp-fail-op": {
@@ -147,6 +150,13 @@
             "componentRef": {
               "name": "comp-exit-handler-1"
             },
+            "inputs": {
+              "parameters": {
+                "pipelineparam--message": {
+                  "componentInputParameter": "message"
+                }
+              }
+            },
             "taskInfo": {
               "name": "exit-handler-1"
             }
@@ -177,12 +187,24 @@
             }
           }
         }
+      },
+      "inputDefinitions": {
+        "parameters": {
+          "message": {
+            "type": "STRING"
+          }
+        }
       }
     },
     "schemaVersion": "2.0.0",
-    "sdkVersion": "kfp-1.6.2"
+    "sdkVersion": "kfp-1.6.3"
   },
   "runtimeConfig": {
-    "gcsOutputDirectory": "dummy_root"
+    "gcsOutputDirectory": "dummy_root",
+    "parameters": {
+      "message": {
+        "stringValue": "Hello World!"
+      }
+    }
   }
 }
diff --git a/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_exit_handler.py b/sdk/python/kfp/v2/compiler_cli_tests/test_data/pipeline_with_exit_handler.py
@@ -33,12 +33,12 @@ def fail_op(message: str):
 
 
 @dsl.pipeline(name='pipeline-with-exit-handler', pipeline_root='dummy_root')
-def my_pipeline():
+def my_pipeline(message: str = 'Hello World!'):
 
   exit_task = print_op('Exit handler has worked!')
 
   with dsl.ExitHandler(exit_task):
-    print_op('Hello World!')
+    print_op(message)
     fail_op('Task failed.')