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