Skip to content

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

Open
MauriceDHanisch wants to merge 1 commit into
It4innovations:mainfrom
MauriceDHanisch:fix/scheduler-solve-timeout
Open

fix(scheduler): bound MILP solve time and gap to prevent server hang#1119
MauriceDHanisch wants to merge 1 commit into
It4innovations:mainfrom
MauriceDHanisch:fix/scheduler-solve-timeout

Conversation

@MauriceDHanisch

Copy link
Copy Markdown

Summary

The global task scheduler's MILP solve (crates/tako/src/internal/scheduler/solver.rs) can take an unbounded amount of time on some workloads, hanging the server. This PR adds a bounded solve path used only by the scheduler:

  • mip_rel_gap (default 10%): accept a solution once provably within X% of optimal, instead of proving exact optimality. Task resource estimates are themselves approximate, so exact optimality isn't buying anything real.
  • time_limit (default 5s): hard wall-clock cap. If HiGHS hits the limit but already has a feasible incumbent, that incumbent is dispatched instead of discarded.
  • Both are tunable via HQ_SCHEDULER_MIP_REL_GAP and HQ_SCHEDULER_MIP_TIME_LIMIT_MS env vars, no new CLI surface.

Crucially, this is scoped to the scheduler only. The same LpSolver/HighsSolver abstraction is also used by the worker's own NUMA/socket resource allocator (worker/resources/groups.rs), which needs an exact, guaranteed-feasible answer, not an approximate one. The change splits solve() (exact, unchanged — used by the allocator) from a new solve_bounded() (gap/timeout-tuned, used by the scheduler only), with a default trait implementation so the coin_cbc/microlp backends need no changes.

Test plan

  • cargo test -p tako --lib — 209 passed, run 5x consecutively for stability (thread-local test overrides instead of process-global env vars, to avoid races across concurrently-running tests)
  • cargo build (full workspace) — clean
  • New coverage: a many-distinct-shapes scheduler test that previously hung now completes in bounded time; a test that a timed-out-but-feasible incumbent still gets dispatched; a test that a genuinely infeasible request still returns None under the bounded solve; a test that the worker's NUMA/socket allocator stays exact regardless of the scheduler's gap/timeout tuning

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.
@spirali

spirali commented Jul 3, 2026

Copy link
Copy Markdown
Collaborator

Can you please share with us your use case when you observe the hanging server?

@MauriceDHanisch

Copy link
Copy Markdown
Author

Repro: 50 tasks, each with a different (cpus, mem) request — 10 cpu values times 5 memory tiers, one task per combination, submitted to 20 workers.

#!/usr/bin/env bash
set -e

python3 - repro.toml <<'PY'
cpu_choices = [4, 8, 12, 16, 20, 24, 32, 40, 48, 64]
mem_choices = [8000, 32000, 64000, 128000, 220000]

with open("repro.toml", "w") as f:
    for cpus in cpu_choices:
        for mem in mem_choices:
            f.write(f'[[task]]\ncommand = ["sleep", "1"]\n'
                     f'[[task.request]]\nresources = {{ "cpus" = "{cpus}", "mem" = "{mem}" }}\n\n')
PY

hq server start &
for i in $(seq 1 20); do
  hq worker start --cpus 64 --resource "mem=sum(220000)" &
done
sleep 3

hq job submit-file --output-mode json repro.toml
# grab the job id from the output above, then:
#   time hq job wait <id>

I built hq at the commit right before this PR and at the commit with this PR, ran the exact same job against both.

Before this PR: it hangs. Not slow, hangs. 2+ minutes, server CPU pegged at 100%, doesn't even acknowledge the submission, no response to any RPC. I had to kill the process.

With this PR: 2.06 seconds.

The reason 50 distinct shapes shows up at all: tasks here get their resource request sized individually instead of picked from a small menu of standard sizes. That's normal for batch/HPC workloads where each task's actual resource need varies continuously. Once you do that, you end up close to one distinct shape per task instead of many tasks sharing a handful of shapes. The scheduler puts one MILP variable per (worker, batch, variant), so cost scales with how many distinct shapes are in flight, not with how many tasks there are. 50 shapes over 20 workers is already enough to blow it up.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants