Bug 1625200 - [taskgraph] Implement an 'All' composite strategy, r=tomprince

I'd like to implement a 'backstop' strategy, such that it will prevent all other
optimizers from removing tasks under certain conditions (e.g every 10th push).

The nicest way to implement this seems to be an 'All' composite strategy
(similar to 'Either' which this patch renames to 'Any'). This means we could
do something like:

All("seta", "backstop")

which means we would only remove tasks if *all* substrategies say to remove
tasks.

Differential Revision: https://phabricator.services.mozilla.com/D68620

--HG--
extra : moz-landing-system : lando

taskcluster/taskgraph/optimize/__init__.py

@@ -14,8 +14,10 @@ See ``taskcluster/docs/optimization.rst`` for more information.
 from __future__ import absolute_import, print_function, unicode_literals
 
 import logging
+from abc import ABCMeta, abstractmethod, abstractproperty
 from collections import defaultdict
 
+import six
 from slugid import nice as slugid
 
 from taskgraph.graph import Graph
@@ -31,6 +33,8 @@ def register_strategy(name, args=()):
     def wrap(cls):
         if name not in registry:
             registry[name] = cls(*args)
+            if not hasattr(registry[name], 'description'):
+                registry[name].description = name
         return cls
     return wrap
 
@@ -255,18 +259,15 @@ class OptimizationStrategy(object):
         return False
 
 
-class Either(OptimizationStrategy):
+@six.add_metaclass(ABCMeta)
-    """Given one or more optimization strategies, remove a task if any of them
+class CompositeStrategy(OptimizationStrategy):
-    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):
         missing = set(substrategies) - set(registry.keys())
         if missing:
             raise TypeError("substrategies aren't registered: {}".format(
                 ",  ".join(sorted(missing))))
 
-        self.description = "-or-".join(substrategies)
         self.substrategies = [registry[sub] for sub in substrategies]
         self.split_args = kwargs.pop('split_args', None)
         if not self.split_args:
@@ -274,25 +275,70 @@ class Either(OptimizationStrategy):
         if kwargs:
             raise TypeError("unexpected keyword args")
 
-    def _for_substrategies(self, arg, fn):
+    @abstractproperty
+    def description(self):
+        """A textual description of the combined substrategies."""
+        pass
+
+    @abstractmethod
+    def reduce(self, results):
+        """Given all substrategy results as a generator, return the overall
+        result."""
+        pass
+
+    def _generate_results(self, fname, task, params, arg):
         for sub, arg in zip(self.substrategies, self.split_args(arg)):
-            rv = fn(sub, arg)
+            yield getattr(sub, fname)(task, params, arg)
+
+    def should_remove_task(self, *args):
+        results = self._generate_results('should_remove_task', *args)
+        return self.reduce(results)
+
+    def should_replace_task(self, *args):
+        results = self._generate_results('should_replace_task', *args)
+        return self.reduce(results)
+
+
+class Any(CompositeStrategy):
+    """Given one or more optimization strategies, remove or replace a task if any of them
+    says to.
+
+    Replacement will use the value returned by the first strategy that says to replace.
+    """
+
+    @property
+    def description(self):
+        return "-or-".join([s.description for s in self.substrategies])
+
+    @classmethod
+    def reduce(cls, results):
+        for rv in results:
             if rv:
                 return rv
         return False
 
-    def should_remove_task(self, task, params, arg):
-        return self._for_substrategies(
-            arg,
-            lambda sub, arg: sub.should_remove_task(task, params, arg))
 
-    def should_replace_task(self, task, params, arg):
+class All(CompositeStrategy):
-        return self._for_substrategies(
+    """Given one or more optimization strategies, remove or replace a task if all of them
-            arg,
+    says to.
-            lambda sub, arg: sub.should_replace_task(task, params, arg))
+
+    Replacement will use the value returned by the first strategy passed in.
+    Note the values used for replacement need not be the same, as long as they
+    all say to replace.
+    """
+    @property
+    def description(self):
+        return "-and-".join([s.description for s in self.substrategies])
+
+    @classmethod
+    def reduce(cls, results):
+        rvs = list(results)
+        if all(rvs):
+            return rvs[0]
+        return False
 
 
-class Alias(Either):
+class Alias(CompositeStrategy):
     """Provides an alias to an existing strategy.
 
     This can be useful to swap strategies in and out without needing to modify
@@ -301,16 +347,23 @@ class Alias(Either):
     def __init__(self, strategy):
         super(Alias, self).__init__(strategy)
 
+    @property
+    def description(self):
+        return self.substrategies[0].description
+
+    def reduce(self, results):
+        return next(results)
+
 
 # Trigger registration in sibling modules.
 import_sibling_modules()
 
 
 # Register composite strategies.
-register_strategy('test', args=('skip-unless-schedules', 'seta'))(Either)
+register_strategy('test', args=('skip-unless-schedules', 'seta'))(Any)
 register_strategy('test-inclusive', args=('skip-unless-schedules',))(Alias)
 register_strategy('test-try', args=('skip-unless-schedules',))(Alias)
-register_strategy('fuzzing-builds', args=('skip-unless-schedules', 'seta'))(Either)
+register_strategy('fuzzing-builds', args=('skip-unless-schedules', 'seta'))(Any)
 
 
 class experimental(object):
@@ -323,12 +376,12 @@ class experimental(object):
     """
 
     relevant_tests = {
-        'test': Either('skip-unless-schedules', 'skip-unless-has-relevant-tests'),
+        'test': Any('skip-unless-schedules', 'skip-unless-has-relevant-tests'),
     }
     """Runs task containing tests in the same directories as modified files."""
 
     seta = {
-        'test': Either('skip-unless-schedules', 'seta'),
+        'test': Any('skip-unless-schedules', 'seta'),
     }
     """Provides a stable history of SETA's performance in the event we make it
     non-default in the future. Only useful as a benchmark."""
@@ -338,26 +391,26 @@ class experimental(object):
         learning to determine which tasks to run."""
 
         all = {
-            'test': Either('skip-unless-schedules', 'bugbug-all'),
+            'test': Any('skip-unless-schedules', 'bugbug-all'),
         }
         """Doesn't limit platforms, medium confidence threshold."""
 
         all_low = {
-            'test': Either('skip-unless-schedules', 'bugbug-all-low'),
+            'test': Any('skip-unless-schedules', 'bugbug-all-low'),
         }
         """Doesn't limit platforms, low confidence threshold."""
 
         all_high = {
-            'test': Either('skip-unless-schedules', 'bugbug-all-high'),
+            'test': Any('skip-unless-schedules', 'bugbug-all-high'),
         }
         """Doesn't limit platforms, high confidence threshold."""
 
         debug = {
-            'test': Either('skip-unless-schedules', 'bugbug-debug'),
+            'test': Any('skip-unless-schedules', 'bugbug-debug'),
         }
         """Restricts tests to debug platforms."""
 
         reduced = {
-            'test': Either('skip-unless-schedules', 'bugbug-reduced'),
+            'test': Any('skip-unless-schedules', 'bugbug-reduced'),
         }
         """Use the reduced set of tasks (and no groups) chosen by bugbug."""

taskcluster/taskgraph/test/test_optimize.py

@@ -6,14 +6,20 @@ from __future__ import absolute_import, print_function, unicode_literals
 
 import unittest
 
+import pytest
 from taskgraph import graph, optimize
-from taskgraph.optimize import OptimizationStrategy
+from taskgraph.optimize import OptimizationStrategy, All, Any
 from taskgraph.taskgraph import TaskGraph
 from taskgraph.task import Task
 from mozunit import main
 from slugid import nice as slugid
 
 
+@pytest.fixture
+def set_monkeypatch(request, monkeypatch):
+    request.cls.monkeypatch = monkeypatch
+
+
 class Remove(OptimizationStrategy):
 
     def should_remove_task(self, task, params, arg):
@@ -26,6 +32,7 @@ class Replace(OptimizationStrategy):
         return taskid
 
 
+@pytest.mark.usefixtures("set_monkeypatch")
 class TestOptimize(unittest.TestCase):
 
     strategies = {
@@ -69,11 +76,11 @@ class TestOptimize(unittest.TestCase):
             ('t3', 't1', 'dep2'),
             ('t2', 't1', 'dep'))
 
-    def assert_remove_tasks(self, graph, exp_removed, do_not_optimize=set()):
+    def assert_remove_tasks(self, graph, exp_removed, do_not_optimize=set(), strategies=None):
-        optimize.registry = self.strategies
+        strategies = strategies or self.strategies
         got_removed = optimize.remove_tasks(
             target_task_graph=graph,
-            optimizations=optimize._get_optimizations(graph, self.strategies),
+            optimizations=optimize._get_optimizations(graph, strategies),
             params={},
             do_not_optimize=do_not_optimize)
         self.assertEqual(got_removed, exp_removed)
@@ -91,6 +98,29 @@ class TestOptimize(unittest.TestCase):
             t3={'remove': None})
         self.assert_remove_tasks(graph, {'t1', 't2', 't3'})
 
+    def test_composite_strategies_any(self):
+        self.monkeypatch.setattr(optimize, 'registry', self.strategies)
+        strategies = self.strategies.copy()
+        strategies['any'] = Any('never', 'remove')
+
+        graph = self.make_triangle(
+            t1={'any': None},
+            t2={'any': None},
+            t3={'any': None})
+
+        self.assert_remove_tasks(graph, {'t1', 't2', 't3'}, strategies=strategies)
+
+    def test_composite_strategies_all(self):
+        self.monkeypatch.setattr(optimize, 'registry', self.strategies)
+        strategies = self.strategies.copy()
+        strategies['all'] = All('never', 'remove')
+
+        graph = self.make_triangle(
+            t1={'all': None},
+            t2={'all': None},
+            t3={'all': None})
+        self.assert_remove_tasks(graph, set(), strategies=strategies)
+
     def test_remove_tasks_blocked(self):
         "Removable tasks that are depended on by non-removable tasks are not removed"
         graph = self.make_triangle(