Implements Query leftovers flag

Author: spirali

crates/tako/src/control.rs

@@ -41,9 +41,11 @@ pub struct WorkerTypeQuery {
 
 #[derive(Debug)]
 pub struct NewWorkerAllocationResponse {
-    /// Array of the same size as number of queries, it returns the number of workers that should
+    /// Array of the same size as the number of queries, it returns the number of workers that should
     /// be spawned for the given query
     pub single_node_workers_per_query: Vec<usize>,
+    /// True iff not all tasks could be executed on resources provided in the query
+    pub single_node_leftovers: bool,
     pub multi_node_allocations: Vec<MultiNodeAllocationResponse>,
 }
 

crates/tako/src/internal/scheduler/query.rs

@@ -22,7 +22,7 @@ pub(crate) fn compute_new_worker_query(
     // Scheduler has to be performed before the query, so there should be no ready_to_assign tasks
     assert!(core.sn_ready_to_assign().is_empty() || !core.has_workers());
 
-    let add_task = |new_loads: &mut [WorkerTypeState], task: &Task| {
+    let add_task = |new_loads: &mut [WorkerTypeState], task: &Task| -> bool {
         let request = &task.configuration.resources;
         for ws in new_loads.iter_mut() {
             if let Some(time_limit) = ws.time_limit {
@@ -38,16 +38,17 @@ pub(crate) fn compute_new_worker_query(
             for load in ws.loads.iter_mut() {
                 if load.have_immediate_resources_for_rqv(request, &ws.w_resources) {
                     load.add_request(task.id, request, &ws.w_resources);
-                    return;
+                    return false;
                 }
             }
             if ws.loads.len() < ws.max as usize {
                 let mut load = WorkerLoad::new(&ws.w_resources);
                 load.add_request(task.id, request, &ws.w_resources);
                 ws.loads.push(load);
-                return;
+                return false;
             }
         }
+        true
     };
 
     /* Make sure that all named resources provided has an Id */
@@ -68,6 +69,7 @@ pub(crate) fn compute_new_worker_query(
         })
         .collect();
 
+    let mut leftover = false;
     for worker in core.get_workers() {
         let mut load = WorkerLoad::new(&worker.resources);
         for task_id in worker.sn_tasks() {
@@ -79,12 +81,12 @@ pub(crate) fn compute_new_worker_query(
                 load.add_request(task.id, request, &worker.resources);
                 continue;
             }
-            add_task(&mut new_loads, task);
+            leftover |= add_task(&mut new_loads, task);
         }
     }
     for task_id in core.sleeping_sn_tasks() {
         let task = core.get_task(*task_id);
-        add_task(&mut new_loads, task);
+        leftover |= add_task(&mut new_loads, task);
     }
 
     // `compute_new_worker_query` should be called immediately after scheduling was performed,
@@ -93,7 +95,7 @@ pub(crate) fn compute_new_worker_query(
     // postponing ready_to_assign. So we have to look also into this array
     for task_id in core.sn_ready_to_assign() {
         let task = core.get_task(*task_id);
-        add_task(&mut new_loads, task);
+        leftover |= add_task(&mut new_loads, task);
     }
 
     let single_node_allocations = new_loads
@@ -140,6 +142,7 @@ pub(crate) fn compute_new_worker_query(
 
     NewWorkerAllocationResponse {
         single_node_workers_per_query: single_node_allocations,
+        single_node_leftovers: leftover,
         multi_node_allocations,
     }
 }

crates/tako/src/internal/tests/integration/test_basic.rs

@@ -139,6 +139,7 @@ async fn test_query_no_output_immediate_call() {
         );
         assert_eq!(msg.single_node_workers_per_query, vec![0]);
         assert!(msg.multi_node_allocations.is_empty());
+        assert!(!msg.single_node_leftovers);
         handler.wait(&ids).await;
     })
     .await;
@@ -164,6 +165,7 @@ async fn test_query_no_output_delayed_call() {
         );
         assert_eq!(msg.single_node_workers_per_query, vec![0]);
         assert!(msg.multi_node_allocations.is_empty());
+        assert!(!msg.single_node_leftovers);
         handler.wait(&ids).await;
     })
     .await;
@@ -192,6 +194,7 @@ async fn test_query_new_workers_delayed_call() {
         );
         assert_eq!(msg.single_node_workers_per_query, vec![1]);
         assert!(msg.multi_node_allocations.is_empty());
+        assert!(!msg.single_node_leftovers);
     })
     .await;
 }
@@ -218,6 +221,7 @@ async fn test_query_new_workers_immediate() {
         );
         assert_eq!(msg.single_node_workers_per_query, vec![1]);
         assert!(msg.multi_node_allocations.is_empty());
+        assert!(!msg.single_node_leftovers);
     })
     .await;
 }

crates/tako/src/internal/tests/test_query.rs

@@ -24,6 +24,7 @@ fn test_query_no_tasks() {
     );
     assert_eq!(r.single_node_workers_per_query, vec![0]);
     assert!(r.multi_node_allocations.is_empty());
+    assert!(!r.single_node_leftovers);
 }
 
 #[test]
@@ -53,6 +54,7 @@ fn test_query_enough_workers() {
     );
     assert_eq!(r.single_node_workers_per_query, vec![0]);
     assert!(r.multi_node_allocations.is_empty());
+    assert!(!r.single_node_leftovers);
 }
 
 #[test]
@@ -91,6 +93,7 @@ fn test_query_no_enough_workers1() {
     );
     assert_eq!(r.single_node_workers_per_query, vec![0, 1]);
     assert!(r.multi_node_allocations.is_empty());
+    assert!(!r.single_node_leftovers);
 }
 
 #[test]
@@ -124,6 +127,7 @@ fn test_query_enough_workers2() {
     );
     assert_eq!(r.single_node_workers_per_query, vec![0, 0]);
     assert!(r.multi_node_allocations.is_empty());
+    assert!(!r.single_node_leftovers);
 }
 
 #[test]
@@ -158,6 +162,7 @@ fn test_query_not_enough_workers3() {
     );
     assert_eq!(r.single_node_workers_per_query, vec![1, 0]);
     assert!(r.multi_node_allocations.is_empty());
+    assert!(!r.single_node_leftovers);
 }
 
 #[test]
@@ -199,6 +204,7 @@ fn test_query_many_workers_needed() {
     );
     assert_eq!(r.single_node_workers_per_query, vec![5, 1, 26]);
     assert!(r.multi_node_allocations.is_empty());
+    assert!(!r.single_node_leftovers);
 }
 
 #[test]
@@ -255,6 +261,7 @@ fn test_query_multi_node_tasks() {
     assert_eq!(r.multi_node_allocations[2].worker_type, 1);
     assert_eq!(r.multi_node_allocations[2].worker_per_allocation, 6);
     assert_eq!(r.multi_node_allocations[2].max_allocations, 10);
+    assert!(!r.single_node_leftovers);
 }
 
 #[test]
@@ -276,6 +283,7 @@ fn test_query_multi_node_time_limit() {
             }],
         );
         assert_eq!(r.multi_node_allocations.len(), allocs);
+        assert!(!r.single_node_leftovers);
     }
 }
 
@@ -466,3 +474,28 @@ fn test_query_min_time1() {
     assert_eq!(r.single_node_workers_per_query, vec![1]);
     assert!(r.multi_node_allocations.is_empty());
 }
+
+#[test]
+fn test_query_leftovers() {
+    for (n, leftovers) in [(1, false), (4, false), (8, false), (9, true), (12, true)] {
+        let mut rt = TestEnv::new();
+
+        rt.new_workers(&[4]);
+        for i in 1..=n {
+            rt.new_task(TaskBuilder::new(i).cpus_compact(1));
+        }
+        rt.schedule();
+
+        let r = compute_new_worker_query(
+            rt.core(),
+            &[WorkerTypeQuery {
+                descriptor: ResourceDescriptor::simple(2),
+                time_limit: None,
+                max_sn_workers: 2,
+                max_workers_per_allocation: 1,
+                min_utilization: 0.0,
+            }],
+        );
+        assert_eq!(r.single_node_leftovers, leftovers);
+    }
+}