kfp-tekton/sdk/python/tests/compiler/testdata/loop_with_conditional_depen...

87 lines
2.2 KiB
Python

# 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.
from kfp import dsl
from kfp.components import load_component_from_text
from kfp_tekton.tekton import CEL_ConditionOp
from kfp_tekton.compiler import TektonCompiler
class Coder:
def empty(self):
return ""
TektonCompiler._get_unique_id_code = Coder.empty
class CelCondition(dsl.Condition):
def __init__(self, pred: str, name: str = None):
super().__init__(CEL_ConditionOp(pred).output == 'true', name)
def PrintOp(name: str, msg: str = None):
if msg is None:
msg = name
print_op = load_component_from_text(
"""
name: %s
inputs:
- {name: input_text, type: String, description: 'Represents an input parameter.'}
outputs:
- {name: output_value, type: String, description: 'Represents an output paramter.'}
implementation:
container:
image: alpine:3.6
command:
- sh
- -c
- |
set -e
echo $0 > $1
- {inputValue: input_text}
- {outputPath: output_value}
""" % (name)
)
return print_op(msg)
@dsl.pipeline("empty-loop")
def condition_1(param: list = ["a", "b", "c"]):
op0 = PrintOp("print-0")
with CelCondition(f'{op0.output} == "print-0"'):
op1 = PrintOp("print-1")
# works fine for a task outside of loop:
op2 = PrintOp("print-2")
op2.after(op1)
# ...but breaks for loop or task inside of a loop:
loop = dsl.ParallelFor(param)
loop2 = dsl.ParallelFor(param)
# both A) and B) yield the same result:
# A)
loop.after(op1)
with loop:
op3 = PrintOp("print-3")
with loop2:
op4 = PrintOp("print-4")
# B)
op4.after(op1)
if __name__ == '__main__':
TektonCompiler().compile(condition_1, __file__.replace('.py', '.yaml'))