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