def test_component_spec_with_placeholder_referencing_nonexisting_input_output(
self):
with self.assertRaisesRegex(
pydantic.ValidationError, 'Argument "input_name=\'input000\'" '
'references non-existing input.'):
structures.ComponentSpec(
name='component_1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
image='alpine',
commands=[
'sh',
'-c',
'set -ex\necho "$0" > "$1"',
structures.InputValuePlaceholder(
input_name='input000'),
structures.OutputPathPlaceholder(
output_name='output1'),
],
)),
inputs={'input1': structures.InputSpec(type='String')},
outputs={'output1': structures.OutputSpec(type='String')},
)
with self.assertRaisesRegex(
pydantic.ValidationError,
'Argument "output_name=\'output000\'" '
'references non-existing output.'):
structures.ComponentSpec(
name='component_1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
image='alpine',
commands=[
'sh',
'-c',
'set -ex\necho "$0" > "$1"',
structures.InputValuePlaceholder(
input_name='input1'),
structures.OutputPathPlaceholder(
output_name='output000'),
],
)),
inputs={'input1': structures.InputSpec(type='String')},
outputs={'output1': structures.OutputSpec(type='String')},
)
def test_create_pipeline_task_valid(self):
expected_component_spec = structures.ComponentSpec(
name='component1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
image='alpine',
commands=['sh', '-c', 'echo "$0" >> "$1"'],
arguments=[
structures.InputValuePlaceholder(input_name='input1'),
structures.OutputPathPlaceholder(
output_name='output1'),
],
)),
inputs={
'input1': structures.InputSpec(type='String'),
},
outputs={
'output1': structures.OutputSpec(type='Artifact'),
},
)
expected_task_spec = structures.TaskSpec(
name='component1',
inputs={'input1': 'value'},
dependent_tasks=[],
component_ref='component1',
)
expected_container_spec = structures.ContainerSpec(
image='alpine',
commands=['sh', '-c', 'echo "$0" >> "$1"'],
arguments=[
"{{$.inputs.parameters['input1']}}",
"{{$.outputs.artifacts['output1'].path}}",
],
)
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
arguments={'input1': 'value'},
)
self.assertEqual(task.task_spec, expected_task_spec)
self.assertEqual(task.component_spec, expected_component_spec)
self.assertEqual(task.container_spec, expected_container_spec)
def test_simple_component_spec_save_to_component_yaml(self):
open_mock = mock.mock_open()
expected_yaml = textwrap.dedent("""\
name: component_1
inputs:
input1: {type: String}
outputs:
output1: {type: String}
implementation:
container:
image: alpine
commands:
- sh
- -c
- 'set -ex
echo "$0" > "$1"'
- {inputValue: input1}
- {outputPath: output1}
""")
with mock.patch("builtins.open", open_mock, create=True):
structures.ComponentSpec(
name='component_1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
image='alpine',
commands=[
'sh',
'-c',
'set -ex\necho "$0" > "$1"',
structures.InputValuePlaceholder(
input_name='input1'),
structures.OutputPathPlaceholder(
output_name='output1'),
],
)),
inputs={
'input1': structures.InputSpec(type='String')
},
outputs={
'output1': structures.OutputSpec(type='String')
},
).save_to_component_yaml('test_save_file.txt')
open_mock.assert_called_with('test_save_file.txt', 'a')
open_mock.return_value.write.assert_called_once_with(expected_yaml)
def test_resolve_concat_placeholder(self):
expected_container_spec = structures.ContainerSpec(
image='alpine',
commands=[
'sh',
'-c',
'echo "$0"',
"{{$.inputs.parameters['input1']}}+{{$.inputs.parameters['input2']}}",
],
)
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML_CONCAT_PLACEHOLDER),
arguments={
'input1': '1',
'input2': '2',
},
)
self.assertEqual(task.container_spec, expected_container_spec)
def test_simple_component_spec_load_from_v2_component_yaml(self):
component_yaml_v2 = textwrap.dedent("""\
name: component_1
inputs:
input1:
type: String
outputs:
output1:
type: String
implementation:
container:
image: alpine
commands:
- sh
- -c
- 'set -ex
echo "$0" > "$1"'
- inputValue: input1
- outputPath: output1
""")
generated_spec = structures.ComponentSpec.load_from_component_yaml(
component_yaml_v2)
expected_spec = structures.ComponentSpec(
name='component_1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
image='alpine',
commands=[
'sh',
'-c',
'set -ex\necho "$0" > "$1"',
structures.InputValuePlaceholder(input_name='input1'),
structures.OutputPathPlaceholder(output_name='output1'),
],
)),
inputs={'input1': structures.InputSpec(type='String')},
outputs={'output1': structures.OutputSpec(type='String')})
self.assertEqual(generated_spec, expected_spec)
class PipelineTaskTest(parameterized.TestCase):
def test_create_pipeline_task_valid(self):
expected_component_spec = structures.ComponentSpec(
name='component1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
image='alpine',
commands=['sh', '-c', 'echo "$0" >> "$1"'],
arguments=[
structures.InputValuePlaceholder(input_name='input1'),
structures.OutputPathPlaceholder(
output_name='output1'),
],
)),
inputs={
'input1': structures.InputSpec(type='String'),
},
outputs={
'output1': structures.OutputSpec(type='Artifact'),
},
)
expected_task_spec = structures.TaskSpec(
name='component1',
inputs={'input1': 'value'},
dependent_tasks=[],
component_ref='component1',
)
expected_container_spec = structures.ContainerSpec(
image='alpine',
commands=['sh', '-c', 'echo "$0" >> "$1"'],
arguments=[
"{{$.inputs.parameters['input1']}}",
"{{$.outputs.artifacts['output1'].path}}",
],
)
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
arguments={'input1': 'value'},
)
self.assertEqual(task.task_spec, expected_task_spec)
self.assertEqual(task.component_spec, expected_component_spec)
self.assertEqual(task.container_spec, expected_container_spec)
def test_create_pipeline_task_invalid_missing_required_input(self):
with self.assertRaisesRegex(ValueError,
'No value provided for input: input1.'):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.
load_from_component_yaml(V2_YAML),
arguments={},
)
def test_create_pipeline_task_invalid_wrong_input(self):
with self.assertRaisesRegex(
ValueError,
'Component "component1" got an unexpected input: input0.'):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.
load_from_component_yaml(V2_YAML),
arguments={
'input1': 'value',
'input0': 'abc',
},
)
@parameterized.parameters(
{
'component_yaml':
V2_YAML_IF_PLACEHOLDER,
'arguments': {
'optional_input_1': 'value'
},
'expected_container_spec':
structures.ContainerSpec(
image='alpine',
commands=['sh', '-c', 'echo "$0" "$1"'],
arguments=[
'input: ',
"{{$.inputs.parameters['optional_input_1']}}",
],
)
},
{
'component_yaml':
V2_YAML_IF_PLACEHOLDER,
'arguments': {},
'expected_container_spec':
structures.ContainerSpec(
image='alpine',
commands=['sh', '-c', 'echo "$0" "$1"'],
arguments=[
'default: ',
'Hello world!',
],
)
},
)
def test_resolve_if_placeholder(
self,
component_yaml: str,
arguments: dict,
expected_container_spec: structures.ContainerSpec,
):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
component_yaml),
arguments=arguments,
)
self.assertEqual(task.container_spec, expected_container_spec)
def test_resolve_concat_placeholder(self):
expected_container_spec = structures.ContainerSpec(
image='alpine',
commands=[
'sh',
'-c',
'echo "$0"',
"{{$.inputs.parameters['input1']}}+{{$.inputs.parameters['input2']}}",
],
)
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML_CONCAT_PLACEHOLDER),
arguments={
'input1': '1',
'input2': '2',
},
)
self.assertEqual(task.container_spec, expected_container_spec)
def test_set_caching_options(self):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
arguments={'input1': 'value'},
)
task.set_caching_options(False)
self.assertEqual(False, task.task_spec.enable_caching)
@parameterized.parameters(
{
'cpu_limit': '123',
'expected_cpu_number': 123,
},
{
'cpu_limit': '123m',
'expected_cpu_number': 0.123,
},
{
'cpu_limit': '123.0',
'expected_cpu_number': 123,
},
{
'cpu_limit': '123.0m',
'expected_cpu_number': 0.123,
},
)
def test_set_valid_cpu_limit(self, cpu_limit: str,
expected_cpu_number: float):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
arguments={'input1': 'value'},
)
task.set_cpu_limit(cpu_limit)
self.assertEqual(expected_cpu_number,
task.container_spec.resources.cpu_limit)
@parameterized.parameters(
{
'gpu_limit': '666',
'expected_gpu_number': 666,
}, )
def test_set_valid_gpu_limit(self, gpu_limit: str,
expected_gpu_number: int):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
arguments={'input1': 'value'},
)
task.set_gpu_limit(gpu_limit)
self.assertEqual(expected_gpu_number,
task.container_spec.resources.accelerator_count)
@parameterized.parameters(
{
'memory': '1E',
'expected_memory_number': 1000000000,
},
{
'memory': '15Ei',
'expected_memory_number': 17293822569.102703,
},
{
'memory': '2P',
'expected_memory_number': 2000000,
},
{
'memory': '25Pi',
'expected_memory_number': 28147497.6710656,
},
{
'memory': '3T',
'expected_memory_number': 3000,
},
{
'memory': '35Ti',
'expected_memory_number': 38482.90697216,
},
{
'memory': '4G',
'expected_memory_number': 4,
},
{
'memory': '45Gi',
'expected_memory_number': 48.31838208,
},
{
'memory': '5M',
'expected_memory_number': 0.005,
},
{
'memory': '55Mi',
'expected_memory_number': 0.05767168,
},
{
'memory': '6K',
'expected_memory_number': 0.000006,
},
{
'memory': '65Ki',
'expected_memory_number': 0.00006656,
},
{
'memory': '7000',
'expected_memory_number': 0.000007,
},
)
def test_set_memory_limit(self, memory: str, expected_memory_number: int):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
arguments={'input1': 'value'},
)
task.set_memory_limit(memory)
self.assertEqual(expected_memory_number,
task.container_spec.resources.memory_limit)
def test_add_node_selector_constraint_type_only(self):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
arguments={'input1': 'value'},
)
task.add_node_selector_constraint('NVIDIA_TESLA_K80')
self.assertEqual(
structures.ResourceSpec(accelerator_type='NVIDIA_TESLA_K80',
accelerator_count=1),
task.container_spec.resources)
def test_add_node_selector_constraint_accelerator_count(self):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
arguments={'input1': 'value'},
)
task.set_gpu_limit('5').add_node_selector_constraint('TPU_V3')
self.assertEqual(
structures.ResourceSpec(accelerator_type='TPU_V3',
accelerator_count=5),
task.container_spec.resources)
def test_set_env_variable(self):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
arguments={'input1': 'value'},
)
task.set_env_variable('env_name', 'env_value')
self.assertEqual({'env_name': 'env_value'}, task.container_spec.env)
def test_set_display_name(self):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
arguments={'input1': 'value'},
)
task.set_display_name('test_name')
self.assertEqual('test_name', task.task_spec.display_name)
def test_component_spec_load_from_v1_component_yaml(self):
component_yaml_v1 = textwrap.dedent("""\
name: Component with 2 inputs and 2 outputs
inputs:
- {name: Input parameter, type: String}
- {name: Input artifact}
outputs:
- {name: Output 1}
- {name: Output 2}
implementation:
container:
image: busybox
command: [sh, -c, '
mkdir -p $(dirname "$2")
mkdir -p $(dirname "$3")
echo "$0" > "$2"
cp "$1" "$3"
'
]
args:
- {inputValue: Input parameter}
- {inputPath: Input artifact}
- {outputPath: Output 1}
- {outputPath: Output 2}
""")
generated_spec = structures.ComponentSpec.load_from_component_yaml(
component_yaml_v1)
expected_spec = structures.ComponentSpec(
name='Component with 2 inputs and 2 outputs',
implementation=structures.Implementation(
container=structures.ContainerSpec(
image='busybox',
commands=[
'sh',
'-c',
(' mkdir -p $(dirname "$2") mkdir -p $(dirname "$3") '
'echo "$0" > "$2" cp "$1" "$3" '),
],
arguments=[
structures.InputValuePlaceholder(
input_name='input_parameter'),
structures.InputPathPlaceholder(
input_name='input_artifact'),
structures.OutputPathPlaceholder(
output_name='output_1'),
structures.OutputPathPlaceholder(
output_name='output_2'),
],
env={},
)),
inputs={
'input_parameter': structures.InputSpec(type='String'),
'input_artifact': structures.InputSpec(type='Artifact')
},
outputs={
'output_1': structures.OutputSpec(type='Artifact'),
'output_2': structures.OutputSpec(type='Artifact'),
})
self.assertEqual(generated_spec, expected_spec)
- --arg1
- {inputValue: optional_input_1}
else: [--arg2, default]
""")
V2_COMPONENT_SPEC_IF_PLACEHOLDER = structures.ComponentSpec(
name='component_if',
implementation=structures.Implementation(
container=structures.ContainerSpec(
image='alpine',
arguments=[
structures.IfPresentPlaceholder(
if_structure=structures.IfPresentPlaceholderStructure(
input_name='optional_input_1',
then=[
'--arg1',
structures.InputValuePlaceholder(
input_name='optional_input_1'),
],
otherwise=[
'--arg2',
'default',
]))
])),
inputs={
'optional_input_1': structures.InputSpec(type='String', default=None)
},
)
V1_YAML_CONCAT_PLACEHOLDER = textwrap.dedent("""\
name: component_concat
inputs:
class TestComponent(base_component.BaseComponent):
def execute(self, *args, **kwargs):
pass
component_op = TestComponent(
component_spec=structures.ComponentSpec(
name='component_1',
implementation=structures.ContainerSpec(
image='alpine',
commands=[
'sh',
'-c',
'set -ex\necho "$0" "$1" "$2" > "$3"',
structures.InputValuePlaceholder(input_name='input1'),
structures.InputValuePlaceholder(input_name='input2'),
structures.InputValuePlaceholder(input_name='input3'),
structures.OutputPathPlaceholder(output_name='output1'),
],
),
inputs={
'input1':
structures.InputSpec(type='String'),
'input2':
structures.InputSpec(type='Integer'),
'input3':
structures.InputSpec(type='Float', default=3.14),
'input4':
structures.InputSpec(type='Optional[Float]', default=None),
},
def create_component_from_func(func: Callable,
base_image: Optional[str] = None,
target_image: Optional[str] = None,
packages_to_install: List[str] = None,
output_component_file: Optional[str] = None,
install_kfp_package: bool = True,
kfp_package_path: Optional[str] = None):
"""Implementation for the @component decorator.
The decorator is defined under component_decorator.py. See the
decorator for the canonical documentation for this function.
"""
packages_to_install = packages_to_install or []
if install_kfp_package and target_image is None:
if kfp_package_path is None:
kfp_package_path = _get_default_kfp_package_path()
packages_to_install.append(kfp_package_path)
packages_to_install_command = _get_packages_to_install_command(
package_list=packages_to_install)
command = []
args = []
if base_image is None:
base_image = _DEFAULT_BASE_IMAGE
component_image = base_image
if target_image:
component_image = target_image
command, args = _get_command_and_args_for_containerized_component(
function_name=func.__name__, )
else:
command, args = _get_command_and_args_for_lightweight_component(
func=func)
component_spec = extract_component_interface(func)
component_spec.implementation = structures.Implementation(
container=structures.ContainerSpec(
image=component_image,
commands=packages_to_install_command + command,
arguments=args,
))
module_path = pathlib.Path(inspect.getsourcefile(func))
module_path.resolve()
component_name = _python_function_name_to_component_name(func.__name__)
component_info = ComponentInfo(name=component_name,
function_name=func.__name__,
func=func,
target_image=target_image,
module_path=module_path,
component_spec=component_spec,
output_component_file=output_component_file,
base_image=base_image)
if REGISTERED_MODULES is not None:
REGISTERED_MODULES[component_name] = component_info
if output_component_file:
component_spec.save_to_component_yaml(output_component_file)
return python_component.PythonComponent(component_spec=component_spec,
python_func=func)
