def test_prune_nodes(): """ Test to make sure that node-pruning works """ node0 = OperatorNode({}, {'name': 'test0', 'type': 'none'}) node1 = OperatorNode({}, { 'name': 'test1', 'type': 'none', 'upstream_dependencies': ['test0'] }) node2 = OperatorNode({}, { 'name': 'test2', 'type': 'none', 'upstream_dependencies': ['test0', 'test1'] }) node3 = OperatorNode({}, { 'name': 'test3', 'type': 'none', 'upstream_dependencies': ['test0', 'test1'] }) graph = _GraphUtil.build_subgraph([node0, node1, node2, node3]) _GraphUtil.prune_nodes(graph, node_selector=lambda node: node.name == 'test1') assert _GraphUtil.upstream_dependency_set(node3, graph) == frozenset([node0]) assert _GraphUtil.downstream_dependency_set(node0, graph) == frozenset( [node2, node3])
def _prune_workflow(self, workflow, is_primary, graph, prune_nodes): """ Produce the pruned workflow. Works by pruning the provided graph using the provided prune_nodes argument, and then reconstructing the workflow to include only the nodes remaining after pruning, with their upstream dependencies set properly to those computed during pruning. :param workflow: the workflow to prune :type workflow: dict :param is_primary: whether or not this workflow is the primary workflow :type is_primary: boolean :param graph: the OperatorGraph associated with this workflow :type graph: OperatorGraph :param prune_nodes: the set of all nodes being pruned from all workflows, specified as (workflow_name, node_name) tuples :type prune_nodes: set<(string, string)> :returns: the pruned workflow :rtype: dict """ prune_node_names = [ node_name for (prune_workflow_name, node_name) in prune_nodes if (is_primary and prune_workflow_name is None) or (workflow['name'] == prune_workflow_name) ] logger.debug( 'Graph for workflow %s (%s) before pruning: %s', workflow['name'], 'primary' if is_primary else 'secondary', graph.graph.nodes()) _GraphUtil.prune_nodes( graph=graph.graph, nodes=[graph.nodes[node_name] for node_name in prune_node_names]) logger.debug( 'Graph for workflow %s (%s) after pruning: %s', workflow['name'], 'primary' if is_primary else 'secondary', graph.graph.nodes()) return self._strip_workflow_nodes(workflow, graph.graph)
def _prune_flow_control_nodes(self, wf, graph): prune_nodes = list(self._flow_control_nodes(wf).keys()) if self.prune_forks: prune_nodes += list(self._fork_nodes(wf).keys()) if self.prune_joins: prune_nodes += list(self._join_nodes(wf).keys()) return _GraphUtil.prune_nodes(graph, nodes=prune_nodes)