variant.py (4971B)
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 import datetime 5 6 import jsone 7 from taskgraph.transforms.base import TransformSequence 8 from taskgraph.util.copy import deepcopy 9 from taskgraph.util.schema import Schema, resolve_keyed_by, validate_schema 10 from taskgraph.util.templates import merge 11 from taskgraph.util.treeherder import join_symbol, split_symbol 12 from voluptuous import Any, Optional, Required 13 14 from gecko_taskgraph.util.chunking import TEST_VARIANTS 15 16 transforms = TransformSequence() 17 18 """List of available test variants defined.""" 19 20 21 variant_description_schema = Schema({ 22 str: { 23 Required("description"): str, 24 Required("suffix"): str, 25 Optional("mozinfo"): str, 26 Required("component"): str, 27 Required("expiration"): str, 28 Optional("when"): {Any("$eval", "$if"): str}, 29 Optional("replace"): {str: object}, 30 Optional("merge"): {str: object}, 31 } 32 }) 33 """variant description schema""" 34 35 36 @transforms.add 37 def split_variants(config, tasks): 38 """Splits test definitions into multiple tasks based on the `variants` key. 39 40 If `variants` are defined, the original task will be yielded along with a 41 copy of the original task for each variant defined in the list. The copies 42 will have the 'unittest_variant' attribute set. 43 """ 44 validate_schema(variant_description_schema, TEST_VARIANTS, "In variants.yml:") 45 46 def find_expired_variants(variants): 47 expired = [] 48 49 # do not expire on esr/beta/release 50 if config.params.get("release_type", "") in [ 51 "release", 52 "beta", 53 ]: 54 return [] 55 56 if "esr" in config.params.get("release_type", ""): 57 return [] 58 59 today = datetime.datetime.today() 60 for variant in variants: 61 expiration = variants[variant]["expiration"] 62 if len(expiration.split("-")) == 1: 63 continue 64 expires_at = datetime.datetime.strptime(expiration, "%Y-%m-%d") 65 if expires_at < today: 66 expired.append(variant) 67 return expired 68 69 def remove_expired(variants, expired): 70 remaining_variants = [] 71 for name in variants: 72 parts = [p for p in name.split("+") if p in expired] 73 if len(parts) > 0: 74 continue 75 76 remaining_variants.append(name) 77 return remaining_variants 78 79 def replace_task_items(task_key, variant_key): 80 for item in variant_key: 81 if isinstance(variant_key[item], dict): 82 task_key[item] = replace_task_items( 83 task_key.get(item, {}), variant_key[item] 84 ) 85 else: 86 task_key[item] = variant_key[item] 87 return task_key 88 89 def apply_variant(variant, task, name): 90 task["description"] = variant["description"].format(**task) 91 92 suffix = f"-{variant['suffix']}" 93 group, symbol = split_symbol(task["treeherder-symbol"]) 94 if group != "?": 95 group += suffix 96 else: 97 symbol += suffix 98 task["treeherder-symbol"] = join_symbol(group, symbol) 99 100 # This will be used to set the label and try-name in 'make_job_description'. 101 task.setdefault("variant-suffix", "") 102 task["variant-suffix"] += suffix 103 104 # Replace and/or merge the configuration. 105 106 # we only want to update the leaf node, the the entire top level dict 107 task = replace_task_items(task, variant.get("replace", {})) 108 109 resolve_keyed_by( 110 task, 111 "mozharness.extra-options", 112 item_name=task["test-name"], 113 enforce_single_match=False, 114 variant=name, 115 ) 116 117 return merge(task, deepcopy(variant.get("merge", {}))) 118 119 expired_variants = find_expired_variants(TEST_VARIANTS) 120 for task in tasks: 121 variants = task.pop("variants", []) 122 variants = remove_expired(variants, expired_variants) 123 124 if task.pop("run-without-variant"): 125 taskv = deepcopy(task) if variants else task 126 taskv["attributes"]["unittest_variant"] = None 127 yield taskv 128 129 for name in variants: 130 # Apply composite variants (joined by '+') in order. 131 parts = name.split("+") 132 taskv = deepcopy(task) 133 for part in parts: 134 variant = TEST_VARIANTS[part] 135 136 # If any variant in a composite fails this check we skip it. 137 if "when" in variant: 138 context = {"task": task} 139 if not jsone.render(variant["when"], context): 140 break 141 142 taskv = apply_variant(variant, taskv, name) 143 else: 144 taskv["attributes"]["unittest_variant"] = name 145 yield taskv