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