Source code
Revision control
Copy as Markdown
Other Tools
# 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
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.main import commands as taskgraph_commands
from mach.decorators import Command, CommandArgument, SubCommand
from mach.util import strtobool
from mozsystemmonitor.resourcemonitor import SystemResourceMonitor
logger = logging.getLogger("taskcluster")
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,
)
# 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",
"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)