From 22e8f8ddf31546b0408eae94685ee618cc43ee71 Mon Sep 17 00:00:00 2001 From: Marco Castelluccio Date: Thu, 9 Nov 2023 15:48:16 +0000 Subject: [PATCH] Bug 1557848 - Don't list all chunks for jobs in mach try fuzzy. r=ahal DONTBUILD Differential Revision: https://phabricator.services.mozilla.com/D157732 --- tools/tryselect/selectors/coverage.py | 4 +-- tools/tryselect/selectors/fuzzy.py | 47 ++++++++++++++++++++++----- tools/tryselect/tasks.py | 25 ++++++++++++-- tools/tryselect/test/test_fuzzy.py | 11 +++++-- tools/tryselect/test/test_tasks.py | 2 +- 5 files changed, 72 insertions(+), 17 deletions(-) diff --git a/tools/tryselect/selectors/coverage.py b/tools/tryselect/selectors/coverage.py index 2244c8810938..2d1362db76bd 100644 --- a/tools/tryselect/selectors/coverage.py +++ b/tools/tryselect/selectors/coverage.py @@ -354,7 +354,7 @@ def filter_tasks_by_chunks(tasks, chunks): platform = PLATFORM_MAP[platform] selected_task = None - for task in tasks: + for task in tasks.keys(): if not task.startswith(platform): continue @@ -405,7 +405,7 @@ def run( return 1 tg = generate_tasks(parameters, full) - all_tasks = tg.tasks.keys() + all_tasks = tg.tasks tasks_by_chunks = filter_tasks_by_chunks(all_tasks, test_chunks) tasks_by_path = filter_tasks_by_paths(all_tasks, test_files) diff --git a/tools/tryselect/selectors/fuzzy.py b/tools/tryselect/selectors/fuzzy.py index cb7d53a2ca5c..8965617098fa 100644 --- a/tools/tryselect/selectors/fuzzy.py +++ b/tools/tryselect/selectors/fuzzy.py @@ -97,6 +97,15 @@ class FuzzyParser(BaseTryParser): "disables this filtering.", }, ], + [ + ["--show-chunk-numbers"], + { + "action": "store_true", + "default": False, + "help": "Chunk numbers are hidden to simplify the selection. This flag " + "makes them appear again.", + }, + ], ] common_groups = ["push", "task", "preset"] task_configs = [ @@ -131,6 +140,7 @@ def run( show_estimates=False, disable_target_task_filter=False, push_to_lando=False, + show_chunk_numbers=False, ): fzf = fzf_bootstrap(update) @@ -143,7 +153,7 @@ def run( tg = generate_tasks( parameters, full=full, disable_target_task_filter=disable_target_task_filter ) - all_tasks = sorted(tg.tasks.keys()) + all_tasks = tg.tasks # graph_Cache created by generate_tasks, recreate the path to that file. cache_dir = os.path.join( @@ -163,9 +173,11 @@ def run( make_trimmed_taskgraph_cache(graph_cache, dep_cache, target_file=target_set) if not full and not disable_target_task_filter: - # Put all_tasks into a list because it's used multiple times, and "filter()" - # returns a consumable iterator. - all_tasks = list(filter(filter_by_uncommon_try_tasks, all_tasks)) + all_tasks = { + task_name: task + for task_name, task in all_tasks.items() + if filter_by_uncommon_try_tasks(task_name) + } if test_paths: all_tasks = filter_tasks_by_paths(all_tasks, test_paths) @@ -214,7 +226,12 @@ def run( if query_arg and query_arg != "INTERACTIVE": cmd.extend(["-f", query_arg]) - query_str, tasks = run_fzf(cmd, sorted(candidate_tasks)) + if not show_chunk_numbers: + fzf_tasks = set(task.chunk_pattern for task in candidate_tasks.values()) + else: + fzf_tasks = set(candidate_tasks.keys()) + + query_str, tasks = run_fzf(cmd, sorted(fzf_tasks)) queries.append(query_str) return set(tasks) @@ -223,11 +240,16 @@ def run( for q in intersect_query or []: if not selected: - tasks = get_tasks(q) - selected |= tasks + selected |= get_tasks(q) else: - tasks = get_tasks(q, selected) - selected &= tasks + selected &= get_tasks( + q, + { + task_name: task + for task_name, task in all_tasks.items() + if task_name in selected or task.chunk_pattern in selected + }, + ) if not queries: selected = get_tasks() @@ -239,6 +261,13 @@ def run( if save_query: return queries + if not show_chunk_numbers: + selected = set( + task_name + for task_name, task in all_tasks.items() + if task.chunk_pattern in selected + ) + # build commit message msg = "Fuzzy" args = ["query={}".format(q) for q in queries] diff --git a/tools/tryselect/tasks.py b/tools/tryselect/tasks.py index bca6b3de1ba8..eec222eacd86 100644 --- a/tools/tryselect/tasks.py +++ b/tools/tryselect/tasks.py @@ -82,6 +82,23 @@ def generate_tasks(params=None, full=False, disable_target_task_filter=False): taskgraph.fast = True generator = TaskGraphGenerator(root_dir=root, parameters=params) + def add_chunk_patterns(tg): + for task_name, task in tg.tasks.items(): + chunks = task.task.get("extra", {}).get("chunks", {}) + if isinstance(chunks, int): + task.chunk_pattern = "{}-*/{}".format( + "-".join(task_name.split("-")[:-1]), chunks + ) + else: + assert isinstance(chunks, dict) + if chunks.get("total", 1) == 1: + task.chunk_pattern = task_name + else: + task.chunk_pattern = "{}-*".format( + "-".join(task_name.split("-")[:-1]) + ) + return tg + cache_dir = os.path.join( get_state_dir(specific_to_topsrcdir=True), "cache", "taskgraph" ) @@ -91,7 +108,7 @@ def generate_tasks(params=None, full=False, disable_target_task_filter=False): invalidate(cache) if os.path.isfile(cache): with open(cache) as fh: - return TaskGraph.from_json(json.load(fh))[1] + return add_chunk_patterns(TaskGraph.from_json(json.load(fh))[1]) if not os.path.isdir(cache_dir): os.makedirs(cache_dir) @@ -112,7 +129,7 @@ def generate_tasks(params=None, full=False, disable_target_task_filter=False): key = cache_key(attr, generator.parameters, disable_target_task_filter) with open(os.path.join(cache_dir, key), "w") as fh: json.dump(tg.to_json(), fh) - return tg + return add_chunk_patterns(tg) # Cache both full_task_set and target_task_set regardless of whether or not # --full was requested. Caching is cheap and can potentially save a lot of @@ -152,7 +169,9 @@ def filter_tasks_by_paths(tasks, paths): def match_task(task): return any(re.search(pattern, task) for pattern in task_regexes) - return filter(match_task, tasks) + return { + task_name: task for task_name, task in tasks.items() if match_task(task_name) + } def resolve_tests_by_suite(paths): diff --git a/tools/tryselect/test/test_fuzzy.py b/tools/tryselect/test/test_fuzzy.py index 330ba99825df..b38b7d919b7a 100644 --- a/tools/tryselect/test/test_fuzzy.py +++ b/tools/tryselect/test/test_fuzzy.py @@ -9,7 +9,8 @@ import pytest @pytest.mark.skipif(os.name == "nt", reason="fzf not installed on host") -def test_query_paths(run_mach, capfd): +@pytest.mark.parametrize("show_chunk_numbers", [True, False]) +def test_query_paths(run_mach, capfd, show_chunk_numbers): cmd = [ "try", "fuzzy", @@ -18,6 +19,9 @@ def test_query_paths(run_mach, capfd): "^test-linux '64-qr/debug-mochitest-chrome-1proc-", "caps/tests/mochitest/test_addonMayLoad.html", ] + if show_chunk_numbers: + cmd.append("--show-chunk-numbers") + assert run_mach(cmd) == 0 output = capfd.readouterr().out @@ -34,8 +38,11 @@ def test_query_paths(run_mach, capfd): @pytest.mark.skipif(os.name == "nt", reason="fzf not installed on host") -def test_query(run_mach, capfd): +@pytest.mark.parametrize("full", [True, False]) +def test_query(run_mach, capfd, full): cmd = ["try", "fuzzy", "--no-push", "-q", "'source-test-python-taskgraph-tests-py3"] + if full: + cmd.append("--full") assert run_mach(cmd) == 0 output = capfd.readouterr().out diff --git a/tools/tryselect/test/test_tasks.py b/tools/tryselect/test/test_tasks.py index cdf8cf84c1ca..2e99c72d8b4e 100644 --- a/tools/tryselect/test/test_tasks.py +++ b/tools/tryselect/test/test_tasks.py @@ -10,7 +10,7 @@ from tryselect.tasks import cache_key, filter_tasks_by_paths, resolve_tests_by_s def test_filter_tasks_by_paths(patch_resolver): - tasks = ["foobar/xpcshell-1", "foobar/mochitest", "foobar/xpcshell"] + tasks = {"foobar/xpcshell-1": {}, "foobar/mochitest": {}, "foobar/xpcshell": {}} patch_resolver(["xpcshell"], {}) assert list(filter_tasks_by_paths(tasks, "dummy")) == []