fix(sdk): Add new logic to cover new tekton runafter condition (kubeflow#1016)

Tomcli · web-flow · commit b219b2ee1efc · 2022-07-28T20:30:55.000Z
* add new logic to cover new tekton runafter condition

* fix loop condition error
diff --git a/sdk/python/kfp_tekton/compiler/compiler.py b/sdk/python/kfp_tekton/compiler/compiler.py
@@ -1394,6 +1394,45 @@ def get_when_task(input_task_when, depended_conditions):
       pipeline_run['spec']['podTemplate']['nodeSelector'] = copy.deepcopy(pipeline_conf.default_pod_node_selector)
     workflow = pipeline_run
 
+    # populate dependend condition for all the runafter tasks
+    def populate_runafter_condition(task):
+      task_runafter = task.get('runAfter')
+      if task_runafter:
+        for t in workflow['spec']['pipelineSpec']['tasks']:
+          if t['name'] in task_runafter:
+            if t.get('when'):
+              task.setdefault('when', [])
+              for when_item in t['when']:
+                if when_item not in task['when']:
+                  add_conditions = True
+                  # Do not add condition if the condition is not in the same graph/pipelineloop
+                  for pipeline_loop in self.loops_pipeline.values():
+                    if task['name'] in pipeline_loop['task_list']:
+                      task_input = re.findall('\$\(tasks.([^ \t\n.:,;{}]+).results.([^ \t\n.:,;{}]+)\)', when_item['input'])
+                      if task_input and task_input[0][0] not in pipeline_loop['task_list']:
+                        add_conditions = False
+                  if add_conditions:
+                    task['when'].append(when_item)
+
+    # search runafter tree logic before populating the condition
+    visited_tasks = {}
+    task_queue = []
+    for task in workflow['spec']['pipelineSpec']['tasks']:
+      task_runafter = task.get('runAfter')
+      if task_runafter:
+        task_queue.append(task)
+    while task_queue:
+      popped_task = task_queue.pop(0)
+      populate_condition = True
+      for queued_task in task_queue:
+        if queued_task['name'] in popped_task['runAfter'] and len(task_queue) != visited_tasks.get(popped_task['name']):
+          visited_tasks[popped_task['name']] = len(task_queue)
+          task_queue.append(popped_task)
+          populate_condition = False
+          break
+      if populate_condition:
+        populate_runafter_condition(popped_task)
+
     return workflow
 
   def _sanitize_and_inject_artifact(self, pipeline: dsl.Pipeline, pipeline_conf=None):
diff --git a/sdk/python/tests/compiler/compiler_tests.py b/sdk/python/tests/compiler/compiler_tests.py
@@ -148,6 +148,13 @@ def test_big_data_multi_volumes_2_workflow(self):
     from .testdata.big_data_multi_volumes_2 import big_data
     self._test_pipeline_workflow(big_data, 'big_data_multi_volumes_2.yaml', skip_noninlined=True)
 
+  def test_condition_depend_workflow(self):
+    """
+    Test compiling a condition depend workflow.
+    """
+    from .testdata.condition_depend import pipeline
+    self._test_pipeline_workflow(pipeline, 'condition_depend.yaml', skip_noninlined=True)
+
   def test_recur_cond_workflow(self):
     """
     Test compiling a recurive condition workflow.
diff --git a/sdk/python/tests/compiler/testdata/condition_depend.py b/sdk/python/tests/compiler/testdata/condition_depend.py
@@ -0,0 +1,66 @@
+# Copyright 2022 kubeflow.org
+#
+# 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.
+
+# This use case was discussed in https://github.com/kubeflow/kfp-tekton/issues/845
+
+from kfp import dsl, components
+
+op1_yaml = '''\
+name: 'my-in-coop1'
+inputs:
+- {name: item, type: Integer}
+- {name: param, type: Integer}
+implementation:
+    container:
+        image: library/bash:4.4.23
+        command: ['sh', '-c']
+        args:
+        - |
+          set -e
+          echo op1 "$0" "$1"
+        - {inputValue: item}
+        - {inputValue: param}
+'''
+
+
+def false() -> str:
+    return 'false'
+
+
+false_op = components.create_component_from_func(
+    false, base_image='python:alpine3.6')
+
+
+@dsl.pipeline(name='pipeline')
+def pipeline(param: int = 10):
+    condition_op = false_op()
+    cond = condition_op.output
+    with dsl.Condition(cond == 'true'):
+        op2_template = components.load_component_from_text(op1_yaml)
+        op2 = op2_template(1, param)
+
+    op3_template = components.load_component_from_text(op1_yaml)
+    op3 = op3_template(1, param)
+    op4_template = components.load_component_from_text(op1_yaml)
+    op4 = op4_template(1, param)
+    op4.after(op2)
+    op3.after(op4)
+
+
+if __name__ == '__main__':
+  from kfp_tekton.compiler import TektonCompiler as Compiler
+  from kfp_tekton.compiler.pipeline_utils import TektonPipelineConf
+  tekton_pipeline_conf = TektonPipelineConf()
+  tekton_pipeline_conf.set_tekton_inline_spec(True)
+  Compiler().compile(pipeline, __file__.replace('.py', '.yaml'), tekton_pipeline_conf=tekton_pipeline_conf)
diff --git a/sdk/python/tests/compiler/testdata/condition_depend.yaml b/sdk/python/tests/compiler/testdata/condition_depend.yaml
@@ -0,0 +1,250 @@
+# Copyright 2021 kubeflow.org
+#
+# 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.
+
+apiVersion: tekton.dev/v1beta1
+kind: PipelineRun
+metadata:
+  name: pipeline
+  annotations:
+    tekton.dev/output_artifacts: '{"false": [{"key": "artifacts/$PIPELINERUN/false/Output.tgz",
+      "name": "false-Output", "path": "/tmp/outputs/Output/data"}]}'
+    tekton.dev/input_artifacts: '{}'
+    tekton.dev/artifact_bucket: mlpipeline
+    tekton.dev/artifact_endpoint: minio-service.kubeflow:9000
+    tekton.dev/artifact_endpoint_scheme: http://
+    tekton.dev/artifact_items: '{"false": [["Output", "$(results.output.path)"]],
+      "my-in-coop1": [], "my-in-coop1-2": [], "my-in-coop1-3": []}'
+    sidecar.istio.io/inject: "false"
+    pipelines.kubeflow.org/big_data_passing_format: $(workspaces.$TASK_NAME.path)/artifacts/$ORIG_PR_NAME/$TASKRUN_NAME/$TASK_PARAM_NAME
+    pipelines.kubeflow.org/pipeline_spec: '{"inputs": [{"default": "10", "name": "param",
+      "optional": true, "type": "Integer"}], "name": "pipeline"}'
+spec:
+  params:
+  - name: param
+    value: '10'
+  pipelineSpec:
+    params:
+    - name: param
+      default: '10'
+    tasks:
+    - name: "false"
+      taskSpec:
+        steps:
+        - name: main
+          args:
+          - '----output-paths'
+          - $(results.output.path)
+          command:
+          - sh
+          - -ec
+          - |
+            program_path=$(mktemp)
+            printf "%s" "$0" > "$program_path"
+            python3 -u "$program_path" "$@"
+          - |
+            def false():
+                return 'false'
+
+            def _serialize_str(str_value: str) -> str:
+                if not isinstance(str_value, str):
+                    raise TypeError('Value "{}" has type "{}" instead of str.'.format(
+                        str(str_value), str(type(str_value))))
+                return str_value
+
+            import argparse
+            _parser = argparse.ArgumentParser(prog='False', description='')
+            _parser.add_argument("----output-paths", dest="_output_paths", type=str, nargs=1)
+            _parsed_args = vars(_parser.parse_args())
+            _output_files = _parsed_args.pop("_output_paths", [])
+
+            _outputs = false(**_parsed_args)
+
+            _outputs = [_outputs]
+
+            _output_serializers = [
+                _serialize_str,
+
+            ]
+
+            import os
+            for idx, output_file in enumerate(_output_files):
+                try:
+                    os.makedirs(os.path.dirname(output_file))
+                except OSError:
+                    pass
+                with open(output_file, 'w') as f:
+                    f.write(_output_serializers[idx](_outputs[idx]))
+          image: python:alpine3.6
+        results:
+        - name: output
+          type: string
+          description: /tmp/outputs/Output/data
+        metadata:
+          labels:
+            pipelines.kubeflow.org/pipelinename: ''
+            pipelines.kubeflow.org/generation: ''
+            pipelines.kubeflow.org/cache_enabled: "true"
+          annotations:
+            pipelines.kubeflow.org/component_spec_digest: '{"name": "False", "outputs":
+              [{"name": "Output", "type": "String"}], "version": "False@sha256=ab6d70d69626be9a388e3d6d92c4eae87095dc005400a41e6613f4b855c694f6"}'
+            tekton.dev/template: ''
+      timeout: 525600m
+    - name: my-in-coop1
+      params:
+      - name: param
+        value: $(params.param)
+      taskSpec:
+        steps:
+        - name: main
+          args:
+          - |
+            set -e
+            echo op1 "$0" "$1"
+          - '1'
+          - $(inputs.params.param)
+          command:
+          - sh
+          - -c
+          image: library/bash:4.4.23
+        params:
+        - name: param
+        metadata:
+          labels:
+            pipelines.kubeflow.org/pipelinename: ''
+            pipelines.kubeflow.org/generation: ''
+            pipelines.kubeflow.org/cache_enabled: "true"
+          annotations:
+            pipelines.kubeflow.org/component_spec_digest: '{"name": "my-in-coop1",
+              "outputs": [], "version": "my-in-coop1@sha256=8ccab3a28a39a406554d964865f2ccb0aed854a43b6de827f613eff2bccd6f8f"}'
+            tekton.dev/template: ''
+      when:
+      - input: $(tasks.condition-1.results.outcome)
+        operator: in
+        values:
+        - "true"
+      timeout: 525600m
+    - name: my-in-coop1-2
+      params:
+      - name: param
+        value: $(params.param)
+      taskSpec:
+        steps:
+        - name: main
+          args:
+          - |
+            set -e
+            echo op1 "$0" "$1"
+          - '1'
+          - $(inputs.params.param)
+          command:
+          - sh
+          - -c
+          image: library/bash:4.4.23
+        params:
+        - name: param
+        metadata:
+          labels:
+            pipelines.kubeflow.org/pipelinename: ''
+            pipelines.kubeflow.org/generation: ''
+            pipelines.kubeflow.org/cache_enabled: "true"
+          annotations:
+            pipelines.kubeflow.org/component_spec_digest: '{"name": "my-in-coop1",
+              "outputs": [], "version": "my-in-coop1@sha256=8ccab3a28a39a406554d964865f2ccb0aed854a43b6de827f613eff2bccd6f8f"}'
+            tekton.dev/template: ''
+      runAfter:
+      - my-in-coop1-3
+      timeout: 525600m
+      when:
+      - input: $(tasks.condition-1.results.outcome)
+        operator: in
+        values:
+        - "true"
+    - name: my-in-coop1-3
+      params:
+      - name: param
+        value: $(params.param)
+      taskSpec:
+        steps:
+        - name: main
+          args:
+          - |
+            set -e
+            echo op1 "$0" "$1"
+          - '1'
+          - $(inputs.params.param)
+          command:
+          - sh
+          - -c
+          image: library/bash:4.4.23
+        params:
+        - name: param
+        metadata:
+          labels:
+            pipelines.kubeflow.org/pipelinename: ''
+            pipelines.kubeflow.org/generation: ''
+            pipelines.kubeflow.org/cache_enabled: "true"
+          annotations:
+            pipelines.kubeflow.org/component_spec_digest: '{"name": "my-in-coop1",
+              "outputs": [], "version": "my-in-coop1@sha256=8ccab3a28a39a406554d964865f2ccb0aed854a43b6de827f613eff2bccd6f8f"}'
+            tekton.dev/template: ''
+      runAfter:
+      - my-in-coop1
+      timeout: 525600m
+      when:
+      - input: $(tasks.condition-1.results.outcome)
+        operator: in
+        values:
+        - "true"
+    - name: condition-1
+      params:
+      - name: operand1
+        value: $(tasks.false.results.output)
+      - name: operand2
+        value: "true"
+      - name: operator
+        value: ==
+      taskSpec:
+        results:
+        - name: outcome
+          type: string
+          description: Conditional task outcome
+        params:
+        - name: operand1
+        - name: operand2
+        - name: operator
+        steps:
+        - name: main
+          command:
+          - sh
+          - -ec
+          - program_path=$(mktemp); printf "%s" "$0" > "$program_path";  python3 -u
+            "$program_path" "$1" "$2"
+          args:
+          - |
+            import sys
+            input1=str.rstrip(sys.argv[1])
+            input2=str.rstrip(sys.argv[2])
+            try:
+              input1=int(input1)
+              input2=int(input2)
+            except:
+              input1=str(input1)
+            outcome="true" if (input1 $(inputs.params.operator) input2) else "false"
+            f = open("/tekton/results/outcome", "w")
+            f.write(outcome)
+            f.close()
+          - $(inputs.params.operand1)
+          - $(inputs.params.operand2)
+          image: python:alpine3.6
+  timeout: 525600m
diff --git a/sdk/python/tests/compiler/testdata/loop_with_conditional_dependency.yaml b/sdk/python/tests/compiler/testdata/loop_with_conditional_dependency.yaml
@@ -147,6 +147,11 @@ spec:
       runAfter:
       - print-1
       timeout: 525600m
+      when:
+      - input: $(tasks.condition-cel.results.outcome)
+        operator: in
+        values:
+        - "true"
     - runAfter:
       - print-1
       name: empty-loop-for-loop-2
diff --git a/sdk/python/tests/compiler/testdata/loop_with_conditional_dependency_noninlined.yaml b/sdk/python/tests/compiler/testdata/loop_with_conditional_dependency_noninlined.yaml
diff --git a/sdk/python/tests/compiler/testdata/multi_nested_loop_condi.yaml b/sdk/python/tests/compiler/testdata/multi_nested_loop_condi.yaml
