From f5b5676618a5f58e6afcdfb447cf6619a96f4e9d Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahalberstadt@mozilla.com>
Date: Mon, 25 May 2020 18:52:12 +0000
Subject: [PATCH] Bug 1633866 - [taskgraph] Implement the ability for tasks to
 dynamically specify chunks, r=gbrown,jmaher

This allows tasks to opt into dynamic chunking. This means rather than define
the chunks ahead of time, in-tree manifest runtime data is used to guess how
many chunks are needed for a reasonable average task runtime (currently 20
min).

Only suites that are chunked in the taskgraph, and therefore whose manifests
are known ahead of time, can opt into this feature.

Differential Revision: https://phabricator.services.mozilla.com/D76497
---
 taskcluster/taskgraph/transforms/tests.py | 58 +++++++++++++++++++++--
 1 file changed, 53 insertions(+), 5 deletions(-)

diff --git a/taskcluster/taskgraph/transforms/tests.py b/taskcluster/taskgraph/transforms/tests.py
index e7e27019b770..6ad61ec744c9 100644
--- a/taskcluster/taskgraph/transforms/tests.py
+++ b/taskcluster/taskgraph/transforms/tests.py
@@ -47,6 +47,7 @@ from taskgraph.util.schema import (
 from taskgraph.util.chunking import (
     chunk_manifests,
     get_manifests,
+    get_runtimes,
     guess_mozinfo_from_task,
 )
 from taskgraph.util.taskcluster import (
@@ -249,6 +250,10 @@ 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."""
+
+
 logger = logging.getLogger(__name__)
 
 transforms = TransformSequence()
@@ -336,7 +341,7 @@ test_description_schema = Schema({
     # test platform is not found, the key 'default' will be tried.
     Required('chunks'): optionally_keyed_by(
         'test-platform',
-        int),
+        Any(int, 'dynamic')),
 
     # the time (with unit) after which this task is deleted; default depends on
     # the branch (see below)
@@ -1362,7 +1367,23 @@ def set_test_manifests(config, tasks):
     """Determine the set of test manifests that should run in this task."""
 
     for task in tasks:
-        if taskgraph.fast or task['suite'] in CHUNK_SUITES_BLACKLIST or 'test-manifests' in task:
+        if task['suite'] in CHUNK_SUITES_BLACKLIST:
+            yield task
+            continue
+
+        if taskgraph.fast:
+            # We want to avoid evaluating manifests when taskgraph.fast is set. But
+            # manifests are required for dynamic chunking. Just set the number of
+            # chunks to one in this case.
+            if task['chunks'] == 'dynamic':
+                task['chunks'] = 1
+            yield task
+            continue
+
+        manifests = task.get('test-manifests')
+        if manifests:
+            if isinstance(manifests, list):
+                task['test-manifests'] = {'active': manifests, 'skipped': []}
             yield task
             continue
 
@@ -1377,6 +1398,36 @@ def set_test_manifests(config, tasks):
         yield task
 
 
+@transforms.add
+def resolve_dynamic_chunks(config, tasks):
+    """Determine how many chunks are needed to handle the given set of manifests."""
+
+    for task in tasks:
+        if task['chunks'] != "dynamic":
+            yield task
+            continue
+
+        if not task.get('test-manifests'):
+            raise Exception(
+                "{} must define 'test-manifests' to use dynamic chunking!".format(
+                    task['test-name']))
+
+        runtimes = {m: r for m, r in get_runtimes(task['test-platform']).items()
+                    if m in task['test-manifests']['active']}
+
+        times = list(runtimes.values())
+        avg = round(sum(times) / len(times), 2) if times else 0
+        total = sum(times)
+
+        # If there are manifests missing from the runtimes data, fill them in
+        # with the average of all present manifests.
+        missing = [m for m in task['test-manifests']['active'] if m not in runtimes]
+        total += avg * len(missing)
+
+        task['chunks'] = int(round(total / DYNAMIC_CHUNK_DURATION)) or 1
+        yield task
+
+
 @transforms.add
 def split_chunks(config, tasks):
     """Based on the 'chunks' key, split tests up into chunks by duplicating
@@ -1391,9 +1442,6 @@ def split_chunks(config, tasks):
         if 'test-manifests' in task:
             suite_definition = TEST_SUITES[task['suite']]
             manifests = task['test-manifests']
-            if isinstance(manifests, list):
-                manifests = {'active': manifests, 'skipped': []}
-
             chunked_manifests = chunk_manifests(
                 suite_definition['build_flavor'],
                 suite_definition.get('kwargs', {}).get('subsuite', 'undefined'),