From 7e9c59f4785a431b86a27e0945ecd5c803c2d743 Mon Sep 17 00:00:00 2001 From: Andrew Halberstadt Date: Sat, 3 Feb 2024 02:29:38 +0000 Subject: [PATCH] Bug 1752677 - Add a task config to set the 'existing_tasks' parameter, r=releng-reviewers,bhearsum,jcristau This adds the `-E/--use-existing-tasks` flag to the `fuzzy` or `chooser` selectors. For example: ./mach try fuzzy -E This queries the Treeherder API to find your latest try push, then any tasks that exist on that push, but aren't explicitly selected this time, will be replaced with those older tasks. You can also run: ./mach try fuzzy -E task-id= Where `` is the id of the Decision task from the push you would like to use as the base. Differential Revision: https://phabricator.services.mozilla.com/D200380 --- python/sites/python-test.txt | 1 + tools/tryselect/push.py | 9 ++ tools/tryselect/selectors/chooser/__init__.py | 1 + tools/tryselect/selectors/fuzzy.py | 1 + tools/tryselect/task_config.py | 82 ++++++++++-- tools/tryselect/test/conftest.py | 7 + tools/tryselect/test/python.toml | 2 + tools/tryselect/test/test_push.py | 54 ++++++++ tools/tryselect/test/test_task_configs.py | 124 +++++++++++++++--- tools/tryselect/util/ssh.py | 24 ++++ 10 files changed, 278 insertions(+), 27 deletions(-) create mode 100644 tools/tryselect/test/test_push.py create mode 100644 tools/tryselect/util/ssh.py diff --git a/python/sites/python-test.txt b/python/sites/python-test.txt index c6cf70643dd7..58ad4f5efbce 100644 --- a/python/sites/python-test.txt +++ b/python/sites/python-test.txt @@ -1,4 +1,5 @@ pypi:pytest==7.0.1 +pypi:pytest-mock==3.12.0 pypi:Flask==2.1.3 # (indirect) avoids dependency on markupsafe >= 2.1.0, which is currently incompatible with glean-parser pypi:MarkupSafe==2.0.1 diff --git a/tools/tryselect/push.py b/tools/tryselect/push.py index 2b30c78e6aff..6b36cb42e194 100644 --- a/tools/tryselect/push.py +++ b/tools/tryselect/push.py @@ -97,8 +97,17 @@ def generate_try_task_config(method, labels, params=None, routes=None): # True). Their dependencies can be optimized though. params.setdefault("optimize_target_tasks", False) + # Remove selected labels from 'existing_tasks' parameter if present + if "existing_tasks" in params: + params["existing_tasks"] = { + label: tid + for label, tid in params["existing_tasks"].items() + if label not in labels + } + try_config = params.setdefault("try_task_config", {}) try_config.setdefault("env", {})["TRY_SELECTOR"] = method + try_config["tasks"] = sorted(labels) if routes: diff --git a/tools/tryselect/selectors/chooser/__init__.py b/tools/tryselect/selectors/chooser/__init__.py index 588d54302054..d6a32e08d032 100644 --- a/tools/tryselect/selectors/chooser/__init__.py +++ b/tools/tryselect/selectors/chooser/__init__.py @@ -31,6 +31,7 @@ class ChooserParser(BaseTryParser): "chemspill-prio", "disable-pgo", "env", + "existing-tasks", "gecko-profile", "path", "pernosco", diff --git a/tools/tryselect/selectors/fuzzy.py b/tools/tryselect/selectors/fuzzy.py index c4387ce2b1d4..e948100a0b01 100644 --- a/tools/tryselect/selectors/fuzzy.py +++ b/tools/tryselect/selectors/fuzzy.py @@ -114,6 +114,7 @@ class FuzzyParser(BaseTryParser): "chemspill-prio", "disable-pgo", "env", + "existing-tasks", "gecko-profile", "path", "pernosco", diff --git a/tools/tryselect/task_config.py b/tools/tryselect/task_config.py index ac1e244b71c3..c3124672a9a0 100644 --- a/tools/tryselect/task_config.py +++ b/tools/tryselect/task_config.py @@ -19,14 +19,17 @@ from contextlib import contextmanager from textwrap import dedent import mozpack.path as mozpath +import requests import six from mozbuild.base import BuildEnvironmentNotFoundException, MozbuildObject from mozversioncontrol import Repository +from taskgraph.util import taskcluster from .tasks import resolve_tests_by_suite +from .util.ssh import get_ssh_user -here = os.path.abspath(os.path.dirname(__file__)) -build = MozbuildObject.from_environment(cwd=here) +here = pathlib.Path(__file__).parent +build = MozbuildObject.from_environment(cwd=str(here)) @contextmanager @@ -155,13 +158,7 @@ class Pernosco(TryConfig): # log in. Prevent people with non-Mozilla addresses from using this # flag so they don't end up consuming time and resources only to # realize they can't actually log in and see the reports. - cmd = ["ssh", "-G", "hg.mozilla.org"] - output = subprocess.check_output( - cmd, universal_newlines=True - ).splitlines() - address = [ - l.rsplit(" ", 1)[-1] for l in output if l.startswith("user") - ][0] + address = get_ssh_user() if not address.endswith("@mozilla.com"): print( dedent( @@ -260,6 +257,70 @@ class Environment(TryConfig): } +class ExistingTasks(ParameterConfig): + TREEHERDER_PUSH_ENDPOINT = ( + "https://treeherder.mozilla.org/api/project/try/push/?count=1&author={user}" + ) + TREEHERDER_PUSH_URL = ( + "https://treeherder.mozilla.org/jobs?repo={branch}&revision={revision}" + ) + + arguments = [ + [ + ["-E", "--use-existing-tasks"], + { + "const": "last_try_push", + "default": None, + "nargs": "?", + "help": """ + Use existing tasks from a previous push. Without args this + uses your most recent try push. You may also specify + `rev=` where is the head revision of the + try push or `task-id=` where is the Decision + task id of the push. This last method even works for non-try + branches. + """, + }, + ] + ] + + def find_decision_task(self, use_existing_tasks): + branch = "try" + if use_existing_tasks == "last_try_push": + # Use existing tasks from user's previous try push. + user = get_ssh_user() + url = self.TREEHERDER_PUSH_ENDPOINT.format(user=user) + res = requests.get(url, headers={"User-Agent": "gecko-mach-try/1.0"}) + res.raise_for_status() + data = res.json() + if data["meta"]["count"] == 0: + raise Exception(f"Could not find a try push for '{user}'!") + revision = data["results"][0]["revision"] + + elif use_existing_tasks.startswith("rev="): + revision = use_existing_tasks[len("rev=") :] + + else: + raise Exception("Unable to parse '{use_existing_tasks}'!") + + url = self.TREEHERDER_PUSH_URL.format(branch=branch, revision=revision) + print(f"Using existing tasks from: {url}") + index_path = f"gecko.v2.{branch}.revision.{revision}.taskgraph.decision" + return taskcluster.find_task_id(index_path) + + def get_parameters(self, use_existing_tasks, **kwargs): + if not use_existing_tasks: + return + + if use_existing_tasks.startswith("task-id="): + tid = use_existing_tasks[len("task-id=") :] + else: + tid = self.find_decision_task(use_existing_tasks) + + label_to_task_id = taskcluster.get_artifact(tid, "public/label-to-taskid.json") + return {"existing_tasks": label_to_task_id} + + class RangeAction(Action): def __init__(self, min, max, *args, **kwargs): self.min = min @@ -418,7 +479,7 @@ class GeckoProfile(TryConfig): gecko_profile_entries, gecko_profile_features, gecko_profile_threads, - **kwargs + **kwargs, ): if profile or not all( s is None for s in (gecko_profile_features, gecko_profile_threads) @@ -548,6 +609,7 @@ all_task_configs = { "chemspill-prio": ChemspillPrio, "disable-pgo": DisablePgo, "env": Environment, + "existing-tasks": ExistingTasks, "gecko-profile": GeckoProfile, "path": Path, "pernosco": Pernosco, diff --git a/tools/tryselect/test/conftest.py b/tools/tryselect/test/conftest.py index ba27480a4064..d9cb7daee3eb 100644 --- a/tools/tryselect/test/conftest.py +++ b/tools/tryselect/test/conftest.py @@ -8,12 +8,19 @@ from unittest.mock import MagicMock import pytest import yaml from moztest.resolve import TestResolver +from responses import RequestsMock from taskgraph.graph import Graph from taskgraph.task import Task from taskgraph.taskgraph import TaskGraph from tryselect import push +@pytest.fixture +def responses(): + with RequestsMock() as rsps: + yield rsps + + @pytest.fixture def tg(request): if not hasattr(request.module, "TASKS"): diff --git a/tools/tryselect/test/python.toml b/tools/tryselect/test/python.toml index c624359b371c..f88156f69b6f 100644 --- a/tools/tryselect/test/python.toml +++ b/tools/tryselect/test/python.toml @@ -20,6 +20,8 @@ subsuite = "try" # shouldn't be run in parallel with those other tests. sequential = true +["test_push.py"] + ["test_release.py"] ["test_scriptworker.py"] diff --git a/tools/tryselect/test/test_push.py b/tools/tryselect/test/test_push.py new file mode 100644 index 000000000000..97f2e047d773 --- /dev/null +++ b/tools/tryselect/test/test_push.py @@ -0,0 +1,54 @@ +import mozunit +import pytest +from tryselect import push + + +@pytest.mark.parametrize( + "method,labels,params,routes,expected", + ( + pytest.param( + "fuzzy", + ["task-foo", "task-bar"], + None, + None, + { + "parameters": { + "optimize_target_tasks": False, + "try_task_config": { + "env": {"TRY_SELECTOR": "fuzzy"}, + "tasks": ["task-bar", "task-foo"], + }, + }, + "version": 2, + }, + id="basic", + ), + pytest.param( + "fuzzy", + ["task-foo"], + {"existing_tasks": {"task-foo": "123", "task-bar": "abc"}}, + None, + { + "parameters": { + "existing_tasks": {"task-bar": "abc"}, + "optimize_target_tasks": False, + "try_task_config": { + "env": {"TRY_SELECTOR": "fuzzy"}, + "tasks": ["task-foo"], + }, + }, + "version": 2, + }, + id="existing_tasks", + ), + ), +) +def test_generate_try_task_config(method, labels, params, routes, expected): + assert ( + push.generate_try_task_config(method, labels, params=params, routes=routes) + == expected + ) + + +if __name__ == "__main__": + mozunit.main() diff --git a/tools/tryselect/test/test_task_configs.py b/tools/tryselect/test/test_task_configs.py index fd646689fbfe..00adbf6228ce 100644 --- a/tools/tryselect/test/test_task_configs.py +++ b/tools/tryselect/test/test_task_configs.py @@ -3,7 +3,6 @@ # file, You can obtain one at http://mozilla.org/MPL/2.0/. import inspect -import subprocess from argparse import ArgumentParser from textwrap import dedent @@ -11,6 +10,9 @@ import mozunit import pytest from tryselect.task_config import Pernosco, all_task_configs +TC_URL = "https://firefox-ci-tc.services.mozilla.com" +TH_URL = "https://treeherder.mozilla.org" + # task configs have a list of tests of the form (input, expected) TASK_CONFIG_TESTS = { "artifact": [ @@ -130,26 +132,25 @@ def test_task_configs(config_patch_resolver, task_config, args, expected): @pytest.fixture -def patch_pernosco_email_check(monkeypatch): - def inner(val): - def fake_check_output(*args, **kwargs): - return val - - monkeypatch.setattr(subprocess, "check_output", fake_check_output) +def patch_ssh_user(mocker): + def inner(user): + mock_stdout = mocker.Mock() + mock_stdout.stdout = dedent( + f""" + key1 foo + user {user} + key2 bar + """ + ) + return mocker.patch( + "tryselect.util.ssh.subprocess.run", return_value=mock_stdout + ) return inner -def test_pernosco(patch_pernosco_email_check): - patch_pernosco_email_check( - dedent( - """ - user foobar@mozilla.com - hostname hg.mozilla.com - """ - ) - ) - +def test_pernosco(patch_ssh_user): + patch_ssh_user("user@mozilla.com") parser = ArgumentParser() cfg = Pernosco() @@ -159,5 +160,94 @@ def test_pernosco(patch_pernosco_email_check): assert params == {"try_task_config": {"env": {"PERNOSCO": "1"}}} +def test_exisiting_tasks(responses, patch_ssh_user): + parser = ArgumentParser() + cfg = all_task_configs["existing-tasks"]() + cfg.add_arguments(parser) + + user = "user@example.com" + rev = "a" * 40 + task_id = "abc" + label_to_taskid = {"task-foo": "123", "task-bar": "456"} + + args = ["--use-existing-tasks"] + args = parser.parse_args(args) + + responses.add( + responses.GET, + f"{TH_URL}/api/project/try/push/?count=1&author={user}", + json={"meta": {"count": 1}, "results": [{"revision": rev}]}, + ) + + responses.add( + responses.GET, + f"{TC_URL}/api/index/v1/task/gecko.v2.try.revision.{rev}.taskgraph.decision", + json={"taskId": task_id}, + ) + + responses.add( + responses.GET, + f"{TC_URL}/api/queue/v1/task/{task_id}/artifacts/public/label-to-taskid.json", + json=label_to_taskid, + ) + + m = patch_ssh_user(user) + params = cfg.get_parameters(**vars(args)) + assert params == {"existing_tasks": label_to_taskid} + + m.assert_called_once_with( + ["ssh", "-G", "hg.mozilla.org"], text=True, check=True, capture_output=True + ) + + +def test_exisiting_tasks_task_id(responses): + parser = ArgumentParser() + cfg = all_task_configs["existing-tasks"]() + cfg.add_arguments(parser) + + task_id = "abc" + label_to_taskid = {"task-foo": "123", "task-bar": "456"} + + args = ["--use-existing-tasks", f"task-id={task_id}"] + args = parser.parse_args(args) + + responses.add( + responses.GET, + f"{TC_URL}/api/queue/v1/task/{task_id}/artifacts/public/label-to-taskid.json", + json=label_to_taskid, + ) + + params = cfg.get_parameters(**vars(args)) + assert params == {"existing_tasks": label_to_taskid} + + +def test_exisiting_tasks_rev(responses): + parser = ArgumentParser() + cfg = all_task_configs["existing-tasks"]() + cfg.add_arguments(parser) + + rev = "aaaaaa" + task_id = "abc" + label_to_taskid = {"task-foo": "123", "task-bar": "456"} + + args = ["--use-existing-tasks", f"rev={rev}"] + args = parser.parse_args(args) + + responses.add( + responses.GET, + f"{TC_URL}/api/index/v1/task/gecko.v2.try.revision.{rev}.taskgraph.decision", + json={"taskId": task_id}, + ) + + responses.add( + responses.GET, + f"{TC_URL}/api/queue/v1/task/{task_id}/artifacts/public/label-to-taskid.json", + json=label_to_taskid, + ) + + params = cfg.get_parameters(**vars(args)) + assert params == {"existing_tasks": label_to_taskid} + + if __name__ == "__main__": mozunit.main() diff --git a/tools/tryselect/util/ssh.py b/tools/tryselect/util/ssh.py new file mode 100644 index 000000000000..7682306bc770 --- /dev/null +++ b/tools/tryselect/util/ssh.py @@ -0,0 +1,24 @@ +# 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 https://mozilla.org/MPL/2.0/. + +import subprocess + + +def get_ssh_user(host="hg.mozilla.org"): + ssh_config = subprocess.run( + ["ssh", "-G", host], + text=True, + check=True, + capture_output=True, + ).stdout + + lines = [l.strip() for l in ssh_config.splitlines()] + for line in lines: + if not line: + continue + key, value = line.split(" ", 1) + if key.lower() == "user": + return value + + raise Exception(f"Could not detect ssh user for '{host}'!")