78 lines
2.8 KiB
Python
78 lines
2.8 KiB
Python
# 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
|