From 4215a0d5b57f0ebc4008f3568d606b16f44e11aa Mon Sep 17 00:00:00 2001 From: MauriceDHanisch Date: Thu, 2 Jul 2026 19:31:53 -0700 Subject: [PATCH] fix(scheduler): bound MILP solve time and gap to prevent server hang The scheduler's per-worker placement solve (run_scheduling_solver) was unbounded: with many distinct task resource shapes across many workers, HiGHS's search can take minutes to hours before returning, freezing the single-threaded server for the entire duration (no heartbeats, no RPCs, no other scheduling). See AI-QChem/QE-NO#6. Add HighsSolver::solve_bounded(), used only by the scheduler's placement solve: - mip_rel_gap (default 10%, HQ_SCHEDULER_MIP_REL_GAP): accept a solution once HiGHS has proven it within this fraction of optimal. Task resource requests are themselves estimates, so exact optimality is false precision; benchmarking showed this alone takes a pathological many-shape instance from a 20s+ timeout to under 2s. - time_limit (default 5s, HQ_SCHEDULER_MIP_TIME_LIMIT_MS): a hard wall-clock backstop. A ReachedTimeLimit result is still dispatched if HiGHS reports a real feasible incumbent (never violates a constraint, just not proven within the gap); otherwise the pass is skipped and the ~20ms scheduler debounce retries. The existing exact solve() is kept unchanged and is still what worker/resources/groups.rs's NUMA/socket resource allocator uses -- it shares the same underlying LpSolver but needs a guaranteed-exact feasible allocation, not a good-enough one. Discovered via a regression: applying mip_rel_gap globally broke test_complex_coupling1, an allocator test with no relation to the scheduler, because both call sites shared one solve() prior to this change. Unit tests default (cfg(test) in config.rs) to exact, unhurried solving, matching every existing test's exact-count assertions; new tests opt into the tuned config via a thread-local override (with_test_solver_config) rather than process-global env vars, so they can't race with unrelated tests running concurrently on other threads. New tests: - test_schedule_many_distinct_shapes_stays_bounded: regression test for the original hang, at the shape/worker scale that reproduces it. - test_schedule_bounded_dispatches_feasible_incumbent_on_time_limit: a too-short time limit still dispatches a partial feasible solution. - test_schedule_bounded_infeasible_returns_none_safely: a genuinely infeasible request stays unassigned, not a panic. - test_allocator_stays_exact_regardless_of_scheduler_gap_tuning: guards the allocator/scheduler solve split found above. --- crates/tako/src/internal/scheduler/solver.rs | 2 +- crates/tako/src/internal/solver/config.rs | 82 ++++++++++++++++ crates/tako/src/internal/solver/highs.rs | 52 +++++++++- crates/tako/src/internal/solver/mod.rs | 38 ++++++++ .../src/internal/tests/test_scheduler_sn.rs | 96 +++++++++++++++++++ .../worker/resources/test_allocator.rs | 61 ++++++++++++ 6 files changed, 328 insertions(+), 3 deletions(-) create mode 100644 crates/tako/src/internal/solver/config.rs 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();