New query format with partial worker types

Author: spirali

crates/hyperqueue/src/server/autoalloc/process.rs

@@ -377,7 +377,7 @@ async fn perform_submits(
         .iter()
         .map(|(_, queue)| create_queue_worker_query(queue))
         .collect();
-    let response = senders.server.new_worker_query(queries, false)?;
+    let response = senders.server.new_worker_query(&queries)?;
     log::debug!("Scheduler query response: {response:?}");
 
     // Merge responses back into a single array
@@ -424,6 +424,7 @@ fn create_queue_worker_query(queue: &AllocationQueue) -> WorkerTypeQuery {
     WorkerTypeQuery {
         // The maximum number of workers that we can provide in this queue
         // TODO: estimate the resources of the queue in a better way
+        partial: false,
         descriptor: ResourceDescriptor::new(vec![ResourceDescriptorItem {
             name: CPU_RESOURCE_NAME.to_string(),
             kind: ResourceDescriptorKind::regular_sockets(1, 1),

crates/tako/src/control.rs

@@ -8,9 +8,7 @@ use tokio::net::TcpListener;
 use tokio::sync::Notify;
 
 use crate::events::EventProcessor;
-use crate::gateway::{
-    MultiNodeAllocationResponse, ResourceRequestVariants, TaskSubmit, WorkerRuntimeInfo,
-};
+use crate::gateway::{MultiNodeAllocationResponse, TaskSubmit, WorkerRuntimeInfo};
 use crate::internal::common::error::DsError;
 use crate::internal::messages::worker::ToWorkerMessage;
 use crate::internal::scheduler::query::compute_new_worker_query;
@@ -28,6 +26,9 @@ use crate::{TaskId, WorkerId};
 
 #[derive(Debug)]
 pub struct WorkerTypeQuery {
+    /// If True indicates that we do not have all information about this worker
+    pub partial: bool,
+    /// Worker resource descriptor
     pub descriptor: ResourceDescriptor,
     /// Worker time limit
     pub time_limit: Option<Duration>,
@@ -47,11 +48,6 @@ pub struct NewWorkerAllocationResponse {
     /// be spawned for the given query
     pub single_node_workers_per_query: Vec<usize>,
     pub multi_node_allocations: Vec<MultiNodeAllocationResponse>,
-    /// Resource requests of tasks that left after resolving the query and could not be assigned to
-    /// any running worker or worker in query response.
-    /// The second value is number of such tasks
-    /// Note: It contains both single node and multi node leftovers
-    pub leftovers: Vec<(ResourceRequestVariants, u32)>,
 }
 
 #[derive(Clone)]
@@ -118,22 +114,17 @@ impl ServerRef {
     */
     pub fn new_worker_query(
         &self,
-        queries: Vec<WorkerTypeQuery>,
-        collect_leftovers: bool,
+        queries: &[WorkerTypeQuery],
     ) -> crate::Result<NewWorkerAllocationResponse> {
-        for query in &queries {
+        for query in queries {
             query.descriptor.validate()?;
         }
         let mut core = self.core_ref.get_mut();
         let mut comm = self.comm_ref.get_mut();
         if comm.get_scheduling_flag() {
             run_scheduling_now(&mut core, &mut comm, Instant::now())
         }
-        Ok(compute_new_worker_query(
-            &mut core,
-            &queries,
-            collect_leftovers,
-        ))
+        Ok(compute_new_worker_query(&mut core, queries))
     }
 
     pub fn try_release_memory(&self) {

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

@@ -1,14 +1,12 @@
-use crate::Map;
 use crate::control::{NewWorkerAllocationResponse, WorkerTypeQuery};
 use crate::gateway::MultiNodeAllocationResponse;
 use crate::internal::server::core::Core;
 use crate::internal::server::task::Task;
 use crate::internal::server::workerload::{WorkerLoad, WorkerResources};
-use crate::resources::ResourceRequestVariants;
-use smallvec::smallvec;
 use std::time::Duration;
 
 struct WorkerTypeState {
+    partial: bool,
     loads: Vec<WorkerLoad>,
     w_resources: WorkerResources,
     time_limit: Option<Duration>,
@@ -19,46 +17,43 @@ struct WorkerTypeState {
 pub(crate) fn compute_new_worker_query(
     core: &mut Core,
     queries: &[WorkerTypeQuery],
-    collect_leftovers: bool,
 ) -> NewWorkerAllocationResponse {
     log::debug!("Compute new worker query: query = {queries:?}");
 
     // 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,
-                    leftovers: &mut Map<ResourceRequestVariants, u32>| {
+    let add_task = |new_loads: &mut [WorkerTypeState], task: &Task| {
         let request = &task.configuration.resources;
         for ws in new_loads.iter_mut() {
-            if let Some(time_limit) = ws.time_limit {
-                if !ws
-                    .w_resources
-                    .is_capable_to_run_with(request, |rq| rq.min_time() <= time_limit)
-                {
+            if !ws.partial {
+                if !ws.w_resources.is_capable_to_run_with(request, |rq| {
+                    ws.time_limit.is_none_or(|t| rq.min_time() <= t)
+                }) {
                     continue;
                 }
-            } else if !ws.w_resources.is_capable_to_run(request) {
-                continue;
-            }
-            for load in ws.loads.iter_mut() {
-                if load.have_immediate_resources_for_rqv(request, &ws.w_resources) {
+                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;
+                    }
+                }
+                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;
                 }
-            }
-            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;
-            }
-        }
-        if collect_leftovers {
-            if let Some(count) = leftovers.get_mut(request) {
-                *count += 1;
             } else {
-                leftovers.insert(request.clone(), 1);
+                if !ws.w_resources.is_lowerbound_for(request, |rq| {
+                    ws.time_limit.is_none_or(|t| rq.min_time() <= t)
+                }) {
+                    continue;
+                }
+                if ws.loads.is_empty() {
+                    let load = WorkerLoad::new(&ws.w_resources);
+                    ws.loads.push(load);
+                }
             }
         }
     };
@@ -74,14 +69,14 @@ pub(crate) fn compute_new_worker_query(
     let mut new_loads: Vec<_> = queries
         .iter()
         .map(|q| WorkerTypeState {
+            partial: q.partial,
             loads: Vec::new(),
             w_resources: WorkerResources::from_description(&q.descriptor, &resource_map),
             time_limit: q.time_limit,
             max: q.max_sn_workers,
         })
         .collect();
 
-    let mut leftovers = Map::new();
     for worker in core.get_workers() {
         let mut load = WorkerLoad::new(&worker.resources);
         for task_id in worker.sn_tasks() {
@@ -93,12 +88,12 @@ pub(crate) fn compute_new_worker_query(
                 load.add_request(task.id, request, &worker.resources);
                 continue;
             }
-            add_task(&mut new_loads, task, &mut leftovers);
+            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, &mut leftovers);
+        add_task(&mut new_loads, task);
     }
 
     // `compute_new_worker_query` should be called immediately after scheduling was performed,
@@ -107,13 +102,16 @@ 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, &mut leftovers);
+        add_task(&mut new_loads, task);
     }
 
     let single_node_allocations = new_loads
         .iter()
         .zip(queries.iter())
         .map(|(ws, q)| {
+            if ws.partial {
+                return ws.loads.len();
+            }
             ws.loads
                 .iter()
                 .map(|load| {
@@ -148,28 +146,13 @@ pub(crate) fn compute_new_worker_query(
                     None
                 }
             });
-            if collect_leftovers && result.is_none() {
-                let request = ResourceRequestVariants::new(smallvec![rq.clone()]);
-                leftovers.insert(request, count);
-            }
             result
         })
         .collect();
     multi_node_allocations.sort_unstable_by_key(|x| (x.worker_type, x.worker_per_allocation));
 
-    let leftovers = if collect_leftovers {
-        let resource_map = core.create_resource_map();
-        leftovers
-            .into_iter()
-            .map(|(v, s)| (v.to_gateway(&resource_map), s))
-            .collect()
-    } else {
-        Vec::new()
-    };
-
     NewWorkerAllocationResponse {
         single_node_workers_per_query: single_node_allocations,
-        leftovers,
         multi_node_allocations,
     }
 }

crates/tako/src/internal/server/workerload.rs

@@ -31,6 +31,10 @@ impl WorkerResources {
             .unwrap_or(ResourceAmount::ZERO)
     }
 
+    pub(crate) fn get_or_none(&self, resource_id: ResourceId) -> Option<ResourceAmount> {
+        self.n_resources.get(resource_id).copied()
+    }
+
     pub(crate) fn from_description(
         resource_desc: &ResourceDescriptor,
         resource_map: &ResourceMap,
@@ -64,12 +68,33 @@ impl WorkerResources {
         })
     }
 
+    pub(crate) fn is_lowerbound_for_request(&self, request: &ResourceRequest) -> bool {
+        request.entries().iter().all(|r| {
+            if let Some(has) = self.get_or_none(r.resource_id) {
+                let ask = r.request.min_amount();
+                ask <= has
+            } else {
+                true
+            }
+        })
+    }
+
     pub(crate) fn is_capable_to_run(&self, rqv: &ResourceRequestVariants) -> bool {
         rqv.requests()
             .iter()
             .any(|rq| self.is_capable_to_run_request(rq))
     }
 
+    pub(crate) fn is_lowerbound_for(
+        &self,
+        rqv: &ResourceRequestVariants,
+        filter_fn: impl Fn(&ResourceRequest) -> bool,
+    ) -> bool {
+        rqv.requests()
+            .iter()
+            .any(|rq| filter_fn(rq) && self.is_lowerbound_for_request(rq))
+    }
+
     pub(crate) fn is_capable_to_run_with(
         &self,
         rqv: &ResourceRequestVariants,

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

@@ -115,12 +115,9 @@ async fn test_task_time_limit_pass() {
 
 fn query_helper(
     handler: &mut ServerHandle,
-    worker_queries: Vec<WorkerTypeQuery>,
+    worker_queries: &[WorkerTypeQuery],
 ) -> NewWorkerAllocationResponse {
-    handler
-        .server_ref
-        .new_worker_query(worker_queries, true)
-        .unwrap()
+    handler.server_ref.new_worker_query(worker_queries).unwrap()
 }
 
 #[tokio::test]
@@ -132,7 +129,8 @@ async fn test_query_no_output_immediate_call() {
             .await;
         let msg = query_helper(
             &mut handler,
-            vec![WorkerTypeQuery {
+            &[WorkerTypeQuery {
+                partial: false,
                 descriptor: ResourceDescriptor::simple(12),
                 time_limit: None,
                 max_sn_workers: 2,
@@ -142,7 +140,6 @@ 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.leftovers.is_empty());
         handler.wait(&ids).await;
     })
     .await;
@@ -158,7 +155,8 @@ async fn test_query_no_output_delayed_call() {
         sleep(Duration::from_secs(1)).await;
         let msg = query_helper(
             &mut handler,
-            vec![WorkerTypeQuery {
+            &[WorkerTypeQuery {
+                partial: false,
                 descriptor: ResourceDescriptor::simple(12),
                 time_limit: None,
                 max_sn_workers: 2,
@@ -168,7 +166,6 @@ 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.leftovers.is_empty());
         handler.wait(&ids).await;
     })
     .await;
@@ -187,7 +184,8 @@ async fn test_query_new_workers_delayed_call() {
         sleep(Duration::from_secs(1)).await;
         let msg = query_helper(
             &mut handler,
-            vec![WorkerTypeQuery {
+            &[WorkerTypeQuery {
+                partial: false,
                 descriptor: ResourceDescriptor::simple(12),
                 time_limit: None,
                 max_sn_workers: 2,
@@ -213,7 +211,8 @@ async fn test_query_new_workers_immediate() {
             .await;
         let msg = query_helper(
             &mut handler,
-            vec![WorkerTypeQuery {
+            &[WorkerTypeQuery {
+                partial: false,
                 descriptor: ResourceDescriptor::simple(12),
                 time_limit: None,
                 max_sn_workers: 2,