diff --git a/crates/tako/src/internal/scheduler/solver.rs b/crates/tako/src/internal/scheduler/solver.rs index c2cc35916..514634af9 100644 --- a/crates/tako/src/internal/scheduler/solver.rs +++ b/crates/tako/src/internal/scheduler/solver.rs @@ -410,7 +410,7 @@ pub(crate) fn run_scheduling_solver( } let mut result = SchedulingSolution::default(); - let Some((solution, _)) = solver.solve() else { + let Some((solution, _)) = solver.solve_bounded() else { return result; }; diff --git a/crates/tako/src/internal/solver/config.rs b/crates/tako/src/internal/solver/config.rs new file mode 100644 index 000000000..8461a5d2d --- /dev/null +++ b/crates/tako/src/internal/solver/config.rs @@ -0,0 +1,82 @@ +use std::time::Duration; + +/// Relative MIP optimality gap: the solver accepts a solution once it is +/// provably within this fraction of the true optimum, instead of always +/// proving exact optimality. Task resource requests are themselves estimates +/// (cpu/memory bucketing), so demanding an exact optimum is false precision; +/// a 10% gap trades a small, usually much smaller in practice, placement +/// suboptimality for a solve that reliably finishes in well under a second on +/// realistic instances instead of stalling the single-threaded server. +pub(crate) fn mip_rel_gap() -> f64 { + // Unit tests assert exact placement counts/priority behavior on small, + // fast-solving instances -- a nonzero default gap there would trade + // correctness-test fidelity for a speedup these tiny instances don't + // need. A test that wants to exercise gap-tuned behavior specifically + // uses with_test_solver_config below (thread-local, not process-global + // env vars, so it can't race with unrelated tests running concurrently + // on other threads). + #[cfg(test)] + if let Some(v) = TEST_REL_GAP_OVERRIDE.with(|c| c.get()) { + return v; + } + #[cfg(test)] + let default = 0.0; + #[cfg(not(test))] + let default = 0.10; + + get_f64_from_env("HQ_SCHEDULER_MIP_REL_GAP").unwrap_or(default) +} + +/// Hard wall-clock cap on a single scheduling solve. The solver is otherwise +/// unbounded and can run for minutes to hours on workloads with many distinct +/// resource shapes, blocking the single-threaded server (no heartbeats, no +/// RPCs, no other scheduling) for the entire duration. 5s clears the steep +/// part of the incumbent-quality cliff observed on realistic and +/// harder-than-realistic synthetic instances while bounding the worst case. +pub(crate) fn mip_time_limit() -> Duration { + // See mip_rel_gap: unit tests need exact, unhurried solves on tiny + // instances, not a production-scale wall-clock bound. + #[cfg(test)] + if let Some(v) = TEST_TIME_LIMIT_OVERRIDE.with(|c| c.get()) { + return v; + } + #[cfg(test)] + let default = Duration::from_secs(60); + #[cfg(not(test))] + let default = Duration::from_secs(5); + + get_duration_from_env("HQ_SCHEDULER_MIP_TIME_LIMIT_MS").unwrap_or(default) +} + +fn get_f64_from_env(key: &str) -> Option { + std::env::var(key).ok().and_then(|value| value.parse::().ok()) +} + +fn get_duration_from_env(key: &str) -> Option { + std::env::var(key) + .ok() + .and_then(|value| value.parse::().ok()) + .map(Duration::from_millis) +} + +#[cfg(test)] +thread_local! { + static TEST_REL_GAP_OVERRIDE: std::cell::Cell> = const { std::cell::Cell::new(None) }; + static TEST_TIME_LIMIT_OVERRIDE: std::cell::Cell> = const { std::cell::Cell::new(None) }; +} + +/// Runs `f` with a scheduler solver config override in effect, for tests +/// that need to exercise the production-tuned (or otherwise non-default) +/// solve_bounded() behavior. Thread-local, not a process-global env var: the +/// Rust test harness runs each #[test] to completion on its own OS thread, +/// so this cannot race with unrelated tests running concurrently on other +/// threads the way a process-global env var would. +#[cfg(test)] +pub(crate) fn with_test_solver_config(rel_gap: f64, time_limit: Duration, f: impl FnOnce() -> R) -> R { + TEST_REL_GAP_OVERRIDE.with(|c| c.set(Some(rel_gap))); + TEST_TIME_LIMIT_OVERRIDE.with(|c| c.set(Some(time_limit))); + let result = f(); + TEST_REL_GAP_OVERRIDE.with(|c| c.set(None)); + TEST_TIME_LIMIT_OVERRIDE.with(|c| c.set(None)); + result +} diff --git a/crates/tako/src/internal/solver/highs.rs b/crates/tako/src/internal/solver/highs.rs index 64e9ad90f..d12b473b5 100644 --- a/crates/tako/src/internal/solver/highs.rs +++ b/crates/tako/src/internal/solver/highs.rs @@ -1,5 +1,6 @@ +use crate::internal::solver::config::{mip_rel_gap, mip_time_limit}; use crate::internal::solver::{ConstraintType, LpInnerSolver, LpSolution}; -use highs::Sense; +use highs::{HighsModelStatus, HighsSolutionStatus, Sense}; pub(crate) struct HighsSolver(highs::RowProblem); @@ -42,15 +43,62 @@ 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(), highs::HighsModelStatus::Optimal) { + 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: accepts a solution once + /// it is provably within `mip_rel_gap` of optimal, and hard-caps wall + /// time at `mip_time_limit`. Task resource requests are themselves + /// estimates, and this solve can otherwise blow up (unboundedly, on the + /// single-threaded server) with many distinct resource shapes, so a + /// bounded solve is the right tradeoff here -- unlike `solve()`, which + /// callers that need a guaranteed-exact answer (the resource allocator) + /// must keep using instead. + fn solve_bounded(self) -> Option<(Self::Solution, f64)> { + let mut model = self.0.optimise(Sense::Maximise); + model.set_option("time_limit", mip_time_limit().as_secs_f64()); + model.set_option("mip_rel_gap", mip_rel_gap()); + let solved_model = model.solve(); + + match solved_model.status() { + // Either an exact optimum, or (since mip_rel_gap is set) a + // solution HiGHS has proven is within the accepted gap of + // optimal. + HighsModelStatus::Optimal => {} + // The hard wall-clock cap fired before the gap could be proven. + // Still dispatch it if it's a real feasible incumbent -- it + // never violates a constraint, it's just not proven close to + // optimal -- so a slow-converging solve degrades scheduling + // quality for one pass instead of blocking the server. + HighsModelStatus::ReachedTimeLimit + if solved_model.primal_solution_status() == HighsSolutionStatus::Feasible => + { + log::warn!( + "Scheduler MILP solve hit the {:?} time limit before proving the {:.0}% \ + optimality gap; dispatching the best incumbent found so far.", + mip_time_limit(), + mip_rel_gap() * 100.0 + ); + } + _ => return None, + } + + let solution = solved_model.get_solution(); + let objective_value = solved_model.objective_value(); + Some((solution, objective_value)) + } } impl LpSolution for highs::Solution { diff --git a/crates/tako/src/internal/solver/mod.rs b/crates/tako/src/internal/solver/mod.rs index 261133200..ec6db9530 100644 --- a/crates/tako/src/internal/solver/mod.rs +++ b/crates/tako/src/internal/solver/mod.rs @@ -1,5 +1,6 @@ #[cfg(all(feature = "coin_cbc", not(feature = "microlp"), not(feature = "highs")))] pub(crate) mod coin_cbc; +pub(crate) mod config; #[cfg(feature = "highs")] pub(crate) mod highs; #[cfg(all(feature = "microlp", not(feature = "highs")))] @@ -38,6 +39,16 @@ pub(crate) trait LpInnerSolver { variables: impl Iterator, ); fn solve(self) -> Option<(Self::Solution, f64)>; + + /// Like `solve`, but allowed to trade exactness for bounded solve time + /// (see `highs::HighsSolver::solve_bounded`). Backends without a tuned + /// implementation fall back to the exact `solve`. + fn solve_bounded(self) -> Option<(Self::Solution, f64)> + where + Self: Sized, + { + self.solve() + } } pub(crate) trait LpSolution { @@ -182,6 +193,28 @@ impl LpSolver { } s } + + #[inline] + pub fn solve_bounded(self) -> Option<(Solution, f64)> { + 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(); + if let Some((s, _)) = &s + && self.verbose + { + println!("==== Solution: ===="); + for (name, _weight, var) in self.variables.iter() { + println!("{} = {}", name, s.get_value(*var)); + } + } + s + } } #[cfg(not(debug_assertions))] @@ -220,6 +253,11 @@ impl LpSolver { pub fn solve(self) -> Option<(Solution, f64)> { self.solver.solve() } + + #[inline] + pub fn solve_bounded(self) -> Option<(Solution, f64)> { + self.solver.solve_bounded() + } } impl LpSolver { diff --git a/crates/tako/src/internal/tests/test_scheduler_sn.rs b/crates/tako/src/internal/tests/test_scheduler_sn.rs index 082b90db1..5c555f2f7 100644 --- a/crates/tako/src/internal/tests/test_scheduler_sn.rs +++ b/crates/tako/src/internal/tests/test_scheduler_sn.rs @@ -1461,3 +1461,99 @@ pub fn test_schedule_min_utilization3() { assert_eq!(ts.iter().filter(|t| rt.task(**t).is_assigned()).count(), 4); assert!(!rt.task(t2).is_assigned()); } + +// Tests below exercise the bounded scheduler solve (solve_bounded / config.rs +// HQ_SCHEDULER_MIP_REL_GAP / HQ_SCHEDULER_MIP_TIME_LIMIT_MS): the scheduler's +// MILP solve is otherwise unbounded and can hang the single-threaded server +// for minutes to hours given many distinct task resource shapes across many +// workers -- see the fix/scheduler-solve-timeout branch. Unit tests default +// (cfg(test) in config.rs) to exact, unhurried solving so unrelated tests +// keep asserting exact placement counts; these tests explicitly opt back +// into the production-tuned config via +// crate::internal::solver::config::with_test_solver_config, which is +// thread-local (not a process-global env var) so it cannot race with +// unrelated tests running concurrently on other threads. +use crate::internal::solver::config::with_test_solver_config; + +#[test] +fn test_schedule_many_distinct_shapes_stays_bounded() { + // Regression test for the original hang: with today's per-worker MILP + // formulation, `distinct shapes x workers` placement variables makes an + // exact solve blow up well past this size. The production defaults + // (10% gap, 5s cap) must keep this fast; if a future change reintroduces + // an unbounded solve on the scheduler's hot path, this test should start + // timing out (or take multiple seconds) instead of passing in + // milliseconds. + with_test_solver_config(0.10, Duration::from_secs(5), || { + let mut rt = TestEnv::new(); + rt.new_named_resource("mem"); + for _ in 0..20 { + rt.new_worker(&WorkerBuilder::new(64).res_sum("mem", 459_000)); + } + // ~60 distinct shapes across ~20 workers, matching the scale that + // hangs an unbounded per-worker MILP in practice. + for i in 0..60u32 { + let cpus = 1 + (i % 60); + rt.new_tasks(2, &TaskBuilder::new().cpus(cpus)); + } + + let start = std::time::Instant::now(); + rt.schedule(); + let elapsed = start.elapsed(); + assert!( + elapsed < Duration::from_secs(10), + "scheduling solve took {elapsed:?}, expected it to stay well \ + under the configured 5s time limit -- this likely means \ + solve_bounded() is no longer being used on the scheduler's \ + hot path" + ); + }); +} + +#[test] +fn test_schedule_bounded_dispatches_feasible_incumbent_on_time_limit() { + // With a time limit far too short to converge, solve_bounded() must + // still dispatch whatever feasible (constraint-respecting) incumbent + // HiGHS found, rather than discarding it -- see the ReachedTimeLimit + + // HighsSolutionStatus::Feasible branch in highs.rs. 50ms and this + // instance size (24 distinct shapes, ~260 tasks, 20 workers) is a known + // operating point from prior benchmarking of hq's real per-worker + // placement weighting: it reliably has a feasible incumbent by 50ms + // without having fully converged (which is also fine for this + // assertion -- either way proves a valid solution wasn't discarded). + with_test_solver_config(0.0, Duration::from_millis(50), || { + let mut rt = TestEnv::new(); + rt.new_named_resource("mem"); + for _ in 0..20 { + rt.new_worker(&WorkerBuilder::new(64).res_sum("mem", 459_000)); + } + for i in 0..24u32 { + let cpus = 1 + (i % 60); + rt.new_tasks(11, &TaskBuilder::new().cpus(cpus)); + } + + rt.schedule(); + let assigned = rt.task_map().tasks().filter(|t| t.is_assigned()).count(); + assert!( + assigned > 0, + "a short time limit should still dispatch a feasible incumbent, \ + not discard the solve entirely" + ); + }); +} + +#[test] +fn test_schedule_bounded_infeasible_returns_none_safely() { + // A genuinely infeasible request (no worker can ever satisfy it) must + // not panic or dispatch anything, regardless of the gap/time-limit + // tuning -- the `_ => return None` branch in solve_bounded(). + with_test_solver_config(0.10, Duration::from_secs(5), || { + let mut rt = TestEnv::new(); + rt.new_worker(&WorkerBuilder::new(4)); + let t = rt.new_task(&TaskBuilder::new().cpus(999)); + + rt.schedule(); + assert!(!rt.task(t).is_assigned()); + }); +} + diff --git a/crates/tako/src/internal/worker/resources/test_allocator.rs b/crates/tako/src/internal/worker/resources/test_allocator.rs index b1dd4ec0c..c641029e3 100644 --- a/crates/tako/src/internal/worker/resources/test_allocator.rs +++ b/crates/tako/src/internal/worker/resources/test_allocator.rs @@ -899,6 +899,67 @@ fn test_coupling3() { assert_eq!(v, vec![1]); } +#[test] +fn test_allocator_stays_exact_regardless_of_scheduler_gap_tuning() { + // Guards against the exact regression found while implementing the + // scheduler's bounded solve (solve_bounded in internal::solver::highs): + // the worker's own NUMA/socket resource allocator below shares the same + // underlying LpSolver but must keep calling the exact `solve()`, never + // the scheduler's gap/time-limit-tuned `solve_bounded()` -- it needs a + // guaranteed-exact feasible allocation, not a good-enough one. Setting a + // coarse scheduler gap here must not affect it; this is verbatim the + // scenario from test_complex_coupling1 below, which is exactly what + // caught the original regression. + // + // with_test_solver_config is thread-local (not a process-global env + // var), so it cannot race with unrelated tests running concurrently on + // other threads. + crate::internal::solver::config::with_test_solver_config( + 0.50, + std::time::Duration::from_millis(1), + || { + let mut coupling = ResourceDescriptorCoupling::default(); + for i in 0..6 { + coupling.add(0, i, 1, i / 2, 256); + coupling.add(1, i / 2, 2, i, 128); + } + let descriptor = ResourceDescriptor::new( + vec![ + ResourceDescriptorItem { + name: "cpus".to_string(), + kind: ResourceDescriptorKind::regular_sockets(6, 2), + }, + ResourceDescriptorItem { + name: "gpus".to_string(), + kind: ResourceDescriptorKind::regular_sockets(3, 1), + }, + ResourceDescriptorItem { + name: "foo".to_string(), + kind: ResourceDescriptorKind::regular_sockets(6, 3), + }, + ], + coupling, + ); + let mut allocator = test_allocator(&descriptor); + + allocator.force_claim_from_groups(0.into(), &[0], 1.into()); + allocator.force_claim_from_groups(2.into(), &[5], 2.into()); + + let rq = ResBuilder::default() + .add_force_compact(0, ResourceAmount::new_units(4)) + .add_force_compact(1, ResourceAmount::new_units(1)) + .add_force_compact(2, ResourceAmount::new_units(5)) + .finish(); + assert!( + allocator.try_allocate(&rq).is_some(), + "the exact NUMA/socket allocator must keep finding a \ + feasible allocation regardless of the scheduler's \ + mip_rel_gap tuning" + ); + }, + ); +} + #[test] fn test_complex_coupling1() { let mut coupling = ResourceDescriptorCoupling::default();