kfp-tekton/sdk/python/tests/compiler/testdata/calc_pipeline.py

107 lines
3.7 KiB
Python

# 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.
import kfp.components as comp
# Define a Python function
def add(a: float, b: float) -> float:
"""Calculates sum of two arguments"""
return a + b
add_op = comp.func_to_container_op(add)
# Advanced function
# Demonstrates imports, helper functions and multiple outputs
from typing import NamedTuple
def my_divmod(dividend: float, divisor: float) -> NamedTuple('MyDivmodOutput', [
('quotient', float), ('remainder', float), ('mlpipeline_ui_metadata', 'UI_metadata'),
('mlpipeline_metrics', 'Metrics')]):
"""Divides two numbers and calculate the quotient and remainder"""
# Pip installs inside a component function.
# NOTE: installs should be placed right at the beginning to avoid upgrading a package
# after it has already been imported and cached by python
import sys, subprocess;
subprocess.run([sys.executable, '-m', 'pip', 'install', 'tensorflow==1.8.0'])
# Imports inside a component function:
import numpy as np
# This function demonstrates how to use nested functions inside a component function:
def divmod_helper(dividend, divisor):
return np.divmod(dividend, divisor)
(quotient, remainder) = divmod_helper(dividend, divisor)
from tensorflow.python.lib.io import file_io # noqa F401
import json
# Exports a sample tensorboard:
metadata = {
'outputs': [{
'type': 'tensorboard',
'source': 'gs://ml-pipeline-dataset/tensorboard-train',
}]
}
# Exports two sample metrics:
metrics = {
'metrics': [{
'name': 'quotient',
'numberValue': float(quotient),
}, {
'name': 'remainder',
'numberValue': float(remainder),
}]}
from collections import namedtuple
divmod_output = namedtuple('MyDivmodOutput', ['quotient', 'remainder', 'mlpipeline_ui_metadata', 'mlpipeline_metrics'])
return divmod_output(quotient, remainder, json.dumps(metadata), json.dumps(metrics))
divmod_op = comp.func_to_container_op(my_divmod, base_image='tensorflow/tensorflow:1.11.0-py3')
import kfp.dsl as dsl
@dsl.pipeline(
name='calculation-pipeline',
description='A toy pipeline that performs arithmetic calculations.'
)
# Currently kfp-tekton doesn't support pass parameter to the pipelinerun yet, so we hard code the number here
def calc_pipeline(
a: float = 7.0,
b: float = 8.0,
c: float = 17.0,
):
# Passing pipeline parameter and a constant value as operation arguments
add_task = add_op(a, 4) # Returns a dsl.ContainerOp class instance.
# Passing a task output reference as operation arguments
# For an operation with a single return value, the output reference can be accessed using
# `task.output` or `task.outputs['output_name']` syntax
divmod_task = divmod_op(add_task.output, b)
# For an operation with a multiple return values, the output references can be accessed using
# `task.outputs['output_name']` syntax
result_task = add_op(divmod_task.outputs['quotient'], c)
if __name__ == '__main__':
from kfp_tekton.compiler import TektonCompiler
TektonCompiler().compile(calc_pipeline, __file__.replace('.py', '.yaml'))