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)