From f01c2c674c1c2055e6d98e3dea66d00770082a0c Mon Sep 17 00:00:00 2001 From: "Kris Wilson (Twitter)" Date: Tue, 5 Apr 2016 20:41:59 -0700 Subject: [PATCH] [engine] Move the visualizer into LocalScheduler. - Migrate the meat of the visualizer into `LocalScheduler.visualize_graph`. - Add basic tests for coverage. Testing Done: https://travis-ci.org/pantsbuild/pants/builds/120774094 + local testing against viz and viz-fs. Bugs closed: 3138, 3141 Reviewed at https://rbcommons.com/s/twitter/r/3649/ --- .../pants/engine/exp/examples/visualizer.py | 75 ++----------------- src/python/pants/engine/exp/scheduler.py | 63 ++++++++++++++++ .../pants_test/engine/exp/test_scheduler.py | 18 +++++ 3 files changed, 86 insertions(+), 70 deletions(-) diff --git a/src/python/pants/engine/exp/examples/visualizer.py b/src/python/pants/engine/exp/examples/visualizer.py index 24d82e5c5a4..8bc3a5de56f 100644 --- a/src/python/pants/engine/exp/examples/visualizer.py +++ b/src/python/pants/engine/exp/examples/visualizer.py @@ -15,80 +15,16 @@ from pants.engine.exp.engine import LocalSerialEngine from pants.engine.exp.examples.planners import setup_json_scheduler from pants.engine.exp.fs import PathGlobs -from pants.engine.exp.nodes import Noop, SelectNode, TaskNode, Throw -from pants.util.contextutil import temporary_file, temporary_file_path - - -def format_type(node): - if type(node) == TaskNode: - return node.func.__name__ - return type(node).__name__ - - -def format_subject(node): - if node.variants: - return '({})@{}'.format(node.subject, ','.join('{}={}'.format(k, v) for k, v in node.variants)) - else: - return '({})'.format(node.subject) - - -def format_product(node): - if type(node) == SelectNode and node.variant_key: - return '{}@{}'.format(node.product.__name__, node.variant_key) - return node.product.__name__ - - -def format_node(node, state): - return '{}:{}:{} == {}'.format(format_product(node), - format_subject(node), - format_type(node), - str(state).replace('"', '\\"')) - - -# NB: there are only 12 colors in `set312`. -colorscheme = 'set312' -max_colors = 12 -colors = {} - - -def format_color(node, node_state): - if type(node_state) is Throw: - return 'tomato' - elif type(node_state) is Noop: - return 'white' - key = node.product - return colors.setdefault(key, (len(colors) % max_colors) + 1) - - -def create_digraph(scheduler, storage, request): - - yield 'digraph plans {' - yield ' node[colorscheme={}];'.format(colorscheme) - yield ' concentrate=true;' - yield ' rankdir=LR;' - - for ((node, node_state), dependency_entries) in scheduler.product_graph.walk(request.roots): - node_str = format_node(node, node_state) - - yield (' "{node}" [style=filled, fillcolor={color}];' - .format(color=format_color(node, node_state), - node=node_str)) - - for (dep, dep_state) in dependency_entries: - yield ' "{}" -> "{}"'.format(node_str, format_node(dep, dep_state)) - - yield '}' +from pants.util.contextutil import temporary_file_path def visualize_execution_graph(scheduler, storage, request): - with temporary_file(cleanup=False, suffix='.dot') as fp: - for line in create_digraph(scheduler, storage, request): - fp.write(line) - fp.write('\n') + with temporary_file_path(cleanup=False, suffix='.dot') as dot_file: + scheduler.visualize_graph_to_file(request.roots, dot_file) + print('dot file saved to: {}'.format(dot_file)) - print('dot file saved to: {}'.format(fp.name)) with temporary_file_path(cleanup=False, suffix='.svg') as image_file: - subprocess.check_call('dot -Tsvg -o{} {}'.format(image_file, fp.name), shell=True) + subprocess.check_call('dot -Tsvg -o{} {}'.format(image_file, dot_file), shell=True) print('svg file saved to: {}'.format(image_file)) binary_util.ui_open(image_file) @@ -119,7 +55,6 @@ def usage(error_message): build_root = args.pop(0) - if not os.path.isdir(build_root): usage('First argument must be a valid build root, {} is not a directory.'.format(build_root)) build_root = os.path.realpath(build_root) diff --git a/src/python/pants/engine/exp/scheduler.py b/src/python/pants/engine/exp/scheduler.py index f12644cf1cc..69a841d1a83 100644 --- a/src/python/pants/engine/exp/scheduler.py +++ b/src/python/pants/engine/exp/scheduler.py @@ -242,6 +242,58 @@ def _walk(entries): for entry in _walk(_filtered_entries(roots)): yield entry + def visualize(self, roots): + """Visualize a graph walk by generating graphviz `dot` output. + + :param iterable roots: An iterable of the root nodes to begin the graph walk from. + """ + viz_colors = {} + viz_color_scheme = 'set312' # NB: There are only 12 colors in `set312`. + viz_max_colors = 12 + + def format_color(node, node_state): + if type(node_state) is Throw: + return 'tomato' + elif type(node_state) is Noop: + return 'white' + return viz_colors.setdefault(node.product, (len(viz_colors) % viz_max_colors) + 1) + + def format_type(node): + return node.func.__name__ if type(node) is TaskNode else type(node).__name__ + + def format_subject(node): + if node.variants: + return '({})@{}'.format(node.subject, + ','.join('{}={}'.format(k, v) for k, v in node.variants)) + else: + return '({})'.format(node.subject) + + def format_product(node): + if type(node) is SelectNode and node.variant_key: + return '{}@{}'.format(node.product.__name__, node.variant_key) + return node.product.__name__ + + def format_node(node, state): + return '{}:{}:{} == {}'.format(format_product(node), + format_subject(node), + format_type(node), + str(state).replace('"', '\\"')) + + yield 'digraph plans {' + yield ' node[colorscheme={}];'.format(viz_color_scheme) + yield ' concentrate=true;' + yield ' rankdir=LR;' + + for ((node, node_state), dependency_entries) in self.walk(roots): + node_str = format_node(node, node_state) + + yield ' "{}" [style=filled, fillcolor={}];'.format(node_str, format_color(node, node_state)) + + for (dep, dep_state) in dependency_entries: + yield ' "{}" -> "{}"'.format(node_str, format_node(dep, dep_state)) + + yield '}' + class ExecutionRequest(datatype('ExecutionRequest', ['roots'])): """Holds the roots for an execution, which might have been requested by a user. @@ -386,6 +438,17 @@ def __init__(self, goals, tasks, storage, project_tree, graph_lock=None, graph_v self._product_graph_lock = graph_lock or threading.RLock() self._step_id = 0 + def visualize_graph_to_file(self, roots, filename): + """Visualize a graph walk by writing graphviz `dot` output to a file. + + :param iterable roots: An iterable of the root nodes to begin the graph walk from. + :param str filename: The filename to output the graphviz output to. + """ + with self._product_graph_lock, open(filename, 'wb') as fh: + for line in self.product_graph.visualize(roots): + fh.write(line) + fh.write('\n') + def _create_step(self, node): """Creates a Step and Promise with the currently available dependencies of the given Node. diff --git a/tests/python/pants_test/engine/exp/test_scheduler.py b/tests/python/pants_test/engine/exp/test_scheduler.py index 84fc079b441..6d6f88cf463 100644 --- a/tests/python/pants_test/engine/exp/test_scheduler.py +++ b/tests/python/pants_test/engine/exp/test_scheduler.py @@ -20,6 +20,7 @@ from pants.engine.exp.nodes import (ConflictingProducersError, DependenciesNode, Return, SelectNode, Throw, Waiting) from pants.engine.exp.scheduler import ProductGraph +from pants.util.contextutil import temporary_dir class SchedulerTest(unittest.TestCase): @@ -244,6 +245,23 @@ def test_sibling_specs(self): # And that an subdirectory address is not. self.assertNotIn(self.managed_guava, root_value) + def test_scheduler_visualize(self): + spec = self.spec_parser.parse_spec('3rdparty/jvm:') + build_request = self.request_specs(['list'], spec) + self.build_and_walk(build_request) + + graphviz_output = '\n'.join(self.scheduler.product_graph.visualize(build_request.roots)) + + with temporary_dir() as td: + output_path = os.path.join(td, 'output.dot') + self.scheduler.visualize_graph_to_file(build_request.roots, output_path) + with open(output_path, 'rb') as fh: + graphviz_disk_output = fh.read().strip() + + self.assertEqual(graphviz_output, graphviz_disk_output) + self.assertIn('digraph', graphviz_output) + self.assertIn(' -> ', graphviz_output) + # TODO: Expand test coverage here. class ProductGraphTest(unittest.TestCase):