Bug 1861075 - filter tasks by --worker-type. r=taskgraph-reviewers,hneiva

Differential Revision: https://phabricator.services.mozilla.com/D203671

taskcluster/gecko_taskgraph/parameters.py

@@ -74,6 +74,10 @@ gecko_parameters_schema = {
             "worker-overrides",
             description="Mapping of worker alias to worker pools to use for those aliases.",
         ): {str: str},
+        Optional(
+            "worker-types",
+            description="List of worker types that we will use to run tasks on.",
+        ): [str],
         Optional("routes"): [str],
     },
     Required("version"): str,

tools/tryselect/selectors/fuzzy.py

@@ -12,7 +12,7 @@ from mach.util import get_state_dir
 
 from ..cli import BaseTryParser
 from ..push import check_working_directory, generate_try_task_config, push_to_try
-from ..tasks import filter_tasks_by_paths, generate_tasks
+from ..tasks import filter_tasks_by_paths, filter_tasks_by_worker_type, generate_tasks
 from ..util.fzf import (
     FZF_NOT_FOUND,
     PREVIEW_SCRIPT,
@@ -182,6 +182,11 @@ def run(
             if filter_by_uncommon_try_tasks(task_name)
         }
 
+    if try_config_params.get("try_task_config", {}).get("worker-types", []):
+        all_tasks = filter_tasks_by_worker_type(all_tasks, try_config_params)
+        if not all_tasks:
+            return 1
+
     if test_paths:
         all_tasks = filter_tasks_by_paths(all_tasks, test_paths)
         if not all_tasks:

tools/tryselect/task_config.py

@@ -80,7 +80,7 @@ class TryConfig(ParameterConfig):
 
     def get_parameters(self, **kwargs):
         result = self.try_config(**kwargs)
-        if result is None:
+        if not result:
             return None
         return {"try_task_config": result}
 
@@ -579,9 +579,18 @@ class WorkerOverrides(TryConfig):
                 ),
             },
         ],
+        [
+            ["--worker-type"],
+            {
+                "action": "append",
+                "dest": "worker_types",
+                "default": [],
+                "help": "Select tasks that only run on the specified worker.",
+            },
+        ],
     ]
 
-    def try_config(self, worker_overrides, worker_suffixes, **kwargs):
+    def try_config(self, worker_overrides, worker_suffixes, worker_types, **kwargs):
         from gecko_taskgraph.util.workertypes import get_worker_type
         from taskgraph.config import load_graph_config
 
@@ -621,8 +630,13 @@ class WorkerOverrides(TryConfig):
                     provisioner=provisioner, worker_type=worker_type, suffix=suffix
                 )
 
+        retVal = {}
+        if worker_types:
+            retVal["worker-types"] = list(overrides.keys()) + worker_types
+
         if overrides:
-            return {"worker-overrides": overrides}
+            retVal["worker-overrides"] = overrides
+        return retVal
 
 
 all_task_configs = {

tools/tryselect/tasks.py

@@ -152,6 +152,17 @@ def generate_tasks(params=None, full=False, disable_target_task_filter=False):
     return tg_target
 
 
+def filter_tasks_by_worker_type(tasks, params):
+    worker_types = params.get("try_task_config", {}).get("worker-types", [])
+    if worker_types:
+        retVal = {}
+        for t in tasks:
+            if tasks[t].task["workerType"] in worker_types:
+                retVal[t] = tasks[t]
+        return retVal
+    return tasks
+
+
 def filter_tasks_by_paths(tasks, paths):
     resolver = TestResolver.from_environment(cwd=here, loader_cls=TestManifestLoader)
     run_suites, run_tests = resolver.resolve_metadata(paths)

tools/tryselect/test/test_tasks.py

@@ -6,7 +6,103 @@ import os
 
 import mozunit
 import pytest
-from tryselect.tasks import cache_key, filter_tasks_by_paths, resolve_tests_by_suite
+from tryselect.tasks import (
+    cache_key,
+    filter_tasks_by_paths,
+    filter_tasks_by_worker_type,
+    resolve_tests_by_suite,
+)
+
+
+class task:
+    def __init__(self, workerType):
+        self.workerType = workerType
+
+    @property
+    def task(self):
+        return {"workerType": self.workerType}
+
+
+@pytest.mark.parametrize(
+    "tasks, params, expected",
+    (
+        pytest.param(
+            {
+                "foobar/xpcshell-1": task("t-unittest-314"),
+                "foobar/mochitest": task("t-unittest-157"),
+                "foobar/xpcshell-gpu": task("t-unittest-314-gpu"),
+                "foobar/xpcshell": task("t-unittest-314"),
+            },
+            {"try_task_config": {"worker-types": ["t-unittest-314"]}},
+            [
+                "foobar/xpcshell-1",
+                "foobar/xpcshell",
+            ],
+            id="single worker",
+        ),
+        pytest.param(
+            {
+                "foobar/xpcshell-1": task("t-unittest-314"),
+                "foobar/mochitest": task("t-unittest-157"),
+                "foobar/xpcshell-gpu": task("t-unittest-314-gpu"),
+                "foobar/xpcshell": task("t-unittest-314"),
+            },
+            {
+                "try_task_config": {
+                    "worker-types": ["t-unittest-314", "t-unittest-314-gpu"]
+                }
+            },
+            [
+                "foobar/xpcshell-1",
+                "foobar/xpcshell-gpu",
+                "foobar/xpcshell",
+            ],
+            id="multiple workers worker",
+        ),
+        pytest.param(
+            {
+                "foobar/xpcshell-1": task("t-unittest-314"),
+                "foobar/mochitest": task("t-unittest-157"),
+                "foobar/xpcshell-gpu": task("t-unittest-314-gpu"),
+                "foobar/xpcshell": task("t-unittest-314"),
+            },
+            {"try_task_config": {"worker-types": ["t-unittest-157"]}},
+            [
+                "foobar/mochitest",
+            ],
+            id="single task",
+        ),
+        pytest.param(
+            {
+                "foobar/xpcshell-1": task("t-unittest-314"),
+                "foobar/mochitest": task("t-unittest-157"),
+                "foobar/xpcshell-gpu": task("t-unittest-314-gpu"),
+                "foobar/xpcshell": task("t-unittest-314"),
+            },
+            {"try_task_config": {"worker-types": []}},
+            [
+                "foobar/xpcshell-1",
+                "foobar/mochitest",
+                "foobar/xpcshell-gpu",
+                "foobar/xpcshell",
+            ],
+            id="no worker",
+        ),
+        pytest.param(
+            {
+                "foobar/xpcshell-1": task("t-unittest-314"),
+                "foobar/mochitest": task("t-unittest-157"),
+                "foobar/xpcshell-gpu": task("t-unittest-314-gpu"),
+                "foobar/xpcshell": task("t-unittest-314"),
+            },
+            {"try_task_config": {"worker-types": ["fake-worker"]}},
+            [],
+            id="invalid worker",
+        ),
+    ),
+)
+def test_filter_tasks_by_worker_type(patch_resolver, tasks, params, expected):
+    assert list(filter_tasks_by_worker_type(tasks, params)) == expected
 
 
 def test_filter_tasks_by_paths(patch_resolver):