SERVER-38115 Consolidate setting of resmoke.py --job to a python script

4 files changed, 211 insertions, 3 deletions

diff --git a/buildscripts/evergreen_gen_fuzzer_tests.py b/buildscripts/evergreen_gen_fuzzer_tests.py
index 71f39271eab..c22df9c972d 100755
--- a/buildscripts/evergreen_gen_fuzzer_tests.py
+++ b/buildscripts/evergreen_gen_fuzzer_tests.py
@@ -28,6 +28,7 @@ ConfigOptions = namedtuple("ConfigOptions", [
     "name",
     "variant",
     "continue_on_failure",
+    "resmoke_jobs_max",
     "should_shuffle",
     "timeout_secs",
     "use_multiversion",
@@ -62,7 +63,7 @@ def _get_config_value(attrib, cmd_line_options, config_file_data, required=False
     return default
 
 
-def _get_config_options(cmd_line_options, config_file):
+def _get_config_options(cmd_line_options, config_file):  # pylint: disable=too-many-locals
     """
     Get the configuration to use.
 
@@ -90,6 +91,8 @@ def _get_config_options(cmd_line_options, config_file):
     variant = _get_config_value("build_variant", cmd_line_options, config_file_data, required=True)
     continue_on_failure = _get_config_value("continue_on_failure", cmd_line_options,
                                             config_file_data, default="false")
+    resmoke_jobs_max = _get_config_value("resmoke_jobs_max", cmd_line_options, config_file_data,
+                                         default="0")
     should_shuffle = _get_config_value("should_shuffle", cmd_line_options, config_file_data,
                                        default="false")
     timeout_secs = _get_config_value("timeout_secs", cmd_line_options, config_file_data,
@@ -98,8 +101,8 @@ def _get_config_options(cmd_line_options, config_file):
                                          default=False)
 
     return ConfigOptions(num_files, num_tasks, resmoke_args, npm_command, jstestfuzz_vars, name,
-                         variant, continue_on_failure, should_shuffle, timeout_secs,
-                         use_multiversion)
+                         variant, continue_on_failure, resmoke_jobs_max, should_shuffle,
+                         timeout_secs, use_multiversion)
 
 
 def _name_task(parent_name, task_index, total_tasks):
@@ -146,6 +149,7 @@ def _generate_evg_tasks(options):
         run_tests_vars = {
             "continue_on_failure": options.continue_on_failure,
             "resmoke_args": options.resmoke_args,
+            "resmoke_jobs_max": options.resmoke_jobs_max,
             "should_shuffle": options.should_shuffle,
             "task_path_suffix": options.use_multiversion,
             "timeout_secs": options.timeout_secs,
@@ -181,6 +185,8 @@ def main():
                         help="Task path suffix for multiversion generated tasks.")
     parser.add_argument("--continue-on-failure", dest="continue_on_failure",
                         help="continue_on_failure value for generated tasks.")
+    parser.add_argument("--resmoke-jobs-max", dest="resmoke_jobs_max",
+                        help="resmoke_jobs_max value for generated tasks.")
     parser.add_argument("--should-shuffle", dest="should_shuffle",
                         help="should_shuffle value for generated tasks.")
     parser.add_argument("--timeout-secs", dest="timeout_secs",
diff --git a/buildscripts/evergreen_resmoke_job_count.py b/buildscripts/evergreen_resmoke_job_count.py
new file mode 100644
index 00000000000..098c65130cc
--- /dev/null
+++ b/buildscripts/evergreen_resmoke_job_count.py
@@ -0,0 +1,98 @@
+#!/usr/bin/env python
+"""Determine the number of resmoke jobs to run."""
+
+from __future__ import division
+from __future__ import print_function
+
+import argparse
+import platform
+import re
+import sys
+
+import psutil
+import yaml
+
+CPU_COUNT = psutil.cpu_count()
+PLATFORM_MACHINE = platform.machine()
+SYS_PLATFORM = sys.platform
+
+VARIANT_TASK_FACTOR_OVERRIDES = {
+    "enterprise-rhel-62-64-bit-inmem": [{"task": "secondary_reads_passthrough", "factor": 0.3}]
+}
+
+TASKS_FACTORS = [{"task": "replica_sets*", "factor": 0.5}, {"task": "sharding.*", "factor": 0.5}]
+
+MACHINE_TASK_FACTOR_OVERRIDES = {"aarch64": TASKS_FACTORS}
+
+PLATFORM_TASK_FACTOR_OVERRIDES = {"win32": TASKS_FACTORS, "cygwin": TASKS_FACTORS}
+
+
+def get_task_factor(task_name, overrides, override_type, factor):
+    """Check for task override and return factor."""
+    for task_override in overrides.get(override_type, []):
+        if re.compile(task_override["task"]).match(task_name):
+            return task_override["factor"]
+    return factor
+
+
+def determine_factor(task_name, variant, factor):
+    """Determine the job factor."""
+    factors = []
+    factors.append(
+        get_task_factor(task_name, MACHINE_TASK_FACTOR_OVERRIDES, PLATFORM_MACHINE, factor))
+    factors.append(get_task_factor(task_name, PLATFORM_TASK_FACTOR_OVERRIDES, SYS_PLATFORM, factor))
+    factors.append(get_task_factor(task_name, VARIANT_TASK_FACTOR_OVERRIDES, variant, factor))
+    return min(factors)
+
+
+def determine_jobs(task_name, variant, jobs_max=0, job_factor=1.0):
+    """Determine the resmoke jobs."""
+    if jobs_max < 0:
+        raise ValueError("The jobs_max must be >= 0.")
+    if job_factor <= 0:
+        raise ValueError("The job_factor must be > 0.")
+    factor = determine_factor(task_name, variant, job_factor)
+    jobs_available = int(round(CPU_COUNT * factor))
+    if jobs_max == 0:
+        return max(1, jobs_available)
+    return min(jobs_max, jobs_available)
+
+
+def output_jobs(jobs, outfile):
+    """Output jobs configuration to the specified location."""
+    output = {"resmoke_jobs": jobs}
+
+    if outfile:
+        with open(outfile, "w") as fh:
+            yaml.dump(output, stream=fh, default_flow_style=False)
+
+    yaml.dump(output, stream=sys.stdout, default_flow_style=False)
+
+
+def main():
+    """Determine the resmoke jobs value a task should use in Evergreen."""
+    parser = argparse.ArgumentParser(description=main.__doc__)
+
+    parser.add_argument("--taskName", dest="task", required=True, help="Task being executed.")
+    parser.add_argument("--buildVariant", dest="variant", required=True,
+                        help="Build variant task is being executed on.")
+    parser.add_argument("--jobFactor", dest="jobs_factor", type=float, default=1.0,
+                        help=("Job factor to use as a mulitplier with the number of CPUs. Defaults"
+                              " to %(default)s."))
+    parser.add_argument("--jobsMax", dest="jobs_max", type=int, default=0,
+                        help=("Maximum number of jobs to use. Specify 0 to indicate the number of"
+                              " jobs is determined by --jobFactor and the number of CPUs. Defaults"
+                              " to %(default)s."))
+    parser.add_argument("--outFile", dest="outfile", help=("File to write configuration to. If"
+                                                           " unspecified no file is generated."))
+
+    options = parser.parse_args()
+
+    jobs = determine_jobs(options.task, options.variant, options.jobs_max, options.jobs_factor)
+    if jobs < CPU_COUNT:
+        print("Reducing number of jobs to run from {} to {}".format(CPU_COUNT, jobs))
+    output_jobs(jobs, options.outfile)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/buildscripts/tests/test_evergreen_gen_fuzzer_tests.py b/buildscripts/tests/test_evergreen_gen_fuzzer_tests.py
index 06e3acf07a2..1956aa97519 100644
--- a/buildscripts/tests/test_evergreen_gen_fuzzer_tests.py
+++ b/buildscripts/tests/test_evergreen_gen_fuzzer_tests.py
@@ -59,6 +59,7 @@ class TestGenerateEvgTasks(unittest.TestCase):
         options.resmoke_args = "resmoke args"
         options.variant = "build variant"
         options.continue_on_failure = "false"
+        options.resmoke_jobs_max = 0
         options.should_shuffle = "false"
         options.timeout_secs = "1800"
 
diff --git a/buildscripts/tests/test_evergreen_resmoke_job_count.py b/buildscripts/tests/test_evergreen_resmoke_job_count.py
new file mode 100644
index 00000000000..f284e122fba
--- /dev/null
+++ b/buildscripts/tests/test_evergreen_resmoke_job_count.py
@@ -0,0 +1,103 @@
+"""Unit tests for the evergreen_resomke_job_count script."""
+
+from __future__ import division
+
+import unittest
+
+import psutil
+
+from buildscripts import evergreen_resmoke_job_count as erjc
+
+# pylint: disable=missing-docstring,no-self-use
+
+
+class DetermineJobsTest(unittest.TestCase):
+    cpu_count = psutil.cpu_count()
+    mytask = "mytask"
+    regex = "regexthatmatches"
+    mytask_factor = 0.5
+    regex_factor = 0.25
+    task_factors = [{"task": mytask, "factor": mytask_factor},
+                    {"task": "regex.*", "factor": regex_factor}]
+
+    def test_determine_jobs_no_matching_task(self):
+        jobs = erjc.determine_jobs("_no_match_", "_no_variant_", 0, 1)
+        self.assertEqual(self.cpu_count, jobs)
+
+    def test_determine_jobs_matching_variant(self):
+        erjc.VARIANT_TASK_FACTOR_OVERRIDES = {"myvariant": self.task_factors}
+        jobs = erjc.determine_jobs(self.mytask, "myvariant", 0, 1)
+        self.assertEqual(int(round(self.cpu_count * self.mytask_factor)), jobs)
+        jobs = erjc.determine_jobs(self.regex, "myvariant", 0, 1)
+        self.assertEqual(int(round(self.cpu_count * self.regex_factor)), jobs)
+
+    def test_determine_factor_matching_variant(self):
+        erjc.VARIANT_TASK_FACTOR_OVERRIDES = {"myvariant": self.task_factors}
+        factor = erjc.determine_factor(self.mytask, "myvariant", 1)
+        self.assertEqual(self.mytask_factor, factor)
+        factor = erjc.determine_factor(self.regex, "myvariant", 1)
+        self.assertEqual(self.regex_factor, factor)
+
+    def test_determine_jobs_matching_machine(self):
+        erjc.PLATFORM_MACHINE = "mymachine"
+        erjc.MACHINE_TASK_FACTOR_OVERRIDES = {"mymachine": self.task_factors}
+        jobs = erjc.determine_jobs(self.mytask, "myvariant", 0, 1)
+        self.assertEqual(int(round(self.cpu_count * self.mytask_factor)), jobs)
+        jobs = erjc.determine_jobs(self.regex, "myvariant", 0, 1)
+        self.assertEqual(int(round(self.cpu_count * self.regex_factor)), jobs)
+
+    def test_determine_factor_matching_machine(self):
+        erjc.PLATFORM_MACHINE = "mymachine"
+        erjc.MACHINE_TASK_FACTOR_OVERRIDES = {"mymachine": self.task_factors}
+        factor = erjc.determine_factor(self.mytask, "myvariant", 1)
+        self.assertEqual(self.mytask_factor, factor)
+        factor = erjc.determine_factor(self.regex, "myvariant", 1)
+        self.assertEqual(self.regex_factor, factor)
+
+    def test_determine_jobs_matching_platform(self):
+        erjc.SYS_PLATFORM = "myplatform"
+        erjc.PLATFORM_TASK_FACTOR_OVERRIDES = {"myplatform": self.task_factors}
+        jobs = erjc.determine_jobs(self.mytask, "myvariant", 0, 1)
+        self.assertEqual(int(round(self.cpu_count * self.mytask_factor)), jobs)
+        jobs = erjc.determine_jobs(self.regex, "myvariant", 0, 1)
+        self.assertEqual(int(round(self.cpu_count * self.regex_factor)), jobs)
+
+    def test_determine_factor_matching_platform(self):
+        erjc.SYS_PLATFORM = "myplatform"
+        erjc.PLATFORM_TASK_FACTOR_OVERRIDES = {"mymachine": self.task_factors}
+        factor = erjc.determine_factor(self.mytask, "myvariant", 1)
+        self.assertEqual(self.mytask_factor, factor)
+        factor = erjc.determine_factor(self.regex, "myvariant", 1)
+        self.assertEqual(self.regex_factor, factor)
+
+    def test_determine_jobs_min_factor(self):
+        erjc.PLATFORM_MACHINE = "mymachine"
+        erjc.SYS_PLATFORM = "myplatform"
+        mytask_factor_min = 0.5
+        regex_factor_min = 0.25
+        task_factors1 = [{"task": "mytask", "factor": mytask_factor_min + .5},
+                         {"task": "regex.*", "factor": regex_factor_min + .5}]
+        task_factors2 = [{"task": "mytask", "factor": mytask_factor_min + .25},
+                         {"task": "regex.*", "factor": regex_factor_min + .25}]
+        task_factors3 = [{"task": "mytask", "factor": mytask_factor_min},
+                         {"task": "regex.*", "factor": regex_factor_min}]
+        erjc.VARIANT_TASK_FACTOR_OVERRIDES = {"myvariant": task_factors1}
+        erjc.MACHINE_TASK_FACTOR_OVERRIDES = {"mymachine": task_factors2}
+        erjc.PLATFORM_TASK_FACTOR_OVERRIDES = {"myplatform": task_factors3}
+        jobs = erjc.determine_jobs(self.mytask, "myvariant", 0, 1)
+        self.assertEqual(int(round(self.cpu_count * mytask_factor_min)), jobs)
+        jobs = erjc.determine_jobs(self.regex, "myvariant", 0, 1)
+        self.assertEqual(int(round(self.cpu_count * regex_factor_min)), jobs)
+
+    def test_determine_jobs_factor(self):
+        factor = 0.4
+        jobs = erjc.determine_jobs("_no_match_", "_no_variant_", 0, factor)
+        self.assertEqual(int(round(self.cpu_count * factor)), jobs)
+
+    def test_determine_jobs_jobs_max(self):
+        jobs_max = 3
+        jobs = erjc.determine_jobs("_no_match_", "_no_variant_", jobs_max, 1)
+        self.assertEqual(min(jobs_max, jobs), jobs)
+        jobs_max = 30
+        jobs = erjc.determine_jobs("_no_match_", "_no_variant_", jobs_max, 1)
+        self.assertEqual(min(jobs_max, jobs), jobs)