Bug 1568277 - [taskgraph] Use a 'register_strategy' decorator in optimize.py r=tomprince
authorAndrew Halberstadt <ahalberstadt@mozilla.com>
Thu, 15 Aug 2019 18:48:21 +0000
changeset 488333 c8f797a197313295516bfa1b05962ea54708e997
parent 488332 9b89f970d46b0c966e7f523171cee853cb47407c
child 488334 cb35fd836621e5d1139149d1ab5485b01c8df781
push id113906
push userncsoregi@mozilla.com
push dateFri, 16 Aug 2019 04:07:24 +0000
treeherdermozilla-inbound@d887276421d3 [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewerstomprince
bugs1568277
milestone70.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1568277 - [taskgraph] Use a 'register_strategy' decorator in optimize.py r=tomprince Differential Revision: https://phabricator.services.mozilla.com/D40202
taskcluster/taskgraph/optimize.py
taskcluster/taskgraph/test/test_generator.py
--- a/taskcluster/taskgraph/optimize.py
+++ b/taskcluster/taskgraph/optimize.py
@@ -12,42 +12,51 @@ See ``taskcluster/docs/optimization.rst`
 """
 
 from __future__ import absolute_import, print_function, unicode_literals
 
 import logging
 import os
 from collections import defaultdict
 
-from .graph import Graph
+from mozbuild.base import MozbuildObject
+from mozbuild.util import memoize
+from slugid import nice as slugid
+
 from . import files_changed
+from .graph import Graph
 from .taskgraph import TaskGraph
+from .util.parameterization import resolve_task_references
 from .util.seta import is_low_value_task
 from .util.taskcluster import find_task_id
-from .util.parameterization import resolve_task_references
-from mozbuild.util import memoize
-from slugid import nice as slugid
-from mozbuild.base import MozbuildObject
 
 logger = logging.getLogger(__name__)
+registry = {}
+
+
+def register_strategy(name, args=()):
+    def wrap(cls):
+        if name not in registry:
+            registry[name] = cls(*args)
+        return cls
+    return wrap
 
 
 def optimize_task_graph(target_task_graph, params, do_not_optimize,
                         existing_tasks=None, strategies=None):
     """
     Perform task optimization, returning a taskgraph and a map from label to
     assigned taskId, including replacement tasks.
     """
     label_to_taskid = {}
     if not existing_tasks:
         existing_tasks = {}
 
     # instantiate the strategies for this optimization process
-    if not strategies:
-        strategies = _make_default_strategies()
+    strategies = strategies or registry.copy()
 
     optimizations = _get_optimizations(target_task_graph, strategies)
 
     removed_tasks = remove_tasks(
         target_task_graph=target_task_graph,
         optimizations=optimizations,
         params=params,
         do_not_optimize=do_not_optimize)
@@ -61,27 +70,16 @@ def optimize_task_graph(target_task_grap
         existing_tasks=existing_tasks,
         removed_tasks=removed_tasks)
 
     return get_subgraph(
             target_task_graph, removed_tasks, replaced_tasks,
             label_to_taskid), label_to_taskid
 
 
-def _make_default_strategies():
-    return {
-        'never': OptimizationStrategy(),  # "never" is the default behavior
-        'index-search': IndexSearch(),
-        'seta': SETA(),
-        'skip-unless-changed': SkipUnlessChanged(),
-        'skip-unless-schedules': SkipUnlessSchedules(),
-        'skip-unless-schedules-or-seta': Either(SkipUnlessSchedules(), SETA()),
-    }
-
-
 def _get_optimizations(target_task_graph, strategies):
     def optimizations(label):
         task = target_task_graph.tasks[label]
         if task.optimization:
             opt_by, arg = task.optimization.items()[0]
             return (opt_by, strategies[opt_by], arg)
         else:
             return ('never', strategies['never'], None)
@@ -238,16 +236,17 @@ def get_subgraph(target_task_graph, remo
         if left in tasks_by_taskid and right in tasks_by_taskid
     )
 
     return TaskGraph(
         tasks_by_taskid,
         Graph(set(tasks_by_taskid), edges_by_taskid))
 
 
+@register_strategy('never')
 class OptimizationStrategy(object):
     def should_remove_task(self, task, params, arg):
         """Determine whether to optimize this task by removing it.  Returns
         True to remove."""
         return False
 
     def should_replace_task(self, task, params, arg):
         """Determine whether to optimize this task by replacing it.  Returns a
@@ -257,17 +256,22 @@ class OptimizationStrategy(object):
 
 
 class Either(OptimizationStrategy):
     """Given one or more optimization strategies, remove a task if any of them
     says to, and replace with a task if any finds a replacement (preferring the
     earliest).  By default, each substrategy gets the same arg, but split_args
     can return a list of args for each strategy, if desired."""
     def __init__(self, *substrategies, **kwargs):
-        self.substrategies = substrategies
+        missing = set(substrategies) - set(registry.keys())
+        if missing:
+            raise TypeError("substrategies aren't registered: {}".format(
+                ",  ".join(sorted(missing))))
+
+        self.substrategies = [registry[sub] for sub in substrategies]
         self.split_args = kwargs.pop('split_args', None)
         if not self.split_args:
             self.split_args = lambda arg: [arg] * len(substrategies)
         if kwargs:
             raise TypeError("unexpected keyword args")
 
     def _for_substrategies(self, arg, fn):
         for sub, arg in zip(self.substrategies, self.split_args(arg)):
@@ -282,16 +286,17 @@ class Either(OptimizationStrategy):
             lambda sub, arg: sub.should_remove_task(task, params, arg))
 
     def should_replace_task(self, task, params, arg):
         return self._for_substrategies(
             arg,
             lambda sub, arg: sub.should_replace_task(task, params, arg))
 
 
+@register_strategy("index-search")
 class IndexSearch(OptimizationStrategy):
 
     # A task with no dependencies remaining after optimization will be replaced
     # if artifacts exist for the corresponding index_paths.
     # Otherwise, we're in one of the following cases:
     # - the task has un-optimized dependencies
     # - the artifacts have expired
     # - some changes altered the index_paths and new artifacts need to be
@@ -309,46 +314,49 @@ class IndexSearch(OptimizationStrategy):
                 return task_id
             except KeyError:
                 # 404 will end up here and go on to the next index path
                 pass
 
         return False
 
 
+@register_strategy('seta')
 class SETA(OptimizationStrategy):
     def should_remove_task(self, task, params, _):
         label = task.label
 
         # we would like to return 'False, None' while it's high_value_task
         # and we wouldn't optimize it. Otherwise, it will return 'True, None'
         if is_low_value_task(label,
                              params.get('project'),
                              params.get('pushlog_id'),
                              params.get('pushdate')):
             # Always optimize away low-value tasks
             return True
         else:
             return False
 
 
+@register_strategy("skip-unless-changed")
 class SkipUnlessChanged(OptimizationStrategy):
     def should_remove_task(self, task, params, file_patterns):
         # pushlog_id == -1 - this is the case when run from a cron.yml job
         if params.get('pushlog_id') == -1:
             return False
 
         changed = files_changed.check(params, file_patterns)
         if not changed:
             logger.debug('no files found matching a pattern in `skip-unless-changed` for ' +
                          task.label)
             return True
         return False
 
 
+@register_strategy("skip-unless-schedules")
 class SkipUnlessSchedules(OptimizationStrategy):
 
     @memoize
     def scheduled_by_push(self, repository, revision):
         changed_files = files_changed.get_changed_files(repository, revision)
 
         mbo = MozbuildObject.from_environment()
         # the decision task has a sparse checkout, so, mozbuild_reader will use
@@ -368,8 +376,12 @@ class SkipUnlessSchedules(OptimizationSt
 
         scheduled = self.scheduled_by_push(params['head_repository'], params['head_rev'])
         conditions = set(conditions)
         # if *any* of the condition components are scheduled, do not optimize
         if conditions & scheduled:
             return False
 
         return True
+
+
+# Register composite strategies.
+register_strategy('skip-unless-schedules-or-seta', args=('skip-unless-schedules', 'seta'))(Either)
--- a/taskcluster/taskgraph/test/test_generator.py
+++ b/taskcluster/taskgraph/test/test_generator.py
@@ -91,22 +91,21 @@ class TestGenerator(unittest.TestCase):
     def maketgg(self, target_tasks=None, kinds=[('_fake', [])], params=None):
         params = params or {}
         FakeKind.loaded_kinds = []
         self.target_tasks = target_tasks or []
 
         def target_tasks_method(full_task_graph, parameters, graph_config):
             return self.target_tasks
 
-        def make_fake_strategies():
-            return {mode: FakeOptimization(mode)
-                    for mode in ('always', 'never', 'even', 'odd')}
+        fake_registry = {mode: FakeOptimization(mode)
+                         for mode in ('always', 'never', 'even', 'odd')}
 
         target_tasks_mod._target_task_methods['test_method'] = target_tasks_method
-        self.patch.setattr(optimize_mod, '_make_default_strategies', make_fake_strategies)
+        self.patch.setattr(optimize_mod, 'registry', fake_registry)
 
         parameters = FakeParameters({
             '_kinds': kinds,
             'target_tasks_method': 'test_method',
             'try_mode': None,
         })
         parameters.update(params)