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

@@ -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)

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]

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):

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

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")) == []