Balancing prefill on heterogenous clusters

CHANGELOG.md

@@ -5,11 +5,13 @@
 ### New features
 
 * The server scheduler now contains a safety limit for computation, configurable via `--scheduler-time-limit` (default: 5s)
+* Better scheduling policy (prefill) for heterogenous clusters
 
 ### Fixes
 
+* Fixed some occasional greedy backfilling in server scheduler + improvements in the reservation algorithm
 * Fixed server crash in a specific situation when an unschedulable high-priority task occurs
-
+* Fixed server crash caused by invalid handling of prefill
 
 ## v0.26.2
 

crates/tako/src/internal/scheduler/batches.rs

@@ -5,8 +5,8 @@ use crate::resources::ResourceRqId;
 use std::cmp::Ordering;
 use std::time::Instant;
 
-const BATCH_PRUNING_MAX_SIZE: usize = 32;
 const BATCH_PRUNING_FIXED_PREFIX: usize = 4;
+const BATCH_PRUNING_GLOBAL_MAX: usize = 64;
 
 #[derive(Debug)]
 #[cfg_attr(test, derive(Eq, PartialEq))]
@@ -169,24 +169,57 @@ pub(crate) fn create_task_batches(
             };
         }
     }
-    batches.retain_mut(|b| {
-        prune_progressive(
-            &mut b.cuts,
-            BATCH_PRUNING_FIXED_PREFIX,
-            BATCH_PRUNING_MAX_SIZE,
-        );
-        b.size > 0
-    });
+    batches.retain_mut(|b| b.size > 0);
+    apply_global_cut_budget(
+        &mut batches,
+        BATCH_PRUNING_FIXED_PREFIX,
+        BATCH_PRUNING_GLOBAL_MAX,
+    );
     batches
 }
 
+fn apply_global_cut_budget(batches: &mut [TaskBatch], prefix: usize, mut budget: usize) {
+    let total: usize = batches.iter().map(|b| b.cuts.len()).sum();
+    if total <= budget {
+        return;
+    }
+    let mut granted = vec![0; batches.len()];
+    for (batch, g) in batches.iter().zip(granted.iter_mut()) {
+        let take = batch.cuts.len().min(prefix);
+        budget = budget.saturating_sub(take);
+        *g += take;
+    }
+
+    while budget > 0 {
+        for (batch, g) in batches.iter().zip(granted.iter_mut()) {
+            if *g < batch.cuts.len() {
+                *g += 1;
+                budget -= 1;
+                if budget == 0 {
+                    break;
+                }
+            }
+        }
+    }
+    for (batch, granted) in batches.iter_mut().zip(granted.iter()) {
+        if *granted < batch.cuts.len() {
+            prune_progressive(&mut batch.cuts, prefix, *granted);
+        }
+    }
+}
+
 fn prune_progressive<T>(vec: &mut Vec<T>, prefix_size: usize, size_limit: usize) {
     let original_len = vec.len();
 
     if original_len <= size_limit {
         return;
     }
 
+    if size_limit <= prefix_size + 1 {
+        vec.truncate(size_limit);
+        return;
+    }
+
     let remaining_slots = size_limit - prefix_size;
 
     let mut indices = Vec::with_capacity(size_limit);
@@ -248,4 +281,58 @@ mod tests {
             ]
         );
     }
+
+    /// A global budget can ask for a limit at or just above the prefix, where the quadratic
+    /// sampler has no slots left to place (`i / (slots - 1)` would be `0.0 / 0.0`).
+    #[test]
+    fn test_prune_progressive_at_prefix_boundary() {
+        for size_limit in 0..=5 {
+            let mut vec = (0..40).collect::<Vec<_>>();
+            prune_progressive(&mut vec, 4, size_limit);
+            assert_eq!(vec, (0..size_limit as i32).collect::<Vec<_>>());
+        }
+
+        let mut vec = (0..40).collect::<Vec<_>>();
+        prune_progressive(&mut vec, 4, 6);
+        assert_eq!(vec, vec![0, 1, 2, 3, 4, 39]);
+    }
+
+    #[test]
+    fn test_global_cut_budget() {
+        fn batch(n_cuts: usize) -> TaskBatch {
+            let mut b = TaskBatch::new(0.into(), 100, false);
+            b.cuts = (0..n_cuts)
+                .map(|i| PriorityCut {
+                    size: i as u32,
+                    blockers: Vec::new(),
+                })
+                .collect();
+            b
+        }
+        let total = |bs: &[TaskBatch]| bs.iter().map(|b| b.cuts.len()).sum::<usize>();
+
+        // Under budget: untouched.
+        let mut batches = vec![batch(3), batch(3)];
+        apply_global_cut_budget(&mut batches, 4, 32);
+        assert_eq!(total(&batches), 6);
+
+        let mut batches: Vec<_> = (0..8).map(|_| batch(3)).collect();
+        apply_global_cut_budget(&mut batches, 4, 8);
+        assert_eq!(total(&batches), 24);
+
+        // Proportional: the bigger batch keeps more, and the budget is spent exactly.
+        let mut batches = vec![batch(60), batch(10), batch(10)];
+        apply_global_cut_budget(&mut batches, 4, 16);
+        assert!(batches[0].cuts.len() > batches[1].cuts.len());
+        assert_eq!(total(&batches), 16);
+
+        let mut batches = vec![batch(60), batch(5), batch(5)];
+        apply_global_cut_budget(&mut batches, 4, 40);
+        assert_eq!(total(&batches), 40);
+
+        let mut batches = vec![batch(30), batch(5), batch(5)];
+        apply_global_cut_budget(&mut batches, 4, 39);
+        assert_eq!(total(&batches), 39);
+        assert_eq!(batches[0].cuts.len(), 29);
+    }
 }

crates/tako/src/internal/scheduler/mapping.rs

@@ -161,10 +161,15 @@ fn process_proactive_filling(core: &mut Core, mapping: &mut WorkerTaskMapping) {
         task_map,
         worker_map,
         task_queues,
-        request_map: _,
+        request_map,
         scheduler_state,
         ..
     } = core.split_mut();
+    let max_prefill = scheduler_state.config.proactive_filling_max as u64;
+    if max_prefill == 0 {
+        // Prefill explicitly disabled.
+        return;
+    }
     let top_priority = task_queues.top_priority();
     for queue in task_queues.iter_mut() {
         if queue.top_priority() != Some(top_priority) {
@@ -176,6 +181,13 @@ fn process_proactive_filling(core: &mut Core, mapping: &mut WorkerTaskMapping) {
         if size == 0 {
             continue;
         }
+        let rqv = request_map.get(queue.resource_rq_id);
+        let max_capacity = worker_map
+            .get_workers()
+            .map(|w| w.resources.task_max_count(rqv))
+            .max()
+            .unwrap_or(1)
+            .max(1) as u64;
         let workers: Vec<_> = worker_map
             .values_mut()
             .filter(|worker| {
@@ -207,28 +219,45 @@ fn process_proactive_filling(core: &mut Core, mapping: &mut WorkerTaskMapping) {
         if workers.is_empty() {
             continue;
         }
-        let prefill_size =
-            (size / workers.len() as u32).min(scheduler_state.config.proactive_filling_max);
-        if prefill_size == 0 {
-            continue;
-        }
-        for worker in workers {
-            let tasks = queue.take_tasks_for_prefill(prefill_size);
-            for task_id in &tasks {
+        let capacities: Vec<u64> = workers
+            .iter()
+            .map(|w| w.resources.task_max_count(rqv).max(1) as u64)
+            .collect();
+        let total_capacity: u64 = capacities.iter().sum();
+        let mut return_back = Vec::new();
+        for (worker, capacity) in workers.into_iter().zip(capacities) {
+            // The shares sum to at most `size`, so the queue entry we are drawing from is
+            // never exhausted before the last worker.
+            let share = size as u64 * capacity / total_capacity;
+            let depth = (max_prefill * capacity / max_capacity).max(1);
+            let prefill_size = share.min(depth);
+            if prefill_size == 0 {
+                continue;
+            }
+            let prefills = &mut mapping.workers.entry(worker.id).or_default().prefills;
+            for _ in 0..prefill_size {
+                let task_id = queue.take_one().unwrap();
                 log::debug!("Prefiling task={task_id} to worker={}", worker.id);
-                let task = task_map.get_task_mut(*task_id);
-                assert!(task.is_waiting());
-                task.state = TaskRuntimeState::Prefilled {
-                    worker_id: worker.id,
-                };
-                worker.insert_prefill_task(*task_id);
+                let task = task_map.get_task_mut(task_id);
+                if task.is_waiting() {
+                    task.state = TaskRuntimeState::Prefilled {
+                        worker_id: worker.id,
+                    };
+                    worker.insert_prefill_task(task_id);
+                    queue.insert_prefill(task_id, top_priority, prefill_size as usize);
+                    prefills.push(task_id);
+                } else {
+                    // This can happen when task is in retracting, and it should be queite rare
+                    log::debug!(
+                        "Task is not in waiting state ({:?}) back to the queue.",
+                        task.state
+                    );
+                    return_back.push(task_id);
+                }
             }
-            mapping
-                .workers
-                .entry(worker.id)
-                .or_default()
-                .prefills
-                .extend(tasks);
+        }
+        for task_id in return_back {
+            queue.return_back(task_id, top_priority);
         }
     }
 }