tools/tryselect/tasks.py

   1 # This Source Code Form is subject to the terms of the Mozilla Public
   2 # License, v. 2.0. If a copy of the MPL was not distributed with this
   3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   4
   5 from __future__ import absolute_import, print_function, unicode_literals
   6
   7 import json
   8 import os
   9 import re
  10 import sys
  11 from collections import defaultdict
  12
  13 from mozboot.util import get_state_dir
  14 from mozbuild.base import MozbuildObject
  15 from mozpack.files import FileFinder
  16 from moztest.resolve import TestResolver, get_suite_definition
  17
  18 import taskgraph
  19 from taskgraph.generator import TaskGraphGenerator
  20 from taskgraph.parameters import (
  21     ParameterMismatch,
  22     parameters_loader,
  23 )
  24 from taskgraph.taskgraph import TaskGraph
  25
  26 here = os.path.abspath(os.path.dirname(__file__))
  27 build = MozbuildObject.from_environment(cwd=here)
  28
  29 PARAMETER_MISMATCH = """
  30 ERROR - The parameters being used to generate tasks differ from those expected
  31 by your working copy:
  32
  33     {}
  34
  35 To fix this, either rebase onto the latest mozilla-central or pass in
  36 -p/--parameters. For more information on how to define parameters, see:
  37 https://firefox-source-docs.mozilla.org/taskcluster/taskcluster/mach.html#parameters
  38 """
  39
  40
  41 def invalidate(cache):
  42     try:
  43         cmod = os.path.getmtime(cache)
  44     except OSError as e:
  45         # File does not exist. We catch OSError rather than use `isfile`
  46         # because the recommended watchman hook could possibly invalidate the
  47         # cache in-between the check to `isfile` and the call to `getmtime`
  48         # below.
  49         if e.errno == 2:
  50             return
  51         raise
  52
  53     tc_dir = os.path.join(build.topsrcdir, 'taskcluster')
  54     tmod = max(os.path.getmtime(os.path.join(tc_dir, p)) for p, _ in FileFinder(tc_dir))
  55
  56     if tmod > cmod:
  57         os.remove(cache)
  58
  59
  60 def generate_tasks(params=None, full=False):
  61     cache_dir = os.path.join(get_state_dir(srcdir=True), 'cache', 'taskgraph')
  62     attr = 'full_task_set' if full else 'target_task_set'
  63     cache = os.path.join(cache_dir, attr)
  64
  65     invalidate(cache)
  66     if os.path.isfile(cache):
  67         with open(cache, 'r') as fh:
  68             return TaskGraph.from_json(json.load(fh))[1]
  69
  70     if not os.path.isdir(cache_dir):
  71         os.makedirs(cache_dir)
  72
  73     print("Task configuration changed, generating {}".format(attr.replace('_', ' ')))
  74
  75     taskgraph.fast = True
  76     cwd = os.getcwd()
  77     os.chdir(build.topsrcdir)
  78
  79     root = os.path.join(build.topsrcdir, 'taskcluster', 'ci')
  80     params = parameters_loader(params, strict=False, overrides={'try_mode': 'try_select'})
  81
  82     # Cache both full_task_set and target_task_set regardless of whether or not
  83     # --full was requested. Caching is cheap and can potentially save a lot of
  84     # time.
  85     generator = TaskGraphGenerator(root_dir=root, parameters=params)
  86
  87     def generate(attr):
  88         try:
  89             tg = getattr(generator, attr)
  90         except ParameterMismatch as e:
  91             print(PARAMETER_MISMATCH.format(e.args[0]))
  92             sys.exit(1)
  93
  94         # write cache
  95         with open(os.path.join(cache_dir, attr), 'w') as fh:
  96             json.dump(tg.to_json(), fh)
  97         return tg
  98
  99     tg_full = generate('full_task_set')
 100     tg_target = generate('target_task_set')
 101     # discard results from these, we only need cache.
 102     if full:
 103         generate('full_task_graph')
 104     generate('target_task_graph')
 105
 106     os.chdir(cwd)
 107     if full:
 108         return tg_full
 109     return tg_target
 110
 111
 112 def filter_tasks_by_paths(tasks, paths):
 113     resolver = TestResolver.from_environment(cwd=here)
 114     run_suites, run_tests = resolver.resolve_metadata(paths)
 115     flavors = set([(t['flavor'], t.get('subsuite')) for t in run_tests])
 116
 117     task_regexes = set()
 118     for flavor, subsuite in flavors:
 119         _, suite = get_suite_definition(flavor, subsuite, strict=True)
 120         if 'task_regex' not in suite:
 121             print("warning: no tasks could be resolved from flavor '{}'{}".format(
 122                     flavor, " and subsuite '{}'".format(subsuite) if subsuite else ""))
 123             continue
 124
 125         task_regexes.update(suite['task_regex'])
 126
 127     def match_task(task):
 128         return any(re.search(pattern, task) for pattern in task_regexes)
 129
 130     return filter(match_task, tasks)
 131
 132
 133 def resolve_tests_by_suite(paths):
 134     resolver = TestResolver.from_environment(cwd=here)
 135     _, run_tests = resolver.resolve_metadata(paths)
 136
 137     suite_to_tests = defaultdict(list)
 138
 139     # A dictionary containing all the input paths that we haven't yet
 140     # assigned to a specific test flavor.
 141     remaining_paths_by_suite = defaultdict(lambda: set(paths))
 142
 143     for test in run_tests:
 144         key, _ = get_suite_definition(test['flavor'], test.get('subsuite'), strict=True)
 145
 146         test_path = test.get('srcdir_relpath')
 147         if test_path is None:
 148             continue
 149         found_path = None
 150         for path in remaining_paths_by_suite[key]:
 151             if test_path.startswith(path):
 152                 found_path = path
 153                 break
 154         if found_path:
 155             suite_to_tests[key].append(found_path)
 156             remaining_paths_by_suite[key].remove(found_path)
 157
 158     return suite_to_tests