Bug 1568277 - [taskgraph] Split optimize strategies out into a separate file r=tomprince

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

--HG--
rename : taskcluster/taskgraph/optimize.py => taskcluster/taskgraph/optimize/__init__.py
extra : moz-landing-system : lando

python/mozbuild/mozbuild/artifact_commands.py

@@ -292,7 +292,7 @@ class PackageFrontend(MachCommandBase):
                          'Do not use --from-build in automation; all dependencies '
                          'should be determined in the decision task.')
                 return 1
-            from taskgraph.optimize import IndexSearch
+            from taskgraph.optimize.strategies import IndexSearch
             from taskgraph.parameters import Parameters
             from taskgraph.generator import load_tasks_for_kind
             params = Parameters(

taskcluster/taskgraph/docker.py

@@ -12,7 +12,7 @@ import tarfile
 from io import BytesIO
 
 from taskgraph.parameters import Parameters
-from taskgraph.optimize import IndexSearch
+from taskgraph.optimize.strategies import IndexSearch
 from taskgraph.util import docker
 from taskgraph.util.taskcluster import (
     get_artifact_url,

taskcluster/taskgraph/optimize/__init__.py

@@ -14,19 +14,14 @@ See ``taskcluster/docs/optimization.rst`` for more information.
 from __future__ import absolute_import, print_function, unicode_literals
 
 import logging
-import os
 from collections import defaultdict
 
-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 taskgraph.graph import Graph
+from taskgraph.taskgraph import TaskGraph
+from taskgraph.util.parameterization import resolve_task_references
+from taskgraph.util.python_path import import_sibling_modules
 
 logger = logging.getLogger(__name__)
 registry = {}
@@ -291,96 +286,8 @@ class Either(OptimizationStrategy):
             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
-    # created.
-    # In every of those cases, we need to run the task to create or refresh
-    # artifacts.
-
-    def should_replace_task(self, task, params, index_paths):
-        "Look for a task with one of the given index paths"
-        for index_path in index_paths:
-            try:
-                task_id = find_task_id(
-                    index_path,
-                    use_proxy=bool(os.environ.get('TASK_ID')))
-                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
-        # a MercurialRevisionFinder with revision '.', which should be the same
-        # as `revision`; in other circumstances, it will use a default reader
-        rdr = mbo.mozbuild_reader(config_mode='empty')
-
-        components = set()
-        for p, m in rdr.files_info(changed_files).items():
-            components |= set(m['SCHEDULES'].components)
-
-        return components
-
-    def should_remove_task(self, task, params, conditions):
-        if params.get('pushlog_id') == -1:
-            return False
-
-        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
+# Trigger registration in sibling modules.
+import_sibling_modules()
 
 
 # Register composite strategies.

taskcluster/taskgraph/optimize/strategies.py

@@ -0,0 +1,110 @@
+# 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 logging
+import os
+
+from mozbuild.base import MozbuildObject
+from mozbuild.util import memoize
+
+from taskgraph import files_changed
+from taskgraph.optimize import register_strategy, OptimizationStrategy
+from taskgraph.util.seta import is_low_value_task
+from taskgraph.util.taskcluster import find_task_id
+
+logger = logging.getLogger(__name__)
+
+
+@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
+    # created.
+    # In every of those cases, we need to run the task to create or refresh
+    # artifacts.
+
+    def should_replace_task(self, task, params, index_paths):
+        "Look for a task with one of the given index paths"
+        for index_path in index_paths:
+            try:
+                task_id = find_task_id(
+                    index_path,
+                    use_proxy=bool(os.environ.get('TASK_ID')))
+                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
+        # a MercurialRevisionFinder with revision '.', which should be the same
+        # as `revision`; in other circumstances, it will use a default reader
+        rdr = mbo.mozbuild_reader(config_mode='empty')
+
+        components = set()
+        for p, m in rdr.files_info(changed_files).items():
+            components |= set(m['SCHEDULES'].components)
+
+        return components
+
+    def should_remove_task(self, task, params, conditions):
+        if params.get('pushlog_id') == -1:
+            return False
+
+        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

taskcluster/taskgraph/test/test_optimize.py

@@ -6,21 +6,21 @@ from __future__ import absolute_import, print_function, unicode_literals
 
 import unittest
 
-from taskgraph import optimize
+from taskgraph import graph, optimize
+from taskgraph.optimize import OptimizationStrategy
 from taskgraph.taskgraph import TaskGraph
-from taskgraph import graph
 from taskgraph.task import Task
 from mozunit import main
 from slugid import nice as slugid
 
 
-class Remove(optimize.OptimizationStrategy):
+class Remove(OptimizationStrategy):
 
     def should_remove_task(self, task, params, arg):
         return True
 
 
-class Replace(optimize.OptimizationStrategy):
+class Replace(OptimizationStrategy):
 
     def should_replace_task(self, task, params, taskid):
         return taskid
@@ -29,7 +29,7 @@ class Replace(optimize.OptimizationStrategy):
 class TestOptimize(unittest.TestCase):
 
     strategies = {
-        'never': optimize.OptimizationStrategy(),
+        'never': OptimizationStrategy(),
         'remove': Remove(),
         'replace': Replace(),
     }