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