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/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..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::{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)) } @@ -237,6 +256,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 +271,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..1b6907ca8 100644 --- a/crates/tako/src/internal/scheduler/main.rs +++ b/crates/tako/src/internal/scheduler/main.rs @@ -16,40 +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); + 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 b33aa6d1c..2ac7be0dd 100644 --- a/crates/tako/src/internal/scheduler/mod.rs +++ b/crates/tako/src/internal/scheduler/mod.rs @@ -8,8 +8,9 @@ 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; 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 539229a0a..719df80e3 100644 --- a/crates/tako/src/internal/scheduler/solver.rs +++ b/crates/tako/src/internal/scheduler/solver.rs @@ -7,10 +7,30 @@ 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, + } + } +} + +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( @@ -410,9 +430,11 @@ pub(crate) fn run_scheduling_solver( } let mut result = SchedulingSolution::default(); - let Some((solution, _)) = solver.solve() 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/highs.rs b/crates/tako/src/internal/solver/highs.rs index 64e9ad90f..46cd0ff97 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::{ConstraintType, LpInnerSolver, LpSolution}; -use highs::Sense; +use highs::{HighsModelStatus, HighsSolutionStatus, Sense}; +use std::time::Duration; pub(crate) struct HighsSolver(highs::RowProblem); @@ -42,15 +43,49 @@ 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: 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", time_limit.as_secs_f64()); + let solved_model = model.solve(); + + let is_optimal = match solved_model.status() { + HighsModelStatus::Optimal => true, + // Time limit fired before optimality was proven. Dispatch the + // incumbent anyway if it's feasible. + HighsModelStatus::ReachedTimeLimit + if solved_model.primal_solution_status() == HighsSolutionStatus::Feasible => + { + log::warn!( + "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(); + Some((solution, is_optimal)) + } } 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..1025ec747 100644 --- a/crates/tako/src/internal/solver/mod.rs +++ b/crates/tako/src/internal/solver/mod.rs @@ -5,6 +5,8 @@ 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; @@ -38,6 +40,17 @@ pub(crate) trait LpInnerSolver { variables: impl Iterator, ); fn solve(self) -> Option<(Self::Solution, f64)>; + + /// Like `solve`, but allowed to trade exactness for a hard wall-clock + /// cap. Returns whether the solution is proven optimal. Backends without + /// a tuned implementation fall back to the exact `solve`. + fn solve_bounded(self, time_limit: Duration) -> Option<(Self::Solution, bool)> + where + Self: Sized, + { + let _ = time_limit; + self.solve().map(|(solution, _)| (solution, true)) + } } pub(crate) trait LpSolution { @@ -182,6 +195,28 @@ impl LpSolver { } s } + + #[inline] + pub fn solve_bounded(self, time_limit: Duration) -> Option<(Solution, bool)> { + if self.verbose { + println!("Weights:"); + for (name, weight, _var) in self.variables.iter() { + if *weight != 0.0 { + println!("{} -> {}", name, weight); + } + } + } + let s = self.solver.solve_bounded(time_limit); + if let Some((s, _)) = &s + && 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 +255,11 @@ impl LpSolver { pub fn solve(self) -> Option<(Solution, f64)> { self.solver.solve() } + + #[inline] + pub fn solve_bounded(self, time_limit: Duration) -> Option<(Solution, bool)> { + self.solver.solve_bounded(time_limit) + } } impl LpSolver { 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 082b90db1..da3627cdd 100644 --- a/crates/tako/src/internal/tests/test_scheduler_sn.rs +++ b/crates/tako/src/internal/tests/test_scheduler_sn.rs @@ -1461,3 +1461,73 @@ 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()); } + +// 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() { + 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_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_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)); + + 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/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; }