DEVPROD-17023 Split large tasks on required variants based on total test runtime

src/evergreen/evg_config.rs

@@ -4,7 +4,7 @@ use std::{collections::HashMap, path::Path, process::Command};
 
 use shrub_rs::models::{project::EvgProject, task::EvgTask, variant::BuildVariant};
 
-const REQUIRED_PREFIX: &str = "!";
+use crate::REQUIRED_PREFIX;
 
 pub trait EvgConfigService: Sync + Send {
     /// Get a map of build variant names to build variant definitions.

src/lib.rs

@@ -58,7 +58,7 @@ mod utils;
 const BURN_IN_TESTS_PREFIX: &str = "burn_in_tests";
 const BURN_IN_TASKS_PREFIX: &str = "burn_in_tasks";
 const BURN_IN_BV_SUFFIX: &str = "generated-by-burn-in-tags";
-const DEFAULT_SUB_TASKS_PER_TASK: usize = 5;
+const REQUIRED_PREFIX: &str = "!";
 
 type GenTaskCollection = HashMap<String, Box<dyn GeneratedSuite>>;
 
@@ -145,6 +145,24 @@ pub struct ExecutionConfiguration<'a> {
     pub burn_in_tests_command: &'a str,
     /// S3 endpoint to get test stats from.
     pub s3_test_stats_endpoint: &'a str,
+    pub subtask_limits: SubtaskLimits,
+}
+
+#[derive(Debug, Clone)]
+pub struct SubtaskLimits {
+    // Ideal total test runtime (in seconds) for individual subtasks on required
+    // variants, used to determine the number of subtasks for tasks on required variants.
+    pub test_runtime_per_required_subtask: f64,
+
+    // Threshold of total test runtime (in seconds) for a required task to be considered
+    // large enough to warrant splitting into more that the default number of tasks.
+    pub large_required_task_runtime_threshold: f64,
+
+    // Default number of subtasks that should be generated for tasks
+    pub default_subtasks_per_task: usize,
+
+    // Maximum number of subtasks that can be generated for tasks
+    pub max_subtasks_per_task: usize,
 }
 
 /// Collection of services needed to execution.
@@ -215,6 +233,7 @@ impl Dependencies {
             multiversion_service,
             fs_service,
             gen_resmoke_config,
+            execution_config.subtask_limits,
         ));
         let gen_task_service = Arc::new(GenerateTasksServiceImpl::new(
             evg_config_service,
@@ -589,7 +608,7 @@ impl GenerateTasksService for GenerateTasksServiceImpl {
             )?;
             Some(
                 self.gen_resmoke_service
-                    .generate_resmoke_task(&params, &build_variant.name)
+                    .generate_resmoke_task(&params, build_variant)
                     .await?,
             )
         };
@@ -1021,7 +1040,7 @@ mod tests {
         async fn generate_resmoke_task(
             &self,
             _params: &ResmokeGenParams,
-            _build_variant: &str,
+            _build_variant: &BuildVariant,
         ) -> Result<Box<dyn GeneratedSuite>> {
             todo!()
         }

src/main.rs

@@ -7,7 +7,7 @@ use std::{
 use anyhow::Result;
 use clap::Parser;
 use mongo_task_generator::{
-    generate_configuration, Dependencies, ExecutionConfiguration, ProjectInfo,
+    generate_configuration, Dependencies, ExecutionConfiguration, ProjectInfo, SubtaskLimits,
 };
 use serde::Deserialize;
 use tracing::{error, event, Level};
@@ -19,6 +19,10 @@ const DEFAULT_RESMOKE_COMMAND: &str = "python buildscripts/resmoke.py";
 const DEFAULT_BURN_IN_TESTS_COMMAND: &str = "python buildscripts/burn_in_tests.py run";
 const DEFAULT_TARGET_DIRECTORY: &str = "generated_resmoke_config";
 const DEFAULT_S3_TEST_STATS_ENDPOINT: &str = "https://mongo-test-stats.s3.amazonaws.com";
+const DEFAULT_MAX_SUBTASKS_PER_TASK: &str = "10";
+const DEFAULT_DEFAULT_SUBTASKS_PER_TASKS: &str = "5";
+const DEFAULT_TEST_RUNTIME_PER_REQUIRED_SUBTASK: &str = "3600";
+const DEFAULT_LARGE_REQUIRED_TASK_RUNTIME_THRESHOLD: &str = "7200";
 
 /// Expansions from evergreen to determine settings for how task should be generated.
 #[derive(Debug, Deserialize)]
@@ -129,6 +133,24 @@ struct Args {
     /// S3 endpoint to get test stats from.
     #[clap(long, default_value = DEFAULT_S3_TEST_STATS_ENDPOINT)]
     s3_test_stats_endpoint: String,
+
+    // Ideal total test runtime (in seconds) for individual subtasks on required
+    // variants, used to determine the number of subtasks for tasks on required variants.
+    #[clap(long, default_value = DEFAULT_TEST_RUNTIME_PER_REQUIRED_SUBTASK)]
+    test_runtime_per_required_subtask: f64,
+
+    // Threshold of total test runtime (in seconds) for a required task to be considered
+    // large enough to warrant splitting into more that the default number of tasks.
+    #[clap(long, default_value = DEFAULT_LARGE_REQUIRED_TASK_RUNTIME_THRESHOLD)]
+    large_required_task_runtime_threshold: f64,
+
+    // Default number of subtasks that should be generated for tasks
+    #[clap(long, default_value = DEFAULT_DEFAULT_SUBTASKS_PER_TASKS)]
+    default_subtasks_per_task: usize,
+
+    // Maximum number of subtasks that can be generated for tasks
+    #[clap(long, default_value = DEFAULT_MAX_SUBTASKS_PER_TASK)]
+    max_subtasks_per_task: usize,
 }
 
 /// Configure logging for the command execution.
@@ -165,6 +187,12 @@ async fn main() {
         include_fully_disabled_feature_tests: args.include_fully_disabled_feature_tests,
         burn_in_tests_command: &args.burn_in_tests_command,
         s3_test_stats_endpoint: &args.s3_test_stats_endpoint,
+        subtask_limits: SubtaskLimits {
+            test_runtime_per_required_subtask: args.test_runtime_per_required_subtask,
+            max_subtasks_per_task: args.max_subtasks_per_task,
+            default_subtasks_per_task: args.default_subtasks_per_task,
+            large_required_task_runtime_threshold: args.large_required_task_runtime_threshold,
+        },
     };
     let deps = Dependencies::new(execution_config).unwrap();
 

src/services/config_extraction.rs

@@ -18,7 +18,6 @@ use crate::{
         multiversion::MultiversionService, resmoke_tasks::ResmokeGenParams,
     },
     utils::task_name::remove_gen_suffix,
-    DEFAULT_SUB_TASKS_PER_TASK,
 };
 
 /// Interface for performing extractions of evergreen project configuration.
@@ -249,13 +248,10 @@ impl ConfigExtractionService for ConfigExtractionServiceImpl {
                 .evg_config_utils
                 .lookup_build_variant_expansion(UNIQUE_GEN_SUFFIX_EXPANSION, variant);
         }
-        let num_tasks = match self
+        let num_tasks = self
             .evg_config_utils
             .get_gen_task_var(task_def, "num_tasks")
-        {
-            Some(str) => str.parse().unwrap(),
-            _ => DEFAULT_SUB_TASKS_PER_TASK,
-        };
+            .map(|str| str.parse().unwrap());
 
         Ok(ResmokeGenParams {
             task_name,

src/task_types/burn_in_tests.rs

@@ -630,7 +630,7 @@ mod tests {
         async fn generate_resmoke_task(
             &self,
             _params: &ResmokeGenParams,
-            _build_variant: &str,
+            _build_variant: &BuildVariant,
         ) -> Result<Box<dyn GeneratedSuite>> {
             todo!()
         }