304 lines
13 KiB
Python
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()
|