From b222af709081237a3728bdb9f8b0a67270b1592b Mon Sep 17 00:00:00 2001 From: Ada Bohm Date: Sat, 8 Aug 2026 09:27:17 +0700 Subject: [PATCH 1/7] Fixed prefill balancing --- CHANGELOG.md | 1 + crates/tako/src/internal/scheduler/mapping.rs | 35 ++++++++--- .../src/internal/tests/test_scheduler_sn.rs | 58 +++++++++++++++++++ 3 files changed, 87 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c17ad77d7..59f67d196 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ### New features * The server scheduler now contains a safety limit for computation, configurable via `--scheduler-time-limit` (default: 5s) +* Better scheduling policy (prefill) for heterogenous clusters ### Fixes diff --git a/crates/tako/src/internal/scheduler/mapping.rs b/crates/tako/src/internal/scheduler/mapping.rs index 306110fc4..9d094cdfe 100644 --- a/crates/tako/src/internal/scheduler/mapping.rs +++ b/crates/tako/src/internal/scheduler/mapping.rs @@ -161,10 +161,15 @@ fn process_proactive_filling(core: &mut Core, mapping: &mut WorkerTaskMapping) { task_map, worker_map, task_queues, - request_map: _, + request_map, scheduler_state, .. } = core.split_mut(); + let max_prefill = scheduler_state.config.proactive_filling_max as u64; + if max_prefill == 0 { + // Prefill explicitly disabled. + return; + } let top_priority = task_queues.top_priority(); for queue in task_queues.iter_mut() { if queue.top_priority() != Some(top_priority) { @@ -176,6 +181,13 @@ fn process_proactive_filling(core: &mut Core, mapping: &mut WorkerTaskMapping) { if size == 0 { continue; } + let rqv = request_map.get(queue.resource_rq_id); + let max_capacity = worker_map + .get_workers() + .map(|w| w.resources.task_max_count(rqv)) + .max() + .unwrap_or(1) + .max(1) as u64; let workers: Vec<_> = worker_map .values_mut() .filter(|worker| { @@ -207,12 +219,21 @@ fn process_proactive_filling(core: &mut Core, mapping: &mut WorkerTaskMapping) { if workers.is_empty() { continue; } - let prefill_size = - (size / workers.len() as u32).min(scheduler_state.config.proactive_filling_max); - if prefill_size == 0 { - continue; - } - for worker in workers { + let capacities: Vec = workers + .iter() + .map(|w| w.resources.task_max_count(rqv).max(1) as u64) + .collect(); + let total_capacity: u64 = capacities.iter().sum(); + + for (worker, capacity) in workers.into_iter().zip(capacities) { + // The shares sum to at most `size`, so the queue entry we are drawing from is + // never exhausted before the last worker. + let share = size as u64 * capacity / total_capacity; + let depth = (max_prefill * capacity / max_capacity).max(1); + let prefill_size = share.min(depth) as u32; + if prefill_size == 0 { + continue; + } let tasks = queue.take_tasks_for_prefill(prefill_size); for task_id in &tasks { log::debug!("Prefiling task={task_id} to worker={}", worker.id); diff --git a/crates/tako/src/internal/tests/test_scheduler_sn.rs b/crates/tako/src/internal/tests/test_scheduler_sn.rs index 5e8b63ba0..61ab2d95b 100644 --- a/crates/tako/src/internal/tests/test_scheduler_sn.rs +++ b/crates/tako/src/internal/tests/test_scheduler_sn.rs @@ -1305,6 +1305,64 @@ fn test_prefill_steal() { rt.sanity_check(); } +/// Prefill has to be weighted by worker capacity, otherwise a small worker is handed as many +/// tasks as a large one and takes proportionally longer to drain them. Prefilled tasks are +/// removed from the global queue and are not reclaimed while regular tasks of the same priority +/// remain, so the small worker ends up sitting on an older job's tail long after every large +/// worker has moved on to newer jobs. +#[test] +fn test_prefill_weighted_by_worker_capacity() { + let mut rt = TestEnv::new(); + rt.set_scheduler_config(SchedulerConfig { + proactive_filling_reserve: 0, + proactive_filling_max: 32, + ..Default::default() + }); + let w_big = rt.new_worker(&WorkerBuilder::new(16)); + let w_small = rt.new_worker(&WorkerBuilder::new(2)); + rt.new_tasks(300, &TaskBuilder::new()); + rt.schedule(); + + // `proactive_filling_max` applies to the largest worker; everyone else is scaled down by + // capacity. An unweighted split would give both workers 32. + let big = prefill_count(&mut rt, w_big); + let small = prefill_count(&mut rt, w_small); + assert_eq!(big, 32); + assert_eq!(small, 4); + + // The property that actually matters: both hold the same *duration* of backlog, i.e. the + // same number of task generations (two each here). + assert_eq!(big / 16, small / 2); + rt.sanity_check(); +} + +/// The prefill depth must be measured against the largest worker in the *cluster*, not the +/// largest one eligible for prefill in this round. A worker is skipped while it still holds +/// prefill of the request, so the eligible set is routinely all-small -- and if the reference +/// capacity is taken from it, the depth springs back to `proactive_filling_max` for a tiny +/// worker, which is the whole bug. +#[test] +fn test_prefill_depth_when_large_worker_is_ineligible() { + let mut rt = TestEnv::new(); + rt.set_scheduler_config(SchedulerConfig { + proactive_filling_reserve: 0, + proactive_filling_max: 32, + ..Default::default() + }); + let w_big = rt.new_worker(&WorkerBuilder::new(16)); + rt.new_tasks(400, &TaskBuilder::new()); + rt.schedule(); + assert_eq!(prefill_count(&mut rt, w_big), 32); + + // w_big now holds prefill of this request, so it is excluded from further prefill and + // only the 2-cpu worker is eligible. + let w_small = rt.new_worker(&WorkerBuilder::new(2)); + rt.schedule(); + assert_eq!(prefill_count(&mut rt, w_big), 32); + assert_eq!(prefill_count(&mut rt, w_small), 4); + rt.sanity_check(); +} + #[test] pub fn test_schedule_running() { let mut rt = TestEnv::new(); From ddd4adea0fbaf26e51df8c0df5ed2b3784ea9b74 Mon Sep 17 00:00:00 2001 From: Ada Bohm Date: Sat, 8 Aug 2026 12:47:16 +0700 Subject: [PATCH 2/7] Added scheduler tests --- .../src/internal/tests/test_scheduler_sn.rs | 34 +++++++++++++++++++ .../src/internal/tests/utils/scheduler.rs | 8 +++++ 2 files changed, 42 insertions(+) diff --git a/crates/tako/src/internal/tests/test_scheduler_sn.rs b/crates/tako/src/internal/tests/test_scheduler_sn.rs index 61ab2d95b..c722eeab8 100644 --- a/crates/tako/src/internal/tests/test_scheduler_sn.rs +++ b/crates/tako/src/internal/tests/test_scheduler_sn.rs @@ -1336,6 +1336,40 @@ fn test_prefill_weighted_by_worker_capacity() { rt.sanity_check(); } +/// A resource weight (`hq submit --weight`) multiplies a request's placement value in the solver +/// objective. It is the supported way to make a request that only a few workers can serve win +/// those workers instead of being crowded out by work that could have run anywhere -- see the +/// `S5W` scenario in `benchmarks/scheduler-fairness`. It must never outrank user priority, so +/// here the heavily weighted 4-cpu request is given the *lower* priority. +#[test] +fn test_priority_is_not_overridden_by_weight() { + let narrow = TaskBuilder::new().cpus(1).user_priority(10); + let wide = TaskBuilder::new().cpus(4).user_priority(0).weight(10.0); + + let mut c = TestCase::new(); + c.n_tasks(10, &narrow); + c.n_tasks(4, &wide); + // All capacity goes to the high-priority 1-cpu tasks despite the 10x weight on the others. + c.w(&WorkerBuilder::new(8)).expect_request(8, &narrow); + c.w(&WorkerBuilder::new(2)).expect_request(2, &narrow); + c.check(); +} + +/// At *equal* priority the weight does steer placement, which is the behaviour that lets a wide +/// request claim the only workers able to run it. Miniature of `S5W`. +#[test] +fn test_weight_prefers_request_at_equal_priority() { + let narrow = TaskBuilder::new().cpus(1); + let wide = TaskBuilder::new().cpus(4).weight(4.0); + + let mut c = TestCase::new(); + c.n_tasks(10, &narrow); + c.n_tasks(2, &wide); + c.w(&WorkerBuilder::new(8)).expect_request(2, &wide); + c.w(&WorkerBuilder::new(2)).expect_request(2, &narrow); + c.check(); +} + /// The prefill depth must be measured against the largest worker in the *cluster*, not the /// largest one eligible for prefill in this round. A worker is skipped while it still holds /// prefill of the request, so the eligible set is routinely all-small -- and if the reference diff --git a/crates/tako/src/internal/tests/utils/scheduler.rs b/crates/tako/src/internal/tests/utils/scheduler.rs index 69ce4a7bc..167cad306 100644 --- a/crates/tako/src/internal/tests/utils/scheduler.rs +++ b/crates/tako/src/internal/tests/utils/scheduler.rs @@ -81,6 +81,14 @@ impl TestCase { self.rt.get_mut().new_tasks_cpus(cpus) } + /// `count` tasks from an explicit builder, for properties `pc_tasks` cannot express + /// (resource weight, variants, ...). + pub fn n_tasks(&mut self, count: usize, builder: &TaskBuilder) -> Vec { + (0..count) + .map(|_| self.rt.get_mut().new_task(builder)) + .collect() + } + // priority + cpu tasks pub fn pc_tasks(&mut self, priority_cpus: &[(i32, u32)]) -> Vec { priority_cpus From 6129ec28a0715ac0d567e10e03d9ba3752ae5aaa Mon Sep 17 00:00:00 2001 From: Ada Bohm Date: Mon, 10 Aug 2026 22:19:32 +0700 Subject: [PATCH 3/7] Fix: Fixed crash in prefil tries to prefill again --- CHANGELOG.md | 2 +- crates/tako/src/internal/scheduler/mapping.rs | 40 ++++++++----- .../tako/src/internal/scheduler/taskqueue.rs | 21 ++++--- .../tako/src/internal/tests/test_reactor.rs | 59 +++++++++++-------- tests/test_resources.py | 22 +++++++ 5 files changed, 92 insertions(+), 52 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 59f67d196..21d3f7ff2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ ### Fixes * Fixed server crash in a specific situation when an unschedulable high-priority task occurs - +* Fixed server crash caused by invalid handling of prefill ## v0.26.2 diff --git a/crates/tako/src/internal/scheduler/mapping.rs b/crates/tako/src/internal/scheduler/mapping.rs index 9d094cdfe..b5be85838 100644 --- a/crates/tako/src/internal/scheduler/mapping.rs +++ b/crates/tako/src/internal/scheduler/mapping.rs @@ -224,32 +224,40 @@ fn process_proactive_filling(core: &mut Core, mapping: &mut WorkerTaskMapping) { .map(|w| w.resources.task_max_count(rqv).max(1) as u64) .collect(); let total_capacity: u64 = capacities.iter().sum(); - + let mut return_back = Vec::new(); for (worker, capacity) in workers.into_iter().zip(capacities) { // The shares sum to at most `size`, so the queue entry we are drawing from is // never exhausted before the last worker. let share = size as u64 * capacity / total_capacity; let depth = (max_prefill * capacity / max_capacity).max(1); - let prefill_size = share.min(depth) as u32; + let prefill_size = share.min(depth); if prefill_size == 0 { continue; } - let tasks = queue.take_tasks_for_prefill(prefill_size); - for task_id in &tasks { + let prefills = &mut mapping.workers.entry(worker.id).or_default().prefills; + for _ in 0..prefill_size { + let task_id = queue.take_one().unwrap(); log::debug!("Prefiling task={task_id} to worker={}", worker.id); - let task = task_map.get_task_mut(*task_id); - assert!(task.is_waiting()); - task.state = TaskRuntimeState::Prefilled { - worker_id: worker.id, - }; - worker.insert_prefill_task(*task_id); + let task = task_map.get_task_mut(task_id); + if task.is_waiting() { + task.state = TaskRuntimeState::Prefilled { + worker_id: worker.id, + }; + worker.insert_prefill_task(task_id); + queue.insert_prefill(task_id, top_priority, prefill_size as usize); + prefills.push(task_id); + } else { + // This can happen when task is in retracting, and it should be queite rare + log::debug!( + "Task is not in waiting state ({:?}) back to the queue.", + task.state + ); + return_back.push(task_id); + } } - mapping - .workers - .entry(worker.id) - .or_default() - .prefills - .extend(tasks); + } + for task_id in return_back { + queue.return_back(task_id, top_priority); } } } diff --git a/crates/tako/src/internal/scheduler/taskqueue.rs b/crates/tako/src/internal/scheduler/taskqueue.rs index e63a14c15..8adfeb76a 100644 --- a/crates/tako/src/internal/scheduler/taskqueue.rs +++ b/crates/tako/src/internal/scheduler/taskqueue.rs @@ -151,6 +151,11 @@ impl TaskQueue { } } + #[inline] + pub fn return_back(&mut self, task_id: TaskId, priority: Priority) { + self.add(task_id, priority); + } + fn add(&mut self, task_id: TaskId, priority: Priority) { match self.queue.entry(Reverse(priority)) { Entry::Vacant(e) => { @@ -301,20 +306,14 @@ impl TaskQueue { } } - pub fn take_tasks_for_prefill(&mut self, mut count: u32) -> Vec { - let entry = self.queue.first_entry().unwrap(); - let mut result = Vec::with_capacity(count as usize); - let priority = entry.key().0; - take_from_entry(entry, &mut count, &mut result); + pub fn insert_prefill(&mut self, task_id: TaskId, priority: Priority, max_prefill: usize) { if let Some(prefill) = &mut self.prefill { - assert_eq!(prefill.0, priority); - for task_id in &result { - prefill.1.insert(*task_id); - } + prefill.1.insert(task_id); } else { - self.prefill = Some((priority, result.iter().copied().collect())) + let mut v = Set::with_capacity(max_prefill); + v.insert(task_id); + self.prefill = Some((priority, v)) } - result } pub fn take_tasks(&mut self, mut count: u32) -> Vec { diff --git a/crates/tako/src/internal/tests/test_reactor.rs b/crates/tako/src/internal/tests/test_reactor.rs index ddb87cf26..cfd169fbe 100644 --- a/crates/tako/src/internal/tests/test_reactor.rs +++ b/crates/tako/src/internal/tests/test_reactor.rs @@ -16,7 +16,6 @@ use crate::internal::tests::utils::sorted_vec; use crate::internal::tests::utils::task::{TaskBuilder, task_running_msg}; use crate::internal::tests::utils::workflows::{submit_example_1, submit_example_3}; use crate::internal::worker::configuration::OverviewConfiguration; -use crate::internal::worker::task::RunningTask; use crate::resources::{ResourceAmount, ResourceDescriptorItem, ResourceIdMap}; use crate::tests::utils::env::{TestComm, TestEnv}; use crate::tests::utils::worker::WorkerBuilder; @@ -772,6 +771,14 @@ fn test_task_reject3() { assert!(rt.task(t2).is_waiting()); } +fn get_prefilled(rt: &mut TestEnv, tasks: &[TaskId]) -> Option { + tasks.iter().find(|t| rt.task(**t).is_prefilled()).copied() +} + +fn get_assigned(rt: &mut TestEnv, tasks: &[TaskId]) -> Option { + tasks.iter().find(|t| rt.task(**t).is_assigned()).copied() +} + fn setup_prefill(rt: &mut TestEnv) -> (WorkerId, TaskId, TaskId) { rt.set_scheduler_config(SchedulerConfig { proactive_filling_reserve: 1, @@ -781,16 +788,8 @@ fn setup_prefill(rt: &mut TestEnv) -> (WorkerId, TaskId, TaskId) { let tasks = rt.new_tasks(3, &TaskBuilder::new()); let w1 = rt.new_worker(&WorkerBuilder::new(1)); rt.schedule(); - let prefilled = tasks - .iter() - .find(|t| rt.task(**t).is_prefilled()) - .copied() - .unwrap(); - let assigned = tasks - .iter() - .find(|t| rt.task(**t).is_assigned()) - .copied() - .unwrap(); + let prefilled = get_prefilled(rt, &tasks).unwrap(); + let assigned = get_assigned(rt, &tasks).unwrap(); (w1, assigned, prefilled) } @@ -825,6 +824,30 @@ fn test_prefill_submit_high_priority() { } } +#[test] +fn test_prefill_retracted_and_prefill_again() { + let mut rt = TestEnv::new(); + rt.set_scheduler_config(SchedulerConfig { + proactive_filling_reserve: 0, + proactive_filling_max: 1, + ..Default::default() + }); + let tasks1 = rt.new_tasks(3, &TaskBuilder::new()); + let w = rt.new_worker(&WorkerBuilder::new(1)); + rt.schedule(); + let p1 = get_prefilled(&mut rt, &tasks1).unwrap(); + let a1 = get_assigned(&mut rt, &tasks1).unwrap(); + let new_task = rt.new_task(&TaskBuilder::new().user_priority(1)); + assert!(matches!( + rt.task(p1).state, + TaskRuntimeState::Retracting { worker_id } if w == worker_id + )); + rt.finish_task(a1, w); + rt.schedule(); + assert!(rt.task(new_task).is_assigned()); + assert!(rt.task(p1).is_retracting()); +} + #[test] fn test_prefill_submit_same_priority() { for cpus in [1, 2] { @@ -878,24 +901,12 @@ fn test_prefill_started_on_same_worker() { assert!(rt.task(t1).is_assigned()); let tasks = rt.new_tasks(2, &TaskBuilder::new()); rt.schedule(); - let prefilled: TaskId = tasks - .iter() - .find(|t| rt.task(**t).is_prefilled()) - .copied() - .unwrap(); - let assigned: TaskId = tasks - .iter() - .find(|t| rt.task(**t).is_assigned()) - .copied() - .unwrap(); + let prefilled: TaskId = get_prefilled(&mut rt, &tasks).unwrap(); let up1 = WorkerTaskUpdate::Finished { task_id: t1 }; let mut comm = TestComm::new(); on_task_update(rt.core(), &mut comm, w1, smallvec![up1]); - rt.schedule(); - assert!(rt.task(prefilled).is_retracting()); - let up2 = WorkerTaskUpdate::Running(task_running_msg(prefilled)); on_task_update(rt.core(), &mut comm, w1, smallvec![up2]); rt.sanity_check(); diff --git a/tests/test_resources.py b/tests/test_resources.py index 01291aa05..511d41bbf 100644 --- a/tests/test_resources.py +++ b/tests/test_resources.py @@ -636,3 +636,25 @@ def test_scheduler_unschedulable_sn_blocker(hq_env: HqEnv): hq_env.check_running_processes() table = hq_env.command(["job", "info", "3"], as_table=True) assert table.get_row_value("State").endswith("WAITING (5)") + + +def test_scheduler_priority_churn(hq_env: HqEnv): + hq_env.start_server() + hq_env.start_workers(1, cpus=4) + hq_env.command(["submit", "--array=0-999", "--stdout=none", "--stderr=none", "--", "sleep", "0.05"]) + wait_for_job_state(hq_env, 1, "RUNNING") + + for priority in range(1, 6): + hq_env.command( + [ + "submit", + f"--priority={priority}", + "--stdout=none", + "--stderr=none", + "--", + "sleep", + "0.05", + ] + ) + time.sleep(0.5) + hq_env.check_running_processes() From a563307fa82b9b4d8ad2bf829836464c116b88c5 Mon Sep 17 00:00:00 2001 From: Ada Bohm Date: Mon, 10 Aug 2026 15:06:30 +0700 Subject: [PATCH 4/7] Reservation test --- tests/test_resources.py | 37 +++++++++++++++++++++++++++++++++++++ tests/test_task.py | 2 +- 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/tests/test_resources.py b/tests/test_resources.py index 511d41bbf..ce364f3df 100644 --- a/tests/test_resources.py +++ b/tests/test_resources.py @@ -658,3 +658,40 @@ def test_scheduler_priority_churn(hq_env: HqEnv): ) time.sleep(0.5) hq_env.check_running_processes() + +def test_scheduler_reservation(hq_env: HqEnv, tmp_path): + hq_env.start_server() + hq_env.start_workers(4, cpus=6) + hq_env.command(["submit", "--array=1-4", "--cpus=4", "--", "sleep", "100"]) + wait_for_job_state(hq_env, 1, "RUNNING") + time.sleep(0.5) + content = [""" +[[task]] +id = 0 +priority = 10 +command = ["sleep", "100"] + +[[task.request]] +resources = { "cpus" = 6 } + """] + for i in range(1, 31): + content.append(f""" +[[task]] +id = {i} +command = ["sleep", "100"] +[[task.request]] +resources = {{ "cpus" = 1 }} +""") + tmp_path.joinpath("job.toml").write_text("\n".join(content)) + hq_env.command(["job", "submit-file", "job.toml"]) + wait_for_job_state(hq_env, 1, "RUNNING") + time.sleep(0.5) + print(hq_env.command(["job", "info", "2"])) + + ts = hq_env.command(["task", "--output-mode=json", "info", "2", "0-30"], as_json=True) + print(ts) + assert ts[0]["state"] == "waiting" + + assert sum(1 if t["state"] == "running" else 0 for t in ts) == 6 + assert len(set(t["worker"] for t in ts if t["state"] == "running")) == 3 + diff --git a/tests/test_task.py b/tests/test_task.py index c76094b16..27493bcd5 100644 --- a/tests/test_task.py +++ b/tests/test_task.py @@ -76,4 +76,4 @@ def test_long_running_task(hq_env: HqEnv): hq_env.start_server() hq_env.start_worker() hq_env.command(["submit", "sleep", "20"]) - wait_for_job_state(hq_env, 1, "FINISHED", timeout_s=30) + wait_for_job_state(hq_env, 1, "FINISHED", timeout_s=30) \ No newline at end of file From 94f2d18c6ad7075505c113b84d58a4c2c860e5e4 Mon Sep 17 00:00:00 2001 From: Ada Bohm Date: Thu, 13 Aug 2026 15:44:07 +0700 Subject: [PATCH 5/7] Fix prefill when raced with running state --- .../tako/src/internal/scheduler/taskqueue.rs | 30 ++++++++++++++++ crates/tako/src/internal/server/reactor.rs | 35 +++++++++++++++---- .../tako/src/internal/tests/test_reactor.rs | 27 ++++++++++++++ tests/test_resources.py | 10 +++--- tests/test_task.py | 2 +- 5 files changed, 92 insertions(+), 12 deletions(-) diff --git a/crates/tako/src/internal/scheduler/taskqueue.rs b/crates/tako/src/internal/scheduler/taskqueue.rs index 8adfeb76a..b58df48e5 100644 --- a/crates/tako/src/internal/scheduler/taskqueue.rs +++ b/crates/tako/src/internal/scheduler/taskqueue.rs @@ -221,6 +221,36 @@ impl TaskQueue { } } + /// Removes the task if it is in the queue and returns whether it was removed. + /// A retracting task may have been already taken out of the queue by the scheduler. + pub fn remove_if_queued(&mut self, task_id: TaskId, priority: Priority) -> bool { + if let Some((p, ts)) = &mut self.prefill + && priority == *p + && ts.remove(&task_id) + { + return true; + } + match self.queue.entry(Reverse(priority)) { + Entry::Vacant(_) => false, + Entry::Occupied(mut e) => match e.get_mut() { + OneOrMoreTaskIds::One(v) => { + if *v != task_id { + return false; + } + e.remove(); + true + } + OneOrMoreTaskIds::More(tasks) => { + let found = tasks.remove(&task_id); + if tasks.is_empty() { + e.remove(); + } + found + } + }, + } + } + pub fn shrink_to_fit(&mut self) { // Do nothing } diff --git a/crates/tako/src/internal/server/reactor.rs b/crates/tako/src/internal/server/reactor.rs index 66f30eaa8..7e34c0843 100644 --- a/crates/tako/src/internal/server/reactor.rs +++ b/crates/tako/src/internal/server/reactor.rs @@ -320,13 +320,20 @@ fn task_running( // By removing redirections first, we unassign the task so we can later assign it back // In theory, we could optimize this special case by doing nothing, but it should be quite rare // So I prefer to keep the code simple. - try_remove_redirection( + if !try_remove_redirection( worker_map, scheduler_state, request_map, task_id, task.resource_rq_id, - ); + ) { + // We have tried to retract the task, but it has already started. + // Without a redirection, the task is still in the ready queue and it has + // to be removed, otherwise the queue keeps an id of an already removed task. + task_queues + .get_mut(task.resource_rq_id) + .remove_if_queued(task.id, task.priority()); + } let rqv = request_map.get(task.resource_rq_id); worker_map .get_worker_mut(worker_id) @@ -546,13 +553,17 @@ fn task_finished( } TaskRuntimeState::Retracting { worker_id: w_id } => { assert_eq!(*w_id, worker_id); - try_remove_redirection( + if !try_remove_redirection( worker_map, scheduler_state, request_map, task_id, task.resource_rq_id, - ); + ) { + task_queues + .get_mut(task.resource_rq_id) + .remove_if_queued(task.id, task.priority()); + } } TaskRuntimeState::Prefilled { .. } | TaskRuntimeState::Waiting { .. } @@ -589,17 +600,23 @@ fn task_finished( true } +/// Returns true if a redirection was found (and removed). +/// A task with a redirection was already taken out of the ready queue by the scheduler, +/// a retracting task without a redirection is still waiting in the ready queue. fn try_remove_redirection( worker_map: &mut WorkerMap, scheduler_state: &mut SchedulerState, request_map: &ResourceRqMap, task_id: TaskId, resource_rq_id: ResourceRqId, -) { +) -> bool { if let Some((worker_id, rv_id)) = scheduler_state.redirects.remove(&task_id) { let worker = worker_map.get_worker_mut(worker_id); let rq = request_map.get(resource_rq_id).get(rv_id); worker.remove_sn_task(task_id, rq); + true + } else { + false } } @@ -653,13 +670,17 @@ fn task_failed( } TaskRuntimeState::Retracting { worker_id: w } => { assert_eq!(worker_id, *w); - try_remove_redirection( + if !try_remove_redirection( worker_map, scheduler_state, request_map, task_id, task.resource_rq_id, - ); + ) { + task_queues + .get_mut(task.resource_rq_id) + .remove_if_queued(task.id, task.priority()); + } } _ => {} } diff --git a/crates/tako/src/internal/tests/test_reactor.rs b/crates/tako/src/internal/tests/test_reactor.rs index cfd169fbe..29c90d524 100644 --- a/crates/tako/src/internal/tests/test_reactor.rs +++ b/crates/tako/src/internal/tests/test_reactor.rs @@ -912,6 +912,33 @@ fn test_prefill_started_on_same_worker() { rt.sanity_check(); } +#[test] +fn test_prefill_retracted_but_already_started() { + let mut rt = TestEnv::new(); + let (w1, t1, t2) = setup_prefill(&mut rt); + let mut comm = TestComm::new(); + + let t3 = TaskId::new(100.into(), 501.into()); + let task3 = TaskBuilder::new().user_priority(10).build(t3, rt.core()); + on_new_tasks(rt.core(), &mut comm, vec![task3]); + comm.check_need_scheduling(); + match &comm.take_worker_msgs(w1, 1)[0] { + ToWorkerMessage::RetractTasks(ts) => assert_eq!(ts.ids, vec![t2]), + _ => panic!("Invalid worker msg"), + } + comm.emptiness_check(); + assert!(rt.task(t2).is_retracting()); + + rt.finish_task(t1, w1); + let up = WorkerTaskUpdate::Running(task_running_msg(t2)); + on_task_update(rt.core(), &mut comm, w1, smallvec![up]); + comm.client.take_task_running(1); + comm.check_need_scheduling(); + comm.emptiness_check(); + rt.finish_task(t2, w1); + rt.schedule(); +} + #[test] fn test_prefill_started() { let mut rt = TestEnv::new(); diff --git a/tests/test_resources.py b/tests/test_resources.py index ce364f3df..a24ce2ddc 100644 --- a/tests/test_resources.py +++ b/tests/test_resources.py @@ -658,14 +658,16 @@ def test_scheduler_priority_churn(hq_env: HqEnv): ) time.sleep(0.5) hq_env.check_running_processes() - + + def test_scheduler_reservation(hq_env: HqEnv, tmp_path): hq_env.start_server() hq_env.start_workers(4, cpus=6) hq_env.command(["submit", "--array=1-4", "--cpus=4", "--", "sleep", "100"]) wait_for_job_state(hq_env, 1, "RUNNING") time.sleep(0.5) - content = [""" + content = [ + """ [[task]] id = 0 priority = 10 @@ -673,7 +675,8 @@ def test_scheduler_reservation(hq_env: HqEnv, tmp_path): [[task.request]] resources = { "cpus" = 6 } - """] + """ + ] for i in range(1, 31): content.append(f""" [[task]] @@ -694,4 +697,3 @@ def test_scheduler_reservation(hq_env: HqEnv, tmp_path): assert sum(1 if t["state"] == "running" else 0 for t in ts) == 6 assert len(set(t["worker"] for t in ts if t["state"] == "running")) == 3 - diff --git a/tests/test_task.py b/tests/test_task.py index 27493bcd5..c76094b16 100644 --- a/tests/test_task.py +++ b/tests/test_task.py @@ -76,4 +76,4 @@ def test_long_running_task(hq_env: HqEnv): hq_env.start_server() hq_env.start_worker() hq_env.command(["submit", "sleep", "20"]) - wait_for_job_state(hq_env, 1, "FINISHED", timeout_s=30) \ No newline at end of file + wait_for_job_state(hq_env, 1, "FINISHED", timeout_s=30) From 294d223382355ce17b5390a719ae071d06d9b4e7 Mon Sep 17 00:00:00 2001 From: Ada Bohm Date: Sat, 15 Aug 2026 11:10:15 +0700 Subject: [PATCH 6/7] Fix gap computation in prefil retract --- crates/tako/src/internal/scheduler/solver.rs | 11 +-- crates/tako/src/internal/server/task.rs | 14 ++++ crates/tako/src/internal/server/worker.rs | 12 +--- .../src/internal/tests/test_scheduler_sn.rs | 69 +++++++++++++++++++ 4 files changed, 93 insertions(+), 13 deletions(-) diff --git a/crates/tako/src/internal/scheduler/solver.rs b/crates/tako/src/internal/scheduler/solver.rs index 719df80e3..905640612 100644 --- a/crates/tako/src/internal/scheduler/solver.rs +++ b/crates/tako/src/internal/scheduler/solver.rs @@ -47,7 +47,7 @@ pub(crate) fn run_scheduling_solver( task_queues: _, request_map, worker_groups, - scheduler_state: scheduler_cache, + scheduler_state, .. } = core.split(); if request_map.is_empty() { @@ -292,13 +292,16 @@ pub(crate) fn run_scheduling_solver( if !w.is_capable_to_run_rqv(blocker_rqv, now) { continue; } - let gap = scheduler_cache.gap_cache.get_gap( + let gap = scheduler_state.gap_cache.get_gap( *blocker_rq_id, batch.resource_rq_id, &w.resources, sn_assignment.assigned_tasks.iter().map(|task_id| { let t = task_map.get_task(*task_id); - (t.resource_rq_id, t.rv_id().unwrap()) + ( + t.resource_rq_id, + t.assigned_placement(&scheduler_state.redirects).unwrap().1, + ) }), request_map, ); @@ -430,7 +433,7 @@ pub(crate) fn run_scheduling_solver( } let mut result = SchedulingSolution::default(); - let Some((solution, is_optimal)) = solver.solve_bounded(scheduler_cache.config.mip_time_limit) + let Some((solution, is_optimal)) = solver.solve_bounded(scheduler_state.config.mip_time_limit) else { return result; }; diff --git a/crates/tako/src/internal/server/task.rs b/crates/tako/src/internal/server/task.rs index 35ccd9fc2..f180cdaef 100644 --- a/crates/tako/src/internal/server/task.rs +++ b/crates/tako/src/internal/server/task.rs @@ -247,6 +247,7 @@ impl Task { } } + #[cfg(test)] pub(crate) fn rv_id(&self) -> Option { match self.state { TaskRuntimeState::Running { rv_id, .. } | TaskRuntimeState::Assigned { rv_id, .. } => { @@ -256,6 +257,19 @@ impl Task { } } + #[inline] + pub(crate) fn assigned_placement( + &self, + redirects: &Map, + ) -> Option<(WorkerId, ResourceVariantId)> { + match self.state { + TaskRuntimeState::Assigned { worker_id, rv_id } + | TaskRuntimeState::Running { worker_id, rv_id } => Some((worker_id, rv_id)), + TaskRuntimeState::Retracting { .. } => redirects.get(&self.id).copied(), + _ => None, + } + } + pub(crate) fn increment_instance_id(&mut self) { self.instance_id = InstanceId::new(self.instance_id.as_num() + 1); } diff --git a/crates/tako/src/internal/server/worker.rs b/crates/tako/src/internal/server/worker.rs index 923714530..9790fc22d 100644 --- a/crates/tako/src/internal/server/worker.rs +++ b/crates/tako/src/internal/server/worker.rs @@ -243,15 +243,9 @@ impl Worker { let mut resources = self.resources.clone(); for task_id in a.assigned_tasks.iter() { let task = task_map.get_task(*task_id); - let (worker_id, rv_id) = match &task.state { - TaskRuntimeState::Assigned { worker_id, rv_id } - | TaskRuntimeState::Running { worker_id, rv_id } => (*worker_id, *rv_id), - TaskRuntimeState::Retracting { .. } => { - let (worker_id, rv_id) = transfers.get(task_id).unwrap(); - (*worker_id, *rv_id) - } - s => panic!("Invalid state {s:?}"), - }; + let (worker_id, rv_id) = task + .assigned_placement(transfers) + .unwrap_or_else(|| panic!("Invalid state {:?}", task.state)); assert_eq!(self.id, worker_id); let rq = request_map.get(task.resource_rq_id).get(rv_id); assert!(resources.is_capable_to_run_request(rq)); diff --git a/crates/tako/src/internal/tests/test_scheduler_sn.rs b/crates/tako/src/internal/tests/test_scheduler_sn.rs index c722eeab8..6d594fe5f 100644 --- a/crates/tako/src/internal/tests/test_scheduler_sn.rs +++ b/crates/tako/src/internal/tests/test_scheduler_sn.rs @@ -1305,6 +1305,75 @@ fn test_prefill_steal() { rt.sanity_check(); } +#[test] +fn test_gap_over_redirected_retracting_task() { + let mut rt = TestEnv::new(); + rt.set_scheduler_config(SchedulerConfig { + proactive_filling_reserve: 3, + proactive_filling_max: 6, + ..Default::default() + }); + + // -- Precondition 1: a redirected retracting task in w2's `assigned_tasks`. + // Same opening as `test_prefill_steal`: w1 prefills, then the larger w2 steals part of that + // prefill, so those tasks become `Retracting` (on w1) while being assigned to w2. + let w1 = rt.new_worker(&WorkerBuilder::new(1)); + let low = rt.new_tasks(9, &TaskBuilder::new()); + let low_rq_id = rt.task(low[0]).resource_rq_id; + rt.schedule(); + assert_eq!(prefill_count(&mut rt, w1), 5, "w1 did not prefill"); + let w2 = rt.new_worker(&WorkerBuilder::new(5)); + rt.schedule(); + + // Asserted rather than assumed: if prefill ever stops producing this state, the test must + // fail loudly instead of passing while reproducing nothing. + assert!( + !rt.core().split().scheduler_state.redirects.is_empty(), + "no redirect was created, so the state under test does not exist" + ); + let retracting_on_w2 = rt + .worker(w2) + .sn_assignment() + .unwrap() + .assigned_tasks + .iter() + .filter(|task_id| rt.task(**task_id).rv_id().is_none()) + .count(); + assert!( + retracting_on_w2 > 0, + "w2 holds no assigned task without an rv_id; the unwrap cannot be reached" + ); + + // -- Precondition 2: spare capacity for the low-priority request. + // The solver skips a batch that has no placement variables, and after the steal both w1 and + // w2 are full -- so without this the low-priority batch, and with it the entire + // priority-condition path, is never visited. A worker added now cannot undo the redirect + // already recorded above. 2 cpus is deliberately too narrow for the blocker below, so this + // worker contributes capacity without becoming a candidate for it. + rt.new_worker(&WorkerBuilder::new(2)); + + // -- Precondition 3: a blocker, so the solver builds a priority condition at all. + // A cut is only emitted when a *second* queue is non-empty at a higher priority, so this + // needs a distinct resource request. 5 cpus fits w2's total width -- w2 is therefore + // `is_capable_to_run_rqv` and is not skipped by the impossible-filter -- but w2 has no room + // for it right now, which is what makes it block. + rt.new_tasks(2, &TaskBuilder::new().cpus(5).user_priority(10)); + + let batches = create_task_batches(rt.core(), std::time::Instant::now(), None); + let low_batch = batches + .iter() + .find(|b| b.resource_rq_id == low_rq_id) + .expect("the low-priority request must still have a batch"); + assert!( + low_batch.cuts.iter().any(|cut| !cut.blockers.is_empty()), + "no blocker cut was produced, so the gap path is never entered: {:?}", + low_batch.cuts + ); + + rt.schedule(); + rt.sanity_check(); +} + /// Prefill has to be weighted by worker capacity, otherwise a small worker is handed as many /// tasks as a large one and takes proportionally longer to drain them. Prefilled tasks are /// removed from the global queue and are not reclaimed while regular tasks of the same priority From 2669a7bd79188c7ebda7d7c476e61c8acf4c22c8 Mon Sep 17 00:00:00 2001 From: Ada Bohm Date: Sat, 15 Aug 2026 20:31:18 +0700 Subject: [PATCH 7/7] Fixed greedy backfilling --- CHANGELOG.md | 1 + crates/tako/src/internal/scheduler/solver.rs | 2 +- .../src/internal/tests/test_scheduler_sn.rs | 62 ++++++++++++++++++- 3 files changed, 63 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 21d3f7ff2..837f6a4ff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ ### Fixes +* Fixed some occasional greedy backfilling in server scheduler * Fixed server crash in a specific situation when an unschedulable high-priority task occurs * Fixed server crash caused by invalid handling of prefill diff --git a/crates/tako/src/internal/scheduler/solver.rs b/crates/tako/src/internal/scheduler/solver.rs index 905640612..39cd1a407 100644 --- a/crates/tako/src/internal/scheduler/solver.rs +++ b/crates/tako/src/internal/scheduler/solver.rs @@ -156,7 +156,7 @@ pub(crate) fn run_scheduling_solver( && worker.is_capable_to_run_rqv(rqv, now) && let Some(a) = worker.sn_assignment() { - let weight = w_idx as f64 / (n_workers * 100) as f64; + let weight = -((n_workers - w_idx) as f64 / (n_workers * 1024) as f64); solver.set_name(|| format!("R{}:{}", worker.id, batch.resource_rq_id)); let v = solver.add_bool_variable(weight); tasks_count_vars diff --git a/crates/tako/src/internal/tests/test_scheduler_sn.rs b/crates/tako/src/internal/tests/test_scheduler_sn.rs index 6d594fe5f..9bfcf4c7f 100644 --- a/crates/tako/src/internal/tests/test_scheduler_sn.rs +++ b/crates/tako/src/internal/tests/test_scheduler_sn.rs @@ -7,7 +7,7 @@ use crate::resources::ResourceRqId; use crate::tests::utils::env::{TestComm, TestEnv}; use crate::tests::utils::task::TaskBuilder; use crate::tests::utils::worker::WorkerBuilder; -use crate::{ResourceVariantId, WorkerId}; +use crate::{ResourceVariantId, TaskId, WorkerId}; use std::time::Duration; #[test] @@ -1672,3 +1672,63 @@ fn test_schedule_bounded_is_optimal_true_when_solve_converges() { assert!(rt.schedule_solution().is_optimal); } + +#[test] +fn test_schedule_reservation_priority() { + let mut c = TestCase::new(); + let ht = c.t(&TaskBuilder::new().cpus(6).user_priority(10)); + c.ts(6, &TaskBuilder::new().cpus(1)); + c.w(&WorkerBuilder::new(6)).running_c(6); + c.w(&WorkerBuilder::new(6)).expect_tasks(&[ht]); + c.check(); +} + +#[test] +fn test_schedule_reservation_used_when_worker_frees_up() { + let mut rt = TestEnv::new(); + let ws = rt.new_workers(4, &WorkerBuilder::new(6)); + let running: Vec<_> = ws + .iter() + .map(|w| rt.new_task_running(&TaskBuilder::new().cpus(4), *w)) + .collect(); + let blocker = rt.new_task(&TaskBuilder::new().cpus(6).user_priority(10)); + let small = rt.new_tasks(30, &TaskBuilder::new().cpus(1)); + + fn on_worker(rt: &TestEnv, task_id: TaskId, worker_id: WorkerId) -> bool { + matches!( + rt.task(task_id).state, + TaskRuntimeState::Assigned { worker_id: w, .. } if w == worker_id + ) + } + + rt.schedule(); + + // One worker is held back for the blocker, the three others take two tasks each. + let reserved_idx = ws + .iter() + .position(|w| !small.iter().any(|t| on_worker(&rt, *t, *w))) + .expect("no worker was reserved for the blocker"); + assert_eq!( + small.iter().filter(|t| rt.task(**t).is_assigned()).count(), + 6 + ); + assert!(rt.task(blocker).is_waiting()); + + // The reserved worker becomes completely free, which is exactly what the blocker waits for. + let reserved = ws[reserved_idx]; + rt.finish_task(running[reserved_idx], reserved); + rt.schedule(); + + assert!( + on_worker(&rt, blocker, reserved), + "reserved worker {reserved} was not used for the blocker, state: {:?}", + rt.task(blocker).state + ); + assert_eq!( + small + .iter() + .filter(|t| on_worker(&rt, **t, reserved)) + .count(), + 0 + ); +}