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