mirror of https://github.com/kubeflow/examples.git
59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
# Copyright 2023 kbthu. All Rights Reserved.
|
|
#
|
|
# 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.
|
|
|
|
import kfp
|
|
from kfp import dsl
|
|
from kfp.components import func_to_container_op
|
|
|
|
def inference():
|
|
return dsl.ContainerOp(
|
|
name='inference',
|
|
image='bash:5.1',
|
|
command=['sh', '-c'],
|
|
arguments=['echo $(( $RANDOM % 10 + 1 ))']
|
|
)
|
|
|
|
@func_to_container_op
|
|
def training() -> str:
|
|
import random
|
|
result = random.choice(['pass', 'fail'])
|
|
return result
|
|
|
|
|
|
@func_to_container_op
|
|
def print_op(message: str):
|
|
print(message)
|
|
|
|
|
|
@dsl.pipeline(
|
|
name='Kubeflow pipeline example',
|
|
description='Demonstrate the flow control of Kubeflow pipeline'
|
|
)
|
|
def dsl_example(stage: dsl.PipelineParam):
|
|
with dsl.Condition(stage == 'training'):
|
|
train_op = training()
|
|
with dsl.Condition(train_op.output == 'pass'):
|
|
print_op('Training pass. Do inference')
|
|
infer_op = inference()
|
|
with dsl.Condition(train_op.output == 'fail'):
|
|
print_op('Training fail. Stop')
|
|
|
|
# inference stage
|
|
with dsl.Condition(stage != 'training'):
|
|
infer_op = inference()
|
|
|
|
if __name__ == '__main__':
|
|
import kfp.compiler as compiler
|
|
compiler.Compiler().compile(dsl_example, __file__ + '.yaml')
|