Updated the "Basic - Recursive loop" sample (#1113)

Modernized the sample pipeline code.
This commit is contained in:
Alexey Volkov 2019-04-18 22:01:55 -07:00 committed by Kubernetes Prow Robot
parent 1e57cd3ce0
commit 07faa08b86
1 changed files with 33 additions and 30 deletions

View File

@ -13,36 +13,38 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import kfp.dsl as dsl
class FlipCoinOp(dsl.ContainerOp): import kfp
from kfp import dsl
def flip_coin_op():
"""Flip a coin and output heads or tails randomly.""" """Flip a coin and output heads or tails randomly."""
return dsl.ContainerOp(
def __init__(self): name='Flip coin',
super(FlipCoinOp, self).__init__(
name='Flip',
image='python:alpine3.6', image='python:alpine3.6',
command=['sh', '-c'], command=['sh', '-c'],
arguments=['python -c "import random; result = \'heads\' if random.randint(0,1) == 0 ' arguments=['python -c "import random; result = \'heads\' if random.randint(0,1) == 0 '
'else \'tails\'; print(result)" | tee /tmp/output'], 'else \'tails\'; print(result)" | tee /tmp/output'],
file_outputs={'output': '/tmp/output'}) file_outputs={'output': '/tmp/output'}
)
class PrintOp(dsl.ContainerOp):
def print_op(msg):
"""Print a message.""" """Print a message."""
return dsl.ContainerOp(
def __init__(self, msg):
super(PrintOp, self).__init__(
name='Print', name='Print',
image='alpine:3.6', image='alpine:3.6',
command=['echo', msg], command=['echo', msg],
) )
# Use the dsl.graph_component to decorate functions that are
# Use the dsl.graph_component to decorate pipeline functions that can be
# recursively called. # recursively called.
@dsl.graph_component @dsl.graph_component
def flip_component(flip_result): def flip_component(flip_result):
print_flip = PrintOp(flip_result) print_flip = print_op(flip_result)
flipA = FlipCoinOp().after(print_flip) flipA = flip_coin_op().after(print_flip)
with dsl.Condition(flipA.output == 'heads'): with dsl.Condition(flipA.output == 'heads'):
# When the flip_component is called recursively, the flipA.output # When the flip_component is called recursively, the flipA.output
# from inside the graph component will be passed to the next flip_component # from inside the graph component will be passed to the next flip_component
@ -50,17 +52,18 @@ def flip_component(flip_result):
# comes from the flipA.output in the flipcoin function. # comes from the flipA.output in the flipcoin function.
flip_component(flipA.output) flip_component(flipA.output)
@dsl.pipeline( @dsl.pipeline(
name='pipeline flip coin', name='Recursive loop pipeline',
description='shows how to use dsl.Condition.' description='Shows how to create recursive loops.'
) )
def flipcoin(): def flipcoin():
flipA = FlipCoinOp() first_flip = flip_coin_op()
flip_loop = flip_component(flipA.output) flip_loop = flip_component(first_flip.output)
# flip_loop is a graph_component with the outputs field # flip_loop is a graph_component with the outputs field
# filled with the returned dictionary. # filled with the returned dictionary.
PrintOp('cool, it is over. %s' % flipA.output).after(flip_loop) print_op('cool, it is over.').after(flip_loop)
if __name__ == '__main__': if __name__ == '__main__':
import kfp.compiler as compiler kfp.compiler.Compiler().compile(flipcoin, __file__ + '.tar.gz')
compiler.Compiler().compile(flipcoin, __file__ + '.tar.gz')