Bug 1643689 - [taskgraph] Allow tasks to override the configured 'test_manifest_loader', r=jmaher

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

taskcluster/taskgraph/transforms/tests.py

@@ -47,9 +47,11 @@ from taskgraph.util.schema import (
 from taskgraph.optimize.schema import OptimizationSchema
 from taskgraph.util.chunking import (
     chunk_manifests,
+    get_manifest_loader,
     get_runtimes,
     guess_mozinfo_from_task,
     manifest_loaders,
+    DefaultLoader,
 )
 from taskgraph.util.taskcluster import (
     get_artifact_path,
@@ -324,7 +326,6 @@ CHUNK_SUITES_BLACKLIST = (
 )
 """These suites will be chunked at test runtime rather than here in the taskgraph."""
 
-
 DYNAMIC_CHUNK_DURATION = 20 * 60  # seconds
 """The approximate time each test chunk should take to run."""
 
@@ -432,6 +433,12 @@ test_description_schema = Schema({
         'test-platform',
         Any(int, 'dynamic')),
 
+
+    # Custom 'test_manifest_loader' to use, overriding the one configured in the
+    # parameters. When 'null', no test chunking will be performed. Can also
+    # be used to disable "manifest scheduling".
+    Optional('test-manifest-loader'): Any(None, *list(manifest_loaders)),
+
     # the time (with unit) after which this task is deleted; default depends on
     # the branch (see below)
     Optional('expires-after'): text_type,
@@ -1496,14 +1503,17 @@ def set_test_verify_chunks(config, tasks):
 def set_test_manifests(config, tasks):
     """Determine the set of test manifests that should run in this task."""
 
-    loader_cls = manifest_loaders[config.params['test_manifest_loader']]
-    loader = loader_cls(config.params)
-
     for task in tasks:
         if task['suite'] in CHUNK_SUITES_BLACKLIST:
             yield task
             continue
 
+        # When a task explicitly requests no 'test_manifest_loader', test
+        # resolving will happen at test runtime rather than in the taskgraph.
+        if 'test-manifest-loader' in task and task['test-manifest-loader'] is None:
+            yield task
+            continue
+
         # Set 'tests_grouped' to "1", so we can differentiate between suites that are
         # chunked at the test runtime and those that are chunked in the taskgraph.
         task.setdefault("tags", {})["tests_grouped"] = "1"
@@ -1526,6 +1536,9 @@ def set_test_manifests(config, tasks):
 
         mozinfo = guess_mozinfo_from_task(task)
 
+        loader_name = task.pop('test-manifest-loader', config.params['test_manifest_loader'])
+        loader = get_manifest_loader(loader_name, config.params)
+
         task['test-manifests'] = loader.get_manifests(
             task['suite'],
             frozenset(mozinfo.items()),
@@ -1535,7 +1548,7 @@ def set_test_manifests(config, tasks):
         # loader, we'll only run some subset of manifests and the hardcoded
         # chunk numbers will no longer be valid. Dynamic chunking should yield
         # better results.
-        if config.params['test_manifest_loader'] != 'default':
+        if not isinstance(loader, DefaultLoader):
             task['chunks'] = "dynamic"
 
         yield task

taskcluster/taskgraph/util/chunking.py

@@ -242,3 +242,16 @@ manifest_loaders = {
     'bugbug': BugbugLoader,
     'default': DefaultLoader,
 }
+
+_loader_cache = {}
+
+
+def get_manifest_loader(name, params):
+    # Ensure we never create more than one instance of the same loader type for
+    # performance reasons.
+    if name in _loader_cache:
+        return _loader_cache[name]
+
+    loader = manifest_loaders[name](dict(params))
+    _loader_cache[name] = loader
+    return loader