diff options
Diffstat (limited to '')
-rw-r--r-- | taskcluster/android_taskgraph/transforms/chunk.py | 78 |
1 files changed, 78 insertions, 0 deletions
diff --git a/taskcluster/android_taskgraph/transforms/chunk.py b/taskcluster/android_taskgraph/transforms/chunk.py new file mode 100644 index 0000000000..bb17a99f4e --- /dev/null +++ b/taskcluster/android_taskgraph/transforms/chunk.py @@ -0,0 +1,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 |