main.py (23334B)
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 import argparse 6 import atexit 7 import logging 8 import os 9 import shutil 10 import subprocess 11 import sys 12 import tempfile 13 import traceback 14 from pathlib import Path 15 from typing import Any 16 17 import appdirs 18 import yaml 19 from taskgraph.main import ( 20 FORMAT_METHODS, 21 argument, 22 command, 23 commands, 24 dump_output, 25 format_kind_graph_mermaid, 26 generate_taskgraph, 27 get_taskgraph_generator, 28 ) 29 30 from gecko_taskgraph import GECKO 31 from gecko_taskgraph.files_changed import get_locally_changed_files 32 33 34 def format_taskgraph_yaml(taskgraph): 35 from taskgraph.util.readonlydict import ReadOnlyDict 36 37 class TGDumper(yaml.SafeDumper): 38 def ignore_aliases(self, data): 39 return True 40 41 def represent_ro_dict(self, data): 42 return self.represent_dict(dict(data)) 43 44 TGDumper.add_representer(ReadOnlyDict, TGDumper.represent_ro_dict) 45 46 return yaml.dump(taskgraph.to_json(), Dumper=TGDumper, default_flow_style=False) 47 48 49 FORMAT_METHODS["yaml"] = format_taskgraph_yaml 50 51 52 @command( 53 "kind-graph", 54 help="Show the kind dependency graph as a Mermaid flowchart diagram.", 55 ) 56 @argument("--root", "-r", help="root of the taskgraph definition relative to topsrcdir") 57 @argument("--quiet", "-q", action="store_true", help="suppress all logging output") 58 @argument( 59 "--verbose", "-v", action="store_true", help="include debug-level logging output" 60 ) 61 @argument( 62 "--parameters", 63 "-p", 64 default=None, 65 help="Parameters to use for the generation. Can be a path to file (.yml or " 66 ".json; see `taskcluster/docs/parameters.rst`), a url, of the form " 67 "`project=mozilla-central` to download latest parameters file for the specified " 68 "project from CI, or of the form `task-id=<decision task id>` to download " 69 "parameters from the specified decision task.", 70 ) 71 @argument( 72 "-o", 73 "--output-file", 74 default=None, 75 help="file path to store generated output.", 76 ) 77 @argument( 78 "-k", 79 "--target-kind", 80 dest="target_kinds", 81 action="append", 82 default=[], 83 help="only return kinds and their dependencies.", 84 ) 85 def show_kinds(options): 86 from taskgraph.parameters import parameters_loader # noqa: PLC0415 87 88 if options.pop("verbose", False): 89 logging.root.setLevel(logging.DEBUG) 90 91 setup_logging() 92 93 target_kinds = options.get("target_kinds", []) 94 parameters = options.get("parameters") 95 if not parameters: 96 parameters = parameters_loader( 97 None, strict=False, overrides={"target-kinds": target_kinds} 98 ) 99 elif isinstance(parameters, str): 100 parameters = parameters_loader( 101 parameters, 102 overrides={"target-kinds": target_kinds}, 103 strict=False, 104 ) 105 elif target_kinds: 106 # Parameters object already exists (from tests) 107 parameters["target-kinds"] = target_kinds 108 109 tgg = get_taskgraph_generator(options.get("root"), parameters) 110 kind_graph = tgg.kind_graph 111 112 output = format_kind_graph_mermaid(kind_graph) 113 114 output_file = options.get("output_file") 115 if output_file: 116 with open(output_file, "w") as fh: 117 print(output, file=fh) 118 print(f"Kind graph written to {output_file}", file=sys.stderr) 119 else: 120 print(output) 121 122 return 0 123 124 125 @command( 126 "tasks", 127 help="Show all tasks in the taskgraph.", 128 defaults={"graph_attr": "full_task_set"}, 129 ) 130 @command( 131 "full", help="Show the full taskgraph.", defaults={"graph_attr": "full_task_graph"} 132 ) 133 @command( 134 "target", 135 help="Show the set of target tasks.", 136 defaults={"graph_attr": "target_task_set"}, 137 ) 138 @command( 139 "target-graph", 140 help="Show the target graph.", 141 defaults={"graph_attr": "target_task_graph"}, 142 ) 143 @command( 144 "optimized", 145 help="Show the optimized graph.", 146 defaults={"graph_attr": "optimized_task_graph"}, 147 ) 148 @command( 149 "morphed", 150 help="Show the morphed graph.", 151 defaults={"graph_attr": "morphed_task_graph"}, 152 ) 153 @argument("--root", "-r", help="root of the taskgraph definition relative to topsrcdir") 154 @argument("--quiet", "-q", action="store_true", help="suppress all logging output") 155 @argument( 156 "--verbose", "-v", action="store_true", help="include debug-level logging output" 157 ) 158 @argument( 159 "--json", 160 "-J", 161 action="store_const", 162 dest="format", 163 const="json", 164 help="Output task graph as a JSON object", 165 ) 166 @argument( 167 "--yaml", 168 "-Y", 169 action="store_const", 170 dest="format", 171 const="yaml", 172 help="Output task graph as a YAML object", 173 ) 174 @argument( 175 "--labels", 176 "-L", 177 action="store_const", 178 dest="format", 179 const="labels", 180 help="Output the label for each task in the task graph (default)", 181 ) 182 @argument( 183 "--parameters", 184 "-p", 185 default=None, 186 action="append", 187 help="Parameters to use for the generation. Can be a path to file (.yml or " 188 ".json; see `taskcluster/docs/parameters.rst`), a directory (containing " 189 "parameters files), a url, of the form `project=mozilla-central` to download " 190 "latest parameters file for the specified project from CI, or of the form " 191 "`task-id=<decision task id>` to download parameters from the specified " 192 "decision task. Can be specified multiple times, in which case multiple " 193 "generations will happen from the same invocation (one per parameters " 194 "specified).", 195 ) 196 @argument( 197 "--force-local-files-changed", 198 default=False, 199 action="store_true", 200 help="Compute the 'files-changed' parameter from local version control, " 201 "even when explicitly using a parameter set that already has it defined. " 202 "Note that this is already the default behaviour when no parameters are " 203 "specified.", 204 ) 205 @argument( 206 "--no-optimize", 207 dest="optimize", 208 action="store_false", 209 default="true", 210 help="do not remove tasks from the graph that are found in the " 211 "index (a.k.a. optimize the graph)", 212 ) 213 @argument( 214 "-o", 215 "--output-file", 216 default=None, 217 help="file path to store generated output.", 218 ) 219 @argument( 220 "--tasks-regex", 221 "--tasks", 222 default=None, 223 help="only return tasks with labels matching this regular expression.", 224 ) 225 @argument( 226 "--exclude-key", 227 default=None, 228 dest="exclude_keys", 229 action="append", 230 help="Exclude the specified key (using dot notation) from the final result. " 231 "This is mainly useful with '--diff' to filter out expected differences.", 232 ) 233 @argument( 234 "-k", 235 "--target-kind", 236 dest="target_kinds", 237 action="append", 238 default=[], 239 help="only return tasks that are of the given kind, or their dependencies.", 240 ) 241 @argument( 242 "-F", 243 "--fast", 244 default=False, 245 action="store_true", 246 help="enable fast task generation for local debugging.", 247 ) 248 @argument( 249 "--diff", 250 const="default", 251 nargs="?", 252 default=None, 253 help="Generate and diff the current taskgraph against another revision. " 254 "Without args the base revision will be used. A revision specifier such as " 255 "the hash or `.~1` (hg) or `HEAD~1` (git) can be used as well.", 256 ) 257 @argument( 258 "-j", 259 "--max-workers", 260 dest="max_workers", 261 default=None, 262 type=int, 263 help="The maximum number of workers to use for parallel operations such as" 264 "when multiple parameters files are passed.", 265 ) 266 def show_taskgraph(options): 267 from mozversioncontrol import get_repository_object as get_repository 268 from taskgraph.parameters import Parameters, parameters_loader 269 270 if options.pop("verbose", False): 271 logging.root.setLevel(logging.DEBUG) 272 273 repo = None 274 cur_ref = None 275 diffdir = None 276 output_file = options["output_file"] 277 278 if options["diff"]: 279 # --root argument is taskgraph's config at <repo>/taskcluster 280 repo_root = os.getcwd() 281 if options["root"]: 282 repo_root = f"{options['root']}/.." 283 repo = get_repository(repo_root) 284 285 if not repo.working_directory_clean(): 286 print( 287 "abort: can't diff taskgraph with dirty working directory", 288 file=sys.stderr, 289 ) 290 return 1 291 292 # We want to return the working directory to the current state 293 # as best we can after we're done. In all known cases, using 294 # branch or bookmark (which are both available on the VCS object) 295 # as `branch` is preferable to a specific revision. 296 cur_ref = repo.branch or repo.head_ref[:12] 297 cur_ref_file = cur_ref.replace("/", "_") 298 299 diffdir = tempfile.mkdtemp() 300 atexit.register( 301 shutil.rmtree, diffdir 302 ) # make sure the directory gets cleaned up 303 options["output_file"] = os.path.join( 304 diffdir, f"{options['graph_attr']}_{cur_ref_file}" 305 ) 306 print(f"Generating {options['graph_attr']} @ {cur_ref}", file=sys.stderr) 307 308 overrides = { 309 "target-kinds": options.get("target_kinds"), 310 } 311 parameters: list[Any[str, Parameters]] = options.pop("parameters") 312 if not parameters: 313 parameters = [ 314 parameters_loader(None, strict=False, overrides=overrides) 315 ] # will use default values 316 317 # This is the default behaviour anyway, so no need to re-compute. 318 options["force_local_files_changed"] = False 319 320 elif options["force_local_files_changed"]: 321 overrides["files-changed"] = sorted(get_locally_changed_files(GECKO)) 322 323 for param in parameters[:]: 324 if isinstance(param, str) and os.path.isdir(param): 325 parameters.remove(param) 326 parameters.extend([ 327 p.as_posix() 328 for p in Path(param).iterdir() 329 if p.suffix in (".yml", ".json") 330 ]) 331 332 logdir = None 333 if len(parameters) > 1: 334 # Log to separate files for each process instead of stderr to 335 # avoid interleaving. 336 basename = os.path.basename(os.getcwd()) 337 logdir = os.path.join(appdirs.user_log_dir("taskgraph"), basename) 338 if not os.path.isdir(logdir): 339 os.makedirs(logdir) 340 else: 341 # Only setup logging if we have a single parameter spec. Otherwise 342 # logging will go to files. This is also used as a hook for Gecko 343 # to setup its `mach` based logging. 344 setup_logging() 345 346 generate_taskgraph(options, parameters, overrides, logdir) 347 348 if options["diff"]: 349 assert diffdir is not None 350 assert repo is not None 351 352 # Reload taskgraph modules to pick up changes and clear global state. 353 for mod in sys.modules.copy(): 354 if mod not in {__name__, "taskgraph.main"} and mod.split(".", 1)[ 355 0 356 ].endswith(("taskgraph", "mozbuild")): 357 del sys.modules[mod] 358 359 # Ensure gecko_taskgraph is ahead of taskcluster_taskgraph in sys.path. 360 # Without this, we may end up validating some things against the wrong 361 # schema. 362 import gecko_taskgraph # noqa 363 364 if options["diff"] == "default": 365 base_ref = repo.base_ref 366 else: 367 base_ref = options["diff"] 368 369 base_ref_file = base_ref.replace("/", "_") 370 try: 371 repo.update(base_ref) 372 base_ref = repo.head_ref[:12] 373 options["output_file"] = os.path.join( 374 diffdir, f"{options['graph_attr']}_{base_ref_file}" 375 ) 376 print(f"Generating {options['graph_attr']} @ {base_ref}", file=sys.stderr) 377 generate_taskgraph(options, parameters, overrides, logdir) 378 finally: 379 repo.update(cur_ref) 380 381 # Generate diff(s) 382 diffcmd = [ 383 "diff", 384 "-U20", 385 "--report-identical-files", 386 f"--label={options['graph_attr']}@{base_ref}", 387 f"--label={options['graph_attr']}@{cur_ref}", 388 ] 389 390 non_fatal_failures = [] 391 for spec in parameters: 392 base_path = os.path.join( 393 diffdir, f"{options['graph_attr']}_{base_ref_file}" 394 ) 395 cur_path = os.path.join(diffdir, f"{options['graph_attr']}_{cur_ref_file}") 396 397 params_name = None 398 if len(parameters) > 1: 399 params_name = Parameters.format_spec(spec) 400 base_path += f"_{params_name}" 401 cur_path += f"_{params_name}" 402 403 # If the base or cur files are missing it means that generation 404 # failed. If one of them failed but not the other, the failure is 405 # likely due to the patch making changes to taskgraph in modules 406 # that don't get reloaded (safe to ignore). If both generations 407 # failed, there's likely a real issue. 408 base_missing = not os.path.isfile(base_path) 409 cur_missing = not os.path.isfile(cur_path) 410 if base_missing != cur_missing: # != is equivalent to XOR for booleans 411 non_fatal_failures.append(os.path.basename(base_path)) 412 continue 413 414 try: 415 # If the output file(s) are missing, this command will raise 416 # CalledProcessError with a returncode > 1. 417 proc = subprocess.run( 418 diffcmd + [base_path, cur_path], 419 capture_output=True, 420 text=True, 421 check=True, 422 ) 423 diff_output = proc.stdout 424 returncode = 0 425 except subprocess.CalledProcessError as e: 426 # returncode 1 simply means diffs were found 427 if e.returncode != 1: 428 print(e.stderr, file=sys.stderr) 429 raise 430 diff_output = e.output 431 returncode = e.returncode 432 433 dump_output( 434 diff_output, 435 # Don't bother saving file if no diffs were found. Log to 436 # console in this case instead. 437 path=None if returncode == 0 else output_file, 438 params_spec=spec if len(parameters) > 1 else None, 439 ) 440 441 if non_fatal_failures: 442 failstr = "\n ".join(sorted(non_fatal_failures)) 443 print( 444 "WARNING: Diff skipped for the following generation{s} " 445 "due to failures:\n {failstr}".format( 446 s="s" if len(non_fatal_failures) > 1 else "", failstr=failstr 447 ), 448 file=sys.stderr, 449 ) 450 451 if options["format"] != "json": 452 print( 453 'If you were expecting differences in task bodies you should pass "-J"\n', 454 file=sys.stderr, 455 ) 456 457 if len(parameters) > 1: 458 print(f"See '{logdir}' for logs", file=sys.stderr) 459 460 461 @command("build-image", help="Build a Docker image") 462 @argument("image_name", help="Name of the image to build") 463 @argument( 464 "-t", "--tag", help="tag that the image should be built as.", metavar="name:tag" 465 ) 466 @argument( 467 "--context-only", 468 help="File name the context tarball should be written to." 469 "with this option it will only build the context.tar.", 470 metavar="context.tar", 471 ) 472 def build_image(args): 473 from gecko_taskgraph.docker import build_context, build_image 474 475 if args["context_only"] is None: 476 build_image(args["image_name"], args["tag"], os.environ) 477 else: 478 build_context(args["image_name"], args["context_only"], os.environ) 479 480 481 @command( 482 "load-image", 483 help="Load a pre-built Docker image. Note that you need to " 484 "have docker installed and running for this to work.", 485 ) 486 @argument( 487 "--task-id", 488 help="Load the image at public/image.tar.zst in this task, rather than searching the index", 489 ) 490 @argument( 491 "-t", 492 "--tag", 493 help="tag that the image should be loaded as. If not " 494 "image will be loaded with tag from the tarball", 495 metavar="name:tag", 496 ) 497 @argument( 498 "image_name", 499 nargs="?", 500 help="Load the image of this name based on the current " 501 "contents of the tree (as built for mozilla-central " 502 "or mozilla-inbound)", 503 ) 504 def load_image(args): 505 from taskgraph.docker import load_image_by_name, load_image_by_task_id 506 507 if not args.get("image_name") and not args.get("task_id"): 508 print("Specify either IMAGE-NAME or TASK-ID") 509 sys.exit(1) 510 try: 511 if args["task_id"]: 512 ok = load_image_by_task_id(args["task_id"], args.get("tag")) 513 else: 514 ok = load_image_by_name(args["image_name"], args.get("tag")) 515 if not ok: 516 sys.exit(1) 517 except Exception: 518 traceback.print_exc() 519 sys.exit(1) 520 521 522 @command("image-digest", help="Print the digest of a docker image.") 523 @argument( 524 "image_name", 525 help="Print the digest of the image of this name based on the current contents of the tree.", 526 ) 527 def image_digest(args): 528 from taskgraph.docker import get_image_digest 529 530 try: 531 digest = get_image_digest(args["image_name"]) 532 print(digest) 533 except Exception: 534 traceback.print_exc() 535 sys.exit(1) 536 537 538 @command("decision", help="Run the decision task") 539 @argument("--root", "-r", help="root of the taskgraph definition relative to topsrcdir") 540 @argument( 541 "--message", 542 required=False, 543 help=argparse.SUPPRESS, 544 ) 545 @argument( 546 "--project", 547 required=True, 548 help="Project to use for creating task graph. Example: --project=try", 549 ) 550 @argument("--pushlog-id", dest="pushlog_id", required=True, default="0") 551 @argument("--pushdate", dest="pushdate", required=True, type=int, default=0) 552 @argument("--owner", required=True, help="email address of who owns this graph") 553 @argument("--level", required=True, help="SCM level of this repository") 554 @argument( 555 "--target-tasks-method", help="method for selecting the target tasks to generate" 556 ) 557 @argument( 558 "--repository-type", 559 required=True, 560 help='Type of repository, either "hg" or "git"', 561 ) 562 @argument("--base-repository", required=True, help='URL for "base" repository to clone') 563 @argument( 564 "--base-ref", default="", help='Reference of the revision in the "base" repository' 565 ) 566 @argument( 567 "--base-rev", 568 default="", 569 help="Taskgraph decides what to do based on the revision range between " 570 "`--base-rev` and `--head-rev`. Value is determined automatically if not provided", 571 ) 572 @argument( 573 "--head-repository", 574 required=True, 575 help='URL for "head" repository to fetch revision from', 576 ) 577 @argument( 578 "--head-ref", required=True, help="Reference (this is same as rev usually for hg)" 579 ) 580 @argument( 581 "--head-rev", required=True, help="Commit revision to use from head repository" 582 ) 583 @argument("--head-tag", help="Tag attached to the revision", default="") 584 @argument( 585 "--tasks-for", required=True, help="the tasks_for value used to generate this task" 586 ) 587 @argument("--try-task-config-file", help="path to try task configuration file") 588 @argument( 589 "--no-verify", 590 dest="verify", 591 default=True, 592 action="store_false", 593 help="Skip graph verifications.", 594 ) 595 def decision(options): 596 from gecko_taskgraph.decision import taskgraph_decision 597 598 taskgraph_decision(options) 599 600 601 @command("action-callback", description="Run action callback used by action tasks") 602 @argument( 603 "--root", 604 "-r", 605 default="taskcluster", 606 help="root of the taskgraph definition relative to topsrcdir", 607 ) 608 def action_callback(options): 609 from taskgraph.util import json 610 611 from gecko_taskgraph.actions import trigger_action_callback 612 from gecko_taskgraph.actions.util import get_parameters 613 614 try: 615 # the target task for this action (or null if it's a group action) 616 task_id = json.loads(os.environ.get("ACTION_TASK_ID", "null")) 617 # the target task group for this action 618 task_group_id = os.environ.get("ACTION_TASK_GROUP_ID", None) 619 input = json.loads(os.environ.get("ACTION_INPUT", "null")) 620 callback = os.environ.get("ACTION_CALLBACK", None) 621 root = options["root"] 622 623 parameters = get_parameters(task_group_id) 624 625 return trigger_action_callback( 626 task_group_id=task_group_id, 627 task_id=task_id, 628 input=input, 629 callback=callback, 630 parameters=parameters, 631 root=root, 632 test=False, 633 ) 634 except Exception: 635 traceback.print_exc() 636 sys.exit(1) 637 638 639 @command("test-action-callback", description="Run an action callback in a testing mode") 640 @argument( 641 "--root", 642 "-r", 643 default="taskcluster", 644 help="root of the taskgraph definition relative to topsrcdir", 645 ) 646 @argument( 647 "--parameters", 648 "-p", 649 default="", 650 help="parameters file (.yml or .json; see `taskcluster/docs/parameters.rst`)`", 651 ) 652 @argument("--task-id", default=None, help="TaskId to which the action applies") 653 @argument( 654 "--task-group-id", default=None, help="TaskGroupId to which the action applies" 655 ) 656 @argument("--input", default=None, help="Action input (.yml or .json)") 657 @argument("callback", default=None, help="Action callback name (Python function name)") 658 def test_action_callback(options): 659 import taskgraph.parameters 660 from taskgraph.config import load_graph_config 661 from taskgraph.util import json, yaml 662 663 import gecko_taskgraph.actions 664 665 def load_data(filename): 666 with open(filename) as f: 667 if filename.endswith(".yml"): 668 return yaml.load_stream(f) 669 if filename.endswith(".json"): 670 return json.load(f) 671 raise Exception(f"unknown filename {filename}") 672 673 try: 674 task_id = options["task_id"] 675 676 if options["input"]: 677 input = load_data(options["input"]) 678 else: 679 input = None 680 681 root = options["root"] 682 graph_config = load_graph_config(root) 683 trust_domain = graph_config["trust-domain"] 684 graph_config.register() 685 686 parameters = taskgraph.parameters.load_parameters_file( 687 options["parameters"], strict=False, trust_domain=trust_domain 688 ) 689 parameters.check() 690 691 return gecko_taskgraph.actions.trigger_action_callback( 692 task_group_id=options["task_group_id"], 693 task_id=task_id, 694 input=input, 695 callback=options["callback"], 696 parameters=parameters, 697 root=root, 698 test=True, 699 ) 700 except Exception: 701 traceback.print_exc() 702 sys.exit(1) 703 704 705 def create_parser(): 706 parser = argparse.ArgumentParser(description="Interact with taskgraph") 707 subparsers = parser.add_subparsers() 708 for _, (func, args, kwargs, defaults) in commands.items(): 709 subparser = subparsers.add_parser(*args, **kwargs) 710 for arg in func.args: 711 subparser.add_argument(*arg[0], **arg[1]) 712 subparser.set_defaults(command=func, **defaults) 713 return parser 714 715 716 def setup_logging(): 717 logging.basicConfig( 718 format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO 719 ) 720 721 722 def main(args=sys.argv[1:]): 723 setup_logging() 724 parser = create_parser() 725 args = parser.parse_args(args) 726 try: 727 args.command(vars(args)) 728 except Exception: 729 traceback.print_exc() 730 sys.exit(1)