pipeline.py

#!/usr/bin/env python3

import json

from kfp.compiler import Compiler
from kfp.dsl import ContainerOp, ExitHandler, pipeline
from kubernetes import client as k8s


@pipeline(name='My pipeline', description='')
def pipeline():

    deploy = demo_op('deploy', is_exit_handler=True)
    with ExitHandler(deploy):
        deps = demo_op('setup dependencies')

        analyze = demo_op('analyze data')
        analyze.after(deps)

        train1 = demo_op('training 1')
        train2 = demo_op('training 2')
        train3 = demo_op('training 3')
        train1.after(analyze)
        train2.after(analyze)
        train3.after(analyze)

        predict = demo_op('predict')
        predict.after(train1)
        predict.after(train2)
        predict.after(train3)

        matrix = demo_op('create confusion-matrix')
        roc = demo_op('create roc')
        matrix.after(predict)
        roc.after(predict)


if __name__ == '__main__':
    Compiler().compile(pipeline)


def markdown_metadata(result: str) -> str:
    return json.dumps({
        'outputs': [{
            'type': 'markdown',
            'source': 'The result: %s' % result,
            'storage': 'inline',
        }]
    })


def demo_op(name: str, is_exit_handler=False) -> ContainerOp:
    op = ContainerOp(name=name,
                     image='alpine:latest',
                     command=['sh', '-c'],
                     arguments=[
                         'echo "Running step $0" && echo "$1" > /output',
                         name,
                         markdown_metadata(name),
                     ],
                     is_exit_handler=is_exit_handler)
    return op