diff --git a/CHANGELOG.md b/CHANGELOG.md index c17ad77d7..59f67d196 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/crates/tako/src/internal/scheduler/mapping.rs b/crates/tako/src/internal/scheduler/mapping.rs index 306110fc4..9d094cdfe 100644 --- a/crates/tako/src/internal/scheduler/mapping.rs +++ b/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 = 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); diff --git a/crates/tako/src/internal/tests/test_scheduler_sn.rs b/crates/tako/src/internal/tests/test_scheduler_sn.rs index 5e8b63ba0..c722eeab8 100644 --- a/crates/tako/src/internal/tests/test_scheduler_sn.rs +++ b/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(); diff --git a/crates/tako/src/internal/tests/utils/scheduler.rs b/crates/tako/src/internal/tests/utils/scheduler.rs index 69ce4a7bc..167cad306 100644 --- a/crates/tako/src/internal/tests/utils/scheduler.rs +++ b/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 { + (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 { priority_cpus