def components(self, components: List[base_node.BaseNode]): deduped_components = set(components) producer_map = {} instances_per_component_type = collections.defaultdict(set) # Fills in producer map. for component in deduped_components: # Guarantees every component of a component type has unique component_id. if component.id in instances_per_component_type[component.type]: raise RuntimeError('Duplicated component_id %s for component type %s' % (component.id, component.type)) instances_per_component_type[component.type].add(component.id) for key, output_channel in component.outputs.items(): assert not producer_map.get( output_channel), '{} produced more than once'.format(output_channel) producer_map[output_channel] = component output_channel.producer_component_id = component.id output_channel.output_key = key # Connects nodes based on producer map. for component in deduped_components: for i in component.inputs.values(): if producer_map.get(i): component.add_upstream_node(producer_map[i]) producer_map[i].add_downstream_node(component) layers = topsort.topsorted_layers( list(deduped_components), get_node_id_fn=lambda c: c.id, get_parent_nodes=lambda c: c.upstream_nodes, get_child_nodes=lambda c: c.downstream_nodes) self._components = [] for layer in layers: for component in layer: self._components.append(component)
def test_topsorted_layers_empty(self): layers = topsort.topsorted_layers( nodes=[], get_node_id_fn=lambda n: n.name, get_parent_nodes=lambda n: [], get_child_nodes=lambda n: []) self.assertEqual([], layers)
def components(self, components: List[base_node.BaseNode]): deduped_components = set(components) node_by_id = {} # TODO(b/202822834): Use better distinction for bound channels. # bound_channels stores the exhaustive list of all nodes' output channels, # which is implicitly *bound* to the single pipeline run, as opposed to # manually constructed channels to fetch artifacts beyond the current # pipeline run. bound_channels = set() # Fills in producer map. for component in deduped_components: # Checks every node has an unique id. if component.id in node_by_id: raise RuntimeError( f'Duplicated node_id {component.id} for component type' f'{component.type}.') node_by_id[component.id] = component for key, output_channel in component.outputs.items(): if (output_channel.producer_component_id is not None and output_channel.producer_component_id != component.id and output_channel.output_key != key): raise AssertionError( f'{output_channel} is produced more than once: ' f'{output_channel.producer_id}[{output_channel.output_key}], ' f'{component.id}[{key}]') output_channel.producer_component_id = component.id output_channel.output_key = key bound_channels.add(output_channel) # Connects nodes based on producer map. for component in deduped_components: channels = list(component.inputs.values()) for exec_property in component.exec_properties.values(): if isinstance(exec_property, ph.ChannelWrappedPlaceholder): channels.append(exec_property.channel) for base_channel in channels: for ch in channel_utils.get_individual_channels(base_channel): if ch not in bound_channels: continue upstream_node = node_by_id.get(ch.producer_component_id) if upstream_node: component.add_upstream_node(upstream_node) upstream_node.add_downstream_node(component) layers = topsort.topsorted_layers( list(deduped_components), get_node_id_fn=lambda c: c.id, get_parent_nodes=lambda c: c.upstream_nodes, get_child_nodes=lambda c: c.downstream_nodes) self._components = [] for layer in layers: for component in layer: self._components.append(component) if self.beam_pipeline_args: for component in self._components: add_beam_pipeline_args_to_component(component, self.beam_pipeline_args)
def test_topsorted_layers_error_if_cycle(self): nodes = [ Node('A', [], ['B', 'E']), Node('B', ['A', 'D'], ['C']), Node('C', ['B'], ['D']), Node('D', ['C'], ['B']), Node('E', ['A'], []) ] node_map = {node.name: node for node in nodes} with self.assertRaisesRegex(topsort.InvalidDAGError, 'Cycle detected.'): topsort.topsorted_layers( nodes, get_node_id_fn=lambda n: n.name, get_parent_nodes=( lambda n: [node_map[name] for name in n.upstream_nodes]), get_child_nodes=( lambda n: [node_map[name] for name in n.downstream_nodes]))
def _topsorted_layers( pipeline: pipeline_pb2.Pipeline ) -> List[List[pipeline_pb2.PipelineNode]]: """Returns pipeline nodes in topologically sorted layers.""" node_by_id = _node_by_id(pipeline) return topsort.topsorted_layers( [node.pipeline_node for node in pipeline.nodes], get_node_id_fn=lambda node: node.node_info.id, get_parent_nodes=( lambda node: [node_by_id[n] for n in node.upstream_nodes]), get_child_nodes=( lambda node: [node_by_id[n] for n in node.downstream_nodes]))
def generate(self) -> List[task_lib.Task]: """Generates tasks for executing the next executable nodes in the pipeline. The returned tasks must have `exec_task` populated. List may be empty if no nodes are ready for execution. Returns: A `list` of tasks to execute. """ layers = topsort.topsorted_layers( [node.pipeline_node for node in self._pipeline.nodes], get_node_id_fn=lambda node: node.node_info.id, get_parent_nodes=( lambda node: [self._node_map[n] for n in node.upstream_nodes]), get_child_nodes=( lambda node: [self._node_map[n] for n in node.downstream_nodes])) result = [] for nodes in layers: # Boolean that's set if there's at least one successfully executed node # in the current layer. executed_nodes = False for node in nodes: if node.node_info.id in self._ignore_node_ids: logging.info( 'Ignoring node for task generation: %s', task_lib.NodeUid.from_pipeline_node( self._pipeline, node)) continue # If a task for the node is already tracked by the task queue, it need # not be considered for generation again. if self._is_task_id_tracked_fn( task_lib.exec_node_task_id_from_pipeline_node( self._pipeline, node)): continue executions = task_gen_utils.get_executions( self._mlmd_handle, node) if (executions and task_gen_utils.is_latest_execution_successful( executions)): executed_nodes = True continue # If all upstream nodes are executed but current node is not executed, # the node is deemed ready for execution. if self._upstream_nodes_executed(node): task = self._generate_task(node) if task: result.append(task) # If there are no executed nodes in the current layer, downstream nodes # need not be checked. if not executed_nodes: break return result
def test_topsorted_layers_ignore_duplicate_parent_node(self): nodes = [ Node('A', [], ['B']), Node('B', ['A', 'A'], []), # Duplicate parent node 'A' ] node_map = {node.name: node for node in nodes} layers = topsort.topsorted_layers( nodes, get_node_id_fn=lambda n: n.name, get_parent_nodes=( lambda n: [node_map[name] for name in n.upstream_nodes]), get_child_nodes=( lambda n: [node_map[name] for name in n.downstream_nodes])) self.assertEqual([['A'], ['B']], [[node.name for node in layer] for layer in layers])
def test_topsorted_layers_ignore_unknown_parent_node(self): nodes = [ Node('A', [], ['B']), Node('B', ['A'], ['C']), Node('C', ['B'], []), ] node_map = {node.name: node for node in nodes} # Exclude node A. Node B now has a parent node 'A' that should be ignored. layers = topsort.topsorted_layers( [node_map['B'], node_map['C']], get_node_id_fn=lambda n: n.name, get_parent_nodes=( lambda n: [node_map[name] for name in n.upstream_nodes]), get_child_nodes=( lambda n: [node_map[name] for name in n.downstream_nodes])) self.assertEqual([['B'], ['C']], [[node.name for node in layer] for layer in layers])
def test_topsorted_layers_DAG(self): nodes = [ Node('A', [], ['B', 'C', 'D']), Node('B', ['A'], []), Node('C', ['A'], ['D']), Node('D', ['A', 'C', 'F'], ['E']), Node('E', ['D'], []), Node('F', [], ['D']) ] node_map = {node.name: node for node in nodes} layers = topsort.topsorted_layers( nodes, get_node_id_fn=lambda n: n.name, get_parent_nodes=( lambda n: [node_map[name] for name in n.upstream_nodes]), get_child_nodes=( lambda n: [node_map[name] for name in n.downstream_nodes])) self.assertEqual([['A', 'F'], ['B', 'C'], ['D'], ['E']], [[node.name for node in layer] for layer in layers])
def generate(self) -> List[task_pb2.Task]: """Generates tasks for executing the next executable nodes in the pipeline. The returned tasks must have `exec_task` populated. List may be empty if no nodes are ready for execution. Returns: A `list` of tasks to execute. """ layers = topsort.topsorted_layers( [node.pipeline_node for node in self._pipeline.nodes], get_node_id_fn=lambda node: node.node_info.id, get_parent_nodes=( lambda node: [self._node_map[n] for n in node.upstream_nodes]), get_child_nodes=( lambda node: [self._node_map[n] for n in node.downstream_nodes])) tasks = [] with self._mlmd_connection as m: # TODO(goutham): Cache executions and/or use TaskQueue so that we don't # have to make MLMD queries for upstream nodes in each iteration. for nodes in layers: # Boolean that's set if there's at least one successfully executed node # in the current layer. executed_nodes = False for node in nodes: executions = task_gen_utils.get_executions(m, node) if (executions and task_gen_utils.is_latest_execution_successful( executions)): executed_nodes = True continue # If all upstream nodes are executed but current node is not executed, # the node is deemed ready for execution. if self._upstream_nodes_executed(m, node): task = self._generate_task(m, node) if task: tasks.append(task) # If there are no executed nodes in the current layer, downstream nodes # need not be checked. if not executed_nodes: break return tasks
def generate(self) -> List[task_lib.Task]: """Generates tasks for executing the next executable nodes in the pipeline. The returned tasks must have `exec_task` populated. List may be empty if no nodes are ready for execution. Returns: A `list` of tasks to execute. """ layers = topsort.topsorted_layers( [node.pipeline_node for node in self._pipeline.nodes], get_node_id_fn=lambda node: node.node_info.id, get_parent_nodes=( lambda node: [self._node_map[n] for n in node.upstream_nodes]), get_child_nodes=( lambda node: [self._node_map[n] for n in node.downstream_nodes])) result = [] for layer_num, nodes in enumerate(layers): # Boolean that's set if there's at least one successfully executed node # in the current layer. completed_node_ids = set() for node in nodes: node_uid = task_lib.NodeUid.from_pipeline_node( self._pipeline, node) node_id = node.node_info.id if self._service_job_manager.is_pure_service_node( self._pipeline_state, node.node_info.id): if not self._upstream_nodes_executed(node): continue service_status = self._service_job_manager.ensure_node_services( self._pipeline_state, node_id) if service_status == service_jobs.ServiceStatus.SUCCESS: logging.info('Service node completed successfully: %s', node_uid) completed_node_ids.add(node_id) elif service_status == service_jobs.ServiceStatus.FAILED: logging.error('Failed service node: %s', node_uid) return [ task_lib.FinalizePipelineTask( pipeline_uid=self._pipeline_state.pipeline_uid, status=status_lib.Status( code=status_lib.Code.ABORTED, message= (f'Aborting pipeline execution due to service ' f'node failure; failed node uid: {node_uid}' ))) ] else: logging.info('Pure service node in progress: %s', node_uid) continue # If a task for the node is already tracked by the task queue, it need # not be considered for generation again. if self._is_task_id_tracked_fn( task_lib.exec_node_task_id_from_pipeline_node( self._pipeline, node)): continue executions = task_gen_utils.get_executions( self._mlmd_handle, node) if (executions and task_gen_utils.is_latest_execution_successful( executions)): completed_node_ids.add(node_id) continue # If all upstream nodes are executed but current node is not executed, # the node is deemed ready for execution. if self._upstream_nodes_executed(node): task = self._generate_task(node) if task_lib.is_finalize_pipeline_task(task): return [task] else: result.append(task) # If there are no completed nodes in the current layer, downstream nodes # need not be checked. if not completed_node_ids: break # If all nodes in the final layer are completed successfully , the # pipeline can be finalized. # TODO(goutham): If there are conditional eval nodes, not all nodes may be # executed in the final layer. Handle this case when conditionals are # supported. if layer_num == len(layers) - 1 and completed_node_ids == set( node.node_info.id for node in nodes): return [ task_lib.FinalizePipelineTask( pipeline_uid=self._pipeline_state.pipeline_uid, status=status_lib.Status(code=status_lib.Code.OK)) ] return result
def generate(self) -> List[task_lib.Task]: """Generates tasks for executing the next executable nodes in the pipeline. The returned tasks must have `exec_task` populated. List may be empty if no nodes are ready for execution. Returns: A `list` of tasks to execute. """ layers = topsort.topsorted_layers( [node.pipeline_node for node in self._pipeline.nodes], get_node_id_fn=lambda node: node.node_info.id, get_parent_nodes=( lambda node: [self._node_map[n] for n in node.upstream_nodes]), get_child_nodes=( lambda node: [self._node_map[n] for n in node.downstream_nodes])) result = [] successful_node_ids = set() for layer_num, layer_nodes in enumerate(layers): for node in layer_nodes: node_uid = task_lib.NodeUid.from_pipeline_node( self._pipeline, node) node_id = node.node_info.id if self._in_successful_nodes_cache(node_uid): successful_node_ids.add(node_id) continue if not self._upstream_nodes_successful(node, successful_node_ids): continue # If this is a pure service node, there is no ExecNodeTask to generate # but we ensure node services and check service status. service_status = self._ensure_node_services_if_pure(node_id) if service_status is not None: if service_status == service_jobs.ServiceStatus.FAILED: return [ self._abort_task( f'service job failed; node uid: {node_uid}') ] if service_status == service_jobs.ServiceStatus.SUCCESS: logging.info('Service node successful: %s', node_uid) successful_node_ids.add(node_id) continue # If a task for the node is already tracked by the task queue, it need # not be considered for generation again but we ensure node services # in case of a mixed service node. if self._is_task_id_tracked_fn( task_lib.exec_node_task_id_from_pipeline_node( self._pipeline, node)): service_status = self._ensure_node_services_if_mixed( node_id) if service_status == service_jobs.ServiceStatus.FAILED: return [ self._abort_task( f'associated service job failed; node uid: {node_uid}' ) ] continue node_executions = task_gen_utils.get_executions( self._mlmd_handle, node) latest_execution = task_gen_utils.get_latest_execution( node_executions) # If the latest execution is successful, we're done. if latest_execution and execution_lib.is_execution_successful( latest_execution): logging.info('Node successful: %s', node_uid) successful_node_ids.add(node_id) continue # If the latest execution failed, the pipeline should be aborted. if latest_execution and not execution_lib.is_execution_active( latest_execution): error_msg_value = latest_execution.custom_properties.get( constants.EXECUTION_ERROR_MSG_KEY) error_msg = data_types_utils.get_metadata_value( error_msg_value) if error_msg_value else '' return [ self._abort_task( f'node failed; node uid: {node_uid}; error: {error_msg}' ) ] # Finally, we are ready to generate an ExecNodeTask for the node. task = self._maybe_generate_task(node, node_executions, successful_node_ids) if task: if task_lib.is_finalize_pipeline_task(task): return [task] else: result.append(task) layer_node_ids = set(node.node_info.id for node in layer_nodes) successful_layer_node_ids = layer_node_ids & successful_node_ids self._update_successful_nodes_cache(successful_layer_node_ids) # If all nodes in the final layer are completed successfully , the # pipeline can be finalized. # TODO(goutham): If there are conditional eval nodes, not all nodes may be # executed in the final layer. Handle this case when conditionals are # supported. if (layer_num == len(layers) - 1 and successful_layer_node_ids == layer_node_ids): return [ task_lib.FinalizePipelineTask( pipeline_uid=self._pipeline_uid, status=status_lib.Status(code=status_lib.Code.OK)) ] return result