def test_pop_input_from_component_spec(self): component_spec = pipeline_spec_pb2.ComponentSpec( executor_label='exec-component1') component_spec.input_definitions.artifacts[ 'input1'].artifact_type.schema_title = 'system.Dataset' component_spec.input_definitions.parameters[ 'input2'].type = pipeline_spec_pb2.PrimitiveType.STRING component_spec.input_definitions.parameters[ 'input3'].type = pipeline_spec_pb2.PrimitiveType.DOUBLE # pop an artifact, and there're other inputs left dsl_component_spec.pop_input_from_component_spec( component_spec, 'input1') expected_dict = { 'inputDefinitions': { 'parameters': { 'input2': { 'type': 'STRING' }, 'input3': { 'type': 'DOUBLE' } } }, 'executorLabel': 'exec-component1' } expected_spec = pipeline_spec_pb2.ComponentSpec() json_format.ParseDict(expected_dict, expected_spec) self.assertEqual(expected_spec, component_spec) # pop an parameter, and there're other inputs left dsl_component_spec.pop_input_from_component_spec( component_spec, 'input2') expected_dict = { 'inputDefinitions': { 'parameters': { 'input3': { 'type': 'DOUBLE' } } }, 'executorLabel': 'exec-component1' } expected_spec = pipeline_spec_pb2.ComponentSpec() json_format.ParseDict(expected_dict, expected_spec) self.assertEqual(expected_spec, component_spec) # pop the last input, expect no inputDefinitions dsl_component_spec.pop_input_from_component_spec( component_spec, 'input3') expected_dict = {'executorLabel': 'exec-component1'} expected_spec = pipeline_spec_pb2.ComponentSpec() json_format.ParseDict(expected_dict, expected_spec) self.assertEqual(expected_spec, component_spec) # pop an input that doesn't exist, expect no-op. dsl_component_spec.pop_input_from_component_spec( component_spec, 'input4') self.assertEqual(expected_spec, component_spec)
def _group_to_dag_spec( self, group: dsl.OpsGroup, inputs: Dict[str, List[Tuple[dsl.PipelineParam, str]]], outputs: Dict[str, List[Tuple[dsl.PipelineParam, str]]], dependencies: Dict[str, List[_GroupOrOp]], pipeline_spec: pipeline_spec_pb2.PipelineSpec, deployment_config: pipeline_spec_pb2.PipelineDeploymentConfig, rootgroup_name: str, ) -> None: """Generate IR spec given an OpsGroup. Args: group: The OpsGroup to generate spec for. inputs: The inputs dictionary. The keys are group/op names and values are lists of tuples (param, producing_op_name). outputs: The outputs dictionary. The keys are group/op names and values are lists of tuples (param, producing_op_name). dependencies: The group dependencies dictionary. The keys are group/op names, and the values are lists of dependent groups/ops. pipeline_spec: The pipeline_spec to update in-place. deployment_config: The deployment_config to hold all executors. rootgroup_name: The name of the group root. Used to determine whether the component spec for the current group should be the root dag. """ group_component_name = dsl_utils.sanitize_component_name(group.name) if group.name == rootgroup_name: group_component_spec = pipeline_spec.root else: group_component_spec = pipeline_spec.components[group_component_name] # Generate task specs and component specs for the dag. subgroups = group.groups + group.ops for subgroup in subgroups: subgroup_task_spec = getattr(subgroup, 'task_spec', pipeline_spec_pb2.PipelineTaskSpec()) subgroup_component_spec = getattr(subgroup, 'component_spec', pipeline_spec_pb2.ComponentSpec()) is_loop_subgroup = (isinstance(group, dsl.ParallelFor)) is_recursive_subgroup = ( isinstance(subgroup, dsl.OpsGroup) and subgroup.recursive_ref) # Special handling for recursive subgroup: use the existing opsgroup name if is_recursive_subgroup: subgroup_key = subgroup.recursive_ref.name else: subgroup_key = subgroup.name subgroup_task_spec.task_info.name = ( subgroup_task_spec.task_info.name or dsl_utils.sanitize_task_name(subgroup_key)) # human_name exists for ops only, and is used to de-dupe component spec. subgroup_component_name = ( subgroup_task_spec.component_ref.name or dsl_utils.sanitize_component_name( getattr(subgroup, 'human_name', subgroup_key))) subgroup_task_spec.component_ref.name = subgroup_component_name if isinstance(subgroup, dsl.OpsGroup) and subgroup.type == 'graph': raise NotImplementedError( 'dsl.graph_component is not yet supported in KFP v2 compiler.') if isinstance(subgroup, dsl.OpsGroup) and subgroup.type == 'exit_handler': raise NotImplementedError( 'dsl.ExitHandler is not yet supported in KFP v2 compiler.') importer_tasks = [] # Add importer node when applicable for input_name in subgroup_task_spec.inputs.artifacts: if not subgroup_task_spec.inputs.artifacts[ input_name].task_output_artifact.producer_task: type_schema = type_utils.get_input_artifact_type_schema( input_name, subgroup._metadata.inputs) importer_name = importer_node.generate_importer_base_name( dependent_task_name=subgroup_task_spec.task_info.name, input_name=input_name) importer_task_spec = importer_node.build_importer_task_spec( importer_name) importer_comp_spec = importer_node.build_importer_component_spec( importer_base_name=importer_name, input_name=input_name, input_type_schema=type_schema) importer_task_name = importer_task_spec.task_info.name importer_comp_name = importer_task_spec.component_ref.name importer_exec_label = importer_comp_spec.executor_label group_component_spec.dag.tasks[importer_task_name].CopyFrom( importer_task_spec) pipeline_spec.components[importer_comp_name].CopyFrom( importer_comp_spec) subgroup_task_spec.inputs.artifacts[ input_name].task_output_artifact.producer_task = ( importer_task_name) subgroup_task_spec.inputs.artifacts[ input_name].task_output_artifact.output_artifact_key = ( importer_node.OUTPUT_KEY) # Retrieve the pre-built importer spec importer_spec = subgroup.importer_specs[input_name] deployment_config.executors[importer_exec_label].importer.CopyFrom( importer_spec) importer_tasks.append(importer_task_name) group_inputs = inputs.get(group.name, []) subgroup_inputs = inputs.get(subgroup.name, []) subgroup_params = [param for param, _ in subgroup_inputs] tasks_in_current_dag = [ dsl_utils.sanitize_task_name(subgroup.name) for subgroup in subgroups ] + importer_tasks is_parent_component_root = group_component_spec == pipeline_spec.root # Additional spec modifications for dsl.ParallelFor's subgroups. if is_loop_subgroup: self._update_loop_specs(group, subgroup, group_component_spec, subgroup_component_spec, subgroup_task_spec) elif isinstance(subgroup, dsl.ContainerOp): dsl_component_spec.update_task_inputs_spec( subgroup_task_spec, group_component_spec.input_definitions, subgroup_params, tasks_in_current_dag, ) if isinstance(subgroup, dsl.OpsGroup) and subgroup.type == 'condition': # "punch the hole", adding inputs needed by its subgroup or tasks. dsl_component_spec.build_component_inputs_spec( component_spec=subgroup_component_spec, pipeline_params=subgroup_params, is_root_component=False, ) dsl_component_spec.build_task_inputs_spec( subgroup_task_spec, subgroup_params, tasks_in_current_dag, is_parent_component_root, ) condition = subgroup.condition operand_values = [] for operand in [condition.operand1, condition.operand2]: operand_values.append(self._resolve_value_or_reference(operand)) condition_string = '{} {} {}'.format(operand_values[0], condition.operator, operand_values[1]) subgroup_task_spec.trigger_policy.CopyFrom( pipeline_spec_pb2.PipelineTaskSpec.TriggerPolicy( condition=condition_string)) # Generate dependencies section for this task. if dependencies.get(subgroup.name, None): group_dependencies = list(dependencies[subgroup.name]) group_dependencies.sort() subgroup_task_spec.dependent_tasks.extend( [dsl_utils.sanitize_task_name(dep) for dep in group_dependencies]) if isinstance(subgroup, dsl.ParallelFor): if subgroup.parallelism is not None: warnings.warn( 'Setting parallelism in ParallelFor is not supported yet.' 'The setting is ignored.') # Remove loop arguments related inputs from parent group component spec. input_names = [param.full_name for param, _ in inputs[subgroup.name]] for input_name in input_names: if _for_loop.LoopArguments.name_is_loop_argument(input_name): dsl_component_spec.pop_input_from_component_spec( group_component_spec, input_name) if subgroup.items_is_pipeline_param: # These loop args are a 'withParam' rather than 'withItems'. # i.e., rather than a static list, they are either the output of # another task or were input as global pipeline parameters. pipeline_param = subgroup.loop_args.items_or_pipeline_param input_parameter_name = pipeline_param.full_name if pipeline_param.op_name: subgroup_task_spec.inputs.parameters[ input_parameter_name].task_output_parameter.producer_task = ( dsl_utils.sanitize_task_name(pipeline_param.op_name)) subgroup_task_spec.inputs.parameters[ input_parameter_name].task_output_parameter.output_parameter_key = ( pipeline_param.name) else: subgroup_task_spec.inputs.parameters[ input_parameter_name].component_input_parameter = ( input_parameter_name) if pipeline_param.op_name is None: # Input parameter is from pipeline func rather than component output. # Correct loop argument input type in the parent component spec. # The loop argument was categorized as an artifact due to its missing # or non-primitive type annotation. But it should always be String # typed, as its value is a serialized JSON string. dsl_component_spec.pop_input_from_component_spec( group_component_spec, input_parameter_name) group_component_spec.input_definitions.parameters[ input_parameter_name].type = pipeline_spec_pb2.PrimitiveType.STRING # Add component spec if not exists if subgroup_component_name not in pipeline_spec.components: pipeline_spec.components[subgroup_component_name].CopyFrom( subgroup_component_spec) # Add task spec group_component_spec.dag.tasks[ subgroup_task_spec.task_info.name].CopyFrom(subgroup_task_spec) # Add executor spec, if applicable. container_spec = getattr(subgroup, 'container_spec', None) if container_spec: if compiler_utils.is_v2_component(subgroup): compiler_utils.refactor_v2_container_spec(container_spec) executor_label = subgroup_component_spec.executor_label if executor_label not in deployment_config.executors: deployment_config.executors[executor_label].container.CopyFrom( container_spec) # Add AIPlatformCustomJobSpec, if applicable. custom_job_spec = getattr(subgroup, 'custom_job_spec', None) if custom_job_spec: executor_label = subgroup_component_spec.executor_label if executor_label not in deployment_config.executors: deployment_config.executors[ executor_label].custom_job.custom_job.update(custom_job_spec) pipeline_spec.deployment_spec.update( json_format.MessageToDict(deployment_config))
def _update_loop_specs( self, group: dsl.OpsGroup, subgroup: _GroupOrOp, group_component_spec: pipeline_spec_pb2.ComponentSpec, subgroup_component_spec: pipeline_spec_pb2.ComponentSpec, subgroup_task_spec: pipeline_spec_pb2.PipelineTaskSpec, ) -> None: """Update IR specs for loop. Args: group: The dsl.ParallelFor OpsGroup. subgroup: One of the subgroups of dsl.ParallelFor. group_component_spec: The component spec of the group to update in place. subgroup_component_spec: The component spec of the subgroup to update. subgroup_task_spec: The task spec of the subgroup to update. """ input_names = [ input_name for input_name in subgroup_task_spec.inputs.parameters ] for input_name in input_names: if subgroup_task_spec.inputs.parameters[input_name].HasField( 'component_input_parameter'): loop_argument_name = subgroup_task_spec.inputs.parameters[ input_name].component_input_parameter else: producer_task_name = dsl_utils.remove_task_name_prefix( subgroup_task_spec.inputs.parameters[input_name] .task_output_parameter.producer_task) producer_task_output_key = subgroup_task_spec.inputs.parameters[ input_name].task_output_parameter.output_parameter_key loop_argument_name = '{}-{}'.format(producer_task_name, producer_task_output_key) # Loop arguments are from dynamic input: pipeline param or task output if _for_loop.LoopArguments.name_is_withparams_loop_argument( loop_argument_name): arg_and_var_name = ( _for_loop.LoopArgumentVariable .parse_loop_args_name_and_this_var_name(loop_argument_name)) # The current IR representation is insufficient for referencing a subvar # which is a key in a list of dictionaries. if arg_and_var_name: raise NotImplementedError( 'Use subvar in dsl.ParallelFor with dynamic loop arguments is not ' 'supported. Got subvar: {}'.format(arg_and_var_name[1])) assert group.items_is_pipeline_param pipeline_param = group.loop_args.items_or_pipeline_param input_parameter_name = pipeline_param.full_name # Correct loop argument input type in the parent component spec. # The loop argument was categorized as an artifact due to its missing # or non-primitive type annotation. But it should always be String # typed, as its value is a serialized JSON string. dsl_component_spec.pop_input_from_component_spec( group_component_spec, input_parameter_name) group_component_spec.input_definitions.parameters[ input_parameter_name].type = pipeline_spec_pb2.PrimitiveType.STRING subgroup_task_spec.inputs.parameters[ input_parameter_name].component_input_parameter = ( input_parameter_name) subgroup_task_spec.parameter_iterator.item_input = input_name subgroup_task_spec.parameter_iterator.items.input_parameter = ( input_parameter_name) # Loop arguments comme from static raw values known at compile time. elif _for_loop.LoopArguments.name_is_withitems_loop_argument( loop_argument_name): # Prepare the raw values, either the whole list or the sliced list based # on subvar_name. subvar_name = None if _for_loop.LoopArgumentVariable.name_is_loop_arguments_variable( loop_argument_name): subvar_name = _for_loop.LoopArgumentVariable.get_subvar_name( loop_argument_name) loop_args = group.loop_args.to_list_for_task_yaml() if subvar_name: raw_values = [loop_arg.get(subvar_name) for loop_arg in loop_args] else: raw_values = loop_args # If the loop iterator component expects `str` or `int` typed items from # the loop argument, make sure the item values are string values. # This is because both integers and strings are assigned to protobuf # [Value.string_value](https://github.com/protocolbuffers/protobuf/blob/133e5e75263be696c06599ab97614a1e1e6d9c66/src/google/protobuf/struct.proto#L70) # Such a conversion is not needed for `float` type. which uses protobuf # [Value.number_value](https://github.com/protocolbuffers/protobuf/blob/133e5e75263be696c06599ab97614a1e1e6d9c66/src/google/protobuf/struct.proto#L68) if subgroup_component_spec.input_definitions.parameters[ input_name].type in [ pipeline_spec_pb2.PrimitiveType.STRING, pipeline_spec_pb2.PrimitiveType.INT ]: raw_values = [str(v) for v in raw_values] if subgroup_component_spec.input_definitions.parameters[ input_name].type == pipeline_spec_pb2.PrimitiveType.INT: warnings.warn( 'The loop iterator component is expecting an `int` value.' 'Consider changing the input type to either `str` or `float`.') subgroup_task_spec.parameter_iterator.item_input = input_name subgroup_task_spec.parameter_iterator.items.raw = json.dumps(raw_values) else: raise AssertionError( 'Unexpected loop argument: {}'.format(loop_argument_name)) # Clean up unused inputs from task spec and parent component spec. dsl_component_spec.pop_input_from_task_spec(subgroup_task_spec, input_name) dsl_component_spec.pop_input_from_component_spec(group_component_spec, loop_argument_name)