tor-browser

mach_commands.py (13620B)

---

# 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/.


import argparse
import json
import logging
import os
import pathlib
import sys
import time
import traceback
from functools import partial

import gecko_taskgraph.main
from gecko_taskgraph.files_changed import get_locally_changed_files
from gecko_taskgraph.main import commands as taskgraph_commands
from mach.decorators import Command, CommandArgument, SubCommand
from mach.util import strtobool
from mozsystemmonitor.resourcemonitor import SystemResourceMonitor

# We're likely going to need the result of get_locally_changed_files, and it
# takes time to finish, so prefetch it as soon as possible.
get_locally_changed_files.preload(os.getcwd())


def setup_logging(command_context, quiet=False, verbose=True):
   """
   Set up Python logging for all loggers, sending results to stderr (so
   that command output can be redirected easily) and adding the typical
   mach timestamp.
   """
   # remove the old terminal handler
   old = command_context.log_manager.replace_terminal_handler(None)

   # re-add it, with level and fh set appropriately
   if not quiet:
       level = logging.DEBUG if verbose else logging.INFO
       command_context.log_manager.add_terminal_logging(
           fh=sys.stderr,
           level=level,
           write_interval=old.formatter.write_interval,
           write_times=old.formatter.write_times,
       )
       logging.getLogger("taskcluster").setLevel(logging.INFO)

   # all of the taskgraph logging is unstructured logging
   command_context.log_manager.enable_unstructured()


def get_taskgraph_command_parser(name):
   """Given a command name, obtain its argument parser.

   Args:
       name (str): Name of the command.

   Returns:
       ArgumentParser: An ArgumentParser instance.
   """
   command = taskgraph_commands[name]
   parser = argparse.ArgumentParser()
   for arg in command.func.args:
       parser.add_argument(*arg[0], **arg[1])

   parser.set_defaults(func=command.func, **command.defaults)
   return parser


def get_taskgraph_decision_parser():
   parser = get_taskgraph_command_parser("decision")

   extra_args = [
       (
           ["--optimize-target-tasks"],
           {
               "type": lambda flag: bool(strtobool(flag)),
               "nargs": "?",
               "const": "true",
               "help": "If specified, this indicates whether the target "
               "tasks are eligible for optimization. Otherwise, the default "
               "for the project is used.",
           },
       ),
       (
           ["--include-push-tasks"],
           {
               "action": "store_true",
               "help": "Whether tasks from the on-push graph should be re-used "
               "in this graph. This allows cron graphs to avoid rebuilding "
               "jobs that were built on-push.",
           },
       ),
       (
           ["--rebuild-kind"],
           {
               "dest": "rebuild_kinds",
               "action": "append",
               "default": argparse.SUPPRESS,
               "help": "Kinds that should not be re-used from the on-push graph.",
           },
       ),
   ]
   for arg in extra_args:
       parser.add_argument(*arg[0], **arg[1])

   return parser


@Command(
   "taskgraph",
   category="ci",
   description="Manipulate TaskCluster task graphs defined in-tree",
)
def taskgraph_command(command_context):
   """The taskgraph subcommands all relate to the generation of task graphs
   for Gecko continuous integration.  A task graph is a set of tasks linked
   by dependencies: for example, a binary must be built before it is tested,
   and that build may further depend on various toolchains, libraries, etc.
   """


@SubCommand(
   "taskgraph",
   "kind-graph",
   description="Generate a graph of the relationship between taskgraph kinds",
   parser=partial(get_taskgraph_command_parser, "kind-graph"),
)
def taskgraph_kind_graph(command_context, **options):
   try:
       setup_logging(command_context)
       return taskgraph_commands["kind-graph"].func(options)
   except Exception:
       traceback.print_exc()
       sys.exit(1)


@SubCommand(
   "taskgraph",
   "tasks",
   description="Show all tasks in the taskgraph",
   parser=partial(get_taskgraph_command_parser, "tasks"),
)
def taskgraph_tasks(command_context, **options):
   return run_show_taskgraph(command_context, **options)


@SubCommand(
   "taskgraph",
   "full",
   description="Show the full taskgraph",
   parser=partial(get_taskgraph_command_parser, "full"),
)
def taskgraph_full(command_context, **options):
   return run_show_taskgraph(command_context, **options)


@SubCommand(
   "taskgraph",
   "target",
   description="Show the target task set",
   parser=partial(get_taskgraph_command_parser, "target"),
)
def taskgraph_target(command_context, **options):
   return run_show_taskgraph(command_context, **options)


@SubCommand(
   "taskgraph",
   "target-graph",
   description="Show the target taskgraph",
   parser=partial(get_taskgraph_command_parser, "target-graph"),
)
def taskgraph_target_graph(command_context, **options):
   return run_show_taskgraph(command_context, **options)


@SubCommand(
   "taskgraph",
   "optimized",
   description="Show the optimized taskgraph",
   parser=partial(get_taskgraph_command_parser, "optimized"),
)
def taskgraph_optimized(command_context, **options):
   return run_show_taskgraph(command_context, **options)


@SubCommand(
   "taskgraph",
   "morphed",
   description="Show the morphed taskgraph",
   parser=partial(get_taskgraph_command_parser, "morphed"),
)
def taskgraph_morphed(command_context, **options):
   return run_show_taskgraph(command_context, **options)


def run_show_taskgraph(command_context, **options):
   # There are cases where we don't want to set up mach logging (e.g logs
   # are being redirected to disk). By monkeypatching the 'setup_logging'
   # function we can let 'taskgraph.main' decide whether or not to log to
   # the terminal.
   gecko_taskgraph.main.setup_logging = partial(
       setup_logging,
       command_context,
       quiet=options["quiet"],
       verbose=options["verbose"],
   )
   show_taskgraph = options.pop("func")
   return show_taskgraph(options)


@SubCommand("taskgraph", "actions", description="Write actions.json to stdout")
@CommandArgument(
   "--root", "-r", help="root of the taskgraph definition relative to topsrcdir"
)
@CommandArgument(
   "--quiet", "-q", action="store_true", help="suppress all logging output"
)
@CommandArgument(
   "--verbose",
   "-v",
   action="store_true",
   help="include debug-level logging output",
)
@CommandArgument(
   "--parameters",
   "-p",
   default="project=mozilla-central",
   help="parameters file (.yml or .json; see `taskcluster/docs/parameters.rst`)`",
)
def taskgraph_actions(command_context, **options):
   import gecko_taskgraph
   import gecko_taskgraph.actions
   from taskgraph.generator import TaskGraphGenerator
   from taskgraph.parameters import parameters_loader

   try:
       setup_logging(
           command_context, quiet=options["quiet"], verbose=options["verbose"]
       )
       parameters = parameters_loader(options["parameters"])

       tgg = TaskGraphGenerator(
           root_dir=options.get("root"),
           parameters=parameters,
       )

       actions = gecko_taskgraph.actions.render_actions_json(
           tgg.parameters,
           tgg.graph_config,
           decision_task_id="DECISION-TASK",
       )
       print(json.dumps(actions, sort_keys=True, indent=2, separators=(",", ": ")))
   except Exception:
       traceback.print_exc()
       sys.exit(1)


@SubCommand(
   "taskgraph",
   "decision",
   description="Run the decision task",
   parser=get_taskgraph_decision_parser,
)
def taskgraph_decision(command_context, **options):
   """Run the decision task: generate a task graph and submit to
   TaskCluster.  This is only meant to be called within decision tasks,
   and requires a great many arguments.  Commands like `mach taskgraph
   optimized` are better suited to use on the command line, and can take
   the parameters file generated by a decision task."""
   try:
       setup_logging(command_context)

       in_automation = os.environ.get("MOZ_AUTOMATION") == "1"
       moz_upload_dir = os.environ.get("MOZ_UPLOAD_DIR")
       if in_automation and moz_upload_dir:
           monitor = SystemResourceMonitor(poll_interval=0.1)
           monitor.start()
       else:
           monitor = None

       try:
           start = time.monotonic()
           ret = taskgraph_commands["decision"].func(options)
           end = time.monotonic()
       finally:
           if monitor is not None:
               monitor.stop()
               upload_dir = pathlib.Path(moz_upload_dir)
               profile_path = upload_dir / "profile_build_resources.json"
               with open(profile_path, "w", encoding="utf-8", newline="\n") as f:
                   to_write = json.dumps(monitor.as_profile(), separators=(",", ":"))
                   f.write(to_write)

       if in_automation:
           perfherder_data = {
               "framework": {"name": "build_metrics"},
               "suites": [
                   {
                       "name": "decision",
                       "value": end - start,
                       "lowerIsBetter": True,
                       "shouldAlert": True,
                       "subtests": [],
                   }
               ],
           }
           print(
               f"PERFHERDER_DATA: {json.dumps(perfherder_data)}",
               file=sys.stderr,
           )

           if moz_upload_dir:
               upload_dir = pathlib.Path(moz_upload_dir)
               out_path = upload_dir / "perfherder-data-decision.json"
               with out_path.open("w", encoding="utf-8") as f:
                   json.dump(perfherder_data, f)

       return ret
   except Exception:
       traceback.print_exc()
       sys.exit(1)


@SubCommand(
   "taskgraph",
   "action-callback",
   description="Run action callback used by action tasks",
   parser=partial(get_taskgraph_command_parser, "action-callback"),
)
def action_callback(command_context, **options):
   setup_logging(command_context)
   taskgraph_commands["action-callback"].func(options)


@SubCommand(
   "taskgraph",
   "test-action-callback",
   description="Run an action callback in a testing mode",
   parser=partial(get_taskgraph_command_parser, "test-action-callback"),
)
def test_action_callback(command_context, **options):
   setup_logging(command_context)

   if not options["parameters"]:
       options["parameters"] = "project=mozilla-central"

   taskgraph_commands["test-action-callback"].func(options)


@SubCommand(
   "taskgraph",
   "load-image",
   description="Load a pre-built Docker image. Note that you need to "
   "have docker installed and running for this to work.",
   parser=partial(get_taskgraph_command_parser, "load-image"),
)
def load_image(command_context, **kwargs):
   setup_logging(command_context)
   taskgraph_commands["load-image"].func(kwargs)


@SubCommand(
   "taskgraph",
   "build-image",
   description="Build a Docker image",
   parser=partial(get_taskgraph_command_parser, "build-image"),
)
def build_image(command_context, **kwargs):
   setup_logging(command_context)
   try:
       taskgraph_commands["build-image"].func(kwargs)
   except Exception:
       traceback.print_exc()
       sys.exit(1)


@SubCommand(
   "taskgraph",
   "image-digest",
   description="Print the digest of the image of this name based on the "
   "current contents of the tree.",
   parser=partial(get_taskgraph_command_parser, "build-image"),
)
def image_digest(command_context, **kwargs):
   setup_logging(command_context)
   taskgraph_commands["image-digest"].func(kwargs)


@SubCommand(
   "taskgraph",
   "load-task",
   description="Loads a pre-built Docker image and drops you into a container with "
   "the same environment variables and run-task setup as the specified task. "
   "The task's payload.command will be replaced with 'bash'. You need to have "
   "docker installed and running for this to work.",
   parser=partial(get_taskgraph_command_parser, "load-task"),
)
def load_task(command_context, **kwargs):
   setup_logging(command_context)
   taskgraph_commands["load-task"].func(kwargs)


@Command(
   "release-history",
   category="ci",
   description="Query balrog for release history used by enable partials generation",
)
@CommandArgument(
   "-b",
   "--branch",
   help="The gecko project branch used in balrog, such as "
   "mozilla-central, release, maple",
)
@CommandArgument(
   "--product", default="Firefox", help="The product identifier, such as 'Firefox'"
)
def generate_partials_builds(command_context, product, branch):
   from gecko_taskgraph.util.partials import populate_release_history

   try:
       import yaml

       release_history = {"release_history": populate_release_history(product, branch)}
       print(
           yaml.safe_dump(
               release_history, allow_unicode=True, default_flow_style=False
           )
       )
   except Exception:
       traceback.print_exc()
       sys.exit(1)