run_task.py (9141B)
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 Support for running jobs that are invoked via the `run-task` script. 6 """ 7 8 import dataclasses 9 import os 10 11 from mozbuild.util import memoize 12 from mozpack import path 13 from taskgraph.transforms.run.common import support_caches 14 from taskgraph.util.schema import Schema 15 from taskgraph.util.yaml import load_yaml 16 from voluptuous import Any, Optional, Required 17 18 from gecko_taskgraph.transforms.job import run_job_using 19 from gecko_taskgraph.transforms.job.common import add_tooltool, support_vcs_checkout 20 from gecko_taskgraph.transforms.task import taskref_or_string 21 22 run_task_schema = Schema({ 23 Required("using"): "run-task", 24 # Use the specified caches. 25 Optional("use-caches"): Any(bool, [str]), 26 # if true (the default), perform a checkout of gecko on the worker 27 Required("checkout"): bool, 28 Optional( 29 "cwd", 30 description="Path to run command in. If a checkout is present, the path " 31 "to the checkout will be interpolated with the key `checkout`", 32 ): str, 33 # The sparse checkout profile to use. Value is the filename relative to 34 # "sparse-profile-prefix" which defaults to "build/sparse-profiles/". 35 Required("sparse-profile"): Any(str, None), 36 # The relative path to the sparse profile. 37 Optional("sparse-profile-prefix"): str, 38 # Whether to use a shallow clone or not, default True (git only). 39 Optional("shallow-clone"): bool, 40 # if true, perform a checkout of a comm-central based branch inside the 41 # gecko checkout 42 Required("comm-checkout"): bool, 43 # The command arguments to pass to the `run-task` script, after the 44 # checkout arguments. If a list, it will be passed directly; otherwise 45 # it will be included in a single argument to `bash -cx`. 46 Required("command"): Any([taskref_or_string], taskref_or_string), 47 # Base work directory used to set up the task. 48 Optional("workdir"): str, 49 # If not false, tooltool downloads will be enabled via relengAPIProxy 50 # for either just public files, or all files. Only supported on 51 # docker-worker. 52 Required("tooltool-downloads"): Any( 53 False, 54 "public", 55 "internal", 56 ), 57 # Whether to run as root. (defaults to False) 58 Optional("run-as-root"): bool, 59 }) 60 61 62 def common_setup(config, job, taskdesc, command): 63 run = job["run"] 64 run_cwd = run.get("cwd") 65 if run["checkout"]: 66 repo_configs = config.repo_configs 67 if len(repo_configs) > 1 and run["checkout"] is True: 68 raise Exception("Must explicitly specify checkouts with multiple repos.") 69 elif run["checkout"] is not True: 70 repo_configs = { 71 repo: dataclasses.replace(repo_configs[repo], **config) 72 for (repo, config) in run["checkout"].items() 73 } 74 75 gecko_path = support_vcs_checkout(config, job, taskdesc, repo_configs) 76 command.append(f"--gecko-checkout={gecko_path}") 77 if config.params["repository_type"] == "git" and run.get("shallow-clone", True): 78 command.append("--gecko-shallow-clone") 79 80 if run_cwd: 81 run_cwd = path.normpath(run_cwd.format(checkout=gecko_path)) 82 83 elif run_cwd and "{checkout}" in run_cwd: 84 raise Exception( 85 "Found `{{checkout}}` interpolation in `cwd` for task {name} " 86 "but the task doesn't have a checkout: {cwd}".format( 87 cwd=run_cwd, name=job.get("name", job.get("label")) 88 ) 89 ) 90 91 if config.params["repository_type"] == "hg" and run["sparse-profile"]: 92 sparse_profile_prefix = run.pop( 93 "sparse-profile-prefix", "build/sparse-profiles" 94 ) 95 sparse_profile_path = path.join(sparse_profile_prefix, run["sparse-profile"]) 96 command.append(f"--gecko-sparse-profile={sparse_profile_path}") 97 98 if run_cwd: 99 command.append(f"--task-cwd={run_cwd}") 100 101 support_caches(config, job, taskdesc) 102 taskdesc["worker"].setdefault("env", {})["MOZ_SCM_LEVEL"] = config.params["level"] 103 104 105 worker_defaults = { 106 "checkout": True, 107 "comm-checkout": False, 108 "sparse-profile": None, 109 "tooltool-downloads": False, 110 "run-as-root": False, 111 } 112 113 114 load_yaml = memoize(load_yaml) 115 116 117 def script_url(config, script): 118 if "MOZ_AUTOMATION" in os.environ and "TASK_ID" not in os.environ: 119 raise Exception("TASK_ID must be defined to use run-task on generic-worker") 120 task_id = os.environ.get("TASK_ID", "<TASK_ID>") 121 tc_url = "http://firefox-ci-tc.services.mozilla.com" 122 return f"{tc_url}/api/queue/v1/task/{task_id}/artifacts/public/{script}" 123 124 125 @run_job_using( 126 "docker-worker", "run-task", schema=run_task_schema, defaults=worker_defaults 127 ) 128 def docker_worker_run_task(config, job, taskdesc): 129 run = job["run"] 130 worker = taskdesc["worker"] = job["worker"] 131 run_task_bin = ( 132 "run-task-git" if config.params["repository_type"] == "git" else "run-task-hg" 133 ) 134 command = [f"/builds/worker/bin/{run_task_bin}"] 135 common_setup(config, job, taskdesc, command) 136 137 if run["tooltool-downloads"]: 138 internal = run["tooltool-downloads"] == "internal" 139 add_tooltool(config, job, taskdesc, internal=internal) 140 141 run_command = run["command"] 142 143 # dict is for the case of `{'task-reference': text_type}`. 144 if isinstance(run_command, (str, dict)): 145 run_command = ["bash", "-cx", run_command] 146 if run["comm-checkout"]: 147 command.append( 148 "--comm-checkout={}/comm".format(taskdesc["worker"]["env"]["GECKO_PATH"]) 149 ) 150 if run["run-as-root"]: 151 command.extend(("--user", "root", "--group", "root")) 152 command.append("--") 153 command.extend(run_command) 154 worker["command"] = command 155 156 157 @run_job_using( 158 "generic-worker", "run-task", schema=run_task_schema, defaults=worker_defaults 159 ) 160 def generic_worker_run_task(config, job, taskdesc): 161 run = job["run"] 162 worker = taskdesc["worker"] = job["worker"] 163 is_win = worker["os"] == "windows" 164 is_mac = worker["os"] == "macosx" 165 is_bitbar = worker["os"] == "linux-bitbar" 166 is_lambda = worker["os"] == "linux-lambda" 167 168 if run["tooltool-downloads"]: 169 internal = run["tooltool-downloads"] == "internal" 170 add_tooltool(config, job, taskdesc, internal=internal) 171 172 if is_win: 173 command = ["C:/mozilla-build/python3/python3.exe", "run-task"] 174 elif is_mac: 175 command = ["/usr/local/bin/python3", "run-task"] 176 else: 177 command = ["./run-task"] 178 179 common_setup(config, job, taskdesc, command) 180 181 worker.setdefault("mounts", []) 182 run_task_bin = ( 183 "run-task-git" if config.params["repository_type"] == "git" else "run-task-hg" 184 ) 185 worker["mounts"].append({ 186 "content": { 187 "url": script_url(config, run_task_bin), 188 }, 189 "file": "./run-task", 190 }) 191 192 if ( 193 job.get("fetches") 194 or job.get("use-uv") 195 or job.get("use-python", "system") != "system" 196 ): 197 worker["mounts"].append({ 198 "content": { 199 "url": script_url(config, "fetch-content"), 200 }, 201 "file": "./fetch-content", 202 }) 203 if run.get("checkout"): 204 worker["mounts"].append({ 205 "content": { 206 "url": script_url(config, "robustcheckout.py"), 207 }, 208 "file": "./robustcheckout.py", 209 }) 210 211 run_command = run["command"] 212 213 # dict is for the case of `{'task-reference': text_type}`. 214 if isinstance(run_command, (str, dict)): 215 if is_win: 216 if isinstance(run_command, dict): 217 for k in run_command.keys(): 218 run_command[k] = f'"{run_command[k]}"' 219 else: 220 run_command = f'"{run_command}"' 221 run_command = ["bash", "-cx", run_command] 222 223 if run["comm-checkout"]: 224 command.append( 225 "--comm-checkout={}/comm".format(taskdesc["worker"]["env"]["GECKO_PATH"]) 226 ) 227 228 if run["run-as-root"]: 229 command.extend(("--user", "root", "--group", "root")) 230 command.append("--") 231 if is_bitbar: 232 # Use the bitbar wrapper script which sets up the device and adb 233 # environment variables 234 command.append("/builds/taskcluster/script.py") 235 elif is_lambda: 236 command.append("/home/ltuser/taskcluster/script.py") 237 command.extend(run_command) 238 239 if is_win: 240 taskref = False 241 for c in command: 242 if isinstance(c, dict): 243 taskref = True 244 245 if taskref: 246 cmd = [] 247 for c in command: 248 if isinstance(c, dict): 249 for v in c.values(): 250 cmd.append(v) 251 else: 252 cmd.append(c) 253 worker["command"] = [{"artifact-reference": " ".join(cmd)}] 254 else: 255 worker["command"] = [" ".join(command)] 256 else: 257 worker["command"] = [ 258 ["chmod", "+x", "run-task"], 259 command, 260 ]