SERVER-53196: Fail to generate tasks if a large distro is required but not specified

(cherry picked from commit f946a64ec3cadae717759479ddeaea368c5a6b1c)
(cherry picked from commit 2e49448cbf41df0c8b75a6ad067459e15f31902a)

5 files changed, 125 insertions, 8 deletions

diff --git a/buildscripts/evergreen_generate_resmoke_tasks.py b/buildscripts/evergreen_generate_resmoke_tasks.py
index 36986b28f2d..1699049a421 100755
--- a/buildscripts/evergreen_generate_resmoke_tasks.py
+++ b/buildscripts/evergreen_generate_resmoke_tasks.py
@@ -15,6 +15,7 @@ import re
 import sys
 from collections import defaultdict
 from distutils.util import strtobool  # pylint: disable=no-name-in-module
+from typing import Set
 
 import click
 import requests
@@ -22,6 +23,7 @@ import structlog
 import yaml
 
 from evergreen.api import RetryingEvergreenApi
+from pydantic import BaseModel
 from shrub.config import Configuration
 from shrub.command import CommandDefinition
 from shrub.operations import CmdTimeoutUpdate
@@ -37,13 +39,15 @@ import buildscripts.resmokelib.suitesconfig as suitesconfig  # pylint: disable=w
 import buildscripts.util.read_config as read_config  # pylint: disable=wrong-import-position
 import buildscripts.util.taskname as taskname  # pylint: disable=wrong-import-position
 import buildscripts.util.testname as testname  # pylint: disable=wrong-import-position
+from buildscripts.util.fileops import read_yaml_file  # pylint: disable=wrong-import-position
 
 LOGGER = structlog.getLogger(__name__)
 
 TEST_SUITE_DIR = os.path.join("buildscripts", "resmokeconfig", "suites")
 CONFIG_DIR = "generated_resmoke_config"
 AVG_SETUP_TIME = int(timedelta(minutes=5).total_seconds())
-CONFIG_FILE = "./.evergreen.yml"
+EVG_CONFIG_FILE = "./.evergreen.yml"
+GENERATE_CONFIG_FILE = "etc/generate_subtasks_config.yml"
 MIN_TIMEOUT_SECONDS = int(timedelta(minutes=5).total_seconds())
 LOOKBACK_DURATION_DAYS = 14
 GEN_SUFFIX = "_gen"
@@ -53,6 +57,21 @@ HEADER_TEMPLATE = """# DO NOT EDIT THIS FILE. All manual edits will be lost.
 # {suite_file}.
 """
 
+NO_LARGE_DISTRO_ERR = """
+***************************************************************************************
+It appears we are trying to generate a task marked as requiring a large distro, but the
+build variant has not specified a large build variant. In order to resolve this error,
+you need to:
+
+(1) add a "large_distro_name" expansion to this build variant ("{build_variant}").
+
+ -- or --
+ 
+(2) add this build variant ("{build_variant}") to the "build_variant_large_distro_exception"
+list in the "etc/generate_subtasks_config.yml" file.
+***************************************************************************************
+"""
+
 REQUIRED_CONFIG_KEYS = {
     "build_variant",
     "fallback_num_sub_suites",
@@ -79,6 +98,22 @@ CONFIG_FORMAT_FN = {
 }
 
 
+class GenerationConfiguration(BaseModel):
+    """Configuration for generating sub-tasks."""
+
+    build_variant_large_distro_exceptions: Set[str]
+
+    @classmethod
+    def from_yaml_file(cls, path: str) -> "GenerationConfiguration":
+        """Read the generation configuration from the given file."""
+        return cls(**read_yaml_file(path))
+
+    @classmethod
+    def default_config(cls) -> "GenerationConfiguration":
+        """Create a default configuration."""
+        return cls(build_variant_large_distro_exceptions=set())
+
+
 class ConfigOptions(object):
     """Retrieve configuration from a config file."""
 
@@ -387,10 +422,14 @@ def should_tasks_be_generated(evg_api, task_id):
 class EvergreenConfigGenerator(object):
     """Generate evergreen configurations."""
 
-    def __init__(self, suites, options, evg_api):
+    # pylint: disable=too-many-instance-attributes
+    def __init__(self, suites, options, evg_api, generate_config=None):
         """Create new EvergreenConfigGenerator object."""
         self.suites = suites
         self.options = options
+        self.gen_config = GenerationConfiguration.default_config()
+        if generate_config:
+            self.gen_config = generate_config
         self.evg_api = evg_api
         self.evg_config = Configuration()
         self.task_specs = []
@@ -398,7 +437,11 @@ class EvergreenConfigGenerator(object):
         self.build_tasks = None
 
     def _set_task_distro(self, task_spec):
-        if self.options.use_large_distro and self.options.large_distro_name:
+        if self.options.use_large_distro:
+            if (self.options.variant not in self.gen_config.build_variant_large_distro_exceptions
+                    and not self.options.large_distro_name):
+                print(NO_LARGE_DISTRO_ERR.format(build_variant=self.options.variant))
+                raise ValueError("Invalid Evergreen Configuration")
             task_spec.distro(self.options.large_distro_name)
 
     def _generate_resmoke_args(self, suite_file):
@@ -646,10 +689,13 @@ class Suite(object):
 class GenerateSubSuites(object):
     """Orchestrate the execution of generate_resmoke_suites."""
 
-    def __init__(self, evergreen_api, config_options):
+    def __init__(self, evergreen_api, config_options, generate_config=None):
         """Initialize the object."""
         self.evergreen_api = evergreen_api
         self.config_options = config_options
+        self.generate_options = GenerationConfiguration.default_config()
+        if generate_config:
+            self.generate_options = generate_config
         self.test_list = []
 
     def calculate_suites(self, start_date, end_date):
@@ -714,7 +760,8 @@ class GenerateSubSuites(object):
 
     def write_evergreen_configuration(self, suites, task):
         """Generate the evergreen configuration for the new suite and write it to disk."""
-        evg_config_gen = EvergreenConfigGenerator(suites, self.config_options, self.evergreen_api)
+        evg_config_gen = EvergreenConfigGenerator(suites, self.config_options, self.evergreen_api,
+                                                  self.generate_options)
         evg_config = evg_config_gen.generate_config()
 
         with open(os.path.join(CONFIG_DIR, task + ".json"), "w") as file_handle:
@@ -746,7 +793,7 @@ class GenerateSubSuites(object):
 @click.command()
 @click.option("--expansion-file", type=str, required=True,
               help="Location of expansions file generated by evergreen.")
-@click.option("--evergreen-config", type=str, default=CONFIG_FILE,
+@click.option("--evergreen-config", type=str, default=EVG_CONFIG_FILE,
               help="Location of evergreen configuration file.")
 @click.option("--verbose", is_flag=True, default=False, help="Enable verbose logging.")
 def main(expansion_file, evergreen_config, verbose):
@@ -761,10 +808,11 @@ def main(expansion_file, evergreen_config, verbose):
     """
     enable_logging(verbose)
     evg_api = RetryingEvergreenApi.get_api(config_file=evergreen_config)
+    generate_config = GenerationConfiguration.from_yaml_file(GENERATE_CONFIG_FILE)
     config_options = ConfigOptions.from_file(expansion_file, REQUIRED_CONFIG_KEYS,
                                              DEFAULT_CONFIG_VALUES, CONFIG_FORMAT_FN)
 
-    GenerateSubSuites(evg_api, config_options).run()
+    GenerateSubSuites(evg_api, config_options, generate_config).run()
 
 
 if __name__ == "__main__":
diff --git a/buildscripts/util/fileops.py b/buildscripts/util/fileops.py
index e56348d8548..b3128d087b9 100644
--- a/buildscripts/util/fileops.py
+++ b/buildscripts/util/fileops.py
@@ -1,6 +1,8 @@
 """Utility to support file operations."""
-
 import os
+from typing import Dict, Any
+
+import yaml
 
 
 def create_empty(path):
@@ -25,3 +27,14 @@ def get_file_handle(path, append_file=False):
     """Open 'path', truncate it if 'append_file' is False, and return file handle."""
     mode = "a+" if append_file else "w"
     return open(path, mode)
+
+
+def read_yaml_file(path: str) -> Dict[str, Any]:
+    """
+    Read the yaml file at the given path and return the contents.
+
+    :param path: Path to file to read.
+    :return: Contents of given file.
+    """
+    with open(path) as file_handle:
+        return yaml.safe_load(file_handle)
diff --git a/etc/evergreen.yml b/etc/evergreen.yml
index 6ea9e9b3047..70abf237ea2 100644
--- a/etc/evergreen.yml
+++ b/etc/evergreen.yml
@@ -7370,6 +7370,7 @@ buildvariants:
     scons_cache_scope: shared
     tooltags: ""
     build_mongoreplay: true
+    large_distro_name: rhel62-large
   tasks:
   - name: compile_all_run_unittests_TG
     distros:
@@ -7405,6 +7406,8 @@ buildvariants:
     tooltags: ""
     build_mongoreplay: true
     target_resmoke_time: 15
+    max_sub_suites: 3
+    large_distro_name: rhel62-large
   tasks:
   - name: compile_all_run_unittests_TG
   - name: .aggregation !.encrypt
@@ -7452,6 +7455,7 @@ buildvariants:
     scons_cache_scope: shared
     tooltags: ""
     build_mongoreplay: true
+    large_distro_name: rhel62-large
   tasks:
   - name: compile_all_run_unittests_TG
     distros:
@@ -7491,6 +7495,7 @@ buildvariants:
     scons_cache_scope: shared
     tooltags: "ssl"
     build_mongoreplay: true
+    large_distro_name: ubuntu1804-build
   tasks:
   - name: compile_all_run_unittests_TG
     distros:
@@ -7548,6 +7553,7 @@ buildvariants:
     scons_cache_scope: shared
     tooltags: "ssl sasl gssapi"
     build_mongoreplay: true
+    large_distro_name: ubuntu1804-build
   tasks:
   - name: compile_all_run_unittests_TG
     distros:
@@ -7595,6 +7601,7 @@ buildvariants:
     scons_cache_scope: shared
     tooltags: "ssl"
     build_mongoreplay: true
+    large_distro_name: ubuntu1604-build
   tasks:
   - name: compile_all_run_unittests_TG
     distros:
@@ -10230,6 +10237,7 @@ buildvariants:
     tooltags: "ssl sasl gssapi"
     scons_cache_scope: shared
     build_mongoreplay: true
+    large_distro_name: debian10-build
   tasks:
   - name: compile_all_run_unittests_TG
     distros:
@@ -10278,6 +10286,7 @@ buildvariants:
     tooltags: "ssl"
     scons_cache_scope: shared
     build_mongoreplay: true
+    large_distro_name: debian92-build
   tasks:
   - name: compile_all_run_unittests_TG
     distros:
@@ -10335,6 +10344,7 @@ buildvariants:
     tooltags: "ssl"
     scons_cache_scope: shared
     build_mongoreplay: true
+    large_distro_name: debian10-build
   tasks:
   - name: compile_all_run_unittests_TG
     distros:
@@ -10875,6 +10885,7 @@ buildvariants:
     multiversion_edition: enterprise
     test_flags: |- # Set the taskExecutorPoolSize for all tests
         --mongosSetParameters="taskExecutorPoolSize: 4"
+    large_distro_name: ubuntu1604-build
   tasks:
   - name: compile_all_run_unittests_TG
     distros:
diff --git a/etc/generate_subtasks_config.yml b/etc/generate_subtasks_config.yml
new file mode 100644
index 00000000000..326067b0ef1
--- /dev/null
+++ b/etc/generate_subtasks_config.yml
@@ -0,0 +1,44 @@
+build_variant_large_distro_exceptions:
+  - amazon
+  - amazon2
+  - debian10
+  - debian92
+  - enterprise-amazon2
+  - enterprise-debian10-64
+  - enterprise-debian92-64
+  - enterprise-linux-64-amazon-ami
+  - enterprise-macos
+  - enterprise-rhel-62-64-bit-coverage
+  - enterprise-rhel-67-s390x
+  - enterprise-rhel-70-64-bit
+  - enterprise-rhel-70-64-bit-no-libunwind
+  - enterprise-rhel-71-ppc64le
+  - enterprise-rhel-71-ppc64le-inmem
+  - enterprise-rhel-72-s390x
+  - enterprise-rhel-72-s390x-inmem
+  - enterprise-rhel-80-64-bit
+  - enterprise-rhel-81-ppc64le
+  - enterprise-rhel-82-arm64
+  - enterprise-suse12-64
+  - enterprise-suse12-s390x
+  - enterprise-suse15-64
+  - enterprise-ubuntu1604-64
+  - enterprise-ubuntu1604-arm64
+  - enterprise-ubuntu1804-arm64
+  - enterprise-ubuntu1804-ppc64le
+  - enterprise-ubuntu1804-s390x
+  - enterprise-ubuntu2004-arm64
+  - hot_backups-rhel-70-64-bit
+  - linux-64-ephemeralForTest
+  - macos
+  - rhel62
+  - rhel70
+  - rhel80
+  - rhel-82-arm64
+  - suse12
+  - suse15
+  - ubi7
+  - ubuntu1604-container
+  - ubuntu1604-debug
+  - ubuntu1804-debug-asan
+  - ubuntu1804-debug-ubsan
diff --git a/etc/pip/components/evergreen.req b/etc/pip/components/evergreen.req
index b5ab0d8b133..fad05a60d9f 100644
--- a/etc/pip/components/evergreen.req
+++ b/etc/pip/components/evergreen.req
@@ -1,4 +1,5 @@
 click ~= 7.0
 GitPython ~= 2.1.11
 psutil
+pydantic ~= 1.7.3
 structlog ~= 19.1.0