Bug 1258497: factor more code out of mach_commands.py; r?ahal draft
authorDustin J. Mitchell <dustin@mozilla.com>
Fri, 06 May 2016 17:20:54 +0000
changeset 364497 d7593c503eb13149c658f0ee28d84db59e9d37f8
parent 364496 a0f01abbd288eaf7111ff2989d7c1979a414f6b7
child 520301 f27cdbed27b52c187bc23f3a605838372aaa77fa
push id17473
push userdmitchell@mozilla.com
push dateFri, 06 May 2016 19:36:39 +0000
reviewersahal
bugs1258497
milestone49.0a1
Bug 1258497: factor more code out of mach_commands.py; r?ahal MozReview-Commit-ID: HLdzpJEkSXt
taskcluster/mach_commands.py
taskcluster/taskgraph/decision.py
taskcluster/taskgraph/target_tasks.py
taskcluster/taskgraph/test/test_decision.py
taskcluster/taskgraph/test/test_target_tasks.py
--- a/taskcluster/mach_commands.py
+++ b/taskcluster/mach_commands.py
@@ -20,21 +20,18 @@ from mach.decorators import (
     CommandArgument,
     CommandProvider,
     Command,
     SubCommand,
 )
 
 from mozbuild.base import MachCommandBase
 
-ROOT = os.path.dirname(os.path.realpath(__file__))
-TOPSRCDIR = os.path.realpath(os.path.join(ROOT, '..'))
 
-
-class TaskGraphSubCommand(SubCommand):
+class ShowTaskGraphSubCommand(SubCommand):
     """A SubCommand with TaskGraph-specific arguments"""
 
     def __call__(self, func):
         after = SubCommand.__call__(self, func)
         args = [
             CommandArgument('--root', '-r', default='taskcluster/ci',
                             help="root of the taskgraph definition relative to topsrcdir"),
             CommandArgument('--parameters', '-p', required=True,
@@ -61,68 +58,55 @@ class MachCommands(MachCommandBase):
         by dependencies: for example, a binary must be built before it is tested,
         and that build may further depend on various toolchains, libraries, etc.
         """
 
     @SubCommand('taskgraph', 'python-tests',
                 'Run the taskgraph unit tests')
     def taskgraph_python_tests(self, **options):
         import unittest
+        import taskgraph.target_tasks
         suite = unittest.defaultTestLoader.discover('taskgraph.test')
         runner = unittest.TextTestRunner(verbosity=2)
         result = runner.run(suite)
         if not result.wasSuccessful:
             sys.exit(1)
 
-    @TaskGraphSubCommand('taskgraph', 'tasks',
+    @ShowTaskGraphSubCommand('taskgraph', 'tasks',
                          "Show all tasks in the taskgraph")
     def taskgraph_tasks(self, **options):
-        import taskgraph.parameters
-        parameters = taskgraph.parameters.load_parameters_file(options)
-        tgg = self.get_taskgraph_generator(options, parameters)
-        self.show_taskgraph(tgg.full_task_set, options)
+        return self.show_taskgraph('full_task_set', False, options)
 
-    @TaskGraphSubCommand('taskgraph', 'full',
+    @ShowTaskGraphSubCommand('taskgraph', 'full',
                          "Show the full taskgraph")
     def taskgraph_full(self, **options):
-        import taskgraph.parameters
-        parameters = taskgraph.parameters.load_parameters_file(options)
-        tgg = self.get_taskgraph_generator(options, parameters)
-        self.show_taskgraph(tgg.full_task_graph, options)
+        return self.show_taskgraph('full_task_graph', False, options)
 
-    @TaskGraphSubCommand('taskgraph', 'target',
+    @ShowTaskGraphSubCommand('taskgraph', 'target',
                          "Show the target task set")
     def taskgraph_target(self, **options):
-        import taskgraph.parameters
-        parameters = taskgraph.parameters.load_parameters_file(options)
-        tgg = self.get_taskgraph_generator(options, parameters)
-        self.set_target_tasks(tgg, parameters)
-        self.show_taskgraph(tgg.target_task_set, options)
+        return self.show_taskgraph('target_task_set', True, options)
 
-    @TaskGraphSubCommand('taskgraph', 'target-graph',
+    @ShowTaskGraphSubCommand('taskgraph', 'target-graph',
                          "Show the target taskgraph (the target set with its dependencies)")
     def taskgraph_target_taskgraph(self, **options):
-        import taskgraph.parameters
-        parameters = taskgraph.parameters.load_parameters_file(options)
-        tgg = self.get_taskgraph_generator(options, parameters)
-        self.set_target_tasks(tgg, parameters)
-        self.show_taskgraph(tgg.target_task_graph, options)
+        return self.show_taskgraph('target_task_graph', True, options)
 
-    @TaskGraphSubCommand('taskgraph', 'optimized',
+    @ShowTaskGraphSubCommand('taskgraph', 'optimized',
                          "Show the optimized taskgraph")
     def taskgraph_optimized(self, **options):
-        import taskgraph.parameters
-        parameters = taskgraph.parameters.load_parameters_file(options)
-        tgg = self.get_taskgraph_generator(options, parameters)
-        self.set_target_tasks(tgg, parameters)
-        self.show_taskgraph(tgg.optimized_task_graph, options)
+        return self.show_taskgraph('optimized_task_graph', True, options)
 
-    @TaskGraphSubCommand('taskgraph', 'decision', textwrap.dedent("""\
-                         Decision task: generate a task graph and submit to TaskCluster.
-                         """))
+    @SubCommand('taskgraph', 'decision', textwrap.dedent("""\
+                 Decision task: generate a task graph and submit to TaskCluster.
+                 This is only meant to be called within decision tasks, and
+                 requires a great many arguments.  Commands like `mach
+                 taskgraph optimized` are better suited to use on the command
+                 line, and can take the parameters file generated by a decision task.
+                 """))
     @CommandArgument('--root', '-r',
         default='taskcluster/ci',
         help="root of the taskgraph definition relative to topsrcdir")
     @CommandArgument('--base-repository',
         required=True,
         help='URL for "base" repository to clone')
     @CommandArgument('--head-repository',
         required=True,
@@ -152,118 +136,34 @@ class MachCommands(MachCommandBase):
     @CommandArgument('--level',
         required=True,
         help='SCM level of this repository')
     @CommandArgument('--target-tasks-method',
         required=False,
         help='Method to use to determine the target task (e.g., `try_option_syntax`); '
              'default is to run the full task graph')
     def taskgraph_decision(self, **options):
-        # create a TaskGraphGenerator instance
+        import taskgraph.decision
+        return taskgraph.decision.taskgraph_decision(self.log, options)
+
+    def show_taskgraph(self, graph_attr, set_target_tasks, options):
+        import taskgraph.parameters
+        import taskgraph.target_tasks
         import taskgraph.generator
-        import taskgraph.create
-        import taskgraph.parameters
 
-        parameters = taskgraph.parameters.get_decision_parameters(options)
+        parameters = taskgraph.parameters.load_parameters_file(options)
+
+        optimization_finder = lambda keys: {} # XXX
 
         tgg = taskgraph.generator.TaskGraphGenerator(
             root_dir=options['root'],
             log=self.log,
             parameters=parameters,
-            optimization_finder=None)  # XXX
-
-        # produce some artifacts
-        def write_artifact(filename, data):
-            self.log(logging.INFO, 'writing-artifact', {
-                'filename': filename,
-            }, 'writing artifact file `{filename}`')
-            if not os.path.isdir('artifacts'):
-                os.mkdir('artifacts')
-            path = os.path.join('artifacts', filename)
-            if filename.endswith('.yml'):
-                import yaml
-                yaml.dump(data, open(path, 'w'), allow_unicode=True, default_flow_style=False)
-            elif filename.endswith('.json'):
-                json.dump(data, open(path, 'w'),
-                          sort_keys=True, indent=2, separators=(',', ': '))
-            else:
-                raise TypeError("Don't know how to write to {}".format(filename))
-
-        # generate the target_tasks list and write it as an artifact in case someone
-        # wants to reproduce this run
-        target_tasks = self.set_target_tasks(tgg, parameters)
-        if target_tasks:
-            write_artifact('target_tasks.json', target_tasks)
-
-        # write out the parameters used to generate this graph
-        write_artifact('parameters.yml', dict(parameters))
-
-        # write out the full graph for reference
-        write_artifact('full-task-graph.json',
-                       self.taskgraph_to_json(tgg.full_task_graph))
-
-        # write out the optimized task graph to describe what will happen
-        write_artifact('task-graph.json',
-                       self.taskgraph_to_json(tgg.optimized_task_graph))
-
-        # actually create the graph
-        taskgraph.create.create_tasks(tgg.optimized_task_graph)
-
-    ##
-    # Target tasks methods
+            optimization_finder=optimization_finder)
 
-    def set_target_tasks(self, tgg, parameters):
-        """If params['target_task_set_method'] is set, use it to determine the
-        target task set, update the task graph with that set, and return it.  Note
-        that as a side-effect, this generates the full task set."""
-        target_tasks_method = parameters.get('target_tasks_method')
-        if target_tasks_method:
-            meth = getattr(self, 'target_tasks_' + target_tasks_method)
-            target_tasks = meth(tgg.full_task_graph, parameters)
-            tgg.set_target_tasks(target_tasks)
-            return target_tasks
-
-    def target_tasks_from_parameters(self, full_task_graph, parameters):
-        """Get the target task set from parameters['target_tasks'].  This is
-        useful for re-running a decision task with the same target set as in an
-        earlier run, by copying `target_tasks.json` into `parameters.yml`."""
-        return parameters['target_tasks']
-
-    def target_tasks_try_option_syntax(self, full_task_graph, parameters):
-        from taskgraph.try_option_syntax import TryOptionSyntax
-        options = TryOptionSyntax(parameters['message'], full_task_graph)
-        return [t.label for t in full_task_graph.tasks.itervalues()
-                if options.task_matches(t.attributes)]
-
-    ##
-    # Utilities
+        if set_target_tasks:
+            taskgraph.target_tasks.set_target_tasks(tgg, parameters)
+        tg = getattr(tgg, graph_attr)
 
-    def get_taskgraph_generator(self, options, parameters):
-        import taskgraph.generator
-        if options['optimize']:
-            optimization_finder = None  # XXX function that searches index
-        else:
-            # null optmization finder
-            optimization_finder = lambda keys: {}
-        tgg = taskgraph.generator.TaskGraphGenerator(
-            root_dir=options['root'],
-            log=self.log,
-            parameters=parameters,
-            optimization_finder=optimization_finder)
-        if 'task_set' in parameters:
-            tgg.set_task_set(parameters['task_set'])
-        return tgg
+        # TODO: optionally output in dot format, select fields, filter, etc.
+        for label in tg.graph.visit_postorder():
+            print(tg.tasks[label])
 
-    def show_taskgraph(self, taskgraph, options):
-        # TODO: optionally output in dot format, select fields, filter, etc.
-        for label in taskgraph.graph.visit_postorder():
-            print(taskgraph.tasks[label])
-
-    def taskgraph_to_json(self, taskgraph):
-        dep_links = taskgraph.graph.links_dict()
-        tasks = taskgraph.tasks
-        def tojson(task):
-            return {
-                'task': task.task,
-                'attributes': task.attributes,
-                'dependencies': list(dep_links[task.label]), # XXX
-            }
-        return {label: tojson(tasks[label]) for label in taskgraph.graph.nodes}
new file mode 100644
--- /dev/null
+++ b/taskcluster/taskgraph/decision.py
@@ -0,0 +1,86 @@
+# -*- coding: utf-8 -*-
+
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from __future__ import absolute_import, print_function, unicode_literals
+
+import os
+import json
+import logging
+
+import taskgraph.generator
+import taskgraph.create
+import taskgraph.parameters
+import taskgraph.target_tasks
+
+ARTIFACTS_DIR = 'artifacts'
+
+
+def taskgraph_decision(log, options):
+    parameters = taskgraph.parameters.get_decision_parameters(options)
+
+    # create a TaskGraphGenerator instance
+    tgg = taskgraph.generator.TaskGraphGenerator(
+        root_dir=options['root'],
+        log=log,
+        parameters=parameters,
+        optimization_finder=None)  # XXX
+
+    # generate the target_tasks list and write it as an artifact in case someone
+    # wants to reproduce this run
+    target_tasks = taskgraph.target_tasks.set_target_tasks(tgg, parameters)
+    if target_tasks:
+        write_artifact('target_tasks.json', target_tasks, log)
+
+    # write out the parameters used to generate this graph
+    write_artifact('parameters.yml', dict(**parameters), log)
+
+    # write out the full graph for reference
+    write_artifact('full-task-graph.json',
+                   taskgraph_to_json(tgg.full_task_graph),
+                   log)
+
+    # write out the optimized task graph to describe what will happen
+    write_artifact('task-graph.json',
+                   taskgraph_to_json(tgg.optimized_task_graph),
+                   log)
+
+    # actually create the graph
+    taskgraph.create.create_tasks(tgg.optimized_task_graph)
+
+
+def taskgraph_to_json(taskgraph):
+    tasks = taskgraph.tasks
+
+    def tojson(task):
+        return {
+            'task': task.task,
+            'attributes': task.attributes,
+            'dependencies': []
+        }
+    rv = {label: tojson(tasks[label]) for label in taskgraph.graph.nodes}
+
+    # add dependencies with one trip through the graph edges
+    for (left, right, name) in taskgraph.graph.edges:
+        rv[left]['dependencies'].append((name, right))
+
+    return rv
+
+
+def write_artifact(filename, data, log):
+    log(logging.INFO, 'writing-artifact', {
+        'filename': filename,
+    }, 'writing artifact file `{filename}`')
+    if not os.path.isdir(ARTIFACTS_DIR):
+        os.mkdir(ARTIFACTS_DIR)
+    path = os.path.join(ARTIFACTS_DIR, filename)
+    if filename.endswith('.yml'):
+        import yaml
+        yaml.safe_dump(data, open(path, 'w'), allow_unicode=True, default_flow_style=False)
+    elif filename.endswith('.json'):
+        json.dump(data, open(path, 'w'),
+                  sort_keys=True, indent=2, separators=(',', ': '))
+    else:
+        raise TypeError("Don't know how to write to {}".format(filename))
new file mode 100644
--- /dev/null
+++ b/taskcluster/taskgraph/target_tasks.py
@@ -0,0 +1,40 @@
+# -*- coding: utf-8 -*-
+
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from __future__ import absolute_import, print_function, unicode_literals
+
+_target_task_methods = {}
+def _target_task(name):
+    def wrap(func):
+        _target_task_methods[name] = func
+        return func
+    return wrap
+
+def set_target_tasks(tgg, parameters):
+    """If params['target_task_set_method'] is set, use it to determine the
+    target task set, update the task graph with that set, and return it.  Note
+    that as a side-effect, this generates the full task set."""
+    target_tasks_method = parameters.get('target_tasks_method')
+    if target_tasks_method:
+        meth = _target_task_methods[target_tasks_method]
+        target_tasks = meth(tgg.full_task_graph, parameters)
+        tgg.set_target_tasks(target_tasks)
+        return target_tasks
+
+@_target_task('from_parameters')
+def target_tasks_from_parameters(full_task_graph, parameters):
+    """Get the target task set from parameters['target_tasks'].  This is
+    useful for re-running a decision task with the same target set as in an
+    earlier run, by copying `target_tasks.json` into `parameters.yml`."""
+    return parameters['target_tasks']
+
+@_target_task('try_option_syntax')
+def target_tasks_try_option_syntax(full_task_graph, parameters):
+    from taskgraph.try_option_syntax import TryOptionSyntax
+    options = TryOptionSyntax(parameters['message'], full_task_graph)
+    return [t.label for t in full_task_graph.tasks.itervalues()
+            if options.task_matches(t.attributes)]
+
new file mode 100644
--- /dev/null
+++ b/taskcluster/taskgraph/test/test_decision.py
@@ -0,0 +1,77 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this file,
+# You can obtain one at http://mozilla.org/MPL/2.0/.
+from __future__ import unicode_literals
+
+import os
+import sys
+import json
+import yaml
+import shutil
+import unittest
+import tempfile
+
+from taskgraph import decision
+from taskgraph.graph import Graph
+from taskgraph.types import Task, TaskGraph
+
+from mozunit import main
+
+class TestDecision(unittest.TestCase):
+
+    def test_taskgraph_to_json(self):
+        tasks = {
+            'a': Task(kind=None, label='a', attributes={'attr': 'a-task'}),
+            'b': Task(kind=None, label='b', task={'task': 'def'}),
+        }
+        graph = Graph(nodes=set('ab'), edges=set([('a', 'b', 'edgelabel')]))
+        taskgraph = TaskGraph(tasks, graph)
+
+        res = decision.taskgraph_to_json(taskgraph)
+
+        self.assertEqual(res, {
+            'a': {
+                'attributes': {'attr': 'a-task'},
+                'task': {},
+                'dependencies': [('edgelabel', 'b')],
+            },
+            'b': {
+                'attributes': {},
+                'task': {'task': 'def'},
+                'dependencies': [],
+            }
+        })
+
+
+    def test_write_artifact_json(self):
+        data = [{'some': 'data'}]
+        tmpdir = tempfile.mkdtemp()
+        try:
+            decision.ARTIFACTS_DIR = os.path.join(tmpdir, "artifacts")
+            decision.write_artifact("artifact.json", data, lambda *args: None)
+            with open(os.path.join(decision.ARTIFACTS_DIR, "artifact.json")) as f:
+                self.assertEqual(json.load(f), data)
+        finally:
+            if os.path.exists(tmpdir):
+                shutil.rmtree(tmpdir)
+            decision.ARTIFACTS_DIR = 'artifacts'
+
+
+    def test_write_artifact_yml(self):
+        data = [{'some': 'data'}]
+        tmpdir = tempfile.mkdtemp()
+        try:
+            decision.ARTIFACTS_DIR = os.path.join(tmpdir, "artifacts")
+            decision.write_artifact("artifact.yml", data, lambda *args: None)
+            with open(os.path.join(decision.ARTIFACTS_DIR, "artifact.yml")) as f:
+                self.assertEqual(yaml.safe_load(f), data)
+        finally:
+            if os.path.exists(tmpdir):
+                shutil.rmtree(tmpdir)
+            decision.ARTIFACTS_DIR = 'artifacts'
+
+
+if __name__ == '__main__':
+    main()
+
+
new file mode 100644
--- /dev/null
+++ b/taskcluster/taskgraph/test/test_target_tasks.py
@@ -0,0 +1,68 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this file,
+# You can obtain one at http://mozilla.org/MPL/2.0/.
+from __future__ import unicode_literals
+
+import os
+import sys
+import json
+import yaml
+import shutil
+import unittest
+import tempfile
+
+from taskgraph import target_tasks
+from taskgraph import try_option_syntax
+from taskgraph.graph import Graph
+from taskgraph.types import Task, TaskGraph
+
+from mozunit import main, MockedOpen
+
+class FakeTGG(object):
+
+    def __init__(self, full_task_graph):
+        self.full_task_graph = full_task_graph
+
+    def set_target_tasks(self, target_tasks):
+        self.got_target_tasks = target_tasks
+
+class FakeTryOptionSyntax(object):
+
+    def __init__(self, message, task_graph):
+        pass
+
+    def task_matches(self, attributes):
+        return 'at-at' in attributes
+
+
+class TestTargetTasks(unittest.TestCase):
+
+    def test_set_target_tasks_no_method(self):
+        # nothing happens..
+        self.assertEqual(target_tasks.set_target_tasks(None, {}), None)
+
+    def test_set_target_tasks_try(self):
+        tasks = {
+            'a': Task(kind=None, label='a'),
+            'b': Task(kind=None, label='b', attributes={'at-at': 'yep'}),
+        }
+        graph = Graph(nodes=set('ab'), edges=set())
+        tg = TaskGraph(tasks, graph)
+        tgg = FakeTGG(tg)
+
+        params = {
+            'message': 'try me',
+            'target_tasks_method': 'try_option_syntax',
+        }
+
+        orig_TryOptionSyntax = try_option_syntax.TryOptionSyntax
+        try:
+            try_option_syntax.TryOptionSyntax = FakeTryOptionSyntax
+            self.assertEqual(target_tasks.set_target_tasks(tgg, params), ['b'])
+            self.assertEqual(tgg.got_target_tasks, ['b'])
+        finally:
+            try_option_syntax.TryOptionSyntax = orig_TryOptionSyntax
+
+if __name__ == '__main__':
+    main()
+