def test_set_display_name(self):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
args={'input1': 'value'},
)
task.set_display_name('test_name')
self.assertEqual('test_name', task.task_spec.display_name)
def test_set_caching_options(self):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
args={'input1': 'value'},
)
task.set_caching_options(False)
self.assertEqual(False, task.task_spec.enable_caching)
def test_set_env_variable(self):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
args={'input1': 'value'},
)
task.set_env_variable('env_name', 'env_value')
self.assertEqual({'env_name': 'env_value'}, task.container_spec.env)
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),
args={},
)
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),
args={'input1': 'value'},
)
task.set_memory_limit(memory)
self.assertEqual(expected_memory_number,
task.container_spec.resources.memory_limit)
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),
args={'input1': 'value'},
)
task.set_gpu_limit(gpu_limit)
self.assertEqual(expected_gpu_number,
task.container_spec.resources.accelerator_count)
def test_add_node_selector_constraint_accelerator_count(self):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
args={'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_add_node_selector_constraint_type_only(self):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
args={'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_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),
args={
'input1': 'value',
'input0': 'abc',
},
)
def test_resolve_if_placeholder(
self,
component_yaml: str,
args: dict,
expected_container_spec: structures.ContainerSpec,
):
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
component_yaml),
args=args,
)
self.assertEqual(task.container_spec, expected_container_spec)
def _create_pipeline_from_component_spec(
self,
component_spec: structures.ComponentSpec,
) -> pipeline_spec_pb2.PipelineSpec:
"""Creates a pipeline instance and constructs the pipeline spec for a
primitive component.
Args:
component_spec: The ComponentSpec to convert to PipelineSpec.
Returns:
A PipelineSpec proto representing the compiled component.
"""
args_dict = {}
for arg_name, input_spec in component_spec.inputs.items():
arg_type = input_spec.type
if not type_utils.is_parameter_type(
arg_type) or type_utils.is_task_final_status_type(arg_type):
raise TypeError(
builder.make_invalid_input_type_error_msg(
arg_name, arg_type))
args_dict[arg_name] = dsl.PipelineParameterChannel(
name=arg_name, channel_type=arg_type)
task = pipeline_task.PipelineTask(component_spec, args_dict)
# instead of constructing a pipeline with pipeline_context.Pipeline,
# just build the single task group
group = tasks_group.TasksGroup(
group_type=tasks_group.TasksGroupType.PIPELINE)
group.tasks.append(task)
pipeline_inputs = component_spec.inputs or {}
# Fill in the default values.
args_list_with_defaults = [
dsl.PipelineParameterChannel(
name=input_name,
channel_type=input_spec.type,
value=input_spec.default,
) for input_name, input_spec in pipeline_inputs.items()
]
group.name = uuid.uuid4().hex
return builder.create_pipeline_spec_for_component(
pipeline_name=component_spec.name,
pipeline_args=args_list_with_defaults,
task_group=group,
)
def test_create_pipeline_task_valid(self):
expected_component_spec = structures.ComponentSpec(
name='component1',
implementation=structures.Implementation(
container=structures.ContainerSpec(
image='alpine',
command=['sh', '-c', 'echo "$0" >> "$1"'],
args=[
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',
command=['sh', '-c', 'echo "$0" >> "$1"'],
args=[
"{{$.inputs.parameters['input1']}}",
"{{$.outputs.artifacts['output1'].path}}",
],
)
task = pipeline_task.PipelineTask(
component_spec=structures.ComponentSpec.load_from_component_yaml(
V2_YAML),
args={'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_resolve_concat_placeholder(self):
expected_container_spec = structures.ContainerSpec(
image='alpine',
command=[
'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),
args={
'input1': '1',
'input2': '2',
},
)
self.assertEqual(task.container_spec, expected_container_spec)
def to_pipeline_spec(self) -> pipeline_spec_pb2.PipelineSpec:
"""Creates a pipeline instance and constructs the pipeline spec for a
single component.
Args:
component_spec: The ComponentSpec to convert to PipelineSpec.
Returns:
A PipelineSpec proto representing the compiled component.
"""
# import here to aviod circular module dependency
from kfp.compiler import pipeline_spec_builder as builder
from kfp.components import pipeline_task
from kfp.components import tasks_group
from kfp.components.types import type_utils
args_dict = {}
pipeline_inputs = self.inputs or {}
for arg_name, input_spec in pipeline_inputs.items():
arg_type = input_spec.type
if not type_utils.is_parameter_type(
arg_type) or type_utils.is_task_final_status_type(
arg_type):
raise TypeError(
builder.make_invalid_input_type_error_msg(
arg_name, arg_type))
args_dict[arg_name] = dsl.PipelineParameterChannel(
name=arg_name, channel_type=arg_type)
task = pipeline_task.PipelineTask(self, args_dict)
# instead of constructing a pipeline with pipeline_context.Pipeline,
# just build the single task group
group = tasks_group.TasksGroup(
group_type=tasks_group.TasksGroupType.PIPELINE)
group.tasks.append(task)
# Fill in the default values.
args_list_with_defaults = [
dsl.PipelineParameterChannel(
name=input_name,
channel_type=input_spec.type,
value=input_spec.default,
) for input_name, input_spec in pipeline_inputs.items()
]
group.name = uuid.uuid4().hex
pipeline_name = self.name
pipeline_args = args_list_with_defaults
task_group = group
builder.validate_pipeline_name(pipeline_name)
pipeline_spec = pipeline_spec_pb2.PipelineSpec()
pipeline_spec.pipeline_info.name = pipeline_name
pipeline_spec.sdk_version = f'kfp-{kfp.__version__}'
# Schema version 2.1.0 is required for kfp-pipeline-spec>0.1.13
pipeline_spec.schema_version = '2.1.0'
pipeline_spec.root.CopyFrom(
builder.build_component_spec_for_group(
pipeline_channels=pipeline_args,
is_root_group=True,
))
deployment_config = pipeline_spec_pb2.PipelineDeploymentConfig()
root_group = task_group
task_name_to_parent_groups, group_name_to_parent_groups = builder.get_parent_groups(
root_group)
def get_inputs(task_group: tasks_group.TasksGroup,
task_name_to_parent_groups):
inputs = collections.defaultdict(set)
if len(task_group.tasks) != 1:
raise ValueError(
f'Error compiling component. Expected one task in task group, got {len(task_group.tasks)}.'
)
only_task = task_group.tasks[0]
if only_task.channel_inputs:
for group_name in task_name_to_parent_groups[only_task.name]:
inputs[group_name].add(
(only_task.channel_inputs[-1], None))
return inputs
inputs = get_inputs(task_group, task_name_to_parent_groups)
builder.build_spec_by_group(
pipeline_spec=pipeline_spec,
deployment_config=deployment_config,
group=root_group,
inputs=inputs,
dependencies={}, # no dependencies for single-component pipeline
rootgroup_name=root_group.name,
task_name_to_parent_groups=task_name_to_parent_groups,
group_name_to_parent_groups=group_name_to_parent_groups,
name_to_for_loop_group=
{}, # no for loop for single-component pipeline
)
return pipeline_spec
