tools/tryselect/task_config.py

   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 """
   6 Templates provide a way of modifying the task definition of selected tasks.
   7 They are added to 'try_task_config.json' and processed by the transforms.
   8 """
   9
  10
  11 import json
  12 import os
  13 import pathlib
  14 import subprocess
  15 import sys
  16 from abc import ABCMeta, abstractmethod, abstractproperty
  17 from argparse import SUPPRESS, Action
  18 from contextlib import contextmanager
  19 from textwrap import dedent
  20
  21 import mozpack.path as mozpath
  22 import requests
  23 import six
  24 from mozbuild.base import BuildEnvironmentNotFoundException, MozbuildObject
  25 from mozversioncontrol import Repository
  26 from taskgraph.util import taskcluster
  27
  28 from .tasks import resolve_tests_by_suite
  29 from .util.ssh import get_ssh_user
  30
  31 here = pathlib.Path(__file__).parent
  32 build = MozbuildObject.from_environment(cwd=str(here))
  33
  34
  35 @contextmanager
  36 def try_config_commit(vcs: Repository, commit_message: str):
  37     """Context manager that creates and removes a try config commit."""
  38     # Add the `try_task_config.json` file if it exists.
  39     try_task_config_path = pathlib.Path(build.topsrcdir) / "try_task_config.json"
  40     if try_task_config_path.exists():
  41         vcs.add_remove_files("try_task_config.json")
  42
  43     try:
  44         # Create a try config commit.
  45         vcs.create_try_commit(commit_message)
  46
  47         yield
  48     finally:
  49         # Revert the try config commit.
  50         vcs.remove_current_commit()
  51
  52
  53 class ParameterConfig:
  54     __metaclass__ = ABCMeta
  55
  56     def __init__(self):
  57         self.dests = set()
  58
  59     def add_arguments(self, parser):
  60         for cli, kwargs in self.arguments:
  61             action = parser.add_argument(*cli, **kwargs)
  62             self.dests.add(action.dest)
  63
  64     @abstractproperty
  65     def arguments(self):
  66         pass
  67
  68     @abstractmethod
  69     def get_parameters(self, **kwargs) -> dict:
  70         pass
  71
  72     def validate(self, **kwargs):
  73         pass
  74
  75
  76 class TryConfig(ParameterConfig):
  77     @abstractmethod
  78     def try_config(self, **kwargs) -> dict:
  79         pass
  80
  81     def get_parameters(self, **kwargs):
  82         result = self.try_config(**kwargs)
  83         if not result:
  84             return None
  85         return {"try_task_config": result}
  86
  87
  88 class Artifact(TryConfig):
  89     arguments = [
  90         [
  91             ["--artifact"],
  92             {"action": "store_true", "help": "Force artifact builds where possible."},
  93         ],
  94         [
  95             ["--no-artifact"],
  96             {
  97                 "action": "store_true",
  98                 "help": "Disable artifact builds even if being used locally.",
  99             },
 100         ],
 101     ]
 102
 103     def add_arguments(self, parser):
 104         group = parser.add_mutually_exclusive_group()
 105         return super().add_arguments(group)
 106
 107     @classmethod
 108     def is_artifact_build(cls):
 109         try:
 110             return build.substs.get("MOZ_ARTIFACT_BUILDS", False)
 111         except BuildEnvironmentNotFoundException:
 112             return False
 113
 114     def try_config(self, artifact, no_artifact, **kwargs):
 115         if artifact:
 116             return {"use-artifact-builds": True, "disable-pgo": True}
 117
 118         if no_artifact:
 119             return
 120
 121         if self.is_artifact_build():
 122             print("Artifact builds enabled, pass --no-artifact to disable")
 123             return {"use-artifact-builds": True, "disable-pgo": True}
 124
 125
 126 class Pernosco(TryConfig):
 127     arguments = [
 128         [
 129             ["--pernosco"],
 130             {
 131                 "action": "store_true",
 132                 "default": None,
 133                 "help": "Opt-in to analysis by the Pernosco debugging service.",
 134             },
 135         ],
 136         [
 137             ["--no-pernosco"],
 138             {
 139                 "dest": "pernosco",
 140                 "action": "store_false",
 141                 "default": None,
 142                 "help": "Opt-out of the Pernosco debugging service (if you are on the include list).",
 143             },
 144         ],
 145     ]
 146
 147     def add_arguments(self, parser):
 148         group = parser.add_mutually_exclusive_group()
 149         return super().add_arguments(group)
 150
 151     def try_config(self, pernosco, **kwargs):
 152         if pernosco is None:
 153             return
 154
 155         if pernosco:
 156             try:
 157                 # The Pernosco service currently requires a Mozilla e-mail address to
 158                 # log in. Prevent people with non-Mozilla addresses from using this
 159                 # flag so they don't end up consuming time and resources only to
 160                 # realize they can't actually log in and see the reports.
 161                 address = get_ssh_user()
 162                 if not address.endswith("@mozilla.com"):
 163                     print(
 164                         dedent(
 165                             """\
 166                         Pernosco requires a Mozilla e-mail address to view its reports. Please
 167                         push to try with an @mozilla.com address to use --pernosco.
 168
 169                             Current user: {}
 170                     """.format(
 171                                 address
 172                             )
 173                         )
 174                     )
 175                     sys.exit(1)
 176
 177             except (subprocess.CalledProcessError, IndexError):
 178                 print("warning: failed to detect current user for 'hg.mozilla.org'")
 179                 print("Pernosco requires a Mozilla e-mail address to view its reports.")
 180                 while True:
 181                     answer = input(
 182                         "Do you have an @mozilla.com address? [Y/n]: "
 183                     ).lower()
 184                     if answer == "n":
 185                         sys.exit(1)
 186                     elif answer == "y":
 187                         break
 188
 189         return {
 190             "env": {
 191                 "PERNOSCO": str(int(pernosco)),
 192             }
 193         }
 194
 195     def validate(self, **kwargs):
 196         try_config = kwargs["try_config_params"].get("try_task_config") or {}
 197         if try_config.get("use-artifact-builds"):
 198             print(
 199                 "Pernosco does not support artifact builds at this time. "
 200                 "Please try again with '--no-artifact'."
 201             )
 202             sys.exit(1)
 203
 204
 205 class Path(TryConfig):
 206     arguments = [
 207         [
 208             ["paths"],
 209             {
 210                 "nargs": "*",
 211                 "default": [],
 212                 "help": "Run tasks containing tests under the specified path(s).",
 213             },
 214         ],
 215     ]
 216
 217     def try_config(self, paths, **kwargs):
 218         if not paths:
 219             return
 220
 221         for p in paths:
 222             if not os.path.exists(p):
 223                 print("error: '{}' is not a valid path.".format(p), file=sys.stderr)
 224                 sys.exit(1)
 225
 226         paths = [
 227             mozpath.relpath(mozpath.join(os.getcwd(), p), build.topsrcdir)
 228             for p in paths
 229         ]
 230         return {
 231             "env": {
 232                 "MOZHARNESS_TEST_PATHS": six.ensure_text(
 233                     json.dumps(resolve_tests_by_suite(paths))
 234                 ),
 235             }
 236         }
 237
 238
 239 class Environment(TryConfig):
 240     arguments = [
 241         [
 242             ["--env"],
 243             {
 244                 "action": "append",
 245                 "default": None,
 246                 "help": "Set an environment variable, of the form FOO=BAR. "
 247                 "Can be passed in multiple times.",
 248             },
 249         ],
 250     ]
 251
 252     def try_config(self, env, **kwargs):
 253         if not env:
 254             return
 255         return {
 256             "env": dict(e.split("=", 1) for e in env),
 257         }
 258
 259
 260 class ExistingTasks(ParameterConfig):
 261     TREEHERDER_PUSH_ENDPOINT = (
 262         "https://treeherder.mozilla.org/api/project/try/push/?count=1&author={user}"
 263     )
 264     TREEHERDER_PUSH_URL = (
 265         "https://treeherder.mozilla.org/jobs?repo={branch}&revision={revision}"
 266     )
 267
 268     arguments = [
 269         [
 270             ["-E", "--use-existing-tasks"],
 271             {
 272                 "const": "last_try_push",
 273                 "default": None,
 274                 "nargs": "?",
 275                 "help": """
 276                     Use existing tasks from a previous push. Without args this
 277                 uses your most recent try push. You may also specify
 278                 `rev=<revision>` where <revision> is the head revision of the
 279                 try push or `task-id=<task id>` where <task id> is the Decision
 280                 task id of the push. This last method even works for non-try
 281                 branches.
 282                 """,
 283             },
 284         ]
 285     ]
 286
 287     def find_decision_task(self, use_existing_tasks):
 288         branch = "try"
 289         if use_existing_tasks == "last_try_push":
 290             # Use existing tasks from user's previous try push.
 291             user = get_ssh_user()
 292             url = self.TREEHERDER_PUSH_ENDPOINT.format(user=user)
 293             res = requests.get(url, headers={"User-Agent": "gecko-mach-try/1.0"})
 294             res.raise_for_status()
 295             data = res.json()
 296             if data["meta"]["count"] == 0:
 297                 raise Exception(f"Could not find a try push for '{user}'!")
 298             revision = data["results"][0]["revision"]
 299
 300         elif use_existing_tasks.startswith("rev="):
 301             revision = use_existing_tasks[len("rev=") :]
 302
 303         else:
 304             raise Exception("Unable to parse '{use_existing_tasks}'!")
 305
 306         url = self.TREEHERDER_PUSH_URL.format(branch=branch, revision=revision)
 307         print(f"Using existing tasks from: {url}")
 308         index_path = f"gecko.v2.{branch}.revision.{revision}.taskgraph.decision"
 309         return taskcluster.find_task_id(index_path)
 310
 311     def get_parameters(self, use_existing_tasks, **kwargs):
 312         if not use_existing_tasks:
 313             return
 314
 315         if use_existing_tasks.startswith("task-id="):
 316             tid = use_existing_tasks[len("task-id=") :]
 317         else:
 318             tid = self.find_decision_task(use_existing_tasks)
 319
 320         label_to_task_id = taskcluster.get_artifact(tid, "public/label-to-taskid.json")
 321         return {"existing_tasks": label_to_task_id}
 322
 323
 324 class RangeAction(Action):
 325     def __init__(self, min, max, *args, **kwargs):
 326         self.min = min
 327         self.max = max
 328         kwargs["metavar"] = "[{}-{}]".format(self.min, self.max)
 329         super().__init__(*args, **kwargs)
 330
 331     def __call__(self, parser, namespace, values, option_string=None):
 332         name = option_string or self.dest
 333         if values < self.min:
 334             parser.error("{} can not be less than {}".format(name, self.min))
 335         if values > self.max:
 336             parser.error("{} can not be more than {}".format(name, self.max))
 337         setattr(namespace, self.dest, values)
 338
 339
 340 class Rebuild(TryConfig):
 341     arguments = [
 342         [
 343             ["--rebuild"],
 344             {
 345                 "action": RangeAction,
 346                 "min": 2,
 347                 "max": 20,
 348                 "default": None,
 349                 "type": int,
 350                 "help": "Rebuild all selected tasks the specified number of times.",
 351             },
 352         ],
 353     ]
 354
 355     def try_config(self, rebuild, **kwargs):
 356         if not rebuild:
 357             return
 358
 359         if (
 360             not kwargs.get("new_test_config", False)
 361             and kwargs.get("full")
 362             and rebuild > 3
 363         ):
 364             print(
 365                 "warning: limiting --rebuild to 3 when using --full. "
 366                 "Use custom push actions to add more."
 367             )
 368             rebuild = 3
 369
 370         return {
 371             "rebuild": rebuild,
 372         }
 373
 374
 375 class Routes(TryConfig):
 376     arguments = [
 377         [
 378             ["--route"],
 379             {
 380                 "action": "append",
 381                 "dest": "routes",
 382                 "help": (
 383                     "Additional route to add to the tasks "
 384                     "(note: these will not be added to the decision task)"
 385                 ),
 386             },
 387         ],
 388     ]
 389
 390     def try_config(self, routes, **kwargs):
 391         if routes:
 392             return {
 393                 "routes": routes,
 394             }
 395
 396
 397 class ChemspillPrio(TryConfig):
 398     arguments = [
 399         [
 400             ["--chemspill-prio"],
 401             {
 402                 "action": "store_true",
 403                 "help": "Run at a higher priority than most try jobs (chemspills only).",
 404             },
 405         ],
 406     ]
 407
 408     def try_config(self, chemspill_prio, **kwargs):
 409         if chemspill_prio:
 410             return {"chemspill-prio": True}
 411
 412
 413 class GeckoProfile(TryConfig):
 414     arguments = [
 415         [
 416             ["--gecko-profile"],
 417             {
 418                 "dest": "profile",
 419                 "action": "store_true",
 420                 "default": False,
 421                 "help": "Create and upload a gecko profile during talos/raptor tasks.",
 422             },
 423         ],
 424         [
 425             ["--gecko-profile-interval"],
 426             {
 427                 "dest": "gecko_profile_interval",
 428                 "type": float,
 429                 "help": "How frequently to take samples (ms)",
 430             },
 431         ],
 432         [
 433             ["--gecko-profile-entries"],
 434             {
 435                 "dest": "gecko_profile_entries",
 436                 "type": int,
 437                 "help": "How many samples to take with the profiler",
 438             },
 439         ],
 440         [
 441             ["--gecko-profile-features"],
 442             {
 443                 "dest": "gecko_profile_features",
 444                 "type": str,
 445                 "default": None,
 446                 "help": "Set the features enabled for the profiler.",
 447             },
 448         ],
 449         [
 450             ["--gecko-profile-threads"],
 451             {
 452                 "dest": "gecko_profile_threads",
 453                 "type": str,
 454                 "help": "Comma-separated list of threads to sample.",
 455             },
 456         ],
 457         # For backwards compatibility
 458         [
 459             ["--talos-profile"],
 460             {
 461                 "dest": "profile",
 462                 "action": "store_true",
 463                 "default": False,
 464                 "help": SUPPRESS,
 465             },
 466         ],
 467         # This is added for consistency with the 'syntax' selector
 468         [
 469             ["--geckoProfile"],
 470             {
 471                 "dest": "profile",
 472                 "action": "store_true",
 473                 "default": False,
 474                 "help": SUPPRESS,
 475             },
 476         ],
 477     ]
 478
 479     def try_config(
 480         self,
 481         profile,
 482         gecko_profile_interval,
 483         gecko_profile_entries,
 484         gecko_profile_features,
 485         gecko_profile_threads,
 486         **kwargs,
 487     ):
 488         if profile or not all(
 489             s is None for s in (gecko_profile_features, gecko_profile_threads)
 490         ):
 491             cfg = {
 492                 "gecko-profile": True,
 493                 "gecko-profile-interval": gecko_profile_interval,
 494                 "gecko-profile-entries": gecko_profile_entries,
 495                 "gecko-profile-features": gecko_profile_features,
 496                 "gecko-profile-threads": gecko_profile_threads,
 497             }
 498             return {key: value for key, value in cfg.items() if value is not None}
 499
 500
 501 class Browsertime(TryConfig):
 502     arguments = [
 503         [
 504             ["--browsertime"],
 505             {
 506                 "action": "store_true",
 507                 "help": "Use browsertime during Raptor tasks.",
 508             },
 509         ],
 510     ]
 511
 512     def try_config(self, browsertime, **kwargs):
 513         if browsertime:
 514             return {
 515                 "browsertime": True,
 516             }
 517
 518
 519 class DisablePgo(TryConfig):
 520     arguments = [
 521         [
 522             ["--disable-pgo"],
 523             {
 524                 "action": "store_true",
 525                 "help": "Don't run PGO builds",
 526             },
 527         ],
 528     ]
 529
 530     def try_config(self, disable_pgo, **kwargs):
 531         if disable_pgo:
 532             return {
 533                 "disable-pgo": True,
 534             }
 535
 536
 537 class NewConfig(TryConfig):
 538     arguments = [
 539         [
 540             ["--new-test-config"],
 541             {
 542                 "action": "store_true",
 543                 "help": "When a test fails (mochitest only) restart the browser and start from the next test",
 544             },
 545         ],
 546     ]
 547
 548     def try_config(self, new_test_config, **kwargs):
 549         if new_test_config:
 550             return {
 551                 "new-test-config": True,
 552             }
 553
 554
 555 class WorkerOverrides(TryConfig):
 556     arguments = [
 557         [
 558             ["--worker-override"],
 559             {
 560                 "action": "append",
 561                 "dest": "worker_overrides",
 562                 "help": (
 563                     "Override the worker pool used for a given taskgraph worker alias. "
 564                     "The argument should be `<alias>=<worker-pool>`. "
 565                     "Can be specified multiple times."
 566                 ),
 567             },
 568         ],
 569         [
 570             ["--worker-suffix"],
 571             {
 572                 "action": "append",
 573                 "dest": "worker_suffixes",
 574                 "help": (
 575                     "Override the worker pool used for a given taskgraph worker alias, "
 576                     "by appending a suffix to the work-pool. "
 577                     "The argument should be `<alias>=<suffix>`. "
 578                     "Can be specified multiple times."
 579                 ),
 580             },
 581         ],
 582         [
 583             ["--worker-type"],
 584             {
 585                 "action": "append",
 586                 "dest": "worker_types",
 587                 "default": [],
 588                 "help": "Select tasks that only run on the specified worker.",
 589             },
 590         ],
 591     ]
 592
 593     def try_config(self, worker_overrides, worker_suffixes, worker_types, **kwargs):
 594         from gecko_taskgraph.util.workertypes import get_worker_type
 595         from taskgraph.config import load_graph_config
 596
 597         overrides = {}
 598         if worker_overrides:
 599             for override in worker_overrides:
 600                 alias, worker_pool = override.split("=", 1)
 601                 if alias in overrides:
 602                     print(
 603                         "Can't override worker alias {alias} more than once. "
 604                         "Already set to use {previous}, but also asked to use {new}.".format(
 605                             alias=alias, previous=overrides[alias], new=worker_pool
 606                         )
 607                     )
 608                     sys.exit(1)
 609                 overrides[alias] = worker_pool
 610
 611         if worker_suffixes:
 612             root = build.topsrcdir
 613             root = os.path.join(root, "taskcluster", "ci")
 614             graph_config = load_graph_config(root)
 615             for worker_suffix in worker_suffixes:
 616                 alias, suffix = worker_suffix.split("=", 1)
 617                 if alias in overrides:
 618                     print(
 619                         "Can't override worker alias {alias} more than once. "
 620                         "Already set to use {previous}, but also asked "
 621                         "to add suffix {suffix}.".format(
 622                             alias=alias, previous=overrides[alias], suffix=suffix
 623                         )
 624                     )
 625                     sys.exit(1)
 626                 provisioner, worker_type = get_worker_type(
 627                     graph_config, worker_type=alias, parameters={"level": "1"}
 628                 )
 629                 overrides[alias] = "{provisioner}/{worker_type}{suffix}".format(
 630                     provisioner=provisioner, worker_type=worker_type, suffix=suffix
 631                 )
 632
 633         retVal = {}
 634         if worker_types:
 635             retVal["worker-types"] = list(overrides.keys()) + worker_types
 636
 637         if overrides:
 638             retVal["worker-overrides"] = overrides
 639         return retVal
 640
 641
 642 all_task_configs = {
 643     "artifact": Artifact,
 644     "browsertime": Browsertime,
 645     "chemspill-prio": ChemspillPrio,
 646     "disable-pgo": DisablePgo,
 647     "env": Environment,
 648     "existing-tasks": ExistingTasks,
 649     "gecko-profile": GeckoProfile,
 650     "new-test-config": NewConfig,
 651     "path": Path,
 652     "pernosco": Pernosco,
 653     "rebuild": Rebuild,
 654     "routes": Routes,
 655     "worker-overrides": WorkerOverrides,
 656 }