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

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

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<f64> {
+    std::env::var(key).ok().and_then(|value| value.parse::<f64>().ok())
+}
+
+fn get_duration_from_env(key: &str) -> Option<Duration> {
+    std::env::var(key)
+        .ok()
+        .and_then(|value| value.parse::<u64>().ok())
+        .map(Duration::from_millis)
+}
+
+#[cfg(test)]
+thread_local! {
+    static TEST_REL_GAP_OVERRIDE: std::cell::Cell<Option<f64>> = const { std::cell::Cell::new(None) };
+    static TEST_TIME_LIMIT_OVERRIDE: std::cell::Cell<Option<Duration>> = 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<R>(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
+}

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 {

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

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());
+    });
+}
+

crates/tako/src/internal/worker/resources/test_allocator.rs

@@ -899,6 +899,67 @@ fn test_coupling3() {
     assert_eq!(v, vec![1]);
 }
 
+#[test]
+fn test_allocator_stays_exact_regardless_of_scheduler_gap_tuning() {
+    // Guards against the exact regression found while implementing the
+    // scheduler's bounded solve (solve_bounded in internal::solver::highs):
+    // the worker's own NUMA/socket resource allocator below shares the same
+    // underlying LpSolver but must keep calling the exact `solve()`, never
+    // the scheduler's gap/time-limit-tuned `solve_bounded()` -- it needs a
+    // guaranteed-exact feasible allocation, not a good-enough one. Setting a
+    // coarse scheduler gap here must not affect it; this is verbatim the
+    // scenario from test_complex_coupling1 below, which is exactly what
+    // caught the original regression.
+    //
+    // with_test_solver_config is thread-local (not a process-global env
+    // var), so it cannot race with unrelated tests running concurrently on
+    // other threads.
+    crate::internal::solver::config::with_test_solver_config(
+        0.50,
+        std::time::Duration::from_millis(1),
+        || {
+            let mut coupling = ResourceDescriptorCoupling::default();
+            for i in 0..6 {
+                coupling.add(0, i, 1, i / 2, 256);
+                coupling.add(1, i / 2, 2, i, 128);
+            }
+            let descriptor = ResourceDescriptor::new(
+                vec![
+                    ResourceDescriptorItem {
+                        name: "cpus".to_string(),
+                        kind: ResourceDescriptorKind::regular_sockets(6, 2),
+                    },
+                    ResourceDescriptorItem {
+                        name: "gpus".to_string(),
+                        kind: ResourceDescriptorKind::regular_sockets(3, 1),
+                    },
+                    ResourceDescriptorItem {
+                        name: "foo".to_string(),
+                        kind: ResourceDescriptorKind::regular_sockets(6, 3),
+                    },
+                ],
+                coupling,
+            );
+            let mut allocator = test_allocator(&descriptor);
+
+            allocator.force_claim_from_groups(0.into(), &[0], 1.into());
+            allocator.force_claim_from_groups(2.into(), &[5], 2.into());
+
+            let rq = ResBuilder::default()
+                .add_force_compact(0, ResourceAmount::new_units(4))
+                .add_force_compact(1, ResourceAmount::new_units(1))
+                .add_force_compact(2, ResourceAmount::new_units(5))
+                .finish();
+            assert!(
+                allocator.try_allocate(&rq).is_some(),
+                "the exact NUMA/socket allocator must keep finding a \
+                 feasible allocation regardless of the scheduler's \
+                 mip_rel_gap tuning"
+            );
+        },
+    );
+}
+
 #[test]
 fn test_complex_coupling1() {
     let mut coupling = ResourceDescriptorCoupling::default();