diff --git a/CHANGELOG.md b/CHANGELOG.md index e46c95abb..e9bbe05dc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,13 @@ * The server scheduler now contains a safety limit for computation, configurable via `--scheduler-time-limit` (default: 5s) * Better scheduling policy (prefill) for heterogenous clusters +* Workers held for a waiting high-priority task are now chosen per priority threshold, so fewer + workers are held back and more lower-priority work can run elsewhere +* More precise computation of gap for low-priority tasks + +### Changes + +* HiGHs solver uses only a single thread ### Fixes @@ -13,6 +20,17 @@ * Fixed server crash in a specific situation when an unschedulable high-priority task occurs * Fixed server crash caused by invalid handling of prefill * Fixed canceling of prefilled tasks +* Fixed multi-node tasks that could be placed on workers from two different worker groups + when worker ids of the groups interleaved +* A held worker is no longer switched to a different worker only because small tasks were placed + into its unusable remainder; the scheduler moves a hold only to a worker that is at least as free + for the waiting task, so held capacity is not repeatedly given away +* Reservations: reservations respect priorities, and they no longer block resources + the waiting task could never use +* Fixed wrong resource amounts in this computation for tasks with resource variants on workers + that lack some resource +* Fixed constraining gap that could lead to priority inversion + ## v0.26.2 diff --git a/Cargo.lock b/Cargo.lock index d833a64e0..4c994fdcf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1528,9 +1528,9 @@ dependencies = [ [[package]] name = "microlp" -version = "0.4.0" +version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "458ed987196f802dc47c69d4c5afcd19002d6c1c5f8f75c76d129bcf2425057a" +checksum = "a8f19803f918039a07f3c32b23716824d8f073543417bdac5b1b8619817e3674" dependencies = [ "log", "sprs", diff --git a/crates/hyperqueue/src/server/autoalloc/process.rs b/crates/hyperqueue/src/server/autoalloc/process.rs index dbd60981c..16d99dcd8 100644 --- a/crates/hyperqueue/src/server/autoalloc/process.rs +++ b/crates/hyperqueue/src/server/autoalloc/process.rs @@ -2502,7 +2502,10 @@ mod tests { let schedule_fn = self.schedule_fn.clone(); let custom_state = self.custom_state.clone(); - Box::pin(async move { (schedule_fn.get())(custom_state.clone(), worker_count).await }) + Box::pin(async move { + let fut = (schedule_fn.get())(custom_state.clone(), worker_count); + fut.await + }) } fn get_status_of_allocations( @@ -2517,8 +2520,8 @@ mod tests { Box::pin(async move { let mut result = Map::default(); for allocation_id in allocation_ids { - let status = - (status_fn.get())(custom_state.clone(), allocation_id.clone()).await; + let fut = (status_fn.get())(custom_state.clone(), allocation_id.clone()); + let status = fut.await; if let Some(status) = status.transpose() { result.insert(allocation_id, status); } @@ -2535,7 +2538,10 @@ mod tests { let custom_state = self.custom_state.clone(); let allocation_id = allocation.id.clone(); - Box::pin(async move { (remove_fn.get())(custom_state.clone(), allocation_id).await }) + Box::pin(async move { + let fut = (remove_fn.get())(custom_state.clone(), allocation_id); + fut.await + }) } } diff --git a/crates/hyperqueue/src/tests/server.rs b/crates/hyperqueue/src/tests/server.rs index b4ca77f01..9157c7b11 100644 --- a/crates/hyperqueue/src/tests/server.rs +++ b/crates/hyperqueue/src/tests/server.rs @@ -100,7 +100,9 @@ where // Propagate the errors if server_error.is_err() && test_error.is_err() { eprintln!("{server_error:?}"); - test_error.expect("Test failed"); + if let Err(error) = test_error { + panic!("Test failed: {error:?}"); + } } else { test_error.expect("Test failed"); server_error.expect("Server failed"); diff --git a/crates/tako/Cargo.toml b/crates/tako/Cargo.toml index f469a7c26..d3ac62c83 100644 --- a/crates/tako/Cargo.toml +++ b/crates/tako/Cargo.toml @@ -39,7 +39,7 @@ tempfile = { workspace = true } itertools = { workspace = true } highs = { version = "2.1", optional = true } -microlp = { version = "0.4", optional = true } +microlp = { version = "0.6", optional = true } coin_cbc = { version = "0.1", optional = true } hashbrown = { version = "0.17", features = ["serde", "inline-more"], default-features = false } diff --git a/crates/tako/src/internal/scheduler/batches.rs b/crates/tako/src/internal/scheduler/batches.rs index 3d9ade529..402806ff0 100644 --- a/crates/tako/src/internal/scheduler/batches.rs +++ b/crates/tako/src/internal/scheduler/batches.rs @@ -98,8 +98,14 @@ pub(crate) fn create_task_batches( let mut current: Vec> = iters.iter_mut().map(|it| it.next()).collect(); let mut unique = None; let mut found = Vec::new(); + // The sweep step at which each batch last consumed tasks. A batch needs a new cut only if + // some other batch consumed since then: otherwise the tasks of other requests above it are + // the same as at its previous cut, so the cut would only repeat that one with a bigger size. + let mut step = 0u32; + let mut last_step: Vec> = vec![None; batches.len()]; loop { + step += 1; found.clear(); let mut highest_p = Priority::new(0); for (idx, c) in current.iter().enumerate() { @@ -121,6 +127,7 @@ pub(crate) fn create_task_batches( let idx = found[0]; let size = current[idx].unwrap().1; if unique == Some(idx) { + last_step[idx] = Some(step); batches[idx].size += size; if batches[idx].size > batches[idx].limit { batches[idx].size = batches[idx].limit; @@ -134,6 +141,14 @@ pub(crate) fn create_task_batches( break; } else { for idx in &found { + let since = last_step[*idx]; + let changed = last_step + .iter() + .enumerate() + .any(|(i, s)| i != *idx && s.is_some_and(|s| since.is_none_or(|l| s >= l))); + if !changed { + continue; + } let size = batches[*idx].size; let higher_priorities: Vec<_> = batches .iter_mut() @@ -153,6 +168,7 @@ pub(crate) fn create_task_batches( } } for idx in &found { + last_step[*idx] = Some(step); batches[*idx].size += current[*idx].unwrap().1; if batches[*idx].size > batches[*idx].limit { batches[*idx].size = batches[*idx].limit; @@ -249,6 +265,26 @@ fn prune_progressive(vec: &mut Vec, prefix_size: usize, size_limit: usize) vec.truncate(size_limit); } +pub(crate) fn trim_to_unblocked(batches: &[TaskBatch]) -> Vec { + batches + .iter() + .filter_map(|batch| { + let size = batch.cuts.first().map_or(batch.size, |cut| cut.size); + if size == 0 { + return None; + } + Some(TaskBatch { + resource_rq_id: batch.resource_rq_id, + cuts: Vec::new(), + size, + limit: batch.limit, + limit_reached: false, + is_blocker: false, + }) + }) + .collect() +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/tako/src/internal/scheduler/gap.rs b/crates/tako/src/internal/scheduler/gap.rs index 063e233f3..f5f72912d 100644 --- a/crates/tako/src/internal/scheduler/gap.rs +++ b/crates/tako/src/internal/scheduler/gap.rs @@ -1,7 +1,9 @@ use crate::internal::common::resources::ResourceId; use crate::internal::server::workerload::WorkerResources; -use crate::internal::solver::{ConstraintType, LpSolver}; -use crate::resources::{ResourceAmount, ResourceRequestVariants, ResourceRqId, ResourceRqMap}; +use crate::internal::solver::{ConstraintType, LpSolution, LpSolver}; +use crate::resources::{ + ResourceAmount, ResourceRequest, ResourceRequestVariants, ResourceRqId, ResourceRqMap, +}; use crate::{Map, ResourceVariantId}; use hashbrown::Equivalent; use std::cell::RefCell; @@ -35,6 +37,7 @@ struct GapCacheInner { } impl GapCache { + #[cfg(test)] pub fn get_gap( &self, high_priority_rq: ResourceRqId, @@ -43,17 +46,50 @@ impl GapCache { assigned_tasks: impl Iterator, resource_rq_map: &ResourceRqMap, ) -> u32 { - let h_rqv = resource_rq_map.get(high_priority_rq); - if h_rqv.is_multi_node() { - return 0; - } let l_rqv = resource_rq_map.get(low_priority_rq); if l_rqv.is_multi_node() { return 0; } + let Some(free) = + self.gap_resources(high_priority_rq, resources, assigned_tasks, resource_rq_map) + else { + return 0; + }; + l_rqv + .requests() + .iter() + .map(|rq| free.task_max_count_for_request(rq)) + .min() + .unwrap_or(0) + } + + /// The gap allowance of a worker for a blocker: resources that tasks of any lower request may + /// jointly consume there without ever denying the blocker (`get_gap` counts how many tasks of + /// one request fit into it). `None` when no gap exists: multi-node blockers, and blockers that + /// ask for all of a resource. + pub fn gap_resources( + &self, + high_priority_rq: ResourceRqId, + resources: &WorkerResources, + assigned_tasks: impl Iterator, + resource_rq_map: &ResourceRqMap, + ) -> Option { + let h_rqv = resource_rq_map.get(high_priority_rq); + if h_rqv.is_multi_node() { + return None; + } + // Callers ask only for workers that can run the blocker; otherwise the gap is meaningless + // and the blocker may name a resource the worker lacks (out of bounds in `remove_multiple`). + debug_assert!( + h_rqv + .requests() + .iter() + .any(|rq| resources.is_capable_to_run_request(rq)), + "gap computed for a worker that cannot run the blocker" + ); let mut free: WorkerResources = if let Some(h_rq) = h_rqv.trivial_request() { if h_rq.entries().iter().any(|r| r.request.amount_is_all()) { - return 0; + return None; } let count = resources.task_max_count_for_request(h_rq); let mut resources = resources.clone(); @@ -78,19 +114,239 @@ impl GapCache { free } }; - for (rq_id, rv_id) in assigned_tasks { - if rq_id != high_priority_rq { - let rq = resource_rq_map.get(rq_id).get(rv_id); - free.remove(rq); + let occupants: Vec<&ResourceRequest> = assigned_tasks + .filter(|(rq_id, _)| *rq_id != high_priority_rq) + .map(|(rq_id, rv_id)| resource_rq_map.get(rq_id).get(rv_id)) + .collect(); + for rq in &occupants { + free.remove(rq); + } + if let Some((resource_id, exact)) = exact_single_resource_gap(h_rqv, resources, &occupants) + { + free.set(resource_id, exact); + } else if let Some(exact) = exact_gap_by_enumeration(h_rqv, resources, &occupants) { + free = exact; + } + Some(free) + } +} + +const MAX_EXACT_GAP_UNITS: u32 = 4096; + +fn exact_single_resource_gap( + h_rqv: &ResourceRequestVariants, + resources: &WorkerResources, + occupants: &[&ResourceRequest], +) -> Option<(ResourceId, ResourceAmount)> { + let h_rq = h_rqv.trivial_request()?; + let entries = h_rq.entries(); + if entries.len() != 1 { + return None; + } + let entry = &entries[0]; + let blocker = entry.request.amount_or_none_if_all()?; + let resource_id = entry.resource_id; + let capacity = resources.get(resource_id); + if blocker.fractions() != 0 || capacity.fractions() != 0 || blocker.units() == 0 { + return None; + } + let r = blocker.units(); + if r > MAX_EXACT_GAP_UNITS { + return None; + } + let capacity = capacity.units(); + + let mut reachable = vec![false; r as usize]; + reachable[0] = true; + for rq in occupants { + let amount = rq.get_amount(resource_id).unwrap_or(ResourceAmount::ZERO); + if amount.fractions() != 0 { + return None; + } + let step = (amount.units() % r) as usize; + if step == 0 { + continue; + } + let previous = reachable.clone(); + for (value, _) in previous.iter().enumerate().filter(|(_, hit)| **hit) { + reachable[(value + step) % r as usize] = true; + } + } + + let gap = (0..r) + .filter(|v| reachable[*v as usize]) + .map(|v| (capacity + r - (v % r)) % r) + .min() + .unwrap_or(0); + Some((resource_id, ResourceAmount::new(gap, 0))) +} + +/// Capacity of one resource that no *mix* of `blockers` can use, whatever subset of the current +/// occupants departs. +/// +/// The gap of `gap_resources` is computed for one blocker at a time, and is safe against that +/// blocker alone. Two blockers that pack better together than either of them does by itself can +/// still be delayed by tasks that respect both gaps separately: on a $12$-cpu worker with a +/// $5$-cpu and a $7$-cpu blocker the two gaps are $2$ and $5$, so two 1-cpu tasks admitted under +/// the first leave $10$ cpus, which no longer holds $5 + 7$. The value computed here bounds what +/// all gap users of a worker may take together, and for a single blocker it equals that blocker's +/// own gap, so nothing changes where only one request blocks. +/// +/// Zero is always safe, and is the answer whenever the exact value is out of reach: fractional +/// amounts, or a capacity above `MAX_EXACT_GAP_UNITS`. A blocker that asks for *all* of the +/// resource is passed in as the whole capacity, which also yields zero. +pub(crate) fn joint_gap_units( + capacity: ResourceAmount, + blockers: &[(ResourceAmount, u32)], + occupants: &[ResourceAmount], +) -> ResourceAmount { + if capacity.fractions() != 0 || capacity.units() > MAX_EXACT_GAP_UNITS { + return ResourceAmount::ZERO; + } + let capacity_units = capacity.units() as usize; + // Each blocker is bounded by how many of its tasks the whole worker could hold: its other + // resources limit the mix just as much as this one. Without that bound a blocker that needs + // a scarce second resource would appear able to fill the worker on its own. + let mut sizes: Vec<(usize, u32)> = Vec::with_capacity(blockers.len()); + for (amount, max_count) in blockers { + if amount.fractions() != 0 { + return ResourceAmount::ZERO; + } + if amount.units() > 0 && *max_count > 0 { + sizes.push((amount.units() as usize, *max_count)); + } + } + if sizes.is_empty() { + // No blocker consumes this resource, so none of it is being withheld from them. + return capacity; + } + + // `reachable[v]`: a mix of blockers, each within its own count, totals exactly `v`. + let mut reachable = vec![false; capacity_units + 1]; + reachable[0] = true; + for (size, max_count) in &sizes { + let size = *size; + if size.saturating_mul(*max_count as usize) >= capacity_units { + // The count never binds here: the plain ascending pass is the unbounded case. + for value in 0..=capacity_units.saturating_sub(size) { + if reachable[value] { + reachable[value + size] = true; + } } + continue; } - l_rqv - .requests() - .iter() - .map(|rq| free.task_max_count_for_request(rq)) - .min() - .unwrap_or(0) + // Bounded: binary splitting turns `max_count` copies into `log(max_count)` items. + let mut remaining = *max_count; + let mut chunk = 1u32; + while remaining > 0 { + let take = chunk.min(remaining) as usize; + let step = size * take; + if step <= capacity_units { + for value in (step..=capacity_units).rev() { + if reachable[value - step] { + reachable[value] = true; + } + } + } + remaining -= chunk.min(remaining); + chunk = chunk.saturating_mul(2); + } + } + // `mix[v]`: the largest reachable total at or below `v`. + let mut mix = vec![0usize; capacity_units + 1]; + for value in 1..=capacity_units { + mix[value] = if reachable[value] { + value + } else { + mix[value - 1] + }; + } + + // Totals the current occupants can still hold once any subset of them has departed. + let mut departed = vec![false; capacity_units + 1]; + departed[0] = true; + for occupant in occupants { + if occupant.fractions() != 0 { + return ResourceAmount::ZERO; + } + let amount = occupant.units() as usize; + if amount == 0 || amount > capacity_units { + continue; + } + for value in (amount..=capacity_units).rev() { + if departed[value - amount] { + departed[value] = true; + } + } + } + + let gap = (0..=capacity_units) + .filter(|held| departed[*held]) + .map(|held| { + let free = capacity_units - held; + free - mix[free] + }) + .min() + .unwrap_or(0); + ResourceAmount::new(gap as u32, 0) +} + +const MAX_GAP_SUB_OCCUPANCIES: u32 = 2048; + +fn exact_gap_by_enumeration( + h_rqv: &ResourceRequestVariants, + resources: &WorkerResources, + occupants: &[&ResourceRequest], +) -> Option { + let h_rq = h_rqv.trivial_request()?; + if h_rq.entries().iter().any(|e| e.request.amount_is_all()) { + return None; + } + let mut groups: Map<&ResourceRequest, u32> = Map::default(); + for rq in occupants { + *groups.entry(*rq).or_default() += 1; } + let groups = groups; + let size = groups.values().try_fold(1u32, |acc, count| { + acc.checked_mul(*count + 1) + .filter(|size| *size <= MAX_GAP_SUB_OCCUPANCIES) + })?; + let mut gap: Option = None; + let mut remaining_counts: Vec = groups.values().copied().collect(); + for _ in 0..size { + let mut remaining = resources.clone(); + for ((rq, _), k) in groups.iter().zip(&remaining_counts) { + if *k > 0 { + remaining.remove_multiple(rq, *k); + } + } + let fit = remaining.task_max_count_for_request(h_rq); + remaining.remove_multiple(h_rq, fit); + if let Some(g) = &mut gap { + for (resource_id, amount) in remaining.iter_all_pairs() { + if amount < g.get(resource_id) { + g.set(resource_id, amount); + } + } + } else { + if h_rq + .entries() + .iter() + .all(|e| remaining.get(e.resource_id).is_zero()) + { + return Some(remaining); + } + gap = Some(remaining); + } + for ((_, count), k) in groups.iter().zip(remaining_counts.iter_mut()) { + if *k > 0 { + *k -= 1; + break; + } + *k = *count; + } + } + gap } fn compute_gap_resources( @@ -107,8 +363,11 @@ fn compute_gap_resources( }; let n_resources = n_unresources + 1; let gap_res: Vec = resources - .iter_pairs() + .iter_all_pairs() .map(|(r_id, r_amount)| { + if r_amount.is_zero() { + return ResourceAmount::ZERO; + } let mut solver = LpSolver::new(false); let mut cst = vec![Vec::new(); n_resources]; let vars: Vec<_> = rqv @@ -137,10 +396,10 @@ fn compute_gap_resources( c.into_iter(), ); } - let Some((_, v)) = solver.solve() else { + let Some(s) = solver.solve(None) else { return ResourceAmount::ZERO; }; - r_amount - ResourceAmount::from_float(v.round() as f32) + r_amount - ResourceAmount::from_float(LpSolution::objective(&s).round() as f32) }) .collect(); WorkerResources::new(gap_res.into()) @@ -149,6 +408,7 @@ fn compute_gap_resources( #[cfg(test)] mod tests { use crate::internal::server::core::CoreSplitMut; + use crate::resources::{CPU_RESOURCE_ID, ResourceAmount}; use std::iter; use crate::tests::utils::env::TestEnv; @@ -169,7 +429,350 @@ mod tests { let res = &worker_map.get_worker(w).resources; scheduler_state .gap_cache - .get_gap(h_rq, l_rq, &res, iter::empty(), request_map) + .get_gap(h_rq, l_rq, res, iter::empty(), request_map) + } + + fn compute_gap_occupied( + rt: &mut TestEnv, + high_task: TaskId, + low_task: TaskId, + w: WorkerId, + occupants: &[TaskId], + ) -> u32 { + let CoreSplitMut { + task_map, + worker_map, + scheduler_state, + request_map, + .. + } = rt.core().split_mut(); + let h_rq = task_map.get_task(high_task).resource_rq_id; + let l_rq = task_map.get_task(low_task).resource_rq_id; + let assigned: Vec<_> = occupants + .iter() + .map(|t| { + let task = task_map.get_task(*t); + (task.resource_rq_id, crate::ResourceVariantId::new(0)) + }) + .collect(); + let res = &worker_map.get_worker(w).resources; + scheduler_state + .gap_cache + .get_gap(h_rq, l_rq, res, assigned.into_iter(), request_map) + } + + #[test] + fn test_gap_is_one() { + let mut rt = TestEnv::new(); + let w = rt.new_worker(&WorkerBuilder::new(12)); + let r_h = rt.new_task_cpus(8); + let r_l = rt.new_task_cpus(4); + assert_eq!(compute_gap(&mut rt, r_h, r_l, w), 1, "empty worker"); + + let occupant = rt.new_task_cpus(2); + assert_eq!( + compute_gap_occupied(&mut rt, r_h, r_l, w, &[occupant]), + 0, + "two cpus occupied" + ); + } + + #[test] + fn test_gap_is_eight() { + let mut rt = TestEnv::new(); + let gpu = rt.new_named_resource("gpus"); + let w = rt.new_worker(&WorkerBuilder::new(12).res_sum("gpus", 4)); + let r_h = rt.new_task(&TaskBuilder::new().cpus(1).add_resource(gpu, 1)); + let r_l = rt.new_task_cpus(1); + assert_eq!(compute_gap(&mut rt, r_h, r_l, w), 8, "empty worker"); + + let occupant = rt.new_task_cpus(2); + assert_eq!( + compute_gap_occupied(&mut rt, r_h, r_l, w, &[occupant]), + 6, + "two cpus occupied" + ); + } + + #[test] + fn test_gap_is_safe_under_future_departures() { + let mut rt = TestEnv::new(); + let w = rt.new_worker(&WorkerBuilder::new(12)); + let unrelated = rt.new_task_cpus(3); + let r_h = rt.new_task_cpus(5); + let r_l = rt.new_task_cpus(1); + assert_eq!( + compute_gap_occupied(&mut rt, r_h, r_l, w, &[unrelated]), + 2, + "must be 12 mod 5 = 2, not the naive 9 - 5 = 4" + ); + } + + /// The case that rules out the tempting shortcut of `min(gap(C), gap(C - occupancy))`. + /// + /// C = 12, r_h = 5, occupants {1, 4}: the occupancy still present at some future point can be + /// 0, 1, 4 or 5, giving `(12 - x) mod 5` of 2, 1, 3, 2. The binding term is the *intermediate* + /// x = 1, so the gap is 1 — both extremes give 2, and granting 2 would over-commit the worker. + #[test] + fn test_gap_binding_term_can_be_an_intermediate_sub_occupancy() { + let mut rt = TestEnv::new(); + let w = rt.new_worker(&WorkerBuilder::new(12)); + let occ_a = rt.new_task_cpus(1); + let occ_b = rt.new_task_cpus(4); + let r_h = rt.new_task_cpus(5); + let r_l = rt.new_task_cpus(1); + assert_eq!( + compute_gap_occupied(&mut rt, r_h, r_l, w, &[occ_a, occ_b]), + 1, + "intermediate sub-occupancy must bind; 2 would over-grant" + ); + } + + /// Cross-check the residue DP against brute-force enumeration of every sub-occupancy, over a + /// spread of capacities, blocker widths and occupant multisets. The result must equal the + /// exact minimum -- and in particular must never exceed it, which is the unsafe direction. + #[test] + fn test_gap_matches_brute_force_over_sub_occupancies() { + for capacity in [6u32, 8, 11, 12, 16] { + for blocker in [2u32, 3, 5, 6, 7] { + if blocker > capacity { + // The gap is only computed for workers that can run the blocker + continue; + } + for occupants in [ + vec![], + vec![1u32], + vec![3], + vec![1, 4], + vec![2, 2], + vec![1, 2, 4], + vec![5, 3, 1], + ] { + if occupants.iter().sum::() > capacity { + continue; + } + // min over subsets that may remain occupied of (C - x) mod r + let mut expected = u32::MAX; + for mask in 0..(1u32 << occupants.len()) { + let x: u32 = occupants + .iter() + .enumerate() + .filter(|(i, _)| mask & (1 << i) != 0) + .map(|(_, a)| *a) + .sum(); + expected = expected.min((capacity - x) % blocker); + } + + let mut rt = TestEnv::new(); + let w = rt.new_worker(&WorkerBuilder::new(capacity)); + let occ: Vec<_> = occupants.iter().map(|a| rt.new_task_cpus(*a)).collect(); + let r_h = rt.new_task_cpus(blocker); + let r_l = rt.new_task_cpus(1); + let got = compute_gap_occupied(&mut rt, r_h, r_l, w, &occ); + assert_eq!( + got, expected, + "C={capacity} r_h={blocker} occupants={occupants:?}" + ); + } + } + } + } + + /// A blocker spanning two resources: the residue table does not apply, enumeration does. + #[test] + fn test_multi_resource_blocker_gap_is_exact() { + let mut rt = TestEnv::new(); + let gpu = rt.new_named_resource("gpus"); + let w = rt.new_worker(&WorkerBuilder::new(12).res_sum("gpus", 4)); + let r_h = rt.new_task(&TaskBuilder::new().cpus(5).add_resource(gpu, 1)); + let r_l = rt.new_task_cpus(1); + let occupant = rt.new_task_cpus(3); + // Occupant running: 9 free cpus fit one blocker, 4 cpus it cannot use. Occupant finished: + // 12 cpus fit two blockers, 2 it cannot use. Exact gap 2; the bound 12 - 10 - 3 gives 0. + assert_eq!(compute_gap_occupied(&mut rt, r_h, r_l, w, &[occupant]), 2); + } + + /// A blocker asking for a large amount of one resource is too big for the residue table, but + /// with only two occupants there are four sub-occupancies, so the gap is still exact. This is + /// the intermediate-sub-occupancy case scaled by 1000: exact 1000, where the conservative bound + /// `C - M(C) - o` would give 0. + #[test] + fn test_large_single_resource_blocker_is_exact_by_enumeration() { + let mut rt = TestEnv::new(); + let w = rt.new_worker(&WorkerBuilder::new(12_000)); + let occ_a = rt.new_task_cpus(1_000); + let occ_b = rt.new_task_cpus(4_000); + let r_h = rt.new_task_cpus(5_000); + let r_l = rt.new_task_cpus(1); + assert_eq!( + compute_gap_occupied(&mut rt, r_h, r_l, w, &[occ_a, occ_b]), + 1_000 + ); + } + + /// Too large for the residue table and too many distinct sub-occupancies to enumerate: the + /// conservative bound is used. The exact gap here would be 922; the bound is 0. + #[test] + fn test_gap_falls_back_to_bound_when_both_exact_methods_are_too_large() { + let mut rt = TestEnv::new(); + let w = rt.new_worker(&WorkerBuilder::new(12_000)); + let mut occupants = vec![rt.new_task_cpus(1_000), rt.new_task_cpus(4_000)]; + // Twelve more distinct shapes: 2^14 sub-occupancies in total. + occupants.extend((1..=12).map(|c| rt.new_task_cpus(c))); + let r_h = rt.new_task_cpus(5_000); + let r_l = rt.new_task_cpus(1); + assert_eq!(compute_gap_occupied(&mut rt, r_h, r_l, w, &occupants), 0); + } + + /// The per-resource minimum of `C - x - fit(C - x) * r_h` over every subset of occupants that + /// may still be running, for a blocker of `blocker.0` cpus + `blocker.1` gpus. Occupants are + /// `(cpus, gpus)`. Returns the allowance `(cpus, gpus)`. + fn brute_force_gap( + capacity: (u32, u32), + blocker: (u32, u32), + occupants: &[(u32, u32)], + ) -> (u32, u32) { + let mut gap = (u32::MAX, u32::MAX); + for mask in 0..(1u32 << occupants.len()) { + let (xc, xg) = occupants + .iter() + .enumerate() + .filter(|(i, _)| mask & (1 << i) != 0) + .fold((0, 0), |(c, g), (_, (oc, og))| (c + oc, g + og)); + let (fc, fg) = (capacity.0 - xc, capacity.1 - xg); + let fit = (fc / blocker.0).min(fg / blocker.1); + gap.0 = gap.0.min(fc - fit * blocker.0); + gap.1 = gap.1.min(fg - fit * blocker.1); + } + gap + } + + /// Cross-check the exact multi-resource computation against brute force over every + /// sub-occupancy, for two-resource blockers. A 1-cpu request reads off the cpu allowance, and + /// a 1-cpu + 1-gpu request the smaller of the two. + #[test] + fn test_multi_resource_gap_matches_brute_force() { + for capacity in [(8u32, 2u32), (12, 4), (16, 4), (20, 4)] { + for blocker in [(4u32, 1u32), (3, 1), (2, 2), (5, 1)] { + for occupants in [ + vec![], + vec![(2u32, 1u32)], + vec![(3, 0)], + vec![(1, 1), (2, 0)], + vec![(2, 1), (2, 1), (1, 0)], + vec![(3, 1), (1, 0), (2, 0)], + ] { + let used = occupants + .iter() + .fold((0, 0), |(c, g), (oc, og)| (c + oc, g + og)); + if used.0 > capacity.0 || used.1 > capacity.1 { + continue; + } + let expected = brute_force_gap(capacity, blocker, &occupants); + + let mut rt = TestEnv::new(); + let gpu = rt.new_named_resource("gpus"); + let w = + rt.new_worker(&WorkerBuilder::new(capacity.0).res_sum("gpus", capacity.1)); + let occ: Vec<_> = occupants + .iter() + .map(|(c, g)| { + let tb = TaskBuilder::new().cpus(*c); + rt.new_task(&if *g > 0 { tb.add_resource(gpu, *g) } else { tb }) + }) + .collect(); + let r_h = rt.new_task( + &TaskBuilder::new() + .cpus(blocker.0) + .add_resource(gpu, blocker.1), + ); + let cpu_only = rt.new_task_cpus(1); + let cpu_and_gpu = rt.new_task(&TaskBuilder::new().cpus(1).add_resource(gpu, 1)); + let context = format!("C={capacity:?} r_h={blocker:?} occupants={occupants:?}"); + assert_eq!( + compute_gap_occupied(&mut rt, r_h, cpu_only, w, &occ), + expected.0, + "cpu allowance, {context}" + ); + assert_eq!( + compute_gap_occupied(&mut rt, r_h, cpu_and_gpu, w, &occ), + expected.0.min(expected.1), + "cpu+gpu allowance, {context}" + ); + } + } + } + } + + /// Where the conservative bound is loose. A 4-cpu + 1-gpu blocker on 20 cpus and 4 gpus can + /// never use more than 16 cpus, and while a 2-cpu + 1-gpu occupant runs it gets only 3 gpus, + /// hence 12 cpus. So 4 cpus are never usable by it, whatever finishes first. The bound + /// `C - M(C) - o = 20 - 16 - 2` charges the occupant's cpus as if they stayed forever and + /// gives only 2. + #[test] + fn test_multi_resource_gap_is_larger_than_the_bound() { + let mut rt = TestEnv::new(); + let gpu = rt.new_named_resource("gpus"); + let w = rt.new_worker(&WorkerBuilder::new(20).res_sum("gpus", 4)); + let occupant = rt.new_task(&TaskBuilder::new().cpus(2).add_resource(gpu, 1)); + let r_h = rt.new_task(&TaskBuilder::new().cpus(4).add_resource(gpu, 1)); + let r_l = rt.new_task_cpus(1); + assert_eq!(compute_gap_occupied(&mut rt, r_h, r_l, w, &[occupant]), 4); + } + + /// Too many distinct sub-occupancies for a multi-resource blocker: the conservative bound is + /// used. Four 1-cpu + 1-gpu occupants and twelve cpu-only occupants of distinct sizes give + /// 5 * 2^12 sub-occupancies. The exact gap would be 200 - 78 - 16 = 106; the bound is + /// 200 - 16 - 82 = 102. + #[test] + fn test_multi_resource_gap_falls_back_to_bound_above_limit() { + let mut rt = TestEnv::new(); + let gpu = rt.new_named_resource("gpus"); + let w = rt.new_worker(&WorkerBuilder::new(200).res_sum("gpus", 4)); + let mut occupants: Vec<_> = (0..4) + .map(|_| rt.new_task(&TaskBuilder::new().cpus(1).add_resource(gpu, 1))) + .collect(); + occupants.extend((1..=12).map(|c| rt.new_task_cpus(c))); + let r_h = rt.new_task(&TaskBuilder::new().cpus(4).add_resource(gpu, 1)); + let r_l = rt.new_task_cpus(1); + assert_eq!(compute_gap_occupied(&mut rt, r_h, r_l, w, &occupants), 102); + } + + /// Multi-variant gaps are stored per resource id. A worker that lacks a resource with a lower + /// id than one it has (here no "mem", id 1, but "gpus", id 2) must still get each value at the + /// right id. + /// + /// Blocker variants: 8 cpus, or 2 cpus + 2 gpus. On 8 cpus + 3 gpus any mix uses at most + /// 8 cpus and 2 gpus, so the gap is 0 cpus and 1 gpu, and nothing for the missing "mem". + #[test] + fn test_multi_variant_gap_keeps_resource_ids_when_worker_lacks_a_resource() { + let mut rt = TestEnv::new(); + let mem = rt.new_named_resource("mem"); + let gpus = rt.new_named_resource("gpus"); + let w = rt.new_worker(&WorkerBuilder::new(8).res_sum("gpus", 3)); + let r_h = rt.new_task( + &TaskBuilder::new() + .cpus(8) + .next_variant() + .cpus(2) + .add_resource(gpus, 2), + ); + let CoreSplitMut { + task_map, + worker_map, + scheduler_state, + request_map, + .. + } = rt.core().split_mut(); + let h_rq = task_map.get_task(r_h).resource_rq_id; + let res = &worker_map.get_worker(w).resources; + let gap = scheduler_state + .gap_cache + .gap_resources(h_rq, res, iter::empty(), request_map) + .unwrap(); + assert_eq!(gap.get(CPU_RESOURCE_ID), ResourceAmount::ZERO, "cpus"); + assert_eq!(gap.get(mem), ResourceAmount::ZERO, "mem"); + assert_eq!(gap.get(gpus), ResourceAmount::new_units(1), "gpus"); } #[test] diff --git a/crates/tako/src/internal/scheduler/solver.rs b/crates/tako/src/internal/scheduler/solver.rs index 91abb24f6..ac107d68e 100644 --- a/crates/tako/src/internal/scheduler/solver.rs +++ b/crates/tako/src/internal/scheduler/solver.rs @@ -1,10 +1,16 @@ -use crate::internal::common::resources::{ResourceId, ResourceRequest, ResourceRequestVariants}; +use crate::internal::common::resources::{ + ResourceAmount, ResourceId, ResourceRequest, ResourceRequestVariants, +}; use crate::internal::scheduler::TaskBatch; +use crate::internal::scheduler::batches::trim_to_unblocked; +use crate::internal::scheduler::gap::joint_gap_units; +use crate::internal::scheduler::state::SchedulerState; use crate::internal::server::core::{Core, CoreSplit}; +use crate::internal::server::taskmap::TaskMap; use crate::internal::server::worker::Worker; use crate::internal::server::workerload::WorkerResources; use crate::internal::solver::{ConstraintType, LpSolution, LpSolver, Variable}; -use crate::resources::{CPU_RESOURCE_ID, ResourceRqId}; +use crate::resources::{CPU_RESOURCE_ID, ResourceRqId, ResourceRqMap}; use crate::{Map, ResourceVariantId, Set, WorkerId}; use thin_vec::ThinVec; @@ -34,16 +40,38 @@ impl SchedulingSolution { } } -struct Candidates { - demand: u32, - candidates: Vec<(f64, WorkerId)>, -} - pub(crate) fn run_scheduling_solver( core: &Core, now: std::time::Instant, task_batches: &[TaskBatch], custom_workers: Option<&[Worker]>, +) -> SchedulingSolution { + let CoreSplit { worker_map, .. } = core.split(); + let workers: Vec<&Worker> = if let Some(ws) = custom_workers { + ws.iter().collect() + } else { + let mut ws = worker_map + .get_workers() + .filter(|w| w.sn_assignment().is_some()) + .collect::>(); + ws.sort_unstable_by_key(|w| w.id); + ws + }; + let solution = run_scheduling_solver_inner(core, now, task_batches, &workers); + if solution.is_optimal || !solution.is_empty() { + return solution; + }; + let task_batches = trim_to_unblocked(task_batches); + let mut solution = run_scheduling_solver_inner(core, now, &task_batches, &workers); + solution.is_optimal = false; + solution +} + +pub fn run_scheduling_solver_inner( + core: &Core, + now: std::time::Instant, + task_batches: &[TaskBatch], + workers: &[&Worker], ) -> SchedulingSolution { let n_resources = core.resource_map().n_resources(); @@ -60,16 +88,8 @@ pub(crate) fn run_scheduling_solver( return SchedulingSolution::default(); } let mut resource_sums = vec![0f64; n_resources]; - let workers: Vec<&Worker> = if let Some(ws) = custom_workers { - ws.iter().collect() - } else { - let mut ws = worker_map - .get_workers() - .filter(|w| w.sn_assignment().is_some()) - .collect::>(); - ws.sort_unstable_by_key(|w| w.id); - ws - }; + + let n_workers = workers.len(); workers.iter().for_each(|worker| { let Some(a) = worker.sn_assignment() else { @@ -87,30 +107,42 @@ pub(crate) fn run_scheduling_solver( }) }); - let n_workers = workers.len(); - let mut solver = LpSolver::new(false); let mut placements: Map<(WorkerId, ResourceRqId, ResourceVariantId), Variable> = Map::new(); let mut tasks_count_vars: Map> = Map::new(); + // Reservation variables by (worker, request). A reservation stands for one task of its request + // set aside on that worker, so it also enters the request's own priority conditions there. + let mut reservations: Map<(WorkerId, ResourceRqId), Variable> = Map::new(); let mut worker_res_constraint = vec![Vec::new(); n_resources]; let mut worker_cpu_constraint_no_reserves: Vec<(Variable, f64)> = Vec::with_capacity(task_batches.len()); - let mut reservation_candidates: Map = Map::new(); - for batch in task_batches { - reservation_candidates.insert( - batch.resource_rq_id, - Candidates { - demand: batch.size, - candidates: Vec::new(), - }, - ); - } + // Held workers are chosen before any variable exists, because a reservation variable is only + // created on a held worker. A reservation elsewhere would count toward its blocker without + // holding anything back: on a worker with no free resources it costs nothing but its small + // penalty, and with no worker held (the blocker fits somewhere right now) it would let a + // lower-priority request take the very capacity the blocker was counted as using. + let reserved = held_workers( + workers, + task_batches, + request_map, + task_map, + scheduler_state, + now, + ); + // The reservation penalty is expressed relative to the placements a reservation can unlock. + // Placement weights shrink as the cluster's free resources grow, so a fixed penalty would + // outweigh every placement on a large cluster and no reservation would ever pay for itself. + let reservation_scale = + placement_weight_lower_bound(workers, task_batches, request_map, &resource_sums); + // Create worker-task placements + let mut worker_reservations: Vec = Vec::new(); for (w_idx, worker) in workers.iter().enumerate() { worker_cpu_constraint_no_reserves.clear(); + worker_reservations.clear(); for batch in task_batches.iter() { let rqv = request_map.get(batch.resource_rq_id); let mut has_variant = false; @@ -133,14 +165,11 @@ pub(crate) fn run_scheduling_solver( ); placements.insert((worker.id, batch.resource_rq_id, v_idx), v); // Insert into worker resource constraints - for (r, amount) in worker.resources.iter_pairs() { + for (r, amount) in worker.resources.iter_nonzero_pairs() { worker_res_constraint[r.as_usize()].push((v, amount.as_f64())); } } - } else if !worker.is_request_blocked(batch.resource_rq_id, v_idx) - && worker.has_time_to_run(rq.min_time(), now) - && worker.have_immediate_resources_for_rq(rq) - { + } else if sn_variant_fits_now(worker, batch.resource_rq_id, v_idx, rq, now) { has_variant = true; set_placement_name(&mut solver, worker.id, batch.resource_rq_id, v_idx); let v = @@ -167,44 +196,55 @@ pub(crate) fn run_scheduling_solver( } } - if has_variant - && !rqv.is_multi_node() - && batch.is_blocker - && let Some(a) = worker.sn_assignment() - { - let demand = &mut reservation_candidates - .get_mut(&batch.resource_rq_id) - .unwrap() - .demand; - *demand = demand.saturating_sub(a.free_resources.task_max_count(rqv)); - } - if !has_variant - && !rqv.is_multi_node() - && batch.is_blocker - && worker.is_capable_to_run_rqv(rqv, now) + && reserved + .get(&batch.resource_rq_id) + .is_some_and(|held| held.order.contains(&worker.id)) && let Some(a) = worker.sn_assignment() { - let weight = -((n_workers - w_idx) as f64 / (n_workers * 1024) as f64); + let weight = + -reservation_scale * (n_workers - w_idx) as f64 / (n_workers * 1024) as f64; solver.set_name(|| format!("R{}:{}", worker.id, batch.resource_rq_id)); let v = solver.add_bool_variable(weight); + reservations.insert((worker.id, batch.resource_rq_id), v); + worker_reservations.push(v); tasks_count_vars .entry(batch.resource_rq_id) .or_default() .push(v); - for (res_id, count) in a.free_resources.iter_pairs() { - worker_res_constraint[res_id.as_usize()].push((v, count.as_f64())); + // A reservation withholds the free resources the blocker could use, but not its gap + // allowance: that capacity is lost to the blocker anyway, and the priority + // conditions already let lower-priority gap tasks use it. + let gap = worker_gap( + worker, + batch.resource_rq_id, + request_map, + task_map, + scheduler_state, + ); + for (res_id, count) in a.free_resources.iter_nonzero_pairs() { + let withheld = gap + .as_ref() + .map_or(count, |g| count.saturating_sub(g.get(res_id))); + if !withheld.is_zero() { + worker_res_constraint[res_id.as_usize()].push((v, withheld.as_f64())); + } } - // The blocker cannot run here now, but this worker could host it once it drains. - // How close it already is decides which workers are held for the blocker below. - reservation_candidates - .get_mut(&batch.resource_rq_id) - .unwrap() - .candidates - .push((fit_ratio(&a.free_resources, rqv), worker.id)); } } + // A worker can later host only one of the blockers it is reserved for, so it may serve + // at most one. The resource constraints do not ensure this: on a worker with no free + // resources a reservation consumes nothing, and any number of them would fit. + if worker_reservations.len() > 1 { + solver.set_name(|| format!("w{}: at most one reservation", worker.id)); + solver.add_constraint( + ConstraintType::Max, + 1.0, + worker_reservations.iter().map(|v| (*v, 1.0)), + ); + } + if worker.configuration.min_utilization > 0.001 { add_min_utilization(&mut solver, worker, &mut worker_cpu_constraint_no_reserves); } @@ -226,30 +266,6 @@ pub(crate) fn run_scheduling_solver( c.clear(); } } - let reserved: Map> = reservation_candidates - .into_iter() - .map( - |( - rq_id, - Candidates { - mut candidates, - demand, - }, - )| { - if demand > 0 { - candidates.sort_unstable_by(|(fit_a, w_a), (fit_b, w_b)| { - fit_b.total_cmp(fit_a).then(w_b.cmp(w_a)) - }); - } - let held = candidates - .into_iter() - .take(demand as usize) - .map(|(_fit, w_id)| w_id) - .collect(); - (rq_id, held) - }, - ) - .collect(); let mut task_counts_per_group: Map<(ResourceRqId, &str), Variable> = Map::new(); let mut temp = Vec::new(); @@ -313,9 +329,20 @@ pub(crate) fn run_scheduling_solver( Some(new_v) }; - let mut zero_cond = Vec::new(); - let mut zero_cond_reserved = Vec::new(); + // Terms of a condition's left-hand side: every placement, minus what sits in a worker's gap. + let mut cond_terms: Vec<(Variable, f64)> = Vec::new(); + let mut cond_terms_reserved: Vec<(Variable, f64)> = Vec::new(); let mut blocked_by_unbounded: Set = Set::new(); + // Gap parts per (worker, blocker, resource): the gap allowance and the terms that share it. + let mut shared_gap: Map)> = Map::new(); + let mut shared_gap_seen: Set<(WorkerId, ResourceRqId, ResourceRqId)> = Set::new(); + // The rows above bound the gap of one blocker at a time. Blockers that pack better together + // than alone need a second bound, on everything their gaps admit on a worker; see + // `joint_gap_units`. Terms per (worker, resource), and the blockers each worker has. + // Keyed per request as well: a request gets one gap variable per blocker, and they stand for + // the same tasks, so the row below must take their maximum rather than their sum. + let mut joint_gap: Map<(WorkerId, ResourceId, ResourceRqId), (f64, Vec)> = Map::new(); + let mut joint_blockers: Map)> = Map::new(); for batch in task_batches.iter() { let Some(task_counts) = tasks_count_vars.get(&batch.resource_rq_id) else { @@ -323,20 +350,18 @@ pub(crate) fn run_scheduling_solver( }; let batch_rqv = request_map.get(batch.resource_rq_id); assert!(!task_counts.is_empty()); - if !batch.limit_reached { - solver.set_name(|| format!("size limit for rq{}", batch.resource_rq_id)); - solver.add_constraint( - ConstraintType::Max, - batch.size as f64, - task_counts.iter().map(|v| (*v, 1.0)), - ) - } + solver.set_name(|| format!("size limit for rq{}", batch.resource_rq_id)); + solver.add_constraint( + ConstraintType::Max, + batch.size as f64, + task_counts.iter().map(|v| (*v, 1.0)), + ); let batch_size = batch.size as f64; blocked_by_unbounded.clear(); for cut in &batch.cuts { for (blocker_rq_id, blocking_size) in &cut.blockers { - zero_cond.clear(); - zero_cond_reserved.clear(); + cond_terms.clear(); + cond_terms_reserved.clear(); let blocker_rqv = request_map.get(*blocker_rq_id); if batch_rqv.is_multi_node() { for (group_name, group) in worker_groups.iter() { @@ -344,20 +369,27 @@ pub(crate) fn run_scheduling_solver( task_counts_per_group.get(&(batch.resource_rq_id, group_name.as_str())) && group.is_capable_to_run(blocker_rqv, now, worker_map) { - zero_cond.push(*v); + cond_terms.push((*v, 1.0)); } } } else { - for w in &workers { + for w in workers { let Some(sn_assignment) = w.sn_assignment() else { continue; }; if !w.is_capable_to_run_rqv(blocker_rqv, now) { continue; } - let gap = scheduler_state.gap_cache.get_gap( + // Every blocker of this worker takes part in the joint bound below, also + // the ones whose own gap is zero: those are the ones that constrain a mix + // of blockers the most. + joint_blockers + .entry(w.id) + .or_insert_with(|| (w, Set::new())) + .1 + .insert(*blocker_rq_id); + let gap_resources = scheduler_state.gap_cache.gap_resources( *blocker_rq_id, - batch.resource_rq_id, &w.resources, sn_assignment.assigned_tasks.iter().map(|task_id| { let t = task_map.get_task(*task_id); @@ -368,106 +400,119 @@ pub(crate) fn run_scheduling_solver( }), request_map, ); + let gap = gap_resources + .as_ref() + .map(|g| gap_count(g, batch_rqv)) + .unwrap_or(0); // A worker held for this blocker gets the unconditional form: with no // blocker-count term there is nothing a reservation variable can discharge. - let is_reserved = reserved - .get(blocker_rq_id) - .is_some_and(|ws| ws.contains(&w.id)); - if gap > 0 { - let vars = batch_rqv.variant_ids().filter_map(|v_id| { - placements.get(&(w.id, batch.resource_rq_id, v_id)).copied() - }); - let cut_size = cut.size as f64; - if let Some(s) = blocking_size - && !is_reserved - && let Some(blocking_v) = get_bvar(&mut solver, *blocker_rq_id, *s) - { - solver.set_name(|| { - format!( - "w{}: if #rq{blocker_rq_id} < {s} then limit #rq{} to {} + {} (gap) where both rqs may run", - w.id, batch.resource_rq_id, cut.size, gap - ) - }); - constraint_extra_var( - &mut solver, - ConstraintType::Max, - cut_size + batch_size + gap as f64, - vars, - blocking_v, - batch_size, - ); - } else if blocking_size.is_none() || is_reserved { - solver.set_name(|| { - format!( - "w{}: limit #rq{} to {} + {} (gap) where it can run with rq{blocker_rq_id}", - w.id, batch.resource_rq_id, gap, cut.size, - ) - }); - solver.add_constraint( - ConstraintType::Max, - cut_size + gap as f64, - vars.into_iter().map(|v| (v, 1.0)), - ); - } - } else { - for v_id in batch_rqv.variant_ids() { - if let Some(var) = - placements.get(&(w.id, batch.resource_rq_id, v_id)) + let is_reserved = blocking_size.is_some_and(|threshold| { + reserved + .get(blocker_rq_id) + .is_some_and(|held| held.holds(w.id, threshold)) + }); + // A reservation counts like a placement of its request: while the blocker is + // unserved on this worker, the request may not set capacity aside here any + // more than it may start a task here. + let reservation = reservations.get(&(w.id, batch.resource_rq_id)).copied(); + let placement_vars = batch_rqv + .variant_ids() + .filter_map(|v_id| { + placements + .get(&(w.id, batch.resource_rq_id, v_id)) + .map(|v| (v_id, *v)) + }) + .collect::>(); + if placement_vars.is_empty() && reservation.is_none() { + continue; + } + let placed = placement_vars + .iter() + .map(|(_, v)| (*v, 1.0)) + .chain(reservation.map(|v| (v, 1.0))); + cond_terms.extend(placed.clone()); + if is_reserved { + cond_terms_reserved.extend(placed); + } + if gap == 0 { + // Nothing can sit in a gap here, so every task placed counts against + // the before-count. No variable and no constraint for this worker. + continue; + } + // Tasks of this request sitting in the worker's gap. They do not count + // against the before-count, so they are subtracted from the condition. + // At most `gap` of them fit, which is a bound of the variable itself. + solver.set_name(|| { + format!( + "w{}: #rq{} within the {gap} gap of rq{blocker_rq_id}", + w.id, batch.resource_rq_id + ) + }); + let gap_var = solver.add_variable(0.0, 0.0, gap as f64); + // Only tasks that really run here may sit in the gap. Without this, a + // request could claim gap usage it does not have, and so shrink the + // condition below and the shared gap constraint for other requests. + solver.set_name(|| { + format!( + "w{}: #rq{} in the gap is at most what runs here", + w.id, batch.resource_rq_id + ) + }); + constraint_extra_var( + &mut solver, + ConstraintType::Min, + 0.0, + placement_vars.iter().map(|(_, v)| *v).chain(reservation), + gap_var, + -1.0, + ); + cond_terms.push((gap_var, -1.0)); + if is_reserved { + cond_terms_reserved.push((gap_var, -1.0)); + } + // All requests blocked by this blocker share the worker's gap, so their + // gap parts are collected per resource. Recorded once per (worker, + // blocker, request): the gap does not depend on the cut. + if let Some(gap_resources) = gap_resources.as_ref() + && shared_gap_seen.insert((w.id, *blocker_rq_id, batch.resource_rq_id)) + { + for (resource_id, gap_amount) in gap_resources.iter_all_pairs() { + if gap_amount >= sn_assignment.free_resources.get(resource_id) { + // The worker's own resource limit is at least as strict. + continue; + } + // Gap tasks are charged at the largest variant amount, so a mix + // of variants can never use more of the gap than counted. + let largest = placement_vars + .iter() + .map(|(v_id, _)| { + batch_rqv + .get(*v_id) + .get_amount(resource_id) + .unwrap_or_else(|| w.resources.get(resource_id)) + }) + .max(); + if let Some(largest) = largest + && !largest.is_zero() { - zero_cond.push(*var); - if is_reserved { - zero_cond_reserved.push(*var); - } + shared_gap + .entry((w.id, *blocker_rq_id, resource_id)) + .or_insert_with(|| (gap_amount, Vec::new())) + .1 + .push((gap_var, largest.as_f64())); + joint_gap + .entry((w.id, resource_id, batch.resource_rq_id)) + .or_insert_with(|| (largest.as_f64(), Vec::new())) + .1 + .push(gap_var); } } } - /*for v_id in batch_rqv.variant_ids() { - if let Some(var) = placements.get(&(w.id, batch.resource_rq_id, v_id)) { - let gap = scheduler_cache.gap_cache.get_gap( - *blocker_rq_id, - batch.resource_rq_id, - v_id, - &w.resources, - request_map, - ); - if gap > 0 { - let cut_size = cut.size as f64; - if let Some(s) = blocking_size { - let blocking_v = get_bvar(&mut solver, *blocker_rq_id, *s); - solver.set_name(|| { - format!( - "w{}: if #rq{blocker_rq_id} < {s} then limit #rq{} to {} + {} (gap) where both rqs may run", - w.id, batch.resource_rq_id, cut.size, gap - ) - }); - solver.add_constraint( - ConstraintType::Max, - cut_size + batch_size + gap as f64, - [(*var, 1.0), (blocking_v, batch_size)].into_iter(), - ); - } else { - solver.set_name(|| { - format!( - "w{}: limit #rq{} to {} + {} (gap) where it can run with rq{blocker_rq_id}", - w.id, batch.resource_rq_id, gap, cut.size, - ) - }); - solver.add_constraint( - ConstraintType::Max, - cut_size + gap as f64, - [(*var, 1.0)].into_iter(), - ); - } - } else { - zero_cond.push(*var); - } - } - }*/ } } // Workers held for the blocker take only what the priority rule allows anyway // (`cut.size`), with no blocker-count escape, so their capacity accumulates. - if !zero_cond_reserved.is_empty() { + if !cond_terms_reserved.is_empty() { solver.set_name(|| { format!( "limit #rq{} to {} on workers held for rq{blocker_rq_id}", @@ -477,12 +522,14 @@ pub(crate) fn run_scheduling_solver( solver.add_constraint( ConstraintType::Max, cut.size as f64, - zero_cond_reserved.iter().map(|v| (*v, 1.0)), + cond_terms_reserved.iter().copied(), ); } - if zero_cond.is_empty() { + if cond_terms.is_empty() { continue; } + // The before-count is one number for the whole cluster: every placement counts + // against it, except what sits in the gap of its worker. if let Some(s) = blocking_size && let Some(blocking_v) = get_bvar(&mut solver, *blocker_rq_id, *s) { @@ -493,15 +540,18 @@ pub(crate) fn run_scheduling_solver( ) }); let cut_size = cut.size as f64; - constraint_extra_var( - &mut solver, + solver.add_constraint( ConstraintType::Max, batch_size + cut_size, - zero_cond.iter().copied(), - blocking_v, - batch_size, + cond_terms + .iter() + .copied() + .chain(std::iter::once((blocking_v, batch_size))), ); - } else if blocking_size.is_none() && !blocked_by_unbounded.contains(blocker_rq_id) { + } else if blocking_size.is_none() + && (cond_terms.iter().any(|(_, coef)| *coef < 0.0) + || !blocked_by_unbounded.contains(blocker_rq_id)) + { blocked_by_unbounded.insert(*blocker_rq_id); solver.set_name(|| { format!( @@ -512,19 +562,134 @@ pub(crate) fn run_scheduling_solver( solver.add_constraint( ConstraintType::Max, cut.size as f64, - zero_cond.iter().map(|v| (*v, 1.0)), + cond_terms.iter().copied(), ); } } } } + // One gap per worker and blocker, shared by every request the blocker blocks. + for ((worker_id, blocker_rq_id, resource_id), (gap_amount, terms)) in shared_gap { + solver.set_name(|| { + format!( + "w{worker_id}: gap of rq{blocker_rq_id} on resource {} is {gap_amount}", + resource_id.as_num() + ) + }); + solver.add_constraint(ConstraintType::Max, gap_amount.as_f64(), terms.into_iter()); + } + + // One term per (worker, resource, request): the largest gap claim that request makes there. + let mut joint_rows: Map<(WorkerId, ResourceId), Vec<(Variable, f64)>> = Map::new(); + for ((worker_id, resource_id, rq_id), (amount, gap_vars)) in joint_gap { + // A request that is itself a blocker on this worker is part of the mix, not a threat to + // it: its tasks are taken in priority order, so a task admitted here is one of the very + // tasks the mix is made of. Only requests from outside the blocker set are bounded. + if joint_blockers + .get(&worker_id) + .is_some_and(|(_, blockers)| blockers.contains(&rq_id)) + { + continue; + } + let term = if let [single] = gap_vars[..] { + single + } else { + solver.set_name(|| format!("w{worker_id}: gap use of rq{rq_id} across its blockers")); + let shared = solver.add_variable(0.0, 0.0, f64::INFINITY); + for gap_var in gap_vars { + constraint_extra_var( + &mut solver, + ConstraintType::Min, + 0.0, + std::iter::once(shared), + gap_var, + -1.0, + ); + } + shared + }; + joint_rows + .entry((worker_id, resource_id)) + .or_default() + .push((term, amount)); + } + + for ((worker_id, resource_id), terms) in joint_rows { + let Some((worker, blockers)) = joint_blockers.get(&worker_id) else { + continue; + }; + if blockers.len() < 2 { + // With one blocker the gap of `gap_resources` is already exact, and it accounts for + // the whole resource vector of every sub-occupancy. This row only exists to catch a + // mix of blockers, and its per-resource view would be the cruder of the two here. + continue; + } + let Some(sn_assignment) = worker.sn_assignment() else { + continue; + }; + let capacity = worker.resources.get(resource_id); + // A blocker is charged at its smallest variant: a variant that packs tighter leaves less + // for the gap, which is the safe direction. A request for *all* of the resource takes the + // whole worker. + let blocker_amounts: Vec<(ResourceAmount, u32)> = blockers + .iter() + .filter_map(|rq_id| { + let rqv = request_map.get(*rq_id); + if rqv.is_multi_node() { + return None; + } + rqv.requests() + .iter() + .map(|rq| { + let amount = rq + .entries() + .iter() + .find(|entry| entry.resource_id == resource_id) + .map(|entry| { + entry + .request + .amount_or_none_if_all() + .unwrap_or_else(|| worker.resources.get(resource_id)) + }) + .unwrap_or(ResourceAmount::ZERO); + (amount, worker.resources.task_max_count_for_request(rq)) + }) + .min_by_key(|(amount, _)| *amount) + }) + .collect(); + let occupant_amounts: Vec = sn_assignment + .assigned_tasks + .iter() + .filter_map(|task_id| { + let task = task_map.get_task(*task_id); + let (_, rv_id) = task.assigned_placement(&scheduler_state.redirects)?; + request_map + .get(task.resource_rq_id) + .get(rv_id) + .get_amount(resource_id) + }) + .collect(); + let joint = joint_gap_units(capacity, &blocker_amounts, &occupant_amounts); + if joint >= sn_assignment.free_resources.get(resource_id) { + // The worker's own capacity row is at least as strict. + continue; + } + solver.set_name(|| { + format!( + "w{worker_id}: all gaps on resource {} are {joint} together", + resource_id.as_num() + ) + }); + solver.add_constraint(ConstraintType::Max, joint.as_f64(), terms.into_iter()); + } + let mut result = SchedulingSolution::default(); - let Some((solution, is_optimal)) = solver.solve_bounded(scheduler_state.config.mip_time_limit) - else { + let Some(solution) = solver.solve(Some(scheduler_state.config.mip_time_limit)) else { + result.is_optimal = false; return result; }; - result.is_optimal = is_optimal; + result.is_optimal = solution.is_optimal(); for batch in task_batches { let resource_rq_id = batch.resource_rq_id; @@ -533,22 +698,26 @@ pub(crate) fn run_scheduling_solver( let v_id = ResourceVariantId::new(0); let n_nodes = rqv.get(v_id).n_nodes() as usize; let mut ws: Vec> = Vec::new(); - for worker in &workers { + // Chunks are cut per worker group: the group constraint makes each group's count a + // multiple of `n_nodes`, but worker ids of different groups may interleave, so cutting + // across all workers in id order could give a task nodes from two groups. + let mut open: Map<&str, ThinVec> = Map::new(); + for worker in workers { if let Some(v) = placements.get(&(worker.id, resource_rq_id, v_id)) { let count = solution.get_value(*v).round() as u32; if count > 0 { - if let Some(last) = ws.last_mut() - && last.len() < n_nodes - { - last.push(worker.id); - } else { - let mut workers = ThinVec::with_capacity(n_nodes); - workers.push(worker.id); - ws.push(workers); + let group = worker.configuration.group.as_str(); + let chunk = open + .entry(group) + .or_insert_with(|| ThinVec::with_capacity(n_nodes)); + chunk.push(worker.id); + if chunk.len() == n_nodes { + ws.push(open.remove(group).unwrap()); } } } } + assert!(open.is_empty()); if !ws.is_empty() { result.mn_workers.insert((resource_rq_id, v_id), ws); } @@ -629,6 +798,210 @@ fn add_min_utilization( worker_res_constraint.pop(); } +/// Whether a single-node variant of a request can be placed on the worker in this round. +/// Shared by placement creation and by `held_workers`, which must agree on it exactly. +fn sn_variant_fits_now( + worker: &Worker, + rq_id: ResourceRqId, + v_idx: ResourceVariantId, + rq: &ResourceRequest, + now: std::time::Instant, +) -> bool { + !worker.is_request_blocked(rq_id, v_idx) + && worker.has_time_to_run(rq.min_time(), now) + && worker.have_immediate_resources_for_rq(rq) +} + +/// The gap allowance of `worker` for blocker `rq_id`, given what currently runs there. +fn worker_gap( + worker: &Worker, + rq_id: ResourceRqId, + request_map: &ResourceRqMap, + task_map: &TaskMap, + scheduler_state: &SchedulerState, +) -> Option { + let a = worker.sn_assignment()?; + scheduler_state.gap_cache.gap_resources( + rq_id, + &worker.resources, + a.assigned_tasks.iter().map(|task_id| { + let t = task_map.get_task(*task_id); + ( + t.resource_rq_id, + t.assigned_placement(&scheduler_state.redirects).unwrap().1, + ) + }), + request_map, + ) +} + +/// Workers held for one single-node blocker, most empty first (ties to the higher worker id). +struct Held { + /// Blocker tasks that the workers' currently free resources can host. + placeable: u32, + /// Held workers for the blocker's largest threshold; any smaller threshold holds a prefix. + order: Vec, +} + +impl Held { + /// Whether a priority condition that asks for `threshold` blocker tasks holds this worker: + /// it needs one held worker per blocker task it counts that cannot be placed now. + fn holds(&self, worker_id: WorkerId, threshold: u32) -> bool { + let needed = threshold.saturating_sub(self.placeable) as usize; + self.order.iter().take(needed).any(|w| *w == worker_id) + } +} + +/// Workers held for each single-node blocker. +/// +/// Holding is sized by the thresholds of the priority conditions that name the blocker, not by +/// its batch: a condition asking for `s` blocker tasks needs `s - placeable` held workers, and +/// blocker tasks below every request they could block need none. Sizing by the whole batch +/// would hold capacity for blocker tasks that the blocked request outranks. +fn held_workers( + workers: &[&Worker], + task_batches: &[TaskBatch], + request_map: &ResourceRqMap, + task_map: &TaskMap, + scheduler_state: &SchedulerState, + now: std::time::Instant, +) -> Map { + let mut max_threshold: Map = Map::new(); + for batch in task_batches { + for cut in &batch.cuts { + for (blocker_rq_id, blocking_size) in &cut.blockers { + if let Some(size) = blocking_size { + let entry = max_threshold.entry(*blocker_rq_id).or_default(); + *entry = (*entry).max(*size); + } + } + } + } + + let mut held = Map::new(); + for batch in task_batches { + let rqv = request_map.get(batch.resource_rq_id); + if !batch.is_blocker || rqv.is_multi_node() { + continue; + } + let Some(threshold) = max_threshold.get(&batch.resource_rq_id).copied() else { + continue; + }; + let mut placeable = 0u32; + let mut candidates: Vec<(f64, WorkerId)> = Vec::new(); + for worker in workers { + let Some(a) = worker.sn_assignment() else { + continue; + }; + let fits_now = rqv.requests_with_ids().any(|(v_idx, rq)| { + sn_variant_fits_now(worker, batch.resource_rq_id, v_idx, rq, now) + }); + if fits_now { + placeable = placeable.saturating_add(a.free_resources.task_max_count(rqv)); + } else if worker.is_capable_to_run_rqv(rqv, now) { + // The blocker cannot run here now, but this worker could host it once it drains. + // Rank by the free resources minus the blocker's gap: gap filling may give the + // gap to lower-priority tasks at any time without delaying the blocker, so it + // does not bring the blocker closer to starting. + // Tasks placed into the gap lower free resources and gap equally, so they never + // move the hold, and a finishing task never lowers the rank. A held worker + // therefore loses its hold only to a worker that offers the blocker more. + let mut usable = a.free_resources.clone(); + if let Some(gap) = worker_gap( + worker, + batch.resource_rq_id, + request_map, + task_map, + scheduler_state, + ) { + for (res_id, amount) in gap.iter_nonzero_pairs() { + usable.set(res_id, usable.get(res_id).saturating_sub(amount)); + } + } + candidates.push((fit_ratio(&usable, rqv), worker.id)); + } + } + let needed = threshold.saturating_sub(placeable) as usize; + if needed == 0 { + continue; + } + candidates.sort_unstable_by(|(fit_a, w_a), (fit_b, w_b)| { + fit_b.total_cmp(fit_a).then(w_b.cmp(w_a)) + }); + held.insert( + batch.resource_rq_id, + Held { + placeable, + order: candidates + .into_iter() + .take(needed) + .map(|(_fit, w_id)| w_id) + .collect(), + }, + ); + } + held +} + +/// A lower bound on the objective weight of any single-node placement this round. +/// +/// `create_sn_var` weighs a placement by its request's share of the free resources, the worker's +/// compaction bias `(n - i_w) / n`, and the request weight. The bias is smallest on the last +/// worker, `1 / n`, and a request for *all* of a resource is valued at the worker's total, which is +/// at least the smallest total any worker has. Only positive weights count; if there are none, +/// nothing can be unlocked and the scale is irrelevant. +fn placement_weight_lower_bound( + workers: &[&Worker], + task_batches: &[TaskBatch], + request_map: &ResourceRqMap, + resource_sums: &[f64], +) -> f64 { + let n_workers = workers.len(); + if n_workers == 0 { + return 1.0; + } + let smallest_total = |r: ResourceId| { + workers + .iter() + .map(|w| w.resources.get(r).as_f64()) + .filter(|amount| *amount > 0.0) + .fold(f64::INFINITY, f64::min) + }; + let mut bound = f64::INFINITY; + for batch in task_batches { + for rq in request_map.get(batch.resource_rq_id).requests() { + if rq.is_multi_node() { + continue; + } + let share: f64 = rq + .entries() + .iter() + .map(|e| { + let global = resource_sums[e.resource_id.as_usize()]; + if global < 0.000001 { + return 0.0; + } + let amount = e + .request + .amount_or_none_if_all() + .map(|a| a.as_f64()) + .unwrap_or_else(|| smallest_total(e.resource_id)); + if amount.is_finite() { + amount / global + } else { + 0.0 + } + }) + .sum(); + let weight = share * rq.weight().as_f64() / n_workers as f64; + if weight > 0.0 { + bound = bound.min(weight); + } + } + } + if bound.is_finite() { bound } else { 1.0 } +} + fn fit_ratio(free: &WorkerResources, rqv: &ResourceRequestVariants) -> f64 { rqv.requests() .iter() @@ -690,7 +1063,7 @@ fn create_mn_var( ) -> Variable { let weight = worker .resources - .iter_pairs() + .iter_nonzero_pairs() .map(|(r, amount)| { let global = resource_sums[r.as_usize()]; if global < 0.000001 { @@ -706,6 +1079,21 @@ fn create_mn_var( solver.add_bool_variable(weight) } +/// How many tasks of `rqv` fit into a worker's gap allowance for some blocker. +/// Worker, blocker request and resource: one shared gap allowance. +type SharedGapKey = (WorkerId, ResourceRqId, ResourceId); + +fn gap_count(gap_resources: &WorkerResources, rqv: &ResourceRequestVariants) -> u32 { + if rqv.is_multi_node() { + return 0; + } + rqv.requests() + .iter() + .map(|rq| gap_resources.task_max_count_for_request(rq)) + .min() + .unwrap_or(0) +} + fn constraint_extra_var( solver: &mut LpSolver, constraint_type: ConstraintType, diff --git a/crates/tako/src/internal/server/workerload.rs b/crates/tako/src/internal/server/workerload.rs index cdcbb0254..888709f03 100644 --- a/crates/tako/src/internal/server/workerload.rs +++ b/crates/tako/src/internal/server/workerload.rs @@ -30,18 +30,20 @@ impl WorkerResources { .unwrap_or(ResourceAmount::ZERO) } - pub(crate) fn iter_pairs(&self) -> impl Iterator { - self.n_resources - .iter() - .copied() + pub(crate) fn iter_all_pairs(&self) -> impl Iterator { + self.iter_amounts() .enumerate() - .filter_map(|(idx, c)| { - if !c.is_zero() { - Some((ResourceId::new(idx as u32), c)) - } else { - None - } - }) + .map(|(idx, c)| (ResourceId::new(idx as u32), c)) + } + + pub(crate) fn iter_nonzero_pairs(&self) -> impl Iterator { + self.iter_amounts().enumerate().filter_map(|(idx, c)| { + if !c.is_zero() { + Some((ResourceId::new(idx as u32), c)) + } else { + None + } + }) } pub(crate) fn iter_amounts(&self) -> impl Iterator { @@ -153,6 +155,12 @@ impl WorkerResources { .sum::() } + /// Overwrite one resource's amount. Used by the gap computation, which derives an exact + /// value for the blocker's resource and keeps the derived-by-subtraction value elsewhere. + pub fn set(&mut self, resource_id: ResourceId, amount: ResourceAmount) { + self.n_resources[resource_id] = amount; + } + pub fn remove(&mut self, rq: &ResourceRequest) { for entry in rq.entries() { if let Some(amount) = entry.request.amount_or_none_if_all() { diff --git a/crates/tako/src/internal/solver/coin_cbc.rs b/crates/tako/src/internal/solver/coin_cbc.rs index fd2f1e76e..08ce84082 100644 --- a/crates/tako/src/internal/solver/coin_cbc.rs +++ b/crates/tako/src/internal/solver/coin_cbc.rs @@ -1,5 +1,6 @@ use crate::internal::solver::{ConstraintType, LpInnerSolver, LpSolution}; use coin_cbc::{Col, Model, Sense}; +use std::time::Duration; pub(crate) struct CoinCbcSolver { model: Model, @@ -60,13 +61,16 @@ impl LpInnerSolver for CoinCbcSolver { } } - fn solve(self) -> Option<(Self::Solution, f64)> { + fn solve(mut self, time_limit: Option) -> Option { + if let Some(time_limit) = time_limit { + self.model + .set_parameter("seconds", &time_limit.as_secs_f64().to_string()); + } let solution = self.model.solve(); if !solution.raw().is_proven_optimal() { return None; } - let obj = solution.raw().obj_value(); - Some((solution, obj)) + Some(solution) } } @@ -77,4 +81,14 @@ impl LpSolution for coin_cbc::Solution { fn get_value(&self, v: Col) -> f64 { self.col(v) } + + #[inline] + fn objective(&self) -> f64 { + self.raw().obj_value() + } + + #[inline] + fn is_optimal(&self) -> bool { + self.raw().is_proven_optimal() + } } diff --git a/crates/tako/src/internal/solver/highs.rs b/crates/tako/src/internal/solver/highs.rs index 46cd0ff97..200c4c92a 100644 --- a/crates/tako/src/internal/solver/highs.rs +++ b/crates/tako/src/internal/solver/highs.rs @@ -1,5 +1,5 @@ use crate::internal::solver::{ConstraintType, LpInnerSolver, LpSolution}; -use highs::{HighsModelStatus, HighsSolutionStatus, Sense}; +use highs::{HighsModelStatus, HighsSolutionStatus, Sense, Solution}; use std::time::Duration; pub(crate) struct HighsSolver(highs::RowProblem); @@ -12,7 +12,7 @@ impl HighsSolver { impl LpInnerSolver for HighsSolver { type Variable = highs::Col; - type Solution = highs::Solution; + type Solution = HighsSolution; #[inline] fn add_variable(&mut self, weight: f64, min: f64, max: f64) -> Self::Variable { @@ -43,34 +43,16 @@ impl LpInnerSolver for HighsSolver { } } - /// Unbounded, exact solve: used by the worker's own NUMA/socket resource - /// allocator (see worker/resources/groups.rs), which relies on finding an - /// exact feasible allocation rather than a merely-good-enough one -- these - /// LPs are tiny (single-worker resource groups), so there is no - /// scheduler-scale performance problem to trade off here. - fn solve(self) -> Option<(Self::Solution, f64)> { - let solved_model = self.0.optimise(Sense::Maximise).solve(); - if !matches!(solved_model.status(), HighsModelStatus::Optimal) { - return None; - } - let solution = solved_model.get_solution(); - let objective_value = solved_model.objective_value(); - Some((solution, objective_value)) - } - - /// Bounded solve for the global task scheduler: hard-caps wall time at - /// `time_limit` instead of always proving exact optimality. Returns - /// whether the solution is proven optimal, since `solve_bounded` may - /// dispatch a merely feasible incumbent on timeout. - fn solve_bounded(self, time_limit: Duration) -> Option<(Self::Solution, bool)> { + fn solve(self, time_limit: Option) -> Option { let mut model = self.0.optimise(Sense::Maximise); - model.set_option("time_limit", time_limit.as_secs_f64()); + if let Some(time_limit) = time_limit { + model.set_option("time_limit", time_limit.as_secs_f64()); + } + model.set_option("threads", 1); let solved_model = model.solve(); - let is_optimal = match solved_model.status() { HighsModelStatus::Optimal => true, - // Time limit fired before optimality was proven. Dispatch the - // incumbent anyway if it's feasible. + HighsModelStatus::ModelEmpty => true, HighsModelStatus::ReachedTimeLimit if solved_model.primal_solution_status() == HighsSolutionStatus::Feasible => { @@ -82,17 +64,37 @@ impl LpInnerSolver for HighsSolver { } _ => return None, }; - let solution = solved_model.get_solution(); - Some((solution, is_optimal)) + + Some(HighsSolution { + solution, + objective: solved_model.objective_value(), + is_optimal, + }) } } -impl LpSolution for highs::Solution { +pub(crate) struct HighsSolution { + solution: Solution, + objective: f64, + is_optimal: bool, +} + +impl LpSolution for HighsSolution { type Variable = highs::Col; #[inline] fn get_value(&self, v: highs::Col) -> f64 { - self[v] + self.solution[v] + } + + #[inline] + fn objective(&self) -> f64 { + self.objective + } + + #[inline] + fn is_optimal(&self) -> bool { + self.is_optimal } } diff --git a/crates/tako/src/internal/solver/microlp.rs b/crates/tako/src/internal/solver/microlp.rs index 92feb2eb8..c4048786a 100644 --- a/crates/tako/src/internal/solver/microlp.rs +++ b/crates/tako/src/internal/solver/microlp.rs @@ -1,5 +1,6 @@ use crate::internal::solver::{ConstraintType, LpInnerSolver, LpSolution}; -use microlp::{ComparisonOp, Solution}; +use microlp::{ComparisonOp, Solution, SolveOutcome}; +use std::time::Duration; pub(crate) struct MicrolpSolver(microlp::Problem); @@ -54,12 +55,14 @@ impl LpInnerSolver for MicrolpSolver { ) } - fn solve(self) -> Option<(Self::Solution, f64)> { - let Ok(solution) = self.0.solve() else { + fn solve(mut self, time_limit: Option) -> Option { + if let Some(time_limit) = time_limit { + self.0.set_time_limit(time_limit); + } + let Ok(SolveOutcome::Solution(solution)) = self.0.solve() else { return None; }; - let objective = solution.objective(); - Some((solution, objective)) + Some(solution) } } @@ -70,4 +73,14 @@ impl LpSolution for microlp::Solution { fn get_value(&self, v: microlp::Variable) -> f64 { self.var_value(v) } + + #[inline] + fn is_optimal(&self) -> bool { + self.status() == microlp::SolutionStatus::Optimal + } + + #[inline] + fn objective(&self) -> f64 { + self.objective() + } } diff --git a/crates/tako/src/internal/solver/mod.rs b/crates/tako/src/internal/solver/mod.rs index 1025ec747..832904f05 100644 --- a/crates/tako/src/internal/solver/mod.rs +++ b/crates/tako/src/internal/solver/mod.rs @@ -39,23 +39,18 @@ pub(crate) trait LpInnerSolver { value: f64, variables: impl Iterator, ); - fn solve(self) -> Option<(Self::Solution, f64)>; /// Like `solve`, but allowed to trade exactness for a hard wall-clock /// cap. Returns whether the solution is proven optimal. Backends without /// a tuned implementation fall back to the exact `solve`. - fn solve_bounded(self, time_limit: Duration) -> Option<(Self::Solution, bool)> - where - Self: Sized, - { - let _ = time_limit; - self.solve().map(|(solution, _)| (solution, true)) - } + fn solve(self, time_limit: Option) -> Option; } pub(crate) trait LpSolution { type Variable: Copy; fn get_value(&self, v: Self::Variable) -> f64; + fn objective(&self) -> f64; + fn is_optimal(&self) -> bool; } pub(crate) struct LpSolver { @@ -175,7 +170,7 @@ impl LpSolver { } #[inline] - pub fn solve(self) -> Option<(Solution, f64)> { + pub fn solve(self, time_limit: Option) -> Option { if self.verbose { println!("Weights:"); for (name, weight, _var) in self.variables.iter() { @@ -184,30 +179,8 @@ impl LpSolver { } } } - let s = self.solver.solve(); - if let Some((s, _)) = &s - && self.verbose - { - println!("==== Solution: ===="); - for (name, _weight, var) in self.variables.iter() { - println!("{} = {}", name, s.get_value(*var)); - } - } - s - } - - #[inline] - pub fn solve_bounded(self, time_limit: Duration) -> Option<(Solution, bool)> { - if self.verbose { - println!("Weights:"); - for (name, weight, _var) in self.variables.iter() { - if *weight != 0.0 { - println!("{} -> {}", name, weight); - } - } - } - let s = self.solver.solve_bounded(time_limit); - if let Some((s, _)) = &s + let s = self.solver.solve(time_limit); + if let Some(s) = &s && self.verbose { println!("==== Solution: ===="); @@ -252,13 +225,8 @@ impl LpSolver { } #[inline] - pub fn solve(self) -> Option<(Solution, f64)> { - self.solver.solve() - } - - #[inline] - pub fn solve_bounded(self, time_limit: Duration) -> Option<(Solution, bool)> { - self.solver.solve_bounded(time_limit) + pub fn solve(self, time_limit: Option) -> Option { + self.solver.solve(time_limit) } } diff --git a/crates/tako/src/internal/tests/mod.rs b/crates/tako/src/internal/tests/mod.rs index 134e62ec2..e6bb17a70 100644 --- a/crates/tako/src/internal/tests/mod.rs +++ b/crates/tako/src/internal/tests/mod.rs @@ -1,4 +1,6 @@ #[cfg(test)] +mod test_min_utilization; +#[cfg(test)] mod test_query; #[cfg(test)] mod test_reactor; diff --git a/crates/tako/src/internal/tests/test_min_utilization.rs b/crates/tako/src/internal/tests/test_min_utilization.rs new file mode 100644 index 000000000..6088d6991 --- /dev/null +++ b/crates/tako/src/internal/tests/test_min_utilization.rs @@ -0,0 +1,142 @@ +use crate::control::WorkerTypeQuery; +use crate::internal::scheduler::query::compute_new_worker_query; +use crate::internal::tests::utils::env::TestEnv; +use crate::internal::tests::utils::task::TaskBuilder; +use crate::internal::tests::utils::worker::WorkerBuilder; +use crate::resources::ResourceDescriptor; + +/// Worker shape. The report used $128$ CPUs; only the ratio matters here. +const WORKER_CPUS: u32 = 8; + +/// Place `n_small` one-CPU tasks and `n_big` whole-worker tasks on one worker that demands full +/// utilization, and report how many of each the round assigned. +fn place( + n_small: usize, + n_big: usize, + weight: f32, + prio_small: i32, + prio_big: i32, +) -> (usize, usize) { + let mut rt = TestEnv::new(); + rt.new_worker(&WorkerBuilder::new(WORKER_CPUS).min_utilization(1.0)); + let small = rt.new_tasks( + n_small, + &TaskBuilder::new() + .cpus(1) + .weight(weight) + .user_priority(prio_small), + ); + let big = rt.new_tasks( + n_big, + &TaskBuilder::new().cpus_all().user_priority(prio_big), + ); + rt.schedule(); + ( + small.iter().filter(|t| rt.task(**t).is_assigned()).count(), + big.iter().filter(|t| rt.task(**t).is_assigned()).count(), + ) +} + +/// How many workers of `WORKER_CPUS` CPUs the scheduler says it could use. +fn workers_wanted(rt: &mut TestEnv, min_utilization: f32, max_sn_workers: u32) -> u32 { + let response = compute_new_worker_query( + rt.core(), + &[WorkerTypeQuery { + partial: false, + descriptor: ResourceDescriptor::simple_cpus(WORKER_CPUS), + time_limit: None, + max_sn_workers, + max_workers_per_allocation: 1, + min_utilization, + }], + ); + response.single_node_workers_per_query[0] +} + +#[test] +fn test_priority_expresses_precedence_and_therefore_blocks() { + // A full batch runs, as it would with any expression of the intent. + assert_eq!(place(8, 1, 1.0, 10, 0), (8, 0)); + // A partial batch cannot use the worker, and its priority forbids the whole-worker task from + // using it either. The worker idles. This is the strict rule, not a defect. + assert_eq!(place(3, 1, 1.0, 10, 0), (0, 0)); +} + +#[test] +fn test_weight_expresses_preference_without_blocking() { + // Same intent, expressed as a weight at equal priority: the one-CPU work still wins the + // worker whenever it can fill it ... + assert_eq!(place(8, 1, 2.0, 0, 0), (8, 0)); + // ... and when it cannot, the whole-worker task keeps the worker busy instead of idling it. + assert_eq!(place(3, 1, 2.0, 0, 0), (0, 1)); + // The preference does not depend on how strongly the weight is set: it is not a priority. + assert_eq!(place(3, 1, 8.0, 0, 0), (0, 1)); +} + +#[test] +fn test_weight_is_a_real_knob() { + // A placement is worth its share of the cluster's resources times its weight. A full batch of + // eight one-CPU tasks of weight `w` is worth `8 * 1/8 * w = w`; the whole-worker task, at the + // default weight, is worth `1`. The knob therefore flips the worker between the two mixes at + // `w = 1`, rather than merely reinforcing a default. + + // Below one, the full batch loses the worker to the whole-worker task ... + assert_eq!(place(8, 1, 0.5, 0, 0), (0, 1)); + assert_eq!(place(8, 1, 0.9, 0, 0), (0, 1)); + // ... and above one it wins it. + assert_eq!(place(8, 1, 1.1, 0, 0), (8, 0)); + assert_eq!(place(8, 1, 2.0, 0, 0), (8, 0)); + + // At exactly one the two mixes are worth the same, and which one an LP backend returns is + // its own tie-break (HiGHS and Cbc differ). Only what both optima share is asserted: the + // worker is filled by one mix or the other, never split or left idle. + let at_tie = place(8, 1, 1.0, 0, 0); + assert!( + at_tie == (8, 0) || at_tie == (0, 1), + "expected one full mix at the tie, got {at_tie:?}" + ); +} + +/// How many workers the allocation queue would ask for, in the same two states. +fn wanted(n_small: usize, n_big: usize, weight: f32, prio_small: i32, prio_big: i32) -> u32 { + let mut rt = TestEnv::new(); + rt.new_tasks( + n_small, + &TaskBuilder::new() + .cpus(1) + .weight(weight) + .user_priority(prio_small), + ); + rt.new_tasks( + n_big, + &TaskBuilder::new().cpus_all().user_priority(prio_big), + ); + workers_wanted(&mut rt, 1.0, 4) +} + +#[test] +fn test_allocation_queue_follows_the_same_two_expressions() { + // Eight one-CPU tasks and one whole-worker task: the one-CPU work can fill a worker, so + // nothing is blocked and both expressions ask for the two workers the queue can use. + assert_eq!(wanted(8, 1, 1.0, 10, 0), 2); + assert_eq!(wanted(8, 1, 2.0, 0, 0), 2); + + // Three one-CPU tasks and one whole-worker task: the expressions part company. The weight + // asks for the one worker the filler can use; the priority asks for nothing, which is the + // idle worker of `test_priority_expresses_precedence_and_therefore_blocks` one round earlier. + assert_eq!(wanted(3, 1, 2.0, 0, 0), 1); + assert_eq!(wanted(3, 1, 1.0, 10, 0), 0); + + // The weight expression tracks the amount of filler, the priority expression stays at zero. + assert_eq!(wanted(3, 2, 2.0, 0, 0), 2); + assert_eq!(wanted(3, 2, 1.0, 10, 0), 0); + + // One-CPU tasks on their own ask only for the workers they can fill, under either expression. + let mut rt = TestEnv::new(); + rt.new_tasks(3, &TaskBuilder::new().cpus(1).weight(2.0)); + assert_eq!(workers_wanted(&mut rt, 1.0, 4), 0); + + let mut rt = TestEnv::new(); + rt.new_tasks(16, &TaskBuilder::new().cpus(1).weight(2.0)); + assert_eq!(workers_wanted(&mut rt, 1.0, 4), 2); +} diff --git a/crates/tako/src/internal/tests/test_query.rs b/crates/tako/src/internal/tests/test_query.rs index 0411432d1..a10aaa31d 100644 --- a/crates/tako/src/internal/tests/test_query.rs +++ b/crates/tako/src/internal/tests/test_query.rs @@ -285,7 +285,7 @@ fn test_query_min_utilization1() { (0.7, 1, 3), ] { let r = compute_new_worker_query( - &mut rt.core(), + rt.core(), &[WorkerTypeQuery { partial: false, descriptor: ResourceDescriptor::simple_cpus(*cpus), @@ -505,7 +505,7 @@ fn test_query_min_time1() { assert!(r.multi_node_allocations.is_empty()); let r = compute_new_worker_query( - &mut rt.core(), + rt.core(), &[WorkerTypeQuery { partial: false, descriptor: descriptor.clone(), @@ -769,3 +769,84 @@ fn test_query_after_task_cancel() { ); assert_eq!(r.single_node_workers_per_query, vec![0]); } + +/// A partial query completes undeclared resources as unbounded sums, which makes their scarcity +/// share nearly zero. The share stays strictly positive, so the probe of \S queries still +/// appears even when every resource a task asks for is undeclared. +#[test] +fn test_query_partial_probes_for_entirely_undeclared_resources() { + let mut rt = TestEnv::new(); + rt.new_named_resource("gpus"); + rt.new_tasks(10, &TaskBuilder::new().cpus(1)); + rt.schedule(); + let r = compute_new_worker_query( + rt.core(), + &[WorkerTypeQuery { + partial: true, + // Declares only gpus; the cpus the tasks ask for are unknown until a worker connects. + descriptor: ResourceDescriptor::new( + vec![ResourceDescriptorItem { + name: "gpus".into(), + kind: ResourceDescriptorKind::simple_indices(8), + }], + Default::default(), + ), + time_limit: None, + max_sn_workers: 3, + max_workers_per_allocation: 1, + min_utilization: 0.0, + }], + ); + assert_eq!( + r.single_node_workers_per_query, + vec![1], + "one worker is asked for as a probe, not none and not all three" + ); +} + +/// An unbounded completion lets one worker absorb any amount of undeclared demand. +/// +/// A partial query fills in the resources the queue has not declared as unbounded sums. Two +/// things follow in the model: such a resource contributes a single unit to the global amount of +/// the objective rather than its numeric value, so a task that needs it keeps an ordinary weight; +/// and a worker whose free amount is unbounded gets no capacity row for it. Work that asks only +/// for undeclared resources therefore never needs a second worker, however much of it waits. +#[test] +fn test_query_partial_absorbs_undeclared_demand_on_one_worker() { + let ask = |n_cpu_only: usize, n_gpu: usize| { + let mut rt = TestEnv::new(); + rt.new_named_resource("gpus"); + if n_cpu_only > 0 { + rt.new_tasks(n_cpu_only, &TaskBuilder::new().cpus(1)); + } + if n_gpu > 0 { + rt.new_tasks(n_gpu, &TaskBuilder::new().cpus(1).add_resource(1, 1)); + } + rt.schedule(); + let r = compute_new_worker_query( + rt.core(), + &[WorkerTypeQuery { + partial: true, + descriptor: ResourceDescriptor::new( + vec![ResourceDescriptorItem { + name: "gpus".into(), + kind: ResourceDescriptorKind::simple_indices(8), + }], + Default::default(), + ), + time_limit: None, + max_sn_workers: 3, + max_workers_per_allocation: 1, + min_utilization: 0.0, + }], + ); + r.single_node_workers_per_query[0] + }; + // Sixteen tasks of one gpu each need two eight-gpu workers. + assert_eq!(ask(0, 16), 2); + // Ten thousand tasks that ask only for the undeclared resource add no worker: the first + // one absorbs all of them. + assert_eq!(ask(10_000, 16), 2); + // On their own they are the probe of \S queries: exactly one worker, not none and not three. + assert_eq!(ask(100, 0), 1); +} diff --git a/crates/tako/src/internal/tests/test_reactor.rs b/crates/tako/src/internal/tests/test_reactor.rs index 29c90d524..edd9da0bb 100644 --- a/crates/tako/src/internal/tests/test_reactor.rs +++ b/crates/tako/src/internal/tests/test_reactor.rs @@ -245,7 +245,7 @@ fn test_assignments_and_finish() { on_task_update(rt.core(), &mut comm, ws[0], smallvec![task_finished(t5)]); - assert!(rt.core().find_task(t5.into()).is_none()); + assert!(rt.core().find_task(t5).is_none()); check_worker_tasks_exact(rt.core(), ws[0], &[t1]); check_worker_tasks_exact(rt.core(), ws[1], &[t2]); check_worker_tasks_exact(rt.core(), ws[2], &[]); @@ -370,7 +370,7 @@ fn test_task_cancel() { rt.assign_task(t2, ws[0]); let mut comm = TestComm::new(); - on_cancel_tasks(rt.core(), &mut comm, &vec![wf[0], wf[1], t1, t2, wf[2]]); + on_cancel_tasks(rt.core(), &mut comm, &[wf[0], wf[1], t1, t2, wf[2]]); let msgs = comm.take_worker_msgs(ws[0], 1); assert!( @@ -441,7 +441,7 @@ fn test_worker_crashing_task() { on_remove_worker( rt.core(), &mut comm, - worker_id.into(), + worker_id, LostWorkerReason::HeartbeatLost, ); let mut lw = comm.client.take_lost_workers(); @@ -488,7 +488,7 @@ fn test_task_mn_fail() { comm.emptiness_check(); assert!(rt.core().find_task(1.into()).is_none()); for w in &ws { - assert!(rt.core().get_worker((*w).into()).mn_assignment().is_none()); + assert!(rt.core().get_worker(*w).mn_assignment().is_none()); } } @@ -605,7 +605,7 @@ fn lost_worker_with_running_and_assign_tasks() { } fn check_worker_tasks_exact(core: &Core, worker_id: WorkerId, tasks: &[TaskId]) { - let worker = core.get_worker(worker_id.into()); + let worker = core.get_worker(worker_id); let sn = worker.sn_assignment().unwrap(); assert_eq!(sn.assigned_tasks.len(), tasks.len()); for task in tasks { @@ -614,7 +614,7 @@ fn check_worker_tasks_exact(core: &Core, worker_id: WorkerId, tasks: &[TaskId]) } fn worker_has_task(core: &Core, worker_id: WorkerId, task_id: TaskId) -> bool { - core.get_worker(worker_id.into()) + core.get_worker(worker_id) .sn_assignment() .unwrap() .assigned_tasks diff --git a/crates/tako/src/internal/tests/test_scheduler_mn.rs b/crates/tako/src/internal/tests/test_scheduler_mn.rs index 3a48cf225..07ff0f845 100644 --- a/crates/tako/src/internal/tests/test_scheduler_mn.rs +++ b/crates/tako/src/internal/tests/test_scheduler_mn.rs @@ -158,8 +158,8 @@ fn test_schedule_mn_reserve() { let ws2 = rt.task(t2).mn_placement().unwrap().to_vec(); for w in &ws { - let s1 = get_worker_status(&ws1, (*w).into()); - let s2 = get_worker_status(&ws2, (*w).into()); + let s1 = get_worker_status(&ws1, *w); + let s2 = get_worker_status(&ws2, *w); let ms = comm.take_worker_msgs(*w, 0); check_worker_status_change(s1, s2, ms.as_slice()); } @@ -171,8 +171,8 @@ fn test_schedule_mn_reserve() { let ws3 = rt.task(t3).mn_placement().unwrap().to_vec(); for w in &ws { - let s1 = get_worker_status(&ws2, (*w).into()); - let s2 = get_worker_status(&ws3, (*w).into()); + let s1 = get_worker_status(&ws2, *w); + let s2 = get_worker_status(&ws3, *w); let ms = comm.take_worker_msgs(*w, 0); check_worker_status_change(s1, s2, ms.as_slice()); } @@ -183,7 +183,7 @@ fn test_schedule_mn_reserve() { let mut comm = rt.schedule(); for w in &ws { - let s = get_worker_status(&ws3, (*w).into()); + let s = get_worker_status(&ws3, *w); let ms = comm.take_worker_msgs(*w, 0); check_worker_status_change(s, WorkerStatus::None, ms.as_slice()); } @@ -354,3 +354,30 @@ fn test_schedule_mn_and_sn4() { assert!(rt.task(t1).is_mn_running()); assert!(rt.task(t2).is_assigned()); } + +#[test] +fn test_mn_task_stays_within_one_group_when_group_ids_interleave() { + // Worker ids of two groups interleave (a, b, a, b), as when two allocations start at the + // same time. Each group can run one 2-node task, and each task must use workers of a single + // group: a task spanning two allocations would not share an interconnect. + let mut rt = TestEnv::new(); + for group in ["a", "b", "a", "b"] { + rt.new_worker(&WorkerBuilder::new(1).group(group)); + } + rt.new_task(&TaskBuilder::new().n_nodes(2)); + rt.new_task(&TaskBuilder::new().n_nodes(2)); + let solution = rt.schedule_solution(); + + let tasks: Vec<_> = solution.mn_workers.values().flatten().collect(); + assert_eq!(tasks.len(), 2); + for workers in tasks { + let groups: Vec<&str> = workers + .iter() + .map(|w| rt.worker(*w).configuration.group.as_str()) + .collect(); + assert!( + groups.iter().all(|g| *g == groups[0]), + "multi-node task {workers:?} spans worker groups {groups:?}" + ); + } +} diff --git a/crates/tako/src/internal/tests/test_scheduler_sn.rs b/crates/tako/src/internal/tests/test_scheduler_sn.rs index a23b137d5..b94ff8a83 100644 --- a/crates/tako/src/internal/tests/test_scheduler_sn.rs +++ b/crates/tako/src/internal/tests/test_scheduler_sn.rs @@ -87,6 +87,60 @@ fn test_task_grouping_blocker() { assert!(!a[1].is_blocker); } +#[test] +fn test_task_grouping_same_level_needs_no_cut() { + // a: 5 4, b: 4. Nothing of b is strictly above a's task at 4, so `a` needs no cut. + let mut rt = TestEnv::new(); + rt.new_workers_cpus(&[5]); + let a5 = rt.new_task(&TaskBuilder::new().user_priority(5)); + rt.new_task(&TaskBuilder::new().user_priority(4)); + rt.new_task(&TaskBuilder::new().cpus(2).user_priority(4)); + let rq_a = rt.task(a5).resource_rq_id; + let now = std::time::Instant::now(); + let a = create_task_batches(rt.core(), now, None); + assert_eq!(a.len(), 2); + assert_eq!(a[0].resource_rq_id, rq_a); + assert_eq!(a[0].size, 2); + assert_eq!(a[0].cuts, vec![]); + assert_eq!( + a[1].cuts, + vec![PriorityCut { + size: 0, + blockers: vec![(rq_a, Some(1))], + }] + ); +} + +#[test] +fn test_task_grouping_no_repeated_cut() { + // c: 9, a: 5 4, b: 4. The tasks of other requests above a's task at 4 are the same as + // above its task at 5 (just c), so `a` gets a single cut. + let mut rt = TestEnv::new(); + rt.new_workers_cpus(&[10]); + let a5 = rt.new_task(&TaskBuilder::new().user_priority(5)); + rt.new_task(&TaskBuilder::new().user_priority(4)); + let b4 = rt.new_task(&TaskBuilder::new().cpus(2).user_priority(4)); + let c9 = rt.new_task(&TaskBuilder::new().cpus(3).user_priority(9)); + let rq_a = rt.task(a5).resource_rq_id; + let rq_b = rt.task(b4).resource_rq_id; + let rq_c = rt.task(c9).resource_rq_id; + let now = std::time::Instant::now(); + let a = create_task_batches(rt.core(), now, None); + assert_eq!(a.len(), 3); + let get = |rq| a.iter().find(|b| b.resource_rq_id == rq).unwrap(); + assert_eq!(get(rq_a).size, 2); + assert_eq!( + get(rq_a).cuts, + vec![PriorityCut { + size: 0, + blockers: vec![(rq_c, Some(1))], + }] + ); + assert_eq!(get(rq_b).cuts.len(), 1); + assert_eq!(get(rq_b).cuts[0].size, 0); + assert_eq!(get(rq_c).cuts, vec![]); +} + #[test] fn test_task_group_saturation() { let mut rt = TestEnv::new(); @@ -805,9 +859,9 @@ fn test_no_deps_scattering_1() { comm.emptiness_check(); rt.core().sanity_check(); - let c1 = if m1.len() > 0 { task_count(&m1[0]) } else { 0 }; - let c2 = if m2.len() > 0 { task_count(&m2[0]) } else { 0 }; - let c3 = if m3.len() > 0 { task_count(&m3[0]) } else { 0 }; + let c1 = m1.first().map_or(0, task_count); + let c2 = m2.first().map_or(0, task_count); + let c3 = m3.first().map_or(0, task_count); assert_eq!(c1, 4); assert_eq!(c2, 0); @@ -1858,6 +1912,385 @@ fn test_schedule_reservation_leaves_other_workers_for_backfill() { ); } +/// Narrow tasks assigned to each of `workers`, in order. +fn narrow_per_worker(rt: &TestEnv, narrow: &[TaskId], workers: &[WorkerId]) -> Vec { + workers + .iter() + .map(|w| { + narrow + .iter() + .filter(|t| { + matches!(rt.task(**t).state, + TaskRuntimeState::Assigned { worker_id, .. } if worker_id == *w) + }) + .count() + }) + .collect() +} + +#[test] +fn test_schedule_placeable_blocker_holds_nothing() { + // A wide blocker that a free worker can host right now needs no held worker, so capable but + // partially occupied workers keep backfilling. + let mut rt = TestEnv::new(); + let ws = rt.new_workers(4, &WorkerBuilder::new(8)); + for w in &ws[1..] { + rt.new_task_running(&TaskBuilder::new().cpus(6), *w); + } + let blocker = rt.new_task(&TaskBuilder::new().cpus(8).user_priority(10)); + let narrow = rt.new_tasks(30, &TaskBuilder::new().cpus(1)); + rt.schedule(); + + assert!( + matches!(rt.task(blocker).state, + TaskRuntimeState::Assigned { worker_id, .. } if worker_id == ws[0]), + "the blocker starts on the free worker" + ); + assert_eq!( + narrow_per_worker(&rt, &narrow, &ws[1..]), + vec![2, 2, 2], + "no worker is held, so every partial worker backfills its 2 free cpus" + ); +} + +#[test] +fn test_schedule_lower_priority_cannot_take_capacity_of_higher() { + // One free worker can host either the higher-priority 4-cpu task or the lower-priority 8-cpu + // task, but not both, so priority must give it to the 4-cpu task. + // + // The saturated worker is what makes this fail. It is capable of the 4-cpu task but has no + // free cpus, so the solver gets a reservation variable for that task there, and reserving a + // worker with nothing free costs no capacity. Setting it discharges the priority condition + // "8-cpu tasks may run only once the 4-cpu task is served", after which the 8-cpu task is the + // more valuable placement. Held workers do not close this: held-worker depth is the number of + // blockers no worker can host *before* the solve, and the free worker is counted as hosting + // the 4-cpu task although the solve then gives it to the 8-cpu one. + // + // Without the saturated worker the 4-cpu task is placed correctly. + let mut rt = TestEnv::new(); + let ws = rt.new_workers(2, &WorkerBuilder::new(8)); + for _ in 0..8 { + rt.new_task_running(&TaskBuilder::new().cpus(1), ws[1]); + } + let high = rt.new_task(&TaskBuilder::new().cpus(4).user_priority(11)); + let wide = rt.new_task(&TaskBuilder::new().cpus(8).user_priority(10)); + rt.schedule(); + + assert!( + matches!(rt.task(high).state, + TaskRuntimeState::Assigned { worker_id, .. } if worker_id == ws[0]), + "the higher-priority task must take the free worker; high = {:?}, wide = {:?}", + rt.task(high).state, + rt.task(wide).state + ); + assert!( + !rt.task(wide).is_assigned(), + "the lower-priority task must wait; wide = {:?}", + rt.task(wide).state + ); +} + +#[test] +fn test_schedule_reservation_cannot_take_capacity_of_higher_blocker() { + // A reservation for a request must obey that request's own priority conditions, exactly as + // its placements do. Here `x` (8 cpus, priority 10) cannot run anywhere now, so the emptiest + // capable worker `w1` is held for it. `h` (4 cpus + a "foo" only `w1` has, priority 20) + // outranks `x` and fits into `w1`'s free cpus. Reserving `w1` for `x` would consume those + // cpus and leave `h` waiting behind a lower-priority request. + // + // The objective alone does not prevent it. `h` is incapable of the other workers, so a + // reservation for `x` releases their freed cpus for narrow work without `h` blocking it, and + // eight workers' worth of backfill outweighs placing `h`. + let mut rt = TestEnv::new(); + rt.new_named_resource("foo"); + let w1 = rt.new_worker(&WorkerBuilder::new(8).res_sum("foo", 1000)); + rt.new_task_running(&TaskBuilder::new().cpus(4), w1); + let others = rt.new_workers(8, &WorkerBuilder::new(8)); + for w in &others { + rt.new_task_running(&TaskBuilder::new().cpus(6), *w); + } + let h = rt.new_task( + &TaskBuilder::new() + .cpus(4) + .add_resource(1, 1) + .user_priority(20), + ); + let x = rt.new_task(&TaskBuilder::new().cpus(8).user_priority(10)); + rt.new_tasks(40, &TaskBuilder::new().cpus(1)); + rt.schedule(); + + assert!( + matches!(rt.task(h).state, + TaskRuntimeState::Assigned { worker_id, .. } if worker_id == w1), + "the higher-priority task must take the free cpus on w1; h = {:?}, x = {:?}", + rt.task(h).state, + rt.task(x).state + ); +} + +#[test] +fn test_schedule_held_worker_not_refilled_when_one_reservation_serves_blocker() { + // Two 8-cpu tasks share one request, but the narrow request's priority lies between them: + // `x_high` (10) > narrow (5) > `x_low` (3). The held count covers the whole batch, so two + // workers are held, while the narrow request is blocked only by `x_high`, i.e. at threshold 1. + // One reservation therefore serves the blocker for the narrow request, and a second held + // worker is left with free capacity. + // + // `w0` is the emptiest worker and should accumulate for `x_high`. If narrow work may enter a + // held worker whenever the blocker is served elsewhere, the solver reserves `w1` instead (a + // reservation on a higher index is cheaper) and refills `w0`, so `x_high` waits for the fuller + // worker to drain. + let mut rt = TestEnv::new(); + let ws = rt.new_workers(2, &WorkerBuilder::new(8)); + let mut running: Vec> = [4, 6] + .iter() + .zip(&ws) + .map(|(n, w)| { + (0..*n) + .map(|_| rt.new_task_running(&TaskBuilder::new().cpus(1), *w)) + .collect() + }) + .collect(); + let x_high = rt.new_task(&TaskBuilder::new().cpus(8).user_priority(10)); + rt.new_tasks(40, &TaskBuilder::new().cpus(1).user_priority(5)); + rt.new_task(&TaskBuilder::new().cpus(8).user_priority(3)); + + // `w0` needs 4 finishes to fit an 8-cpu task, `w1` needs 6. + for _tick in 0..=4 { + rt.schedule(); + if rt.task(x_high).is_assigned() { + return; + } + for (idx, w) in ws.iter().enumerate() { + if let Some(task_id) = running[idx].pop() { + rt.finish_task(task_id, *w); + } + } + } + let free: Vec = ws + .iter() + .map(|w| { + let a = rt.worker(*w).sn_assignment().unwrap(); + format!("w{w}: {:?}", a.free_resources.get(ResourceId::new(0))) + }) + .collect(); + panic!( + "the priority-10 task did not start once the emptiest worker drained ({})", + free.join(", ") + ); +} + +#[test] +fn test_schedule_no_worker_held_for_blocker_tasks_the_blocked_request_outranks() { + // `x_high` (10) > narrow (5) > `x_low` (3), both x tasks 8 cpus and unplaceable now. Narrow is + // blocked only by `x_high`, so only one worker is needed for it: the emptiest, `w0`. Holding a + // second worker for `x_low` would keep narrow off it although narrow outranks `x_low`. + let mut rt = TestEnv::new(); + let ws = rt.new_workers(2, &WorkerBuilder::new(8)); + for (n, w) in [4, 6].iter().zip(&ws) { + for _ in 0..*n { + rt.new_task_running(&TaskBuilder::new().cpus(1), *w); + } + } + rt.new_task(&TaskBuilder::new().cpus(8).user_priority(10)); + let narrow = rt.new_tasks(40, &TaskBuilder::new().cpus(1).user_priority(5)); + rt.new_task(&TaskBuilder::new().cpus(8).user_priority(3)); + rt.schedule(); + + assert_eq!( + narrow_per_worker(&rt, &narrow, &ws), + vec![0, 2], + "w0 is held for x_high; w1 must be free for narrow work" + ); +} + +#[test] +fn test_schedule_held_worker_not_refilled_at_a_shallower_threshold() { + // Two lower requests block on the same 8-cpu request at different thresholds: + // `x_a` (10) > `l1` (9) > `x_b` (8) > `l2` (7). `l2` needs both x tasks served, so two workers + // are held and both carry a reservation variable; `l1` needs only `x_a`, i.e. one. A single + // reservation on the fuller `w1` serves `x_a` for `l1`, so the emptiest `w0` must stay closed to + // `l1` by the held-worker constraint itself, or `l1` refills it and `x_a` waits for `w1`. + let mut rt = TestEnv::new(); + let ws = rt.new_workers(2, &WorkerBuilder::new(8)); + let mut running: Vec> = [4, 6] + .iter() + .zip(&ws) + .map(|(n, w)| { + (0..*n) + .map(|_| rt.new_task_running(&TaskBuilder::new().cpus(1), *w)) + .collect() + }) + .collect(); + let x_a = rt.new_task(&TaskBuilder::new().cpus(8).user_priority(10)); + rt.new_tasks(40, &TaskBuilder::new().cpus(1).user_priority(9)); + rt.new_task(&TaskBuilder::new().cpus(8).user_priority(8)); + rt.new_tasks(40, &TaskBuilder::new().cpus(2).user_priority(7)); + + for _tick in 0..=4 { + rt.schedule(); + if rt.task(x_a).is_assigned() { + return; + } + for (idx, w) in ws.iter().enumerate() { + if let Some(task_id) = running[idx].pop() { + rt.finish_task(task_id, *w); + } + } + } + let free: Vec = ws + .iter() + .map(|w| { + let a = rt.worker(*w).sn_assignment().unwrap(); + format!("w{w}: {:?}", a.free_resources.get(ResourceId::new(0))) + }) + .collect(); + panic!( + "the priority-10 task did not start once the emptiest worker drained ({})", + free.join(", ") + ); +} + +#[test] +fn test_schedule_reservation_pays_on_a_large_cluster() { + // Placement weights scale with 1 / (free cpus in the cluster), so on a large cluster each + // placement is worth very little. The reservation penalty must scale with them: a reservation + // that unlocks placements has to pay for itself regardless of cluster size. + // + // `x` (8 cpus + 1 "foo", priority 10) fits nowhere now. `w0` is held for it. `w1` is capable + // but not held, so narrow work there is blocked until `x` is served, which only a reservation + // on `w0` can do. The `big` workers contribute 10 000 free cpus but have no "foo": `x` never + // runs there, so narrow work fills them unconditionally and more narrow tasks wait than they + // can take. The cpus are split over ten workers because one worker's share of a batch is + // capped at `MAX_TASK_PER_WORKER`. + let mut rt = TestEnv::new(); + rt.new_named_resource("foo"); + let big = rt.new_workers(10, &WorkerBuilder::new(1000)); + let w0 = rt.new_worker(&WorkerBuilder::new(8).res_sum("foo", 1)); + let w1 = rt.new_worker(&WorkerBuilder::new(8).res_sum("foo", 1)); + rt.new_task_running(&TaskBuilder::new().cpus(4), w0); + rt.new_task_running(&TaskBuilder::new().cpus(6), w1); + rt.new_task( + &TaskBuilder::new() + .cpus(8) + .add_resource(1, 1) + .user_priority(10), + ); + let narrow = rt.new_tasks(10_020, &TaskBuilder::new().cpus(1).user_priority(5)); + rt.schedule(); + + assert_eq!( + narrow_per_worker(&rt, &narrow, &big), + vec![1000; 10], + "narrow work must fill the workers without \"foo\"" + ); + assert_eq!( + narrow_per_worker(&rt, &narrow, &[w0, w1]), + vec![0, 2], + "a reservation on w0 must release w1 even when placements are worth little" + ); +} + +#[test] +fn test_schedule_one_worker_cannot_be_reserved_for_two_blockers() { + // `x` and `y` are different requests at equal priority, both needing the single "foo" of a + // worker, and neither fits anywhere now. `w0` has 7 free cpus but its "foo" is taken; `w1` is + // fully occupied. Neither covers any part of either request, so both are held on the higher + // id, `w1`, which has no free resources. Reservations there consume nothing, so without a + // per-worker limit `x` and `y` are both served by one worker that can later host only one + // of them, and narrow work fills `w0` (seven tasks). With the limit, the one not reserved for is + // held back on `w0` by its own condition, and `w0` offers no gap either: `y` alone could never + // use more than 7 of its 8 cpus, but `x` can use all eight, and a gap is only capacity that no + // waiting blocker can use. + let mut rt = TestEnv::new(); + rt.new_named_resource("foo"); + let w0 = rt.new_worker(&WorkerBuilder::new(8).res_sum("foo", 1)); + let w1 = rt.new_worker(&WorkerBuilder::new(8).res_sum("foo", 1)); + rt.new_task_running(&TaskBuilder::new().cpus(1).add_resource(1, 1), w0); + rt.new_task_running(&TaskBuilder::new().cpus(8).add_resource(1, 1), w1); + rt.new_task( + &TaskBuilder::new() + .cpus(8) + .add_resource(1, 1) + .user_priority(10), + ); + rt.new_task( + &TaskBuilder::new() + .cpus(7) + .add_resource(1, 1) + .user_priority(10), + ); + let narrow = rt.new_tasks(20, &TaskBuilder::new().cpus(1).user_priority(5)); + rt.schedule(); + + assert_eq!( + narrow_per_worker(&rt, &narrow, &[w0, w1]), + vec![0, 0], + "one worker can be reserved for at most one blocker, so one of them stays unserved" + ); +} + +#[test] +fn test_schedule_reservation_leaves_gap_allowance() { + // A reservation withholds the worker's free resources for the blocker, but part of them is + // capacity the blocker can never use: its gap allowance. The priority condition already lets + // gap tasks run there, and the reservation must not take that capacity away. + // + // Two workers of 16 cpus and 4 gpus each run four 1-cpu + 1-gpu tasks, so every gpu is taken + // and the blocker (2 cpus + 1 gpu) fits nowhere. A full packing of the blocker uses 8 cpus + // and all 4 gpus, so 8 cpus per worker are never usable by it. While k of the running tasks + // remain it fits 4 - k times and cannot use 8 + k cpus, so the gap is 8 cpus (the bound + // `C - M(C) - o` would charge the running tasks' 4 cpus and give only 4). `w1` is held (tie, + // higher id); a reservation there serves the blocker and releases `w0`, but it must consume + // only 12 - 8 = 4 of `w1`'s free cpus, leaving room for 8 gap tasks. + let mut rt = TestEnv::new(); + rt.new_named_resource("gpu"); + let ws = rt.new_workers(2, &WorkerBuilder::new(16).res_sum("gpu", 4)); + for w in &ws { + for _ in 0..4 { + rt.new_task_running(&TaskBuilder::new().cpus(1).add_resource(1, 1), *w); + } + } + rt.new_task( + &TaskBuilder::new() + .cpus(2) + .add_resource(1, 1) + .user_priority(10), + ); + let narrow = rt.new_tasks(40, &TaskBuilder::new().cpus(1).user_priority(5)); + rt.schedule(); + + assert_eq!( + narrow_per_worker(&rt, &narrow, &ws), + vec![12, 8], + "w0 is released by the reservation; the reserved w1 keeps its 8-cpu gap" + ); +} + +#[test] +fn test_schedule_reservation_example() { + // ten 8-cpu workers, each running a single + // 4-cpu task, one 6-cpu task at priority 2 and a hundred 1-cpu tasks at priority 1. The 6-cpu + // task fits nowhere, so one worker is held and reserved for it, which releases the other nine + // (4 cpus each). On the reserved worker the reservation withholds only what the 6-cpu task + // could use: 8 mod 6 = 2 cpus are its gap and stay available. + // + // The single 4-cpu occupant matters: four 1-cpu occupants could leave 6 cpus free at an + // intermediate point, and the gap would be 0. + let mut rt = TestEnv::new(); + let ws = rt.new_workers(10, &WorkerBuilder::new(8)); + for w in &ws { + rt.new_task_running(&TaskBuilder::new().cpus(4), *w); + } + rt.new_task(&TaskBuilder::new().cpus(6).user_priority(2)); + let narrow = rt.new_tasks(100, &TaskBuilder::new().cpus(1).user_priority(1)); + rt.schedule(); + + let mut expected = vec![4; 9]; + expected.push(2); + assert_eq!(narrow_per_worker(&rt, &narrow, &ws), expected); +} + #[test] fn test_schedule_blockers_hold_all_needed_workers() { const N_WORKERS: usize = 4; @@ -2034,3 +2467,214 @@ fn test_reservation_one_assigned() { .sum::(); assert_eq!(filler_scheduled, 6); } + +#[test] +fn test_schedule_gap_tasks_do_not_move_the_hold() { + // The blocker needs 6 cpus. `w0` runs one 4-cpu task: 4 cpus free, 2 of them a gap the blocker + // can never use. `w1` runs one 5-cpu task: 3 cpus free, 2 of them a gap. `w0` is held and its + // gap is filled with 2 narrow tasks. Then 3 narrow tasks finish on `w1`. By plain free cpus + // `w1` (3) now looks better than `w0` (2), but the blocker can use only 1 of them, against 2 on + // `w0`. The hold must stay on `w0`: filling a gap never makes a worker less useful to the + // blocker, so it must not cost the worker its hold (and then its drained capacity). + let mut rt = TestEnv::new(); + let w0 = rt.new_worker(&WorkerBuilder::new(8)); + let w1 = rt.new_worker(&WorkerBuilder::new(8)); + rt.new_task_running(&TaskBuilder::new().cpus(4), w0); + rt.new_task_running(&TaskBuilder::new().cpus(5), w1); + rt.new_task(&TaskBuilder::new().cpus(6).user_priority(10)); + let narrow = rt.new_tasks(40, &TaskBuilder::new().cpus(1).user_priority(5)); + rt.schedule(); + assert_eq!(narrow_per_worker(&rt, &narrow, &[w0, w1]), vec![2, 3]); + + let finished: Vec<_> = narrow + .iter() + .copied() + .filter(|t| { + matches!(rt.task(*t).state, + TaskRuntimeState::Assigned { worker_id, .. } if worker_id == w1) + }) + .collect(); + for t in &finished { + rt.finish_task(*t, w1); + } + let live: Vec<_> = narrow + .iter() + .copied() + .filter(|t| !finished.contains(t)) + .collect(); + rt.schedule(); + assert_eq!( + narrow_per_worker(&rt, &live, &[w0, w1]), + vec![2, 3], + "w0 keeps its hold; w1 is refilled" + ); +} + +#[test] +fn test_schedule_blocker_count_allowance_is_not_multiplied_per_worker() { + let mut rt = TestEnv::new(); + let ws = rt.new_workers(3, &WorkerBuilder::new(12)); + for w in &ws { + rt.new_task_running(&TaskBuilder::new().cpus(5), *w); + } + rt.new_task(&TaskBuilder::new().cpus(8).user_priority(2)); + let above = rt.new_tasks(3, &TaskBuilder::new().cpus(1).user_priority(3)); + let below = rt.new_tasks(40, &TaskBuilder::new().cpus(1).user_priority(1)); + + rt.schedule(); + + let per_worker: Vec = narrow_per_worker(&rt, &above, &ws) + .iter() + .zip(narrow_per_worker(&rt, &below, &ws)) + .map(|(a, b)| a + b) + .collect(); + assert!( + per_worker.iter().any(|placed| *placed <= 4), + "every worker took more than its 4-cpu gap ({per_worker:?}), \ + so the blocker cannot start anywhere once the 5-cpu tasks finish" + ); +} + +#[test] +fn test_schedule_two_requests_share_one_gap_allowance() { + let mut rt = TestEnv::new(); + let foo = rt.new_named_resource("foo"); + let w = rt.new_worker(&WorkerBuilder::new(12).res_sum("foo", 10)); + rt.new_task_running(&TaskBuilder::new().cpus(5), w); + rt.new_task(&TaskBuilder::new().cpus(8).user_priority(10)); + let plain = rt.new_tasks(20, &TaskBuilder::new().cpus(1).user_priority(5)); + let with_foo = rt.new_tasks( + 20, + &TaskBuilder::new() + .cpus(1) + .add_resource(foo, 1) + .user_priority(5), + ); + + rt.schedule(); + + let plain_placed = narrow_per_worker(&rt, &plain, &[w])[0]; + let foo_placed = narrow_per_worker(&rt, &with_foo, &[w])[0]; + assert!( + plain_placed + foo_placed <= 4, + "worker runs {} lower-priority cpus ({plain_placed} plain + {foo_placed} with foo), \ + but the blocker leaves a gap of 4 there", + plain_placed + foo_placed + ); +} + +#[test] +fn test_schedule_soft_rejected_worker_is_still_limited_by_the_blocker() { + let mut rt = TestEnv::new(); + let w = rt.new_worker(&WorkerBuilder::new(12)); + rt.new_task_running(&TaskBuilder::new().cpus(5), w); + let blocker = rt.new_task(&TaskBuilder::new().cpus(8).user_priority(10)); + let narrow = rt.new_tasks(20, &TaskBuilder::new().cpus(1).user_priority(1)); + + let blocker_rq = rt.task(blocker).resource_rq_id; + rt.core() + .get_worker_mut(w) + .block_request(blocker_rq, ResourceVariantId::new(0)); + + rt.schedule(); + + assert_eq!( + narrow_per_worker(&rt, &narrow, &[w]), + vec![4], + "a soft-rejected worker keeps its gap limit" + ); +} + +/// The excess of a request over a worker's gap must never exceed what the request actually runs +/// there. The excess lowers the shared gap constraint, so excess that no task uses would buy other +/// requests room beyond the gap, paid from an allowance spent on nothing. +/// +/// `w0` has 12 cpus with a 5-cpu task running: 7 free, and a gap of 4 for the 8-cpu blocker at +/// priority 5. The donor request (2-cpu tasks) has three tasks above the blocker, so it carries an +/// allowance of 3, and they are placed on `w1`, which is too small for the blocker. The donor +/// therefore runs nothing on `w0` while holding an allowance there. Its own tasks may use that +/// allowance and outrank the blocker, but the other two requests have no allowance at all: they +/// share `w0`'s gap and must stay within it together, 4 cpus rather than 4 each. +#[test] +fn test_schedule_unused_allowance_does_not_widen_the_shared_gap() { + let mut rt = TestEnv::new(); + let foo = rt.new_named_resource("foo"); + let w0 = rt.new_worker(&WorkerBuilder::new(12).res_sum("foo", 10)); + let _w1 = rt.new_worker(&WorkerBuilder::new(6)); + rt.new_task_running(&TaskBuilder::new().cpus(5), w0); + rt.new_tasks(3, &TaskBuilder::new().cpus(2).user_priority(9)); + rt.new_tasks(10, &TaskBuilder::new().cpus(2).user_priority(1)); + rt.new_task(&TaskBuilder::new().cpus(8).user_priority(5)); + let b1 = rt.new_tasks(20, &TaskBuilder::new().cpus(1).user_priority(1)); + let b2 = rt.new_tasks( + 20, + &TaskBuilder::new() + .cpus(1) + .add_resource(foo, 1) + .user_priority(1), + ); + + rt.schedule(); + + let b1_cpus = narrow_per_worker(&rt, &b1, &[w0])[0]; + let b2_cpus = narrow_per_worker(&rt, &b2, &[w0])[0]; + assert!( + b1_cpus + b2_cpus <= 4, + "requests without an allowance use {} cpus of w0 ({b1_cpus} + {b2_cpus}), \ + but they share a gap of 4", + b1_cpus + b2_cpus + ); +} + +/// Gap filling is bounded by what *all* the blockers of a worker could use together, not by what +/// each of them could use alone. On a 12-cpu worker with a 5-cpu and a 7-cpu blocker the two +/// gaps are 2 and 5, but 5 + 7 fills the worker exactly, so there is nothing to fill. +#[test] +fn test_schedule_gap_is_bounded_by_the_mix_of_blockers() { + let mut rt = TestEnv::new(); + rt.new_worker(&WorkerBuilder::new(12)); + let wide = rt.new_tasks(2, &TaskBuilder::new().cpus(5).user_priority(10)); + let wider = rt.new_tasks(2, &TaskBuilder::new().cpus(7).user_priority(10)); + let narrow = rt.new_tasks(8, &TaskBuilder::new().cpus(1).user_priority(0)); + rt.schedule(); + let placed = |ts: &[TaskId]| ts.iter().filter(|t| rt.task(**t).is_assigned()).count(); + assert_eq!(placed(&wide), 1); + assert_eq!(placed(&wider), 1); + assert_eq!( + placed(&narrow), + 0, + "the blockers pack into the worker exactly, so no narrow task may take a cpu" + ); +} + +/// The same worker with one cpu more does leave a gap of one, and exactly one narrow task fits. +#[test] +fn test_schedule_mix_of_blockers_can_still_leave_a_gap() { + let mut rt = TestEnv::new(); + rt.new_worker(&WorkerBuilder::new(13)); + let wide = rt.new_tasks(2, &TaskBuilder::new().cpus(5).user_priority(10)); + let wider = rt.new_tasks(2, &TaskBuilder::new().cpus(7).user_priority(10)); + let narrow = rt.new_tasks(8, &TaskBuilder::new().cpus(1).user_priority(0)); + rt.schedule(); + let placed = |ts: &[TaskId]| ts.iter().filter(|t| rt.task(**t).is_assigned()).count(); + assert_eq!(placed(&wide), 1); + assert_eq!(placed(&wider), 1); + assert_eq!(placed(&narrow), 1); +} + +/// A single blocker keeps its own gap: the joint bound only exists for a mix. +#[test] +fn test_schedule_single_blocker_keeps_its_gap() { + let mut rt = TestEnv::new(); + rt.new_worker(&WorkerBuilder::new(12)); + let wider = rt.new_tasks(3, &TaskBuilder::new().cpus(7).user_priority(10)); + let narrow = rt.new_tasks(8, &TaskBuilder::new().cpus(1).user_priority(0)); + rt.schedule(); + let placed = |ts: &[TaskId]| ts.iter().filter(|t| rt.task(**t).is_assigned()).count(); + assert_eq!(placed(&wider), 1); + assert_eq!( + placed(&narrow), + 5, + "12 - 7 = 5 cpus no 7-cpu task can ever use" + ); +} diff --git a/crates/tako/src/internal/tests/utils/env.rs b/crates/tako/src/internal/tests/utils/env.rs index 934db2b5c..83fb394bf 100644 --- a/crates/tako/src/internal/tests/utils/env.rs +++ b/crates/tako/src/internal/tests/utils/env.rs @@ -142,7 +142,7 @@ impl TestEnv { let worker_id = WorkerId::new(self.worker_id_counter); self.worker_id_counter += 1; let resource_id_map = self.core.create_resource_map(); - let worker = builder.build(worker_id, &&resource_id_map, Instant::now()); + let worker = builder.build(worker_id, &resource_id_map, Instant::now()); on_new_worker(&mut self.core, &mut TestComm::default(), worker); worker_id } @@ -163,10 +163,7 @@ impl TestEnv { pub fn check_worker_tasks(&self, worker_id: WorkerId, tasks: &[TaskId]) { let ids: Vec = self.worker_tasks(worker_id).iter().copied().collect(); - assert_eq!( - ids, - utils::sorted_vec(tasks.iter().map(|&id| id.into()).collect()) - ); + assert_eq!(ids, utils::sorted_vec(tasks.to_vec())); } pub fn start_and_finish_task(&mut self, task_id: TaskId, worker_id: WorkerId) { @@ -203,7 +200,7 @@ impl TestEnv { panic!("Task {} is not waiting", task_id); } } - let w = worker_map.get_worker_mut(worker_id.into()); + let w = worker_map.get_worker_mut(worker_id); let rq = request_map .get(task.resource_rq_id) .get(ResourceVariantId::new(0)); @@ -256,7 +253,7 @@ impl TestEnv { pub fn schedule_mapping(&mut self) -> WorkerTaskMapping { let batches = create_task_batches(&mut self.core, self.now, None); - let solution = run_scheduling_solver(&mut self.core, self.now, &batches, None); + let solution = run_scheduling_solver(&self.core, self.now, &batches, None); create_task_mapping(&mut self.core, solution) } @@ -264,7 +261,7 @@ impl TestEnv { /// tests can inspect `SchedulingSolution::is_optimal` directly. pub fn schedule_solution(&mut self) -> SchedulingSolution { let batches = create_task_batches(&mut self.core, self.now, None); - run_scheduling_solver(&mut self.core, self.now, &batches, None) + run_scheduling_solver(&self.core, self.now, &batches, None) } } diff --git a/crates/tako/src/internal/tests/utils/scheduler.rs b/crates/tako/src/internal/tests/utils/scheduler.rs index 167cad306..261e09399 100644 --- a/crates/tako/src/internal/tests/utils/scheduler.rs +++ b/crates/tako/src/internal/tests/utils/scheduler.rs @@ -55,7 +55,6 @@ impl TestCase { mapping .workers .get(&worker.get_worker_id()) - .as_deref() .map_or(&[], |up| up.assigned.as_slice()), ); } @@ -108,7 +107,7 @@ pub fn normalize_workers(mapping: &mut WorkerTaskMapping, workers: &[WorkerId]) .map(|w| mapping.workers.get(w).cloned().unwrap_or_default()) .collect(); tasks.sort_unstable(); - for (w, up) in workers.iter().zip(tasks.into_iter()) { + for (w, up) in workers.iter().zip(tasks) { mapping.workers.insert(*w, up); } } @@ -137,7 +136,7 @@ impl TestWorker { } pub fn get_eq_class(&self) -> Option { - self.eq_class.clone() + self.eq_class } pub fn get_worker_id(&self) -> WorkerId { diff --git a/crates/tako/src/internal/tests/utils/task.rs b/crates/tako/src/internal/tests/utils/task.rs index 37ec900d7..f2d1a7fe5 100644 --- a/crates/tako/src/internal/tests/utils/task.rs +++ b/crates/tako/src/internal/tests/utils/task.rs @@ -124,7 +124,7 @@ impl TaskBuilder { pub fn build(&self, task_id: TaskId, core: &mut Core) -> Task { let rq_id = self.build_resource_rq_id(core); Task::new( - task_id.into(), + task_id, rq_id, self.task_deps.iter().copied().collect(), None, @@ -165,7 +165,7 @@ impl TaskBuilder { pub fn task_running_msg(task_id: TaskId) -> TaskRunningMsg { TaskRunningMsg { - task_id: task_id, + task_id, rv_id: ResourceVariantId::new(0), context: Default::default(), } diff --git a/crates/tako/src/internal/worker/resources/groups.rs b/crates/tako/src/internal/worker/resources/groups.rs index 3c95b73d7..5074e8908 100644 --- a/crates/tako/src/internal/worker/resources/groups.rs +++ b/crates/tako/src/internal/worker/resources/groups.rs @@ -139,7 +139,7 @@ pub fn group_solver( [(v2, 1.0), (v3, -1.0)].into_iter(), ); } - let (solution, objective_value): (_, _) = solver.solve()?; + let solution = solver.solve(None)?; Some(( vars.iter() .map(|var_group| { @@ -150,6 +150,6 @@ pub fn group_solver( .collect() }) .collect(), - objective_value, + solution.objective(), )) }