summaryrefslogtreecommitdiffstats
path: root/taskcluster/android_taskgraph/transforms/chunk.py
blob: bb17a99f4eb7c604548d223ae965296b2ef38222 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
# 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 copy import deepcopy

from taskgraph import MAX_DEPENDENCIES
from taskgraph.transforms.base import TransformSequence
from taskgraph.util.treeherder import add_suffix

# XXX Docker images may be added after this transform, so we allow one more dep to be added
MAX_NUMBER_OF_DEPS = MAX_DEPENDENCIES - 1

transforms = TransformSequence()


def build_task_definition(orig_task, deps, soft_deps, count):
    task = deepcopy(orig_task)
    task["dependencies"] = {label: label for label in deps}
    task["soft-dependencies"] = list(soft_deps)
    task["name"] = "{}-{}".format(orig_task["name"], count)
    if "treeherder" in task:
        task["treeherder"]["symbol"] = add_suffix(
            task["treeherder"]["symbol"], f"-{count}"
        )

    task["attributes"]["is_final_chunked_task"] = False
    return task


def get_chunked_label(config, chunked_task):
    return "{}-{}".format(config.kind, chunked_task["name"])


@transforms.add
def add_dependencies(config, tasks):
    for task in tasks:
        count = 1
        soft_deps = set()
        regular_deps = set()
        chunked_labels = set()

        soft_dep_labels = list(task.pop("soft-dependencies", []))
        regular_dep_labels = list(task.get("dependencies", {}).keys())
        # sort for deterministic chunking
        all_dep_labels = sorted(set(soft_dep_labels + regular_dep_labels))

        for dep_label in all_dep_labels:
            if dep_label in regular_dep_labels:
                regular_deps.add(dep_label)
            else:
                soft_deps.add(dep_label)

            if len(regular_deps) + len(soft_deps) == MAX_NUMBER_OF_DEPS:
                chunked_task = build_task_definition(
                    task, regular_deps, soft_deps, count
                )
                chunked_label = get_chunked_label(config, chunked_task)
                chunked_labels.add(chunked_label)
                yield chunked_task
                soft_deps.clear()
                regular_deps.clear()
                count += 1

        if regular_deps or soft_deps:
            chunked_task = build_task_definition(task, regular_deps, soft_deps, count)
            chunked_label = get_chunked_label(config, chunked_task)
            chunked_labels.add(chunked_label)
            yield chunked_task

        task["dependencies"] = {label: label for label in chunked_labels}
        # Chunk yields a last task that doesn't have a number appended to it.
        # It helps configuring Github which waits on a single label.
        # Setting this attribute also enables multi_dep to select the right
        # task to depend on.
        task["attributes"]["is_final_chunked_task"] = True
        yield task