Limit using load instead of shard number

Author: rdettai

quickwit/quickwit-control-plane/src/indexing_scheduler/mod.rs

@@ -49,6 +49,9 @@ pub(crate) const MIN_DURATION_BETWEEN_SCHEDULING: Duration =
         Duration::from_secs(30)
     };
 
+/// That's 80% of a pipeline capacity
+const MAX_LOAD_PER_PIPELINE: CpuCapacity = CpuCapacity::from_cpu_millis(3_200);
+
 #[derive(Debug, Clone, Default, Serialize)]
 pub struct IndexingSchedulerState {
     pub num_applied_physical_indexing_plan: usize,
@@ -221,7 +224,7 @@ fn get_sources_to_schedule(model: &ControlPlaneModel) -> Vec<SourceToSchedule> {
                     source_type: SourceToScheduleType::NonSharded {
                         num_pipelines: source_config.num_pipelines.get() as u32,
                         // FIXME
-                        load_per_pipeline: NonZeroU32::new(PIPELINE_FULL_CAPACITY.cpu_millis())
+                        load_per_pipeline: NonZeroU32::new(MAX_LOAD_PER_PIPELINE.cpu_millis())
                             .unwrap(),
                     },
                     params_fingerprint,

quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs

@@ -26,6 +26,7 @@ use scheduling_logic_model::{IndexerOrd, SourceOrd};
 use tracing::{error, warn};
 
 use crate::indexing_plan::PhysicalIndexingPlan;
+use crate::indexing_scheduler::MAX_LOAD_PER_PIPELINE;
 use crate::indexing_scheduler::scheduling::scheduling_logic_model::{
     IndexerAssignment, SchedulingProblem, SchedulingSolution,
 };
@@ -44,9 +45,6 @@ use crate::model::ShardLocations;
 /// of 30%. Which translates into an overall load of 60%.
 const CPU_PER_PIPELINE_LOAD_LOWER_THRESHOLD: CpuCapacity = CpuCapacity::from_cpu_millis(1_200);
 
-/// That's 80% of a period
-const MAX_LOAD_PER_PIPELINE: CpuCapacity = CpuCapacity::from_cpu_millis(3_200);
-
 fn populate_problem(
     source: &SourceToSchedule,
     problem: &mut SchedulingProblem,
@@ -748,7 +746,7 @@ mod tests {
 
     use fnv::FnvHashMap;
     use itertools::Itertools;
-    use quickwit_proto::indexing::{CpuCapacity, IndexingTask, mcpu};
+    use quickwit_proto::indexing::{CpuCapacity, IndexingTask, PIPELINE_FULL_CAPACITY, mcpu};
     use quickwit_proto::types::{IndexUid, NodeId, PipelineUid, ShardId, SourceUid};
     use rand::seq::SliceRandom;
 
@@ -939,6 +937,71 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_build_physical_plan_with_pipeline_limit() {
+        let indexer1 = "indexer1".to_string();
+        let indexer2 = "indexer2".to_string();
+        let source_uid0 = source_id();
+        let source_uid1 = source_id();
+        let source_0 = SourceToSchedule {
+            source_uid: source_uid0.clone(),
+            source_type: SourceToScheduleType::Sharded {
+                shard_ids: (0..16).map(ShardId::from).collect(),
+                load_per_shard: NonZeroU32::new(800).unwrap(),
+            },
+            params_fingerprint: 0,
+        };
+        let source_1 = SourceToSchedule {
+            source_uid: source_uid1.clone(),
+            source_type: SourceToScheduleType::NonSharded {
+                num_pipelines: 4,
+                load_per_pipeline: NonZeroU32::new(PIPELINE_FULL_CAPACITY.cpu_millis()).unwrap(),
+            },
+            params_fingerprint: 0,
+        };
+        let mut indexer_id_to_cpu_capacities = FnvHashMap::default();
+        indexer_id_to_cpu_capacities.insert(indexer1.clone(), mcpu(16_000));
+        indexer_id_to_cpu_capacities.insert(indexer2.clone(), mcpu(16_000));
+        let shard_locations = ShardLocations::default();
+        let indexing_plan = build_physical_indexing_plan(
+            &[source_0, source_1],
+            &indexer_id_to_cpu_capacities,
+            None,
+            &shard_locations,
+        );
+        assert_eq!(indexing_plan.indexing_tasks_per_indexer().len(), 2);
+
+        let node1_plan = indexing_plan.indexer(&indexer1).unwrap();
+        let node2_plan = indexing_plan.indexer(&indexer2).unwrap();
+
+        println!("node1_plan: {node1_plan:#?}");
+        println!("node2_plan: {node2_plan:#?}");
+
+        let source_0_on_node1 = node1_plan
+            .iter()
+            .filter(|task| task.source_id == source_uid0.source_id)
+            .count();
+        let source_0_on_node2 = node2_plan
+            .iter()
+            .filter(|task| task.source_id == source_uid0.source_id)
+            .count();
+        assert!(source_0_on_node1 <= 3);
+        assert!(source_0_on_node2 <= 3);
+        assert_eq!(source_0_on_node1 + source_0_on_node2, 4);
+
+        let source_1_on_node1 = node1_plan
+            .iter()
+            .filter(|task| task.source_id == source_uid1.source_id)
+            .count();
+        let source_1_on_node2 = node2_plan
+            .iter()
+            .filter(|task| task.source_id == source_uid1.source_id)
+            .count();
+        assert!(source_1_on_node1 <= 3);
+        assert!(source_1_on_node2 <= 3);
+        assert_eq!(source_1_on_node1 + source_1_on_node2, 4);
+    }
+
     fn make_indexing_tasks(
         source_uid: &SourceUid,
         shards: &[(PipelineUid, &[ShardId])],

quickwit/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs

@@ -20,6 +20,7 @@ use itertools::Itertools;
 use quickwit_proto::indexing::CpuCapacity;
 
 use super::scheduling_logic_model::*;
+use crate::indexing_scheduler::MAX_LOAD_PER_PIPELINE;
 use crate::indexing_scheduler::scheduling::inflate_node_capacities_if_necessary;
 
 // ------------------------------------------------------------------------------------
@@ -355,9 +356,12 @@ fn place_unassigned_shards_single_source(
     solution: &mut SchedulingSolution,
 ) -> Result<(), NotEnoughCapacity> {
     let mut num_shards = source.num_shards;
-    // To ensure that merges can keep up, try not to assign more than 3
-    // shards per indexer for a source (except if there aren't enough nodes)
-    let limit_num_shards_per_indexer_per_source = 3.max(num_shards.div_ceil(num_indexers as u32));
+    // To ensure that merges can keep up, we try not to assign more than 3
+    // pipelines per indexer for a source (except if there aren't enough nodes).
+    let target_limit_num_shards_per_indexer_per_source =
+        3 * MAX_LOAD_PER_PIPELINE.cpu_millis() / source.load_per_shard.get();
+    let limit_num_shards_per_indexer_per_source = target_limit_num_shards_per_indexer_per_source
+        .max(num_shards.div_ceil(num_indexers as u32));
     while num_shards > 0 {
         let Some((indexer_ord, available_capacity)) = indexer_with_capacities.next() else {
             return Err(NotEnoughCapacity);
@@ -608,6 +612,21 @@ mod tests {
         problem.add_source(4, NonZeroU32::new(1_000).unwrap());
         problem.add_source(4, NonZeroU32::new(1_000).unwrap());
         problem.inc_affinity(0, 1);
+        problem.inc_affinity(1, 0);
+        let mut solution = problem.new_solution();
+        place_unassigned_shards_with_affinity(&problem, &mut solution);
+        assert_eq!(solution.indexer_assignments[0].num_shards(1), 4);
+        assert_eq!(solution.indexer_assignments[1].num_shards(0), 4);
+    }
+
+    #[test]
+    fn test_limit_placement_to_three_pipelines() {
+        let mut problem =
+            SchedulingProblem::with_indexer_cpu_capacities(vec![mcpu(16_000), mcpu(16_000)]);
+        let max_load_per_pipeline = NonZeroU32::new(MAX_LOAD_PER_PIPELINE.cpu_millis()).unwrap();
+        problem.add_source(4, max_load_per_pipeline);
+        problem.add_source(4, max_load_per_pipeline);
+        problem.inc_affinity(0, 1);
         problem.inc_affinity(0, 1);
         problem.inc_affinity(0, 0);
         problem.inc_affinity(1, 0);