Balancing prefill on heterogenous clusters

CHANGELOG.md

@@ -5,6 +5,7 @@
 ### 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
 

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,12 +219,21 @@ 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 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();
+
+        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) as u32;
+            if prefill_size == 0 {
+                continue;
+            }
             let tasks = queue.take_tasks_for_prefill(prefill_size);
             for task_id in &tasks {
                 log::debug!("Prefiling task={task_id} to worker={}", worker.id);

crates/tako/src/internal/tests/test_scheduler_sn.rs

@@ -1305,6 +1305,98 @@ fn test_prefill_steal() {
     rt.sanity_check();
 }
 
+/// Prefill has to be weighted by worker capacity, otherwise a small worker is handed as many
+/// tasks as a large one and takes proportionally longer to drain them. Prefilled tasks are
+/// removed from the global queue and are not reclaimed while regular tasks of the same priority
+/// remain, so the small worker ends up sitting on an older job's tail long after every large
+/// worker has moved on to newer jobs.
+#[test]
+fn test_prefill_weighted_by_worker_capacity() {
+    let mut rt = TestEnv::new();
+    rt.set_scheduler_config(SchedulerConfig {
+        proactive_filling_reserve: 0,
+        proactive_filling_max: 32,
+        ..Default::default()
+    });
+    let w_big = rt.new_worker(&WorkerBuilder::new(16));
+    let w_small = rt.new_worker(&WorkerBuilder::new(2));
+    rt.new_tasks(300, &TaskBuilder::new());
+    rt.schedule();
+
+    // `proactive_filling_max` applies to the largest worker; everyone else is scaled down by
+    // capacity. An unweighted split would give both workers 32.
+    let big = prefill_count(&mut rt, w_big);
+    let small = prefill_count(&mut rt, w_small);
+    assert_eq!(big, 32);
+    assert_eq!(small, 4);
+
+    // The property that actually matters: both hold the same *duration* of backlog, i.e. the
+    // same number of task generations (two each here).
+    assert_eq!(big / 16, small / 2);
+    rt.sanity_check();
+}
+
+/// A resource weight (`hq submit --weight`) multiplies a request's placement value in the solver
+/// objective. It is the supported way to make a request that only a few workers can serve win
+/// those workers instead of being crowded out by work that could have run anywhere -- see the
+/// `S5W` scenario in `benchmarks/scheduler-fairness`. It must never outrank user priority, so
+/// here the heavily weighted 4-cpu request is given the *lower* priority.
+#[test]
+fn test_priority_is_not_overridden_by_weight() {
+    let narrow = TaskBuilder::new().cpus(1).user_priority(10);
+    let wide = TaskBuilder::new().cpus(4).user_priority(0).weight(10.0);
+
+    let mut c = TestCase::new();
+    c.n_tasks(10, &narrow);
+    c.n_tasks(4, &wide);
+    // All capacity goes to the high-priority 1-cpu tasks despite the 10x weight on the others.
+    c.w(&WorkerBuilder::new(8)).expect_request(8, &narrow);
+    c.w(&WorkerBuilder::new(2)).expect_request(2, &narrow);
+    c.check();
+}
+
+/// At *equal* priority the weight does steer placement, which is the behaviour that lets a wide
+/// request claim the only workers able to run it. Miniature of `S5W`.
+#[test]
+fn test_weight_prefers_request_at_equal_priority() {
+    let narrow = TaskBuilder::new().cpus(1);
+    let wide = TaskBuilder::new().cpus(4).weight(4.0);
+
+    let mut c = TestCase::new();
+    c.n_tasks(10, &narrow);
+    c.n_tasks(2, &wide);
+    c.w(&WorkerBuilder::new(8)).expect_request(2, &wide);
+    c.w(&WorkerBuilder::new(2)).expect_request(2, &narrow);
+    c.check();
+}
+
+/// The prefill depth must be measured against the largest worker in the *cluster*, not the
+/// largest one eligible for prefill in this round. A worker is skipped while it still holds
+/// prefill of the request, so the eligible set is routinely all-small -- and if the reference
+/// capacity is taken from it, the depth springs back to `proactive_filling_max` for a tiny
+/// worker, which is the whole bug.
+#[test]
+fn test_prefill_depth_when_large_worker_is_ineligible() {
+    let mut rt = TestEnv::new();
+    rt.set_scheduler_config(SchedulerConfig {
+        proactive_filling_reserve: 0,
+        proactive_filling_max: 32,
+        ..Default::default()
+    });
+    let w_big = rt.new_worker(&WorkerBuilder::new(16));
+    rt.new_tasks(400, &TaskBuilder::new());
+    rt.schedule();
+    assert_eq!(prefill_count(&mut rt, w_big), 32);
+
+    // w_big now holds prefill of this request, so it is excluded from further prefill and
+    // only the 2-cpu worker is eligible.
+    let w_small = rt.new_worker(&WorkerBuilder::new(2));
+    rt.schedule();
+    assert_eq!(prefill_count(&mut rt, w_big), 32);
+    assert_eq!(prefill_count(&mut rt, w_small), 4);
+    rt.sanity_check();
+}
+
 #[test]
 pub fn test_schedule_running() {
     let mut rt = TestEnv::new();

crates/tako/src/internal/tests/utils/scheduler.rs

@@ -81,6 +81,14 @@ impl TestCase {
         self.rt.get_mut().new_tasks_cpus(cpus)
     }
 
+    /// `count` tasks from an explicit builder, for properties `pc_tasks` cannot express
+    /// (resource weight, variants, ...).
+    pub fn n_tasks(&mut self, count: usize, builder: &TaskBuilder) -> Vec<TaskId> {
+        (0..count)
+            .map(|_| self.rt.get_mut().new_task(builder))
+            .collect()
+    }
+
     // priority + cpu tasks
     pub fn pc_tasks(&mut self, priority_cpus: &[(i32, u32)]) -> Vec<TaskId> {
         priority_cpus