feat: add target_parallelism_per_worker field for scaling (risingwavelabs#11945)

Author: shanicky

proto/meta.proto

@@ -431,6 +431,7 @@ message GetReschedulePlanRequest {
     repeated uint32 include_worker_ids = 1;
     repeated uint32 exclude_worker_ids = 2;
     optional uint32 target_parallelism = 3;
+    optional uint32 target_parallelism_per_worker = 4;
   }
 
   message StableResizePolicy {

src/ctl/src/cmd_impl/scale/resize.rs

@@ -120,6 +120,7 @@ pub async fn resize(context: &CtlContext, resize: ScaleResizeCommands) -> anyhow
         exclude_workers,
         include_workers,
         target_parallelism,
+        target_parallelism_per_worker,
         generate,
         output,
         yes,
@@ -132,8 +133,15 @@ pub async fn resize(context: &CtlContext, resize: ScaleResizeCommands) -> anyhow
         let include_worker_ids =
             worker_input_to_worker_ids(include_workers.unwrap_or_default(), true);
 
-        if let Some(target) = target_parallelism && target == 0 {
-            fail!("Target parallelism must be greater than 0");
+        match (target_parallelism, target_parallelism_per_worker) {
+            (Some(_), Some(_)) => {
+                fail!("Cannot specify both target parallelism and target parallelism per worker")
+            }
+            (_, Some(_)) if include_worker_ids.is_empty() => {
+                fail!("Cannot specify target parallelism per worker without including any worker")
+            }
+            (Some(target), _) if target == 0 => fail!("Target parallelism must be greater than 0"),
+            _ => {}
         }
 
         for worker_id in exclude_worker_ids.iter().chain(include_worker_ids.iter()) {
@@ -161,6 +169,7 @@ pub async fn resize(context: &CtlContext, resize: ScaleResizeCommands) -> anyhow
             include_worker_ids,
             exclude_worker_ids,
             target_parallelism,
+            target_parallelism_per_worker,
         }
     };
 

src/ctl/src/lib.rs

@@ -288,6 +288,15 @@ pub struct ScaleResizeCommands {
     #[clap(long)]
     target_parallelism: Option<u32>,
 
+    /// The target parallelism per worker, conflicts with `target_parallelism`, requires
+    /// `include_workers` to be set.
+    #[clap(
+        long,
+        requires = "include_workers",
+        conflicts_with = "target_parallelism"
+    )]
+    target_parallelism_per_worker: Option<u32>,
+
     /// Will generate a plan supported by the `reschedule` command and save it to the provided path
     /// by the `--output`.
     #[clap(long, default_value_t = false)]

src/meta/src/stream/scale.rs

@@ -1689,6 +1689,7 @@ where
             include_worker_ids: BTreeSet<WorkerId>,
             exclude_worker_ids: BTreeSet<WorkerId>,
             target_parallelism: Option<usize>,
+            target_parallelism_per_worker: Option<usize>,
         }
 
         let mut fragment_worker_changes: HashMap<_, _> = fragment_worker_changes
@@ -1700,6 +1701,9 @@ where
                         include_worker_ids: changes.include_worker_ids.into_iter().collect(),
                         exclude_worker_ids: changes.exclude_worker_ids.into_iter().collect(),
                         target_parallelism: changes.target_parallelism.map(|p| p as usize),
+                        target_parallelism_per_worker: changes
+                            .target_parallelism_per_worker
+                            .map(|p| p as usize),
                     },
                 )
             })
@@ -1718,6 +1722,7 @@ where
                 include_worker_ids,
                 exclude_worker_ids,
                 target_parallelism,
+                target_parallelism_per_worker,
             },
         ) in fragment_worker_changes
         {
@@ -1757,19 +1762,69 @@ where
                 })
                 .collect();
 
+            let include_worker_parallel_unit_ids = include_worker_ids
+                .iter()
+                .flat_map(|worker_id| worker_parallel_units.get(worker_id).unwrap())
+                .cloned()
+                .collect_vec();
+
+            let exclude_worker_parallel_unit_ids = exclude_worker_ids
+                .iter()
+                .flat_map(|worker_id| worker_parallel_units.get(worker_id).unwrap())
+                .cloned()
+                .collect_vec();
+
+            fn refilter_parallel_unit_id_by_target_parallelism(
+                worker_parallel_units: &HashMap<u32, HashSet<ParallelUnitId>>,
+                include_worker_ids: &BTreeSet<WorkerId>,
+                include_worker_parallel_unit_ids: &[ParallelUnitId],
+                target_parallel_unit_ids: &mut BTreeSet<ParallelUnitId>,
+                target_parallelism_per_worker: usize,
+            ) {
+                let limited_worker_parallel_unit_ids = include_worker_ids
+                    .iter()
+                    .flat_map(|worker_id| {
+                        worker_parallel_units
+                            .get(worker_id)
+                            .cloned()
+                            .unwrap()
+                            .into_iter()
+                            .sorted()
+                            .take(target_parallelism_per_worker)
+                    })
+                    .collect_vec();
+
+                // remove all the parallel units in the limited workers
+                target_parallel_unit_ids
+                    .retain(|id| !include_worker_parallel_unit_ids.contains(id));
+
+                // then we re-add the limited parallel units from the limited workers
+                target_parallel_unit_ids.extend(limited_worker_parallel_unit_ids.into_iter());
+            }
+
             match fragment.get_distribution_type().unwrap() {
                 FragmentDistributionType::Unspecified => unreachable!(),
                 FragmentDistributionType::Single => {
                     let single_parallel_unit_id =
                         fragment_parallel_unit_ids.iter().exactly_one().unwrap();
 
-                    let target_parallel_unit_ids: BTreeSet<_> = worker_parallel_units
+                    let mut target_parallel_unit_ids: BTreeSet<_> = worker_parallel_units
                         .keys()
                         .filter(|id| !unschedulable_worker_ids.contains(*id))
                         .filter(|id| !exclude_worker_ids.contains(*id))
                         .flat_map(|id| worker_parallel_units.get(id).cloned().unwrap())
                         .collect();
 
+                    if let Some(target_parallelism_per_worker) = target_parallelism_per_worker {
+                        refilter_parallel_unit_id_by_target_parallelism(
+                            &worker_parallel_units,
+                            &include_worker_ids,
+                            &include_worker_parallel_unit_ids,
+                            &mut target_parallel_unit_ids,
+                            target_parallelism_per_worker,
+                        );
+                    }
+
                     if target_parallel_unit_ids.is_empty() {
                         bail!(
                             "No schedulable ParallelUnits available for single distribution fragment {}",
@@ -1796,18 +1851,6 @@ where
                     }
                 }
                 FragmentDistributionType::Hash => {
-                    let include_worker_parallel_unit_ids = include_worker_ids
-                        .iter()
-                        .flat_map(|worker_id| worker_parallel_units.get(worker_id).unwrap())
-                        .cloned()
-                        .collect_vec();
-
-                    let exclude_worker_parallel_unit_ids = exclude_worker_ids
-                        .iter()
-                        .flat_map(|worker_id| worker_parallel_units.get(worker_id).unwrap())
-                        .cloned()
-                        .collect_vec();
-
                     let mut target_parallel_unit_ids: BTreeSet<_> =
                         fragment_parallel_unit_ids.clone();
                     target_parallel_unit_ids.extend(include_worker_parallel_unit_ids.iter());
@@ -1821,15 +1864,30 @@ where
                         );
                     }
 
-                    if let Some(target_parallelism) = target_parallelism {
-                        if target_parallel_unit_ids.len() < target_parallelism {
-                            bail!("Target parallelism {} is greater than schedulable ParallelUnits {}", target_parallelism, target_parallel_unit_ids.len());
+                    match (target_parallelism, target_parallelism_per_worker) {
+                        (Some(_), Some(_)) => {
+                            bail!("Cannot specify both target parallelism and target parallelism per worker");
                         }
+                        (Some(target_parallelism), _) => {
+                            if target_parallel_unit_ids.len() < target_parallelism {
+                                bail!("Target parallelism {} is greater than schedulable ParallelUnits {}", target_parallelism, target_parallel_unit_ids.len());
+                            }
 
-                        target_parallel_unit_ids = target_parallel_unit_ids
-                            .into_iter()
-                            .take(target_parallelism)
-                            .collect();
+                            target_parallel_unit_ids = target_parallel_unit_ids
+                                .into_iter()
+                                .take(target_parallelism)
+                                .collect();
+                        }
+                        (_, Some(target_parallelism_per_worker)) => {
+                            refilter_parallel_unit_id_by_target_parallelism(
+                                &worker_parallel_units,
+                                &include_worker_ids,
+                                &include_worker_parallel_unit_ids,
+                                &mut target_parallel_unit_ids,
+                                target_parallelism_per_worker,
+                            );
+                        }
+                        _ => {}
                     }
 
                     let to_expand_parallel_units = target_parallel_unit_ids

src/tests/simulation/tests/integration_tests/scale/plan.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::collections::HashMap;
+use std::default::Default;
 
 use anyhow::Result;
 use itertools::Itertools;
@@ -65,7 +66,7 @@ async fn test_resize_normal() -> Result<()> {
                 WorkerChanges {
                     include_worker_ids: vec![],
                     exclude_worker_ids: removed_workers,
-                    target_parallelism: None,
+                    ..Default::default()
                 },
             )]),
         }))
@@ -148,7 +149,7 @@ async fn test_resize_single() -> Result<()> {
                 WorkerChanges {
                     include_worker_ids: vec![],
                     exclude_worker_ids: vec![prev_worker.id],
-                    target_parallelism: None,
+                    ..Default::default()
                 },
             )]),
         }))
@@ -223,15 +224,15 @@ async fn test_resize_single_failed() -> Result<()> {
                     WorkerChanges {
                         include_worker_ids: vec![],
                         exclude_worker_ids: vec![worker_a.id],
-                        target_parallelism: None,
+                        ..Default::default()
                     },
                 ),
                 (
                     downstream_fragment_id,
                     WorkerChanges {
                         include_worker_ids: vec![],
                         exclude_worker_ids: vec![worker_b.id],
-                        target_parallelism: None,
+                        ..Default::default()
                     },
                 ),
             ]),
@@ -302,7 +303,7 @@ join mv5 on mv1.v = mv5.v;",
                 WorkerChanges {
                     include_worker_ids: vec![],
                     exclude_worker_ids: removed_worker_ids,
-                    target_parallelism: None,
+                    ..Default::default()
                 },
             )]),
         }))