From 1a86206dfc0bc9b749b904e0d7f84a694622907e Mon Sep 17 00:00:00 2001 From: MauriceDHanisch Date: Thu, 2 Jul 2026 19:31:53 -0700 Subject: [PATCH 1/3] 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 539229a0a..ffcc7a60e 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 9e8b9a611..b2643bc37 100644 --- a/crates/tako/src/internal/worker/resources/test_allocator.rs +++ b/crates/tako/src/internal/worker/resources/test_allocator.rs @@ -897,6 +897,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(); From fb44115e2afbd10a743df9d5c8a8d4d9ca4cc8e8 Mon Sep 17 00:00:00 2001 From: MauriceDHanisch Date: Tue, 21 Jul 2026 13:44:25 -0700 Subject: [PATCH 2/3] fix(scheduler): drop mip_rel_gap, move solve bound into SchedulerConfig Per review on PR #1119: bound the scheduler's MILP solve by wall-clock time only, not by relative gap, since some downstream logic assumes an optimal placement. The time limit moves from a thread-local/env-var read into SchedulerConfig::mip_time_limit, threaded from `hq server start --scheduler-time-limit` through ServerConfig, server_start, CoreRef, and Core down to the solve call site. solve_bounded now returns (Solution, is_optimal) instead of (Solution, objective_value). SchedulingSolution gains an is_optimal field so downstream consumers (e.g. proactive filling) can react to a non-optimal placement. --- .../hyperqueue/src/client/commands/server.rs | 9 ++ crates/hyperqueue/src/server/bootstrap.rs | 6 + crates/hyperqueue/src/tests/server.rs | 1 + crates/pyhq/src/cluster/server.rs | 1 + crates/tako/src/control.rs | 4 +- crates/tako/src/internal/scheduler/main.rs | 3 + crates/tako/src/internal/scheduler/mod.rs | 3 +- crates/tako/src/internal/scheduler/solver.rs | 19 ++- crates/tako/src/internal/scheduler/state.rs | 19 +++ crates/tako/src/internal/server/core.rs | 8 +- crates/tako/src/internal/solver/config.rs | 82 ----------- crates/tako/src/internal/solver/highs.rs | 45 +++--- crates/tako/src/internal/solver/mod.rs | 22 +-- .../tests/integration/utils/server.rs | 1 + .../src/internal/tests/test_scheduler_sn.rs | 135 +++++++----------- crates/tako/src/internal/tests/utils/env.rs | 11 +- .../worker/resources/test_allocator.rs | 61 -------- crates/tako/src/lib.rs | 1 + 18 files changed, 162 insertions(+), 269 deletions(-) delete mode 100644 crates/tako/src/internal/solver/config.rs diff --git a/crates/hyperqueue/src/client/commands/server.rs b/crates/hyperqueue/src/client/commands/server.rs index 38ef24601..7bec927ea 100644 --- a/crates/hyperqueue/src/client/commands/server.rs +++ b/crates/hyperqueue/src/client/commands/server.rs @@ -130,6 +130,14 @@ pub struct ServerStartOpts { /// USE AT YOUR OWN RISK. #[arg(long)] disable_worker_authentication_and_encryption: bool, + + #[arg( + long, + value_parser = parse_hms_or_human_time, + default_value = "5s", + help = duration_doc!("Maximum time the scheduler's placement solve may run per round.") + )] + scheduler_time_limit: Duration, } #[derive(Parser)] @@ -222,6 +230,7 @@ async fn start_server(gsettings: &GlobalSettings, opts: ServerStartOpts) -> anyh } }), server_uid: access_file.as_ref().map(|a| a.server_uid().to_string()), + scheduler_mip_time_limit: opts.scheduler_time_limit, }; init_hq_server(gsettings, server_cfg).await diff --git a/crates/hyperqueue/src/server/bootstrap.rs b/crates/hyperqueue/src/server/bootstrap.rs index 1e75892c7..ad9c589f6 100644 --- a/crates/hyperqueue/src/server/bootstrap.rs +++ b/crates/hyperqueue/src/server/bootstrap.rs @@ -48,6 +48,8 @@ pub struct ServerConfig { pub worker_secret_key: Option>, pub client_secret_key: Option>, pub server_uid: Option, + /// See `tako::server::SchedulerConfig::mip_time_limit`. + pub scheduler_mip_time_limit: Duration, } /// This function initializes the HQ server. @@ -204,6 +206,10 @@ pub async fn initialize_server( None, state_ref.get().server_info().server_uid.clone(), worker_id_initial_value, + tako::server::SchedulerConfig { + mip_time_limit: server_cfg.scheduler_mip_time_limit, + ..Default::default() + }, )?; let (autoalloc_service, autoalloc_process) = create_autoalloc_service(server_ref.clone(), queue_id_initial_value, events.clone()); diff --git a/crates/hyperqueue/src/tests/server.rs b/crates/hyperqueue/src/tests/server.rs index 04750e1d3..b4ca77f01 100644 --- a/crates/hyperqueue/src/tests/server.rs +++ b/crates/hyperqueue/src/tests/server.rs @@ -57,6 +57,7 @@ where worker_secret_key: None, client_secret_key: None, server_uid: None, + scheduler_mip_time_limit: Duration::from_secs(5), }; let (fut, notify, _state, _senders) = initialize_server(&gsettings, server_cfg, 1.into(), 1, None) diff --git a/crates/pyhq/src/cluster/server.rs b/crates/pyhq/src/cluster/server.rs index f8cc8ad3c..1b0a80fe6 100644 --- a/crates/pyhq/src/cluster/server.rs +++ b/crates/pyhq/src/cluster/server.rs @@ -40,6 +40,7 @@ impl RunningServer { worker_secret_key: None, server_uid: None, journal_flush_period: Duration::from_secs(30), + scheduler_mip_time_limit: Duration::from_secs(5), }; let main_future = async move { diff --git a/crates/tako/src/control.rs b/crates/tako/src/control.rs index 82f774a82..e7a3359a3 100644 --- a/crates/tako/src/control.rs +++ b/crates/tako/src/control.rs @@ -16,7 +16,7 @@ use crate::internal::common::error::DsError; use crate::internal::common::resources::ResourceRqId; use crate::internal::messages::worker::ToWorkerMessage; use crate::internal::scheduler::query::compute_new_worker_query; -use crate::internal::scheduler::{run_scheduling, scheduler_loop}; +use crate::internal::scheduler::{SchedulerConfig, run_scheduling, scheduler_loop}; use crate::internal::server::client::handle_new_tasks; use crate::internal::server::comm::{Comm, CommSenderRef}; use crate::internal::server::core::{CoreRef, CustomConnectionHandler}; @@ -237,6 +237,7 @@ pub fn server_start( custom_conn_handler: Option, server_uid: String, worker_id_initial_value: WorkerId, + scheduler_config: SchedulerConfig, ) -> crate::Result<(ServerRef, impl Future>)> { let listener_port = listener.local_addr()?.port(); @@ -251,6 +252,7 @@ pub fn server_start( custom_conn_handler, server_uid, worker_id_initial_value, + scheduler_config, ); let connections = crate::internal::server::rpc::connection_initiator( listener, diff --git a/crates/tako/src/internal/scheduler/main.rs b/crates/tako/src/internal/scheduler/main.rs index f2dc7a286..a1192d2c8 100644 --- a/crates/tako/src/internal/scheduler/main.rs +++ b/crates/tako/src/internal/scheduler/main.rs @@ -40,6 +40,9 @@ pub(crate) async fn scheduler_loop( pub(crate) fn run_scheduling_inner(core: &mut Core, comm: &mut impl Comm, now: Instant) { let batches = create_task_batches(core, now, None); let solution = run_scheduling_solver(core, now, &batches, None); + if !solution.is_optimal { + log::debug!("Scheduler dispatched a non-optimal placement this round"); + } let mapping = create_task_mapping(core, solution); //mapping.dump(); mapping.send_messages(core, comm); diff --git a/crates/tako/src/internal/scheduler/mod.rs b/crates/tako/src/internal/scheduler/mod.rs index b33aa6d1c..1f0a16172 100644 --- a/crates/tako/src/internal/scheduler/mod.rs +++ b/crates/tako/src/internal/scheduler/mod.rs @@ -10,6 +10,7 @@ mod taskqueue; pub(crate) use batches::{TaskBatch, create_task_batches}; pub(crate) use main::{run_scheduling, scheduler_loop}; pub(crate) use solver::run_scheduling_solver; +pub use state::SchedulerConfig; pub(crate) use state::SchedulerState; pub(crate) use taskqueue::TaskQueues; @@ -20,7 +21,7 @@ pub(crate) use main::run_scheduling_inner; pub(crate) use mapping::{WorkerTaskMapping, create_task_mapping}; #[cfg(test)] -pub(crate) use state::SchedulerConfig; +pub(crate) use solver::SchedulingSolution; #[cfg(test)] pub(crate) use batches::PriorityCut; diff --git a/crates/tako/src/internal/scheduler/solver.rs b/crates/tako/src/internal/scheduler/solver.rs index ffcc7a60e..c63295140 100644 --- a/crates/tako/src/internal/scheduler/solver.rs +++ b/crates/tako/src/internal/scheduler/solver.rs @@ -7,10 +7,23 @@ use crate::resources::{CPU_RESOURCE_ID, ResourceRqId}; use crate::{Map, ResourceVariantId, Set, WorkerId}; use thin_vec::ThinVec; -#[derive(Default, Debug)] +#[derive(Debug)] pub(crate) struct SchedulingSolution { pub(crate) sn_counts: Map<(ResourceRqId, ResourceVariantId), Map>, pub(crate) mn_workers: Map<(ResourceRqId, ResourceVariantId), Vec>>, + /// `false` if the MILP solve hit its time limit before proving + /// optimality (see `SchedulerConfig::mip_time_limit`). + pub(crate) is_optimal: bool, +} + +impl Default for SchedulingSolution { + fn default() -> Self { + SchedulingSolution { + sn_counts: Map::new(), + mn_workers: Map::new(), + is_optimal: true, + } + } } pub(crate) fn run_scheduling_solver( @@ -410,9 +423,11 @@ pub(crate) fn run_scheduling_solver( } let mut result = SchedulingSolution::default(); - let Some((solution, _)) = solver.solve_bounded() else { + let Some((solution, is_optimal)) = solver.solve_bounded(scheduler_cache.config.mip_time_limit) + else { return result; }; + result.is_optimal = is_optimal; for batch in task_batches { let resource_rq_id = batch.resource_rq_id; diff --git a/crates/tako/src/internal/scheduler/state.rs b/crates/tako/src/internal/scheduler/state.rs index 68974c1aa..1e6816a54 100644 --- a/crates/tako/src/internal/scheduler/state.rs +++ b/crates/tako/src/internal/scheduler/state.rs @@ -1,5 +1,6 @@ use crate::internal::scheduler::gap::GapCache; use crate::{Map, ResourceVariantId, TaskId, WorkerId}; +use std::time::Duration; pub struct SchedulerConfig { /// The number of tasks that are never prefilled (wrt a given resource request) @@ -9,6 +10,10 @@ pub struct SchedulerConfig { /// TODO: Maybe we can choose it dynamically wrt. resources of workers /// but small tens looks reasonable pub proactive_filling_max: u32, + /// Hard wall-clock cap on a single scheduler MILP solve. An emergency + /// backstop, not a tuning knob: on expiry, the best incumbent found so + /// far is dispatched and marked non-optimal. + pub mip_time_limit: Duration, } impl Default for SchedulerConfig { @@ -16,10 +21,24 @@ impl Default for SchedulerConfig { SchedulerConfig { proactive_filling_reserve: 16, proactive_filling_max: 40, + mip_time_limit: default_mip_time_limit(), } } } +// Unit tests assert exact placement counts on small instances, so default to +// a generous limit. Tests exercising the bounded solve set mip_time_limit +// explicitly via TestEnv::set_scheduler_config. +#[cfg(test)] +fn default_mip_time_limit() -> Duration { + Duration::from_secs(60) +} + +#[cfg(not(test))] +fn default_mip_time_limit() -> Duration { + Duration::from_secs(5) +} + #[derive(Default)] pub(crate) struct SchedulerState { pub gap_cache: GapCache, diff --git a/crates/tako/src/internal/server/core.rs b/crates/tako/src/internal/server/core.rs index c122d15e1..eaabc144b 100644 --- a/crates/tako/src/internal/server/core.rs +++ b/crates/tako/src/internal/server/core.rs @@ -6,7 +6,7 @@ use crate::internal::common::resources::map::{ }; use crate::internal::common::resources::{ResourceId, ResourceRequestVariants, ResourceRqId}; use crate::internal::common::{Set, WrappedRcRefCell}; -use crate::internal::scheduler::{SchedulerState, TaskQueues}; +use crate::internal::scheduler::{SchedulerConfig, SchedulerState, TaskQueues}; use crate::internal::server::rpc::ConnectionDescriptor; use crate::internal::server::task::{Task, TaskRuntimeState}; use crate::internal::server::taskmap::TaskMap; @@ -64,6 +64,7 @@ pub(crate) struct Core { pub(crate) type CoreRef = WrappedRcRefCell; impl CoreRef { + #[allow(clippy::too_many_arguments)] pub fn new( worker_listen_port: u16, secret_key: Option>, @@ -71,6 +72,7 @@ impl CoreRef { custom_conn_handler: Option, server_uid: String, worker_id_initial_value: WorkerId, + scheduler_config: SchedulerConfig, ) -> Self { CoreRef::wrap(Core { worker_listen_port, @@ -79,6 +81,10 @@ impl CoreRef { custom_conn_handler, server_uid, worker_id_counter: worker_id_initial_value.as_num(), + scheduler_state: SchedulerState { + config: scheduler_config, + ..Default::default() + }, ..Default::default() }) } diff --git a/crates/tako/src/internal/solver/config.rs b/crates/tako/src/internal/solver/config.rs deleted file mode 100644 index 8461a5d2d..000000000 --- a/crates/tako/src/internal/solver/config.rs +++ /dev/null @@ -1,82 +0,0 @@ -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 d12b473b5..46cd0ff97 100644 --- a/crates/tako/src/internal/solver/highs.rs +++ b/crates/tako/src/internal/solver/highs.rs @@ -1,6 +1,6 @@ -use crate::internal::solver::config::{mip_rel_gap, mip_time_limit}; use crate::internal::solver::{ConstraintType, LpInnerSolver, LpSolution}; use highs::{HighsModelStatus, HighsSolutionStatus, Sense}; +use std::time::Duration; pub(crate) struct HighsSolver(highs::RowProblem); @@ -58,46 +58,33 @@ impl LpInnerSolver for HighsSolver { 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)> { + /// 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)> { 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()); + model.set_option("time_limit", time_limit.as_secs_f64()); 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. + 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::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 + "Scheduler MILP solve hit the {time_limit:?} time limit before proving \ + optimality; dispatching the best incumbent found so far." ); + false } _ => return None, - } + }; let solution = solved_model.get_solution(); - let objective_value = solved_model.objective_value(); - Some((solution, objective_value)) + Some((solution, is_optimal)) } } diff --git a/crates/tako/src/internal/solver/mod.rs b/crates/tako/src/internal/solver/mod.rs index ec6db9530..1025ec747 100644 --- a/crates/tako/src/internal/solver/mod.rs +++ b/crates/tako/src/internal/solver/mod.rs @@ -1,11 +1,12 @@ #[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")))] pub(crate) mod microlp; +use std::time::Duration; + #[cfg(feature = "highs")] pub(crate) type LpInnerSolverImpl = highs::HighsSolver; @@ -40,14 +41,15 @@ pub(crate) trait LpInnerSolver { ); 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)> + /// 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, { - self.solve() + let _ = time_limit; + self.solve().map(|(solution, _)| (solution, true)) } } @@ -195,7 +197,7 @@ impl LpSolver { } #[inline] - pub fn solve_bounded(self) -> Option<(Solution, f64)> { + pub fn solve_bounded(self, time_limit: Duration) -> Option<(Solution, bool)> { if self.verbose { println!("Weights:"); for (name, weight, _var) in self.variables.iter() { @@ -204,7 +206,7 @@ impl LpSolver { } } } - let s = self.solver.solve_bounded(); + let s = self.solver.solve_bounded(time_limit); if let Some((s, _)) = &s && self.verbose { @@ -255,8 +257,8 @@ impl LpSolver { } #[inline] - pub fn solve_bounded(self) -> Option<(Solution, f64)> { - self.solver.solve_bounded() + pub fn solve_bounded(self, time_limit: Duration) -> Option<(Solution, bool)> { + self.solver.solve_bounded(time_limit) } } diff --git a/crates/tako/src/internal/tests/integration/utils/server.rs b/crates/tako/src/internal/tests/integration/utils/server.rs index 58a669430..3ec43275c 100644 --- a/crates/tako/src/internal/tests/integration/utils/server.rs +++ b/crates/tako/src/internal/tests/integration/utils/server.rs @@ -278,6 +278,7 @@ async fn create_handle( None, "testuid".to_string(), 1.into(), + Default::default(), ) .expect("Could not start server"); diff --git a/crates/tako/src/internal/tests/test_scheduler_sn.rs b/crates/tako/src/internal/tests/test_scheduler_sn.rs index 5c555f2f7..2e84efe19 100644 --- a/crates/tako/src/internal/tests/test_scheduler_sn.rs +++ b/crates/tako/src/internal/tests/test_scheduler_sn.rs @@ -1462,98 +1462,73 @@ pub fn test_schedule_min_utilization3() { 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; +// Bounded-solve tests: unit tests default to a generous mip_time_limit +// (SchedulerConfig::default), so these opt into a short one explicitly. #[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" - ); + let mut rt = TestEnv::new(); + rt.set_scheduler_config(SchedulerConfig { + mip_time_limit: Duration::from_secs(5), + ..Default::default() }); + rt.new_named_resource("mem"); + for _ in 0..20 { + rt.new_worker(&WorkerBuilder::new(64).res_sum("mem", 459_000)); + } + for i in 0..60u32 { + rt.new_tasks(2, &TaskBuilder::new().cpus(1 + (i % 60))); + } + + let start = std::time::Instant::now(); + rt.schedule(); + assert!(start.elapsed() < Duration::from_secs(10)); } #[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" - ); +fn test_schedule_bounded_infeasible_returns_none_safely() { + let mut rt = TestEnv::new(); + rt.set_scheduler_config(SchedulerConfig { + mip_time_limit: Duration::from_secs(5), + ..Default::default() }); + rt.new_worker(&WorkerBuilder::new(4)); + let t = rt.new_task(&TaskBuilder::new().cpus(999)); + + rt.schedule(); + assert!(!rt.task(t).is_assigned()); } #[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)); +fn test_schedule_bounded_is_optimal_true_when_solve_converges() { + let mut rt = TestEnv::new(); + rt.set_scheduler_config(SchedulerConfig { + mip_time_limit: Duration::from_secs(5), + ..Default::default() + }); + rt.new_worker(&WorkerBuilder::new(4)); + rt.new_tasks(2, &TaskBuilder::new().cpus(1)); - rt.schedule(); - assert!(!rt.task(t).is_assigned()); + assert!(rt.schedule_solution().is_optimal); +} + +#[test] +fn test_schedule_bounded_is_optimal_false_on_time_limit() { + let mut rt = TestEnv::new(); + rt.set_scheduler_config(SchedulerConfig { + mip_time_limit: Duration::from_millis(50), + ..Default::default() }); + 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 { + rt.new_tasks(11, &TaskBuilder::new().cpus(1 + (i % 60))); + } + + let solution = rt.schedule_solution(); + assert!(!solution.sn_counts.is_empty()); + assert!(!solution.is_optimal); } diff --git a/crates/tako/src/internal/tests/utils/env.rs b/crates/tako/src/internal/tests/utils/env.rs index be84d7e0b..934db2b5c 100644 --- a/crates/tako/src/internal/tests/utils/env.rs +++ b/crates/tako/src/internal/tests/utils/env.rs @@ -8,8 +8,8 @@ use crate::internal::messages::worker::{ TaskRunningMsg, ToWorkerMessage, WorkerOverview, WorkerTaskUpdate, }; use crate::internal::scheduler::{ - SchedulerConfig, WorkerTaskMapping, create_task_batches, create_task_mapping, - run_scheduling_inner, run_scheduling_solver, + SchedulerConfig, SchedulingSolution, WorkerTaskMapping, create_task_batches, + create_task_mapping, run_scheduling_inner, run_scheduling_solver, }; use crate::internal::server::comm::Comm; use crate::internal::server::core::{Core, CoreSplitMut}; @@ -259,6 +259,13 @@ impl TestEnv { let solution = run_scheduling_solver(&mut self.core, self.now, &batches, None); create_task_mapping(&mut self.core, solution) } + + /// Runs the scheduler's placement solve without applying its result, so + /// 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) + } } #[derive(Default, Debug)] diff --git a/crates/tako/src/internal/worker/resources/test_allocator.rs b/crates/tako/src/internal/worker/resources/test_allocator.rs index b2643bc37..9e8b9a611 100644 --- a/crates/tako/src/internal/worker/resources/test_allocator.rs +++ b/crates/tako/src/internal/worker/resources/test_allocator.rs @@ -897,67 +897,6 @@ 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(); diff --git a/crates/tako/src/lib.rs b/crates/tako/src/lib.rs index 4d62245f1..a32758051 100644 --- a/crates/tako/src/lib.rs +++ b/crates/tako/src/lib.rs @@ -52,6 +52,7 @@ pub mod resources { pub mod server { pub use crate::control::server_start; + pub use crate::internal::scheduler::SchedulerConfig; pub use crate::internal::server::explain::{TaskExplainItem, TaskExplanation}; pub use crate::internal::server::rpc::ConnectionDescriptor; } From 1371faa359f456654d7d41b768edd5ef7b0979a8 Mon Sep 17 00:00:00 2001 From: Ada Bohm Date: Wed, 29 Jul 2026 17:05:35 +0700 Subject: [PATCH 3/3] Followup changes when scheduler timeouts --- CHANGELOG.md | 5 ++ crates/tako/src/control.rs | 23 +++++++- crates/tako/src/internal/scheduler/main.rs | 52 ++++++++++++++----- crates/tako/src/internal/scheduler/mod.rs | 2 +- crates/tako/src/internal/scheduler/solver.rs | 7 +++ .../src/internal/tests/test_scheduler_sn.rs | 1 - 6 files changed, 73 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 26fd61b90..c17ad77d7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,10 +2,15 @@ ## Unreleased +### New features + +* The server scheduler now contains a safety limit for computation, configurable via `--scheduler-time-limit` (default: 5s) + ### Fixes * Fixed server crash in a specific situation when an unschedulable high-priority task occurs + ## v0.26.2 ### Fixes diff --git a/crates/tako/src/control.rs b/crates/tako/src/control.rs index e7a3359a3..1e0ebf0bb 100644 --- a/crates/tako/src/control.rs +++ b/crates/tako/src/control.rs @@ -16,7 +16,9 @@ use crate::internal::common::error::DsError; use crate::internal::common::resources::ResourceRqId; use crate::internal::messages::worker::ToWorkerMessage; use crate::internal::scheduler::query::compute_new_worker_query; -use crate::internal::scheduler::{SchedulerConfig, run_scheduling, scheduler_loop}; +use crate::internal::scheduler::{ + SchedulerConfig, SchedulerResult, run_scheduling, scheduler_loop, +}; use crate::internal::server::client::handle_new_tasks; use crate::internal::server::comm::{Comm, CommSenderRef}; use crate::internal::server::core::{CoreRef, CustomConnectionHandler}; @@ -130,7 +132,24 @@ impl ServerRef { let mut core = self.core_ref.get_mut(); let mut comm = self.comm_ref.get_mut(); if comm.get_scheduling_flag() { - run_scheduling(&mut core, &mut comm, Instant::now()) + // Try to finish scheduling, at most 3 times + if (0..3) + .find_map( + |_| match run_scheduling(&mut core, &mut comm, Instant::now()) { + SchedulerResult::Done => Some(true), + SchedulerResult::NoProgress => Some(false), + SchedulerResult::NeedMoreCompute => None, + }, + ) + .unwrap_or(false) + { + comm.reset_scheduling_flag(); + } else { + return Ok(NewWorkerAllocationResponse { + single_node_workers_per_query: Vec::new(), + multi_node_allocations: Vec::new(), + }); + } } Ok(compute_new_worker_query(&mut core, queries)) } diff --git a/crates/tako/src/internal/scheduler/main.rs b/crates/tako/src/internal/scheduler/main.rs index a1192d2c8..1b6907ca8 100644 --- a/crates/tako/src/internal/scheduler/main.rs +++ b/crates/tako/src/internal/scheduler/main.rs @@ -16,43 +16,69 @@ pub(crate) async fn scheduler_loop( let mut last_schedule = Instant::now().checked_sub(minimum_delay * 2).unwrap(); loop { scheduler_wakeup.notified().await; + let mut now = Instant::now(); if !comm_ref.get().get_scheduling_flag() { - last_schedule = Instant::now(); + last_schedule = now; continue; } - let mut now = Instant::now(); let since_last_schedule = now - last_schedule; if minimum_delay > since_last_schedule { sleep(minimum_delay - since_last_schedule).await; now = Instant::now(); } - let mut comm = comm_ref.get_mut(); - if !comm.get_scheduling_flag() { + if !comm_ref.get_mut().get_scheduling_flag() { last_schedule = now; continue; } - let mut core = core_ref.get_mut(); - run_scheduling(&mut core, &mut comm, now); + while matches!( + run_scheduling(&mut core_ref.get_mut(), &mut comm_ref.get_mut(), now), + SchedulerResult::NeedMoreCompute + ) { + sleep(minimum_delay).await; + } + comm_ref.get_mut().reset_scheduling_flag(); last_schedule = Instant::now(); } } -pub(crate) fn run_scheduling_inner(core: &mut Core, comm: &mut impl Comm, now: Instant) { +pub(crate) enum SchedulerResult { + Done, + NeedMoreCompute, + NoProgress, +} + +pub(crate) fn run_scheduling_inner( + core: &mut Core, + comm: &mut impl Comm, + now: Instant, +) -> SchedulerResult { let batches = create_task_batches(core, now, None); let solution = run_scheduling_solver(core, now, &batches, None); - if !solution.is_optimal { - log::debug!("Scheduler dispatched a non-optimal placement this round"); - } + let need_more_compute = if !solution.is_optimal { + if solution.is_empty() { + log::error!("Scheduler made no progress within given time limit"); + SchedulerResult::NoProgress + } else { + log::debug!("Scheduler dispatched a non-optimal placement this round"); + SchedulerResult::NeedMoreCompute + } + } else { + SchedulerResult::Done + }; let mapping = create_task_mapping(core, solution); //mapping.dump(); mapping.send_messages(core, comm); + need_more_compute } -pub(crate) fn run_scheduling(core: &mut Core, comm: &mut CommSender, now: Instant) { +pub(crate) fn run_scheduling( + core: &mut Core, + comm: &mut CommSender, + now: Instant, +) -> SchedulerResult { trace_time!( "scheduler", "run_scheduling_inner", run_scheduling_inner(core, comm, now) - ); - comm.reset_scheduling_flag(); + ) } diff --git a/crates/tako/src/internal/scheduler/mod.rs b/crates/tako/src/internal/scheduler/mod.rs index 1f0a16172..2ac7be0dd 100644 --- a/crates/tako/src/internal/scheduler/mod.rs +++ b/crates/tako/src/internal/scheduler/mod.rs @@ -8,7 +8,7 @@ mod state; mod taskqueue; pub(crate) use batches::{TaskBatch, create_task_batches}; -pub(crate) use main::{run_scheduling, scheduler_loop}; +pub(crate) use main::{SchedulerResult, run_scheduling, scheduler_loop}; pub(crate) use solver::run_scheduling_solver; pub use state::SchedulerConfig; pub(crate) use state::SchedulerState; diff --git a/crates/tako/src/internal/scheduler/solver.rs b/crates/tako/src/internal/scheduler/solver.rs index c63295140..719df80e3 100644 --- a/crates/tako/src/internal/scheduler/solver.rs +++ b/crates/tako/src/internal/scheduler/solver.rs @@ -26,6 +26,13 @@ impl Default for SchedulingSolution { } } +impl SchedulingSolution { + pub fn is_empty(&self) -> bool { + self.sn_counts.values().all(|v| v.is_empty()) + && self.mn_workers.values().all(|v| v.is_empty()) + } +} + pub(crate) fn run_scheduling_solver( core: &Core, now: std::time::Instant, diff --git a/crates/tako/src/internal/tests/test_scheduler_sn.rs b/crates/tako/src/internal/tests/test_scheduler_sn.rs index 2e84efe19..da3627cdd 100644 --- a/crates/tako/src/internal/tests/test_scheduler_sn.rs +++ b/crates/tako/src/internal/tests/test_scheduler_sn.rs @@ -1531,4 +1531,3 @@ fn test_schedule_bounded_is_optimal_false_on_time_limit() { assert!(!solution.sn_counts.is_empty()); assert!(!solution.is_optimal); } -