new_worker_query is no longer async

Author: spirali

crates/tako/src/control.rs

@@ -1,17 +1,17 @@
 use std::future::Future;
 use std::rc::Rc;
 use std::sync::Arc;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 use orion::aead::SecretKey;
 use tokio::net::TcpListener;
-use tokio::sync::{Notify, oneshot};
+use tokio::sync::Notify;
 
 use crate::events::EventProcessor;
 use crate::gateway::{MultiNodeAllocationResponse, TaskSubmit, WorkerRuntimeInfo};
 use crate::internal::messages::worker::ToWorkerMessage;
 use crate::internal::scheduler::query::compute_new_worker_query;
-use crate::internal::scheduler::state::scheduler_loop;
+use crate::internal::scheduler::state::{run_scheduling_now, scheduler_loop};
 use crate::internal::server::client::handle_new_tasks;
 use crate::internal::server::comm::{Comm, CommSenderRef};
 use crate::internal::server::core::{CoreRef, CustomConnectionHandler};
@@ -99,26 +99,24 @@ impl ServerRef {
       we can get in one allocation at most.
       This is used for planning multi-node tasks.
 
+      Note: This call may immediately call the full scheduler procedure.
+      This should not bother the user of the call, except it is probably not a good
+      idea to call this function often (several times per second) as it may bypass
+      a scheduler time limitations.
     */
     pub fn new_worker_query(
         &self,
         queries: Vec<WorkerTypeQuery>,
-    ) -> crate::Result<oneshot::Receiver<NewWorkerAllocationResponse>> {
+    ) -> crate::Result<NewWorkerAllocationResponse> {
         for query in &queries {
             query.descriptor.validate()?;
         }
-        let (sx, rx) = tokio::sync::oneshot::channel();
-        if self.comm_ref.get().get_scheduling_flag() {
-            self.comm_ref
-                .get_mut()
-                .add_after_scheduling_callback(Box::new(move |core| {
-                    let _ = sx.send(compute_new_worker_query(core, &queries));
-                }));
-        } else {
-            let mut core = self.core_ref.get_mut();
-            let _ = sx.send(compute_new_worker_query(&mut core, &queries));
-        };
-        Ok(rx)
+        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))
     }
 
     pub fn try_release_memory(&self) {

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

@@ -12,13 +12,16 @@ struct WorkerTypeState {
     max: u32,
 }
 
-/* Read the documentation of `new_worker_query`` in control.rs */
+/// Read the documentation of `new_worker_query`` in control.rs
 pub(crate) fn compute_new_worker_query(
     core: &mut Core,
     queries: &[WorkerTypeQuery],
 ) -> 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| {
         let request = &task.configuration.resources;
         for ws in new_loads.iter_mut() {

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

@@ -8,7 +8,7 @@ use tokio::time::sleep;
 use crate::internal::common::Map;
 use crate::internal::messages::worker::{TaskIdsMsg, ToWorkerMessage};
 use crate::internal::scheduler::multinode::MultiNodeAllocator;
-use crate::internal::server::comm::{Comm, CommSenderRef};
+use crate::internal::server::comm::{Comm, CommSender, CommSenderRef};
 use crate::internal::server::core::{Core, CoreRef};
 use crate::internal::server::dataobjmap::DataObjectMap;
 use crate::internal::server::task::{Task, TaskRuntimeState};
@@ -43,22 +43,33 @@ pub(crate) async fn scheduler_loop(
     let mut last_schedule = Instant::now().checked_sub(minimum_delay * 2).unwrap();
     loop {
         scheduler_wakeup.notified().await;
+        if !comm_ref.get().get_scheduling_flag() {
+            last_schedule = Instant::now();
+            continue;
+        }
         let mut now = Instant::now();
         let since_last_schedule = now - last_schedule;
         if minimum_delay > since_last_schedule {
             sleep(minimum_delay - since_last_schedule).await;
             now = Instant::now();
         }
         let mut comm = comm_ref.get_mut();
-        let mut state = SchedulerState::new(now);
+        if !comm.get_scheduling_flag() {
+            last_schedule = now;
+            continue;
+        }
         let mut core = core_ref.get_mut();
-        state.run_scheduling(&mut core, &mut *comm);
-        comm.reset_scheduling_flag();
+        run_scheduling_now(&mut core, &mut comm, now);
         last_schedule = Instant::now();
-        comm.call_after_scheduling_callbacks(&mut core);
     }
 }
 
+pub fn run_scheduling_now(core: &mut Core, comm: &mut CommSender, now: Instant) {
+    let mut state = SchedulerState::new(now);
+    state.run_scheduling(core, &mut *comm);
+    comm.reset_scheduling_flag();
+}
+
 impl SchedulerState {
     pub fn new(now: std::time::Instant) -> Self {
         SchedulerState {

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

@@ -8,7 +8,6 @@ use crate::WorkerId;
 use crate::events::EventProcessor;
 use crate::internal::common::{Map, WrappedRcRefCell};
 use crate::internal::messages::worker::ToWorkerMessage;
-use crate::internal::server::core::Core;
 use crate::internal::transfer::auth::serialize;
 
 pub trait Comm {
@@ -19,14 +18,11 @@ pub trait Comm {
     fn client(&mut self) -> &mut dyn EventProcessor;
 }
 
-type SchedulingCallback = Box<dyn FnOnce(&mut Core)>;
-
 pub struct CommSender {
     workers: Map<WorkerId, UnboundedSender<Bytes>>,
     need_scheduling: bool,
     scheduler_wakeup: Rc<Notify>,
     client_events: Option<Box<dyn EventProcessor>>,
-    after_scheduling_callbacks: Vec<SchedulingCallback>,
     panic_on_worker_lost: bool,
 }
 
@@ -38,7 +34,6 @@ impl CommSenderRef {
             workers: Default::default(),
             scheduler_wakeup,
             client_events: None,
-            after_scheduling_callbacks: Vec::new(),
             need_scheduling: false,
             panic_on_worker_lost,
         })
@@ -70,19 +65,6 @@ impl CommSender {
     pub fn get_scheduling_flag(&self) -> bool {
         self.need_scheduling
     }
-
-    pub fn add_after_scheduling_callback(&mut self, callback: SchedulingCallback) {
-        self.after_scheduling_callbacks.push(callback)
-    }
-
-    pub fn call_after_scheduling_callbacks(&mut self, core: &mut Core) {
-        if !self.after_scheduling_callbacks.is_empty() {
-            log::debug!("Running after scheduling callbacks");
-            self.after_scheduling_callbacks
-                .drain(..)
-                .for_each(|x| x(core))
-        }
-    }
 }
 
 impl Comm for CommSender {

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

@@ -113,12 +113,11 @@ async fn test_task_time_limit_pass() {
     .await;
 }
 
-async fn query_helper(
+fn query_helper(
     handler: &mut ServerHandle,
     worker_queries: Vec<WorkerTypeQuery>,
 ) -> NewWorkerAllocationResponse {
-    let result = handler.server_ref.new_worker_query(worker_queries).unwrap();
-    result.await.unwrap()
+    handler.server_ref.new_worker_query(worker_queries).unwrap()
 }
 
 #[tokio::test]
@@ -137,8 +136,7 @@ async fn test_query_no_output_immediate_call() {
                 max_worker_per_allocation: 2,
                 min_utilization: 0.0,
             }],
-        )
-        .await;
+        );
         assert_eq!(msg.single_node_allocations, vec![0]);
         assert!(msg.multi_node_allocations.is_empty());
         handler.wait(&ids).await;
@@ -163,8 +161,7 @@ async fn test_query_no_output_delayed_call() {
                 max_worker_per_allocation: 2,
                 min_utilization: 0.0,
             }],
-        )
-        .await;
+        );
         assert_eq!(msg.single_node_allocations, vec![0]);
         assert!(msg.multi_node_allocations.is_empty());
         handler.wait(&ids).await;
@@ -192,8 +189,7 @@ async fn test_query_new_workers_delayed_call() {
                 max_worker_per_allocation: 2,
                 min_utilization: 0.0,
             }],
-        )
-        .await;
+        );
         assert_eq!(msg.single_node_allocations, vec![1]);
         assert!(msg.multi_node_allocations.is_empty());
     })
@@ -219,8 +215,7 @@ async fn test_query_new_workers_immediate() {
                 max_worker_per_allocation: 2,
                 min_utilization: 0.0,
             }],
-        )
-        .await;
+        );
         assert_eq!(msg.single_node_allocations, vec![1]);
         assert!(msg.multi_node_allocations.is_empty());
     })