pipelines/sdk/python/kfp/compiler/pipeline_spec_builder_test.py

304 lines
13 KiB
Python

# Copyright 2021 The Kubeflow Authors
#
# 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.
"""Tests for kfp.compiler.pipeline_spec_builder."""
import os
import tempfile
import unittest
from absl.testing import parameterized
from google.protobuf import json_format
from google.protobuf import struct_pb2
from kfp.compiler import pipeline_spec_builder
from kfp.components import pipeline_channel
from kfp.pipeline_spec import pipeline_spec_pb2
import yaml
class PipelineSpecBuilderTest(parameterized.TestCase):
def setUp(self):
self.maxDiff = None
@parameterized.parameters(
{
'channel':
pipeline_channel.PipelineParameterChannel(
name='output1', task_name='task1', channel_type='String'),
'expected':
'pipelinechannel--task1-output1',
},
{
'channel':
pipeline_channel.PipelineArtifactChannel(
name='output1', task_name='task1', channel_type='Artifact'),
'expected':
'pipelinechannel--task1-output1',
},
{
'channel':
pipeline_channel.PipelineParameterChannel(
name='param1', channel_type='String'),
'expected':
'pipelinechannel--param1',
},
)
def test_additional_input_name_for_pipeline_channel(self, channel,
expected):
self.assertEqual(
expected,
pipeline_spec_builder._additional_input_name_for_pipeline_channel(
channel))
@parameterized.parameters(
{
'parameter_type': pipeline_spec_pb2.ParameterType.NUMBER_INTEGER,
'default_value': None,
'expected': struct_pb2.Value(),
},
{
'parameter_type': pipeline_spec_pb2.ParameterType.NUMBER_INTEGER,
'default_value': 1,
'expected': struct_pb2.Value(number_value=1),
},
{
'parameter_type': pipeline_spec_pb2.ParameterType.NUMBER_DOUBLE,
'default_value': 1.2,
'expected': struct_pb2.Value(number_value=1.2),
},
{
'parameter_type': pipeline_spec_pb2.ParameterType.STRING,
'default_value': 'text',
'expected': struct_pb2.Value(string_value='text'),
},
{
'parameter_type': pipeline_spec_pb2.ParameterType.BOOLEAN,
'default_value': True,
'expected': struct_pb2.Value(bool_value=True),
},
{
'parameter_type': pipeline_spec_pb2.ParameterType.BOOLEAN,
'default_value': False,
'expected': struct_pb2.Value(bool_value=False),
},
{
'parameter_type':
pipeline_spec_pb2.ParameterType.STRUCT,
'default_value': {
'a': 1,
'b': 2,
},
'expected':
struct_pb2.Value(
struct_value=struct_pb2.Struct(
fields={
'a': struct_pb2.Value(number_value=1),
'b': struct_pb2.Value(number_value=2),
})),
},
{
'parameter_type':
pipeline_spec_pb2.ParameterType.LIST,
'default_value': ['a', 'b'],
'expected':
struct_pb2.Value(
list_value=struct_pb2.ListValue(values=[
struct_pb2.Value(string_value='a'),
struct_pb2.Value(string_value='b'),
])),
},
{
'parameter_type':
pipeline_spec_pb2.ParameterType.LIST,
'default_value': [{
'a': 1,
'b': 2
}, {
'a': 10,
'b': 20
}],
'expected':
struct_pb2.Value(
list_value=struct_pb2.ListValue(values=[
struct_pb2.Value(
struct_value=struct_pb2.Struct(
fields={
'a': struct_pb2.Value(number_value=1),
'b': struct_pb2.Value(number_value=2),
})),
struct_pb2.Value(
struct_value=struct_pb2.Struct(
fields={
'a': struct_pb2.Value(number_value=10),
'b': struct_pb2.Value(number_value=20),
})),
])),
},
)
def test_fill_in_component_input_default_value(self, parameter_type,
default_value, expected):
component_spec = pipeline_spec_pb2.ComponentSpec(
input_definitions=pipeline_spec_pb2.ComponentInputsSpec(
parameters={
'input1':
pipeline_spec_pb2.ComponentInputsSpec.ParameterSpec(
parameter_type=parameter_type)
}))
pipeline_spec_builder._fill_in_component_input_default_value(
component_spec=component_spec,
input_name='input1',
default_value=default_value)
self.assertEqual(
expected,
component_spec.input_definitions.parameters['input1'].default_value,
)
def test_merge_deployment_spec_and_component_spec(self):
main_deployment_config = pipeline_spec_pb2.PipelineDeploymentConfig()
main_deployment_config.executors['exec-1'].CopyFrom(
pipeline_spec_pb2.PipelineDeploymentConfig.ExecutorSpec(
container=pipeline_spec_pb2.PipelineDeploymentConfig
.PipelineContainerSpec(image='img-1', command=['cmd'])))
main_pipeline_spec = pipeline_spec_pb2.PipelineSpec()
main_pipeline_spec.components['comp-1'].CopyFrom(
pipeline_spec_pb2.ComponentSpec(executor_label='exec-1'))
sub_deployment_config = pipeline_spec_pb2.PipelineDeploymentConfig()
sub_deployment_config.executors['exec-1'].CopyFrom(
pipeline_spec_pb2.PipelineDeploymentConfig.ExecutorSpec(
container=pipeline_spec_pb2.PipelineDeploymentConfig
.PipelineContainerSpec(image='img-a', command=['cmd'])))
sub_deployment_config.executors['exec-2'].CopyFrom(
pipeline_spec_pb2.PipelineDeploymentConfig.ExecutorSpec(
container=pipeline_spec_pb2.PipelineDeploymentConfig
.PipelineContainerSpec(image='img-b', command=['cmd'])))
sub_pipeline_spec = pipeline_spec_pb2.PipelineSpec()
sub_pipeline_spec.deployment_spec.update(
json_format.MessageToDict(sub_deployment_config))
sub_pipeline_spec.components['comp-1'].CopyFrom(
pipeline_spec_pb2.ComponentSpec(executor_label='exec-1'))
sub_pipeline_spec.components['comp-2'].CopyFrom(
pipeline_spec_pb2.ComponentSpec(executor_label='exec-2'))
sub_pipeline_spec.root.CopyFrom(
pipeline_spec_pb2.ComponentSpec(dag=pipeline_spec_pb2.DagSpec()))
sub_pipeline_spec.root.dag.tasks['task-1'].CopyFrom(
pipeline_spec_pb2.PipelineTaskSpec(
component_ref=pipeline_spec_pb2.ComponentRef(name='comp-1')))
sub_pipeline_spec.root.dag.tasks['task-2'].CopyFrom(
pipeline_spec_pb2.PipelineTaskSpec(
component_ref=pipeline_spec_pb2.ComponentRef(name='comp-2')))
expected_sub_pipeline_spec = pipeline_spec_pb2.PipelineSpec.FromString(
sub_pipeline_spec.SerializeToString())
expected_main_deployment_config = pipeline_spec_pb2.PipelineDeploymentConfig(
)
expected_main_deployment_config.executors['exec-1'].CopyFrom(
pipeline_spec_pb2.PipelineDeploymentConfig.ExecutorSpec(
container=pipeline_spec_pb2.PipelineDeploymentConfig
.PipelineContainerSpec(image='img-1', command=['cmd'])))
expected_main_deployment_config.executors['exec-1-2'].CopyFrom(
pipeline_spec_pb2.PipelineDeploymentConfig.ExecutorSpec(
container=pipeline_spec_pb2.PipelineDeploymentConfig
.PipelineContainerSpec(image='img-a', command=['cmd'])))
expected_main_deployment_config.executors['exec-2'].CopyFrom(
pipeline_spec_pb2.PipelineDeploymentConfig.ExecutorSpec(
container=pipeline_spec_pb2.PipelineDeploymentConfig
.PipelineContainerSpec(image='img-b', command=['cmd'])))
expected_main_pipeline_spec = pipeline_spec_pb2.PipelineSpec()
expected_main_pipeline_spec.components['comp-1'].CopyFrom(
pipeline_spec_pb2.ComponentSpec(executor_label='exec-1'))
expected_main_pipeline_spec.components['comp-1-2'].CopyFrom(
pipeline_spec_pb2.ComponentSpec(executor_label='exec-1-2'))
expected_main_pipeline_spec.components['comp-2'].CopyFrom(
pipeline_spec_pb2.ComponentSpec(executor_label='exec-2'))
expected_main_pipeline_spec.components['inner-pipeline'].CopyFrom(
pipeline_spec_pb2.ComponentSpec(dag=pipeline_spec_pb2.DagSpec()))
expected_main_pipeline_spec.components['inner-pipeline'].dag.tasks[
'task-1'].CopyFrom(
pipeline_spec_pb2.PipelineTaskSpec(
component_ref=pipeline_spec_pb2.ComponentRef(
name='comp-1-2')))
expected_main_pipeline_spec.components['inner-pipeline'].dag.tasks[
'task-2'].CopyFrom(
pipeline_spec_pb2.PipelineTaskSpec(
component_ref=pipeline_spec_pb2.ComponentRef(
name='comp-2')))
pipeline_spec_builder.merge_deployment_spec_and_component_spec(
main_pipeline_spec=main_pipeline_spec,
main_deployment_config=main_deployment_config,
sub_pipeline_spec=sub_pipeline_spec,
sub_pipeline_component_name='inner-pipeline',
)
self.assertEqual(sub_pipeline_spec, expected_sub_pipeline_spec)
self.assertEqual(main_pipeline_spec, expected_main_pipeline_spec)
self.assertEqual(main_deployment_config,
expected_main_deployment_config)
def pipeline_spec_from_file(filepath: str) -> str:
with open(filepath, 'r') as f:
dictionary = yaml.safe_load(f)
return json_format.ParseDict(dictionary, pipeline_spec_pb2.PipelineSpec())
class TestWriteIrToFile(unittest.TestCase):
@classmethod
def setUpClass(cls) -> None:
pipeline_spec = pipeline_spec_pb2.PipelineSpec()
pipeline_spec.pipeline_info.name = 'pipeline-name'
cls.pipeline_spec = pipeline_spec
def test_yaml(self):
with tempfile.TemporaryDirectory() as tempdir:
temp_filepath = os.path.join(tempdir, 'output.yaml')
pipeline_spec_builder.write_pipeline_spec_to_file(
self.pipeline_spec, temp_filepath)
actual = pipeline_spec_from_file(temp_filepath)
self.assertEqual(actual, self.pipeline_spec)
def test_yml(self):
with tempfile.TemporaryDirectory() as tempdir:
temp_filepath = os.path.join(tempdir, 'output.yml')
pipeline_spec_builder.write_pipeline_spec_to_file(
self.pipeline_spec, temp_filepath)
actual = pipeline_spec_from_file(temp_filepath)
self.assertEqual(actual, self.pipeline_spec)
def test_json(self):
with tempfile.TemporaryDirectory() as tempdir, self.assertWarnsRegex(
DeprecationWarning, r'Compiling to JSON is deprecated'):
temp_filepath = os.path.join(tempdir, 'output.json')
pipeline_spec_builder.write_pipeline_spec_to_file(
self.pipeline_spec, temp_filepath)
actual = pipeline_spec_from_file(temp_filepath)
self.assertEqual(actual, self.pipeline_spec)
def test_incorrect_extension(self):
with tempfile.TemporaryDirectory() as tempdir, self.assertRaisesRegex(
ValueError, r'should end with "\.yaml"\.'):
temp_filepath = os.path.join(tempdir, 'output.txt')
pipeline_spec_builder.write_pipeline_spec_to_file(
self.pipeline_spec, temp_filepath)
if __name__ == '__main__':
unittest.main()