taskcluster/android_taskgraph/transforms/chunk.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 from copy import deepcopy
   7
   8 from taskgraph import MAX_DEPENDENCIES
   9 from taskgraph.transforms.base import TransformSequence
  10 from taskgraph.util.treeherder import add_suffix
  11
  12 # XXX Docker images may be added after this transform, so we allow one more dep to be added
  13 MAX_NUMBER_OF_DEPS = MAX_DEPENDENCIES - 1
  14
  15 transforms = TransformSequence()
  16
  17
  18 def build_task_definition(orig_task, deps, soft_deps, count):
  19     task = deepcopy(orig_task)
  20     task["dependencies"] = {label: label for label in deps}
  21     task["soft-dependencies"] = list(soft_deps)
  22     task["name"] = "{}-{}".format(orig_task["name"], count)
  23     if "treeherder" in task:
  24         task["treeherder"]["symbol"] = add_suffix(
  25             task["treeherder"]["symbol"], f"-{count}"
  26         )
  27
  28     task["attributes"]["is_final_chunked_task"] = False
  29     return task
  30
  31
  32 def get_chunked_label(config, chunked_task):
  33     return "{}-{}".format(config.kind, chunked_task["name"])
  34
  35
  36 @transforms.add
  37 def add_dependencies(config, tasks):
  38     for task in tasks:
  39         count = 1
  40         soft_deps = set()
  41         regular_deps = set()
  42         chunked_labels = set()
  43
  44         soft_dep_labels = list(task.pop("soft-dependencies", []))
  45         regular_dep_labels = list(task.get("dependencies", {}).keys())
  46         # sort for deterministic chunking
  47         all_dep_labels = sorted(set(soft_dep_labels + regular_dep_labels))
  48
  49         for dep_label in all_dep_labels:
  50             if dep_label in regular_dep_labels:
  51                 regular_deps.add(dep_label)
  52             else:
  53                 soft_deps.add(dep_label)
  54
  55             if len(regular_deps) + len(soft_deps) == MAX_NUMBER_OF_DEPS:
  56                 chunked_task = build_task_definition(
  57                     task, regular_deps, soft_deps, count
  58                 )
  59                 chunked_label = get_chunked_label(config, chunked_task)
  60                 chunked_labels.add(chunked_label)
  61                 yield chunked_task
  62                 soft_deps.clear()
  63                 regular_deps.clear()
  64                 count += 1
  65
  66         if regular_deps or soft_deps:
  67             chunked_task = build_task_definition(task, regular_deps, soft_deps, count)
  68             chunked_label = get_chunked_label(config, chunked_task)
  69             chunked_labels.add(chunked_label)
  70             yield chunked_task
  71
  72         task["dependencies"] = {label: label for label in chunked_labels}
  73         # Chunk yields a last task that doesn't have a number appended to it.
  74         # It helps configuring Github which waits on a single label.
  75         # Setting this attribute also enables multi_dep to select the right
  76         # task to depend on.
  77         task["attributes"]["is_final_chunked_task"] = True
  78         yield task