From 475d2dcc7ae8efcbfd1aa18c784cee25f3cbc716 Mon Sep 17 00:00:00 2001 From: Pranav Sivaraman Date: Tue, 13 Oct 2026 13:58:14 -0400 Subject: [PATCH] tako: avoid panic when worker group is missing in multi-node solver Signed-off-by: Pranav Sivaraman --- crates/tako/src/internal/scheduler/solver.rs | 3 +-- crates/tako/src/internal/tests/test_query.rs | 26 ++++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/crates/tako/src/internal/scheduler/solver.rs b/crates/tako/src/internal/scheduler/solver.rs index ac107d68e..87d487c03 100644 --- a/crates/tako/src/internal/scheduler/solver.rs +++ b/crates/tako/src/internal/scheduler/solver.rs @@ -151,8 +151,7 @@ pub fn run_scheduling_solver_inner( if worker.is_free() && worker_groups .get(&worker.configuration.group) - .unwrap() - .is_capable_to_run_rq(rq, now, worker_map) + .is_some_and(|g| g.is_capable_to_run_rq(rq, now, worker_map)) { set_placement_name(&mut solver, worker.id, batch.resource_rq_id, v_idx); let v = create_mn_var( diff --git a/crates/tako/src/internal/tests/test_query.rs b/crates/tako/src/internal/tests/test_query.rs index a10aaa31d..2ceafe820 100644 --- a/crates/tako/src/internal/tests/test_query.rs +++ b/crates/tako/src/internal/tests/test_query.rs @@ -850,3 +850,29 @@ fn test_query_partial_absorbs_undeclared_demand_on_one_worker() { // On their own they are the probe of \S queries: exactly one worker, not none and not three. assert_eq!(ask(100, 0), 1); } + +#[test] +fn test_query_multi_node_task_with_free_workers_before_scheduling() { + // Fake workers created by the query belong to groups that are not registered in the core. + // The solver must not panic when it considers them for a multi-node task. + let mut rt = TestEnv::new(); + rt.new_workers_cpus(&[1, 1]); + rt.new_task(&TaskBuilder::new().n_nodes(2)); + + let r = compute_new_worker_query( + rt.core(), + &[WorkerTypeQuery { + partial: false, + descriptor: ResourceDescriptor::simple_cpus(1), + time_limit: None, + max_sn_workers: 2, + max_workers_per_allocation: 2, + min_utilization: 0.0, + }], + ); + assert_eq!(r.single_node_workers_per_query, vec![0]); + assert_eq!(r.multi_node_allocations.len(), 1); + assert_eq!(r.multi_node_allocations[0].worker_type, 0); + assert_eq!(r.multi_node_allocations[0].worker_per_allocation, 2); + assert_eq!(r.multi_node_allocations[0].max_allocations, 1); +}