fix(scheduler): bound MILP solve time and gap to prevent server hang

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

crates/hyperqueue/src/server/bootstrap.rs

@@ -48,6 +48,8 @@ pub struct ServerConfig {
     pub worker_secret_key: Option<Arc<SecretKey>>,
     pub client_secret_key: Option<Arc<SecretKey>>,
     pub server_uid: Option<String>,
+    /// 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());

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)

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 {

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<CustomConnectionHandler>,
     server_uid: String,
     worker_id_initial_value: WorkerId,
+    scheduler_config: SchedulerConfig,
 ) -> crate::Result<(ServerRef, impl Future<Output = crate::Result<()>>)> {
     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,

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);

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;

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<WorkerId, u32>>,
     pub(crate) mn_workers: Map<(ResourceRqId, ResourceVariantId), Vec<ThinVec<WorkerId>>>,
+    /// `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() 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;

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,17 +10,35 @@ 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 {
     fn default() -> Self {
         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,

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,13 +64,15 @@ pub(crate) struct Core {
 pub(crate) type CoreRef = WrappedRcRefCell<Core>;
 
 impl CoreRef {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         worker_listen_port: u16,
         secret_key: Option<Arc<SecretKey>>,
         idle_timeout: Option<Duration>,
         custom_conn_handler: Option<CustomConnectionHandler>,
         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()
         })
     }

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 {

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<Item = (Self::Variable, f64)>,
     );
     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 {

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");